spdlog/details/tcp_client-windows.h

0001 // Copyright(c) 2015-present, Gabi Melman & spdlog contributors.
0002 // Distributed under the MIT License (http://opensource.org/licenses/MIT)
0003
0004 #pragma once
0005
0006 #define WIN32_LEAN_AND_MEAN
0007 // tcp client helper
0008 #include <spdlog/common.h>
0009 #include <spdlog/details/os.h>
0010
0011 #include <winsock2.h>
0012 #include <windows.h>
0013 #include <ws2tcpip.h>
0014 #include <stdlib.h>
0015 #include <stdio.h>
0016 #include <string>
0017
0018 #pragma comment(lib, "Ws2_32.lib")
0019 #pragma comment(lib, "Mswsock.lib")
0020 #pragma comment(lib, "AdvApi32.lib")
0021
0022 namespace spdlog {
0023 namespace details {
0024 class tcp_client
0025 {
0026     SOCKET socket_ = INVALID_SOCKET;
0027
0028     static void init_winsock_()
0029     {
0030         WSADATA wsaData;
0031         auto rv = WSAStartup(MAKEWORD(2, 2), &wsaData);
0032         if (rv != 0)
0033         {
0034             throw_winsock_error_("WSAStartup failed", ::WSAGetLastError());
0035         }
0036     }
0037
0038     static void throw_winsock_error_(const std::string &msg, int last_error)
0039     {
0040         char buf[512];
0041         ::FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, NULL, last_error,
0042             MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), buf, (sizeof(buf) / sizeof(char)), NULL);
0043
0044         throw_spdlog_ex(fmt_lib::format("tcp_sink - {}: {}", msg, buf));
0045     }
0046
0047 public:
0048     tcp_client()
0049     {
0050         init_winsock_();
0051     }
0052
0053     ~tcp_client()
0054     {
0055         close();
0056         ::WSACleanup();
0057     }
0058
0059     bool is_connected() const
0060     {
0061         return socket_ != INVALID_SOCKET;
0062     }
0063
0064     void close()
0065     {
0066         ::closesocket(socket_);
0067         socket_ = INVALID_SOCKET;
0068     }
0069
0070     SOCKET fd() const
0071     {
0072         return socket_;
0073     }
0074
0075     // try to connect or throw on failure
0076     void connect(const std::string &host, int port)
0077     {
0078         if (is_connected())
0079         {
0080             close();
0081         }
0082         struct addrinfo hints
0083         {};
0084         ZeroMemory(&hints, sizeof(hints));
0085
0086         hints.ai_family = AF_INET;       // IPv4
0087         hints.ai_socktype = SOCK_STREAM; // TCP
0088         hints.ai_flags = AI_NUMERICSERV; // port passed as as numeric value
0089         hints.ai_protocol = 0;
0090
0091         auto port_str = std::to_string(port);
0092         struct addrinfo *addrinfo_result;
0093         auto rv = ::getaddrinfo(host.c_str(), port_str.c_str(), &hints, &addrinfo_result);
0094         int last_error = 0;
0095         if (rv != 0)
0096         {
0097             last_error = ::WSAGetLastError();
0098             WSACleanup();
0099             throw_winsock_error_("getaddrinfo failed", last_error);
0100         }
0101
0102         // Try each address until we successfully connect(2).
0103
0104         for (auto *rp = addrinfo_result; rp != nullptr; rp = rp->ai_next)
0105         {
0106             socket_ = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
0107             if (socket_ == INVALID_SOCKET)
0108             {
0109                 last_error = ::WSAGetLastError();
0110                 WSACleanup();
0111                 continue;
0112             }
0113             if (::connect(socket_, rp->ai_addr, (int)rp->ai_addrlen) == 0)
0114             {
0115                 break;
0116             }
0117             else
0118             {
0119                 last_error = ::WSAGetLastError();
0120                 close();
0121             }
0122         }
0123         ::freeaddrinfo(addrinfo_result);
0124         if (socket_ == INVALID_SOCKET)
0125         {
0126             WSACleanup();
0127             throw_winsock_error_("connect failed", last_error);
0128         }
0129
0130         // set TCP_NODELAY
0131         int enable_flag = 1;
0132         ::setsockopt(socket_, IPPROTO_TCP, TCP_NODELAY, reinterpret_cast<char *>(&enable_flag), sizeof(enable_flag));
0133     }
0134
0135     // Send exactly n_bytes of the given data.
0136     // On error close the connection and throw.
0137     void send(const char *data, size_t n_bytes)
0138     {
0139         size_t bytes_sent = 0;
0140         while (bytes_sent < n_bytes)
0141         {
0142             const int send_flags = 0;
0143             auto write_result = ::send(socket_, data + bytes_sent, (int)(n_bytes - bytes_sent), send_flags);
0144             if (write_result == SOCKET_ERROR)
0145             {
0146                 int last_error = ::WSAGetLastError();
0147                 close();
0148                 throw_winsock_error_("send failed", last_error);
0149             }
0150
0151             if (write_result == 0) // (probably should not happen but in any case..)
0152             {
0153                 break;
0154             }
0155             bytes_sent += static_cast<size_t>(write_result);
0156         }
0157     }
0158 };
0159 } // namespace details
0160 } // namespace spdlog