tcp_client-windows.h 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160
  1. // Copyright(c) 2015-present, Gabi Melman & spdlog contributors.
  2. // Distributed under the MIT License (http://opensource.org/licenses/MIT)
  3. #pragma once
  4. #define WIN32_LEAN_AND_MEAN
  5. // tcp client helper
  6. #include <spdlog/common.h>
  7. #include <spdlog/details/os.h>
  8. #include <winsock2.h>
  9. #include <windows.h>
  10. #include <ws2tcpip.h>
  11. #include <stdlib.h>
  12. #include <stdio.h>
  13. #include <string>
  14. #pragma comment(lib, "Ws2_32.lib")
  15. #pragma comment(lib, "Mswsock.lib")
  16. #pragma comment(lib, "AdvApi32.lib")
  17. namespace spdlog {
  18. namespace details {
  19. class tcp_client
  20. {
  21. SOCKET socket_ = INVALID_SOCKET;
  22. static void init_winsock_()
  23. {
  24. WSADATA wsaData;
  25. auto rv = WSAStartup(MAKEWORD(2, 2), &wsaData);
  26. if (rv != 0)
  27. {
  28. throw_winsock_error_("WSAStartup failed", ::WSAGetLastError());
  29. }
  30. }
  31. static void throw_winsock_error_(const std::string &msg, int last_error)
  32. {
  33. char buf[512];
  34. ::FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, NULL, last_error,
  35. MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), buf, (sizeof(buf) / sizeof(char)), NULL);
  36. throw_spdlog_ex(fmt_lib::format("tcp_sink - {}: {}", msg, buf));
  37. }
  38. public:
  39. tcp_client()
  40. {
  41. init_winsock_();
  42. }
  43. ~tcp_client()
  44. {
  45. close();
  46. ::WSACleanup();
  47. }
  48. bool is_connected() const
  49. {
  50. return socket_ != INVALID_SOCKET;
  51. }
  52. void close()
  53. {
  54. ::closesocket(socket_);
  55. socket_ = INVALID_SOCKET;
  56. }
  57. SOCKET fd() const
  58. {
  59. return socket_;
  60. }
  61. // try to connect or throw on failure
  62. void connect(const std::string &host, int port)
  63. {
  64. if (is_connected())
  65. {
  66. close();
  67. }
  68. struct addrinfo hints
  69. {};
  70. ZeroMemory(&hints, sizeof(hints));
  71. hints.ai_family = AF_UNSPEC; // To work with IPv4, IPv6, and so on
  72. hints.ai_socktype = SOCK_STREAM; // TCP
  73. hints.ai_flags = AI_NUMERICSERV; // port passed as as numeric value
  74. hints.ai_protocol = 0;
  75. auto port_str = std::to_string(port);
  76. struct addrinfo *addrinfo_result;
  77. auto rv = ::getaddrinfo(host.c_str(), port_str.c_str(), &hints, &addrinfo_result);
  78. int last_error = 0;
  79. if (rv != 0)
  80. {
  81. last_error = ::WSAGetLastError();
  82. WSACleanup();
  83. throw_winsock_error_("getaddrinfo failed", last_error);
  84. }
  85. // Try each address until we successfully connect(2).
  86. for (auto *rp = addrinfo_result; rp != nullptr; rp = rp->ai_next)
  87. {
  88. socket_ = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
  89. if (socket_ == INVALID_SOCKET)
  90. {
  91. last_error = ::WSAGetLastError();
  92. WSACleanup();
  93. continue;
  94. }
  95. if (::connect(socket_, rp->ai_addr, (int)rp->ai_addrlen) == 0)
  96. {
  97. break;
  98. }
  99. else
  100. {
  101. last_error = ::WSAGetLastError();
  102. close();
  103. }
  104. }
  105. ::freeaddrinfo(addrinfo_result);
  106. if (socket_ == INVALID_SOCKET)
  107. {
  108. WSACleanup();
  109. throw_winsock_error_("connect failed", last_error);
  110. }
  111. // set TCP_NODELAY
  112. int enable_flag = 1;
  113. ::setsockopt(socket_, IPPROTO_TCP, TCP_NODELAY, reinterpret_cast<char *>(&enable_flag), sizeof(enable_flag));
  114. }
  115. // Send exactly n_bytes of the given data.
  116. // On error close the connection and throw.
  117. void send(const char *data, size_t n_bytes)
  118. {
  119. size_t bytes_sent = 0;
  120. while (bytes_sent < n_bytes)
  121. {
  122. const int send_flags = 0;
  123. auto write_result = ::send(socket_, data + bytes_sent, (int)(n_bytes - bytes_sent), send_flags);
  124. if (write_result == SOCKET_ERROR)
  125. {
  126. int last_error = ::WSAGetLastError();
  127. close();
  128. throw_winsock_error_("send failed", last_error);
  129. }
  130. if (write_result == 0) // (probably should not happen but in any case..)
  131. {
  132. break;
  133. }
  134. bytes_sent += static_cast<size_t>(write_result);
  135. }
  136. }
  137. };
  138. } // namespace details
  139. } // namespace spdlog