tcp_client.h 3.8 KB

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