tcp_client.h 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146
  1. // Copyright(c) 2015-present, Gabi Melman & spdlog contributors.
  2. // Distributed under the MIT License (http://opensource.org/licenses/MIT)
  3. #pragma once
  4. #ifdef _WIN32
  5. # error include tcp_client-windows.h instead
  6. #endif
  7. // tcp client helper
  8. #include <spdlog/common.h>
  9. #include <spdlog/details/os.h>
  10. #include <sys/socket.h>
  11. #include <arpa/inet.h>
  12. #include <unistd.h>
  13. #include <netdb.h>
  14. #include <netinet/tcp.h>
  15. #include <string>
  16. namespace spdlog {
  17. namespace details {
  18. class tcp_client
  19. {
  20. int socket_ = -1;
  21. public:
  22. bool is_connected() const
  23. {
  24. return socket_ != -1;
  25. }
  26. void close()
  27. {
  28. if (is_connected())
  29. {
  30. ::close(socket_);
  31. socket_ = -1;
  32. }
  33. }
  34. int fd() const
  35. {
  36. return socket_;
  37. }
  38. ~tcp_client()
  39. {
  40. close();
  41. }
  42. // try to connect or throw on failure
  43. void connect(const std::string &host, int port)
  44. {
  45. close();
  46. struct addrinfo hints
  47. {};
  48. memset(&hints, 0, sizeof(struct addrinfo));
  49. hints.ai_family = AF_INET; // IPv4
  50. hints.ai_socktype = SOCK_STREAM; // TCP
  51. hints.ai_flags = AI_NUMERICSERV; // port passed as as numeric value
  52. hints.ai_protocol = 0;
  53. auto port_str = std::to_string(port);
  54. struct addrinfo *addrinfo_result;
  55. auto rv = ::getaddrinfo(host.c_str(), port_str.c_str(), &hints, &addrinfo_result);
  56. if (rv != 0)
  57. {
  58. throw_spdlog_ex(fmt_lib::format("::getaddrinfo failed: {}", gai_strerror(rv)));
  59. }
  60. // Try each address until we successfully connect(2).
  61. int last_errno = 0;
  62. for (auto *rp = addrinfo_result; rp != nullptr; rp = rp->ai_next)
  63. {
  64. #if defined(SOCK_CLOEXEC)
  65. const int flags = SOCK_CLOEXEC;
  66. #else
  67. const int flags = 0;
  68. #endif
  69. socket_ = ::socket(rp->ai_family, rp->ai_socktype | flags, rp->ai_protocol);
  70. if (socket_ == -1)
  71. {
  72. last_errno = errno;
  73. continue;
  74. }
  75. rv = ::connect(socket_, rp->ai_addr, rp->ai_addrlen);
  76. if (rv == 0)
  77. {
  78. break;
  79. }
  80. last_errno = errno;
  81. ::close(socket_);
  82. socket_ = -1;
  83. }
  84. ::freeaddrinfo(addrinfo_result);
  85. if (socket_ == -1)
  86. {
  87. throw_spdlog_ex("::connect failed", last_errno);
  88. }
  89. // set TCP_NODELAY
  90. int enable_flag = 1;
  91. ::setsockopt(socket_, IPPROTO_TCP, TCP_NODELAY, reinterpret_cast<char *>(&enable_flag), sizeof(enable_flag));
  92. // prevent sigpipe on systems where MSG_NOSIGNAL is not available
  93. #if defined(SO_NOSIGPIPE) && !defined(MSG_NOSIGNAL)
  94. ::setsockopt(socket_, SOL_SOCKET, SO_NOSIGPIPE, reinterpret_cast<char *>(&enable_flag), sizeof(enable_flag));
  95. #endif
  96. #if !defined(SO_NOSIGPIPE) && !defined(MSG_NOSIGNAL)
  97. # error "tcp_sink would raise SIGPIPE since niether SO_NOSIGPIPE nor MSG_NOSIGNAL are available"
  98. #endif
  99. }
  100. // Send exactly n_bytes of the given data.
  101. // On error close the connection and throw.
  102. void send(const char *data, size_t n_bytes)
  103. {
  104. size_t bytes_sent = 0;
  105. while (bytes_sent < n_bytes)
  106. {
  107. #if defined(MSG_NOSIGNAL)
  108. const int send_flags = MSG_NOSIGNAL;
  109. #else
  110. const int send_flags = 0;
  111. #endif
  112. auto write_result = ::send(socket_, data + bytes_sent, n_bytes - bytes_sent, send_flags);
  113. if (write_result < 0)
  114. {
  115. close();
  116. throw_spdlog_ex("write(2) failed", errno);
  117. }
  118. if (write_result == 0) // (probably should not happen but in any case..)
  119. {
  120. break;
  121. }
  122. bytes_sent += static_cast<size_t>(write_result);
  123. }
  124. }
  125. };
  126. } // namespace details
  127. } // namespace spdlog