deps/spdlog/include/spdlog/details/tcp_client-windows.h

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