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