llvm/lib/Support/raw_socket_stream.cpp

   1 //===-- llvm/Support/raw_socket_stream.cpp - Socket streams --*- C++ -*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file contains raw_ostream implementations for streams to communicate
  10 // via UNIX sockets
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/Support/raw_socket_stream.h"
  15 #include "llvm/Config/config.h"
  16 #include "llvm/Support/Error.h"
  17
  18 #ifndef _WIN32
  19 #include <sys/socket.h>
  20 #include <sys/un.h>
  21 #else
  22 #include "llvm/Support/Windows/WindowsSupport.h"
  23 // winsock2.h must be included before afunix.h. Briefly turn off clang-format to
  24 // avoid error.
  25 // clang-format off
  26 #include <winsock2.h>
  27 #include <afunix.h>
  28 // clang-format on
  29 #include <io.h>
  30 #endif // _WIN32
  31
  32 #if defined(HAVE_UNISTD_H)
  33 #include <unistd.h>
  34 #endif
  35
  36 using namespace llvm;
  37
  38 #ifdef _WIN32
  39 WSABalancer::WSABalancer() {
  40   WSADATA WsaData;
  41   ::memset(&WsaData, 0, sizeof(WsaData));
  42   if (WSAStartup(MAKEWORD(2, 2), &WsaData) != 0) {
  43     llvm::report_fatal_error("WSAStartup failed");
  44   }
  45 }
  46
  47 WSABalancer::~WSABalancer() { WSACleanup(); }
  48
  49 #endif // _WIN32
  50
  51 static std::error_code getLastSocketErrorCode() {
  52 #ifdef _WIN32
  53   return std::error_code(::WSAGetLastError(), std::system_category());
  54 #else
  55   return std::error_code(errno, std::system_category());
  56 #endif
  57 }
  58
  59 ListeningSocket::ListeningSocket(int SocketFD, StringRef SocketPath)
  60     : FD(SocketFD), SocketPath(SocketPath) {}
  61
  62 ListeningSocket::ListeningSocket(ListeningSocket &&LS)
  63     : FD(LS.FD), SocketPath(LS.SocketPath) {
  64   LS.FD = -1;
  65 }
  66
  67 Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath,
  68                                                       int MaxBacklog) {
  69
  70 #ifdef _WIN32
  71   WSABalancer _;
  72   SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
  73   if (MaybeWinsocket == INVALID_SOCKET) {
  74 #else
  75   int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
  76   if (MaybeWinsocket == -1) {
  77 #endif
  78     return llvm::make_error<StringError>(getLastSocketErrorCode(),
  79                                          "socket create failed");
  80   }
  81
  82   struct sockaddr_un Addr;
  83   memset(&Addr, 0, sizeof(Addr));
  84   Addr.sun_family = AF_UNIX;
  85   strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
  86
  87   if (bind(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1) {
  88     std::error_code Err = getLastSocketErrorCode();
  89     if (Err == std::errc::address_in_use)
  90       ::close(MaybeWinsocket);
  91     return llvm::make_error<StringError>(Err, "Bind error");
  92   }
  93   if (listen(MaybeWinsocket, MaxBacklog) == -1) {
  94     return llvm::make_error<StringError>(getLastSocketErrorCode(),
  95                                          "Listen error");
  96   }
  97   int UnixSocket;
  98 #ifdef _WIN32
  99   UnixSocket = _open_osfhandle(MaybeWinsocket, 0);
 100 #else
 101   UnixSocket = MaybeWinsocket;
 102 #endif // _WIN32
 103   return ListeningSocket{UnixSocket, SocketPath};
 104 }
 105
 106 Expected<std::unique_ptr<raw_socket_stream>> ListeningSocket::accept() {
 107   int AcceptFD;
 108 #ifdef _WIN32
 109   SOCKET WinServerSock = _get_osfhandle(FD);
 110   SOCKET WinAcceptSock = ::accept(WinServerSock, NULL, NULL);
 111   AcceptFD = _open_osfhandle(WinAcceptSock, 0);
 112 #else
 113   AcceptFD = ::accept(FD, NULL, NULL);
 114 #endif //_WIN32
 115   if (AcceptFD == -1)
 116     return llvm::make_error<StringError>(getLastSocketErrorCode(),
 117                                          "Accept failed");
 118   return std::make_unique<raw_socket_stream>(AcceptFD);
 119 }
 120
 121 ListeningSocket::~ListeningSocket() {
 122   if (FD == -1)
 123     return;
 124   ::close(FD);
 125   unlink(SocketPath.c_str());
 126 }
 127
 128 static Expected<int> GetSocketFD(StringRef SocketPath) {
 129 #ifdef _WIN32
 130   SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
 131   if (MaybeWinsocket == INVALID_SOCKET) {
 132 #else
 133   int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
 134   if (MaybeWinsocket == -1) {
 135 #endif // _WIN32
 136     return llvm::make_error<StringError>(getLastSocketErrorCode(),
 137                                          "Create socket failed");
 138   }
 139
 140   struct sockaddr_un Addr;
 141   memset(&Addr, 0, sizeof(Addr));
 142   Addr.sun_family = AF_UNIX;
 143   strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
 144
 145   int status = connect(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr));
 146   if (status == -1) {
 147     return llvm::make_error<StringError>(getLastSocketErrorCode(),
 148                                          "Connect socket failed");
 149   }
 150 #ifdef _WIN32
 151   return _open_osfhandle(MaybeWinsocket, 0);
 152 #else
 153   return MaybeWinsocket;
 154 #endif // _WIN32
 155 }
 156
 157 raw_socket_stream::raw_socket_stream(int SocketFD)
 158     : raw_fd_stream(SocketFD, true) {}
 159
 160 Expected<std::unique_ptr<raw_socket_stream>>
 161 raw_socket_stream::createConnectedUnix(StringRef SocketPath) {
 162 #ifdef _WIN32
 163   WSABalancer _;
 164 #endif // _WIN32
 165   Expected<int> FD = GetSocketFD(SocketPath);
 166   if (!FD)
 167     return FD.takeError();
 168   return std::make_unique<raw_socket_stream>(*FD);
 169 }
 170
 171 raw_socket_stream::~raw_socket_stream() {}
 172
 173 //===----------------------------------------------------------------------===//
 174 //  raw_string_ostream
 175 //===----------------------------------------------------------------------===//
 176
 177 void raw_string_ostream::write_impl(const char *Ptr, size_t Size) {
 178   OS.append(Ptr, Size);
 179 }