tcp_client-windows.h 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  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 <stdio.h>
  9. #include <stdlib.h>
  10. #include <string>
  11. #include <windows.h>
  12. #include <winsock2.h>
  13. #include <ws2tcpip.h>
  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. SOCKET socket_ = INVALID_SOCKET;
  21. static void init_winsock_() {
  22. WSADATA wsaData;
  23. auto rv = WSAStartup(MAKEWORD(2, 2), &wsaData);
  24. if (rv != 0) {
  25. throw_winsock_error_("WSAStartup failed", ::WSAGetLastError());
  26. }
  27. }
  28. static void throw_winsock_error_(const std::string &msg, int last_error) {
  29. char buf[512];
  30. ::FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, NULL,
  31. last_error, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), buf,
  32. (sizeof(buf) / sizeof(char)), NULL);
  33. throw_spdlog_ex(fmt_lib::format("tcp_sink - {}: {}", msg, buf));
  34. }
  35. public:
  36. tcp_client() { init_winsock_(); }
  37. ~tcp_client() {
  38. close();
  39. ::WSACleanup();
  40. }
  41. bool is_connected() const { return socket_ != INVALID_SOCKET; }
  42. void close() {
  43. ::closesocket(socket_);
  44. socket_ = INVALID_SOCKET;
  45. }
  46. SOCKET fd() const { return socket_; }
  47. // try to connect or throw on failure
  48. void connect(const std::string &host, int port) {
  49. if (is_connected()) {
  50. close();
  51. }
  52. struct addrinfo hints {};
  53. ZeroMemory(&hints, sizeof(hints));
  54. hints.ai_family = AF_UNSPEC; // To work with IPv4, IPv6, and so on
  55. hints.ai_socktype = SOCK_STREAM; // TCP
  56. hints.ai_flags = AI_NUMERICSERV; // port passed as as numeric value
  57. hints.ai_protocol = 0;
  58. auto port_str = std::to_string(port);
  59. struct addrinfo *addrinfo_result;
  60. auto rv = ::getaddrinfo(host.c_str(), port_str.c_str(), &hints, &addrinfo_result);
  61. int last_error = 0;
  62. if (rv != 0) {
  63. last_error = ::WSAGetLastError();
  64. WSACleanup();
  65. throw_winsock_error_("getaddrinfo failed", last_error);
  66. }
  67. // Try each address until we successfully connect(2).
  68. for (auto *rp = addrinfo_result; rp != nullptr; rp = rp->ai_next) {
  69. socket_ = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
  70. if (socket_ == INVALID_SOCKET) {
  71. last_error = ::WSAGetLastError();
  72. WSACleanup();
  73. continue;
  74. }
  75. if (::connect(socket_, rp->ai_addr, (int)rp->ai_addrlen) == 0) {
  76. break;
  77. } else {
  78. last_error = ::WSAGetLastError();
  79. close();
  80. }
  81. }
  82. ::freeaddrinfo(addrinfo_result);
  83. if (socket_ == INVALID_SOCKET) {
  84. WSACleanup();
  85. throw_winsock_error_("connect failed", last_error);
  86. }
  87. // set TCP_NODELAY
  88. int enable_flag = 1;
  89. ::setsockopt(socket_, IPPROTO_TCP, TCP_NODELAY, reinterpret_cast<char *>(&enable_flag),
  90. sizeof(enable_flag));
  91. }
  92. // Send exactly n_bytes of the given data.
  93. // On error close the connection and throw.
  94. void send(const char *data, size_t n_bytes) {
  95. size_t bytes_sent = 0;
  96. while (bytes_sent < n_bytes) {
  97. const int send_flags = 0;
  98. auto write_result =
  99. ::send(socket_, data + bytes_sent, (int)(n_bytes - bytes_sent), send_flags);
  100. if (write_result == SOCKET_ERROR) {
  101. int last_error = ::WSAGetLastError();
  102. close();
  103. throw_winsock_error_("send failed", last_error);
  104. }
  105. if (write_result == 0) // (probably should not happen but in any case..)
  106. {
  107. break;
  108. }
  109. bytes_sent += static_cast<size_t>(write_result);
  110. }
  111. }
  112. };
  113. } // namespace details
  114. } // namespace spdlog