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 #include "llvm/Support/FileSystem.h"
  18
  19 #include <atomic>
  20 #include <fcntl.h>
  21 #include <functional>
  22
  23 #ifndef _WIN32
  24 #include <poll.h>
  25 #include <sys/socket.h>
  26 #include <sys/un.h>
  27 #else
  28 #include "llvm/Support/Windows/WindowsSupport.h"
  29 // winsock2.h must be included before afunix.h. Briefly turn off clang-format to
  30 // avoid error.
  31 // clang-format off
  32 #include <winsock2.h>
  33 #include <afunix.h>
  34 // clang-format on
  35 #include <io.h>
  36 #endif // _WIN32
  37
  38 #if defined(HAVE_UNISTD_H)
  39 #include <unistd.h>
  40 #endif
  41
  42 using namespace llvm;
  43
  44 #ifdef _WIN32
  45 WSABalancer::WSABalancer() {
  46   WSADATA WsaData;
  47   ::memset(&WsaData, 0, sizeof(WsaData));
  48   if (WSAStartup(MAKEWORD(2, 2), &WsaData) != 0) {
  49     llvm::report_fatal_error("WSAStartup failed");
  50   }
  51 }
  52
  53 WSABalancer::~WSABalancer() { WSACleanup(); }
  54 #endif // _WIN32
  55
  56 static std::error_code getLastSocketErrorCode() {
  57 #ifdef _WIN32
  58   return std::error_code(::WSAGetLastError(), std::system_category());
  59 #else
  60   return errnoAsErrorCode();
  61 #endif
  62 }
  63
  64 static sockaddr_un setSocketAddr(StringRef SocketPath) {
  65   struct sockaddr_un Addr;
  66   memset(&Addr, 0, sizeof(Addr));
  67   Addr.sun_family = AF_UNIX;
  68   strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
  69   return Addr;
  70 }
  71
  72 static Expected<int> getSocketFD(StringRef SocketPath) {
  73 #ifdef _WIN32
  74   SOCKET Socket = socket(AF_UNIX, SOCK_STREAM, 0);
  75   if (Socket == INVALID_SOCKET) {
  76 #else
  77   int Socket = socket(AF_UNIX, SOCK_STREAM, 0);
  78   if (Socket == -1) {
  79 #endif // _WIN32
  80     return llvm::make_error<StringError>(getLastSocketErrorCode(),
  81                                          "Create socket failed");
  82   }
  83
  84   struct sockaddr_un Addr = setSocketAddr(SocketPath);
  85   if (::connect(Socket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1)
  86     return llvm::make_error<StringError>(getLastSocketErrorCode(),
  87                                          "Connect socket failed");
  88
  89 #ifdef _WIN32
  90   return _open_osfhandle(Socket, 0);
  91 #else
  92   return Socket;
  93 #endif // _WIN32
  94 }
  95
  96 ListeningSocket::ListeningSocket(int SocketFD, StringRef SocketPath,
  97                                  int PipeFD[2])
  98     : FD(SocketFD), SocketPath(SocketPath), PipeFD{PipeFD[0], PipeFD[1]} {}
  99
 100 ListeningSocket::ListeningSocket(ListeningSocket &&LS)
 101     : FD(LS.FD.load()), SocketPath(LS.SocketPath),
 102       PipeFD{LS.PipeFD[0], LS.PipeFD[1]} {
 103
 104   LS.FD = -1;
 105   LS.SocketPath.clear();
 106   LS.PipeFD[0] = -1;
 107   LS.PipeFD[1] = -1;
 108 }
 109
 110 Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath,
 111                                                       int MaxBacklog) {
 112
 113   // Handle instances where the target socket address already exists and
 114   // differentiate between a preexisting file with and without a bound socket
 115   //
 116   // ::bind will return std::errc:address_in_use if a file at the socket address
 117   // already exists (e.g., the file was not properly unlinked due to a crash)
 118   // even if another socket has not yet binded to that address
 119   if (llvm::sys::fs::exists(SocketPath)) {
 120     Expected<int> MaybeFD = getSocketFD(SocketPath);
 121     if (!MaybeFD) {
 122
 123       // Regardless of the error, notify the caller that a file already exists
 124       // at the desired socket address and that there is no bound socket at that
 125       // address. The file must be removed before ::bind can use the address
 126       consumeError(MaybeFD.takeError());
 127       return llvm::make_error<StringError>(
 128           std::make_error_code(std::errc::file_exists),
 129           "Socket address unavailable");
 130     }
 131     ::close(std::move(*MaybeFD));
 132
 133     // Notify caller that the provided socket address already has a bound socket
 134     return llvm::make_error<StringError>(
 135         std::make_error_code(std::errc::address_in_use),
 136         "Socket address unavailable");
 137   }
 138
 139 #ifdef _WIN32
 140   WSABalancer _;
 141   SOCKET Socket = socket(AF_UNIX, SOCK_STREAM, 0);
 142   if (Socket == INVALID_SOCKET)
 143 #else
 144   int Socket = socket(AF_UNIX, SOCK_STREAM, 0);
 145   if (Socket == -1)
 146 #endif
 147     return llvm::make_error<StringError>(getLastSocketErrorCode(),
 148                                          "socket create failed");
 149
 150   struct sockaddr_un Addr = setSocketAddr(SocketPath);
 151   if (::bind(Socket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1) {
 152     // Grab error code from call to ::bind before calling ::close
 153     std::error_code EC = getLastSocketErrorCode();
 154     ::close(Socket);
 155     return llvm::make_error<StringError>(EC, "Bind error");
 156   }
 157
 158   // Mark socket as passive so incoming connections can be accepted
 159   if (::listen(Socket, MaxBacklog) == -1)
 160     return llvm::make_error<StringError>(getLastSocketErrorCode(),
 161                                          "Listen error");
 162
 163   int PipeFD[2];
 164 #ifdef _WIN32
 165   // Reserve 1 byte for the pipe and use default textmode
 166   if (::_pipe(PipeFD, 1, 0) == -1)
 167 #else
 168   if (::pipe(PipeFD) == -1)
 169 #endif // _WIN32
 170     return llvm::make_error<StringError>(getLastSocketErrorCode(),
 171                                          "pipe failed");
 172
 173 #ifdef _WIN32
 174   return ListeningSocket{_open_osfhandle(Socket, 0), SocketPath, PipeFD};
 175 #else
 176   return ListeningSocket{Socket, SocketPath, PipeFD};
 177 #endif // _WIN32
 178 }
 179
 180 // If a file descriptor being monitored by ::poll is closed by another thread,
 181 // the result is unspecified. In the case ::poll does not unblock and return,
 182 // when ActiveFD is closed, you can provide another file descriptor via CancelFD
 183 // that when written to will cause poll to return. Typically CancelFD is the
 184 // read end of a unidirectional pipe.
 185 //
 186 // Timeout should be -1 to block indefinitly
 187 //
 188 // getActiveFD is a callback to handle ActiveFD's of std::atomic<int> and int
 189 static std::error_code
 190 manageTimeout(const std::chrono::milliseconds &Timeout,
 191               const std::function<int()> &getActiveFD,
 192               const std::optional<int> &CancelFD = std::nullopt) {
 193   struct pollfd FD[2];
 194   FD[0].events = POLLIN;
 195 #ifdef _WIN32
 196   SOCKET WinServerSock = _get_osfhandle(getActiveFD());
 197   FD[0].fd = WinServerSock;
 198 #else
 199   FD[0].fd = getActiveFD();
 200 #endif
 201   uint8_t FDCount = 1;
 202   if (CancelFD.has_value()) {
 203     FD[1].events = POLLIN;
 204     FD[1].fd = CancelFD.value();
 205     FDCount++;
 206   }
 207
 208   // Keep track of how much time has passed in case ::poll or WSAPoll are
 209   // interupted by a signal and need to be recalled
 210   auto Start = std::chrono::steady_clock::now();
 211   auto RemainingTimeout = Timeout;
 212   int PollStatus = 0;
 213   do {
 214     // If Timeout is -1 then poll should block and RemainingTimeout does not
 215     // need to be recalculated
 216     if (PollStatus != 0 && Timeout != std::chrono::milliseconds(-1)) {
 217       auto TotalElapsedTime =
 218           std::chrono::duration_cast<std::chrono::milliseconds>(
 219               std::chrono::steady_clock::now() - Start);
 220
 221       if (TotalElapsedTime >= Timeout)
 222         return std::make_error_code(std::errc::operation_would_block);
 223
 224       RemainingTimeout = Timeout - TotalElapsedTime;
 225     }
 226 #ifdef _WIN32
 227     PollStatus = WSAPoll(FD, FDCount, RemainingTimeout.count());
 228   } while (PollStatus == SOCKET_ERROR &&
 229            getLastSocketErrorCode() == std::errc::interrupted);
 230 #else
 231     PollStatus = ::poll(FD, FDCount, RemainingTimeout.count());
 232   } while (PollStatus == -1 &&
 233            getLastSocketErrorCode() == std::errc::interrupted);
 234 #endif
 235
 236   // If ActiveFD equals -1 or CancelFD has data to be read then the operation
 237   // has been canceled by another thread
 238   if (getActiveFD() == -1 || (CancelFD.has_value() && FD[1].revents & POLLIN))
 239     return std::make_error_code(std::errc::operation_canceled);
 240 #if _WIN32
 241   if (PollStatus == SOCKET_ERROR)
 242 #else
 243   if (PollStatus == -1)
 244 #endif
 245     return getLastSocketErrorCode();
 246   if (PollStatus == 0)
 247     return std::make_error_code(std::errc::timed_out);
 248   if (FD[0].revents & POLLNVAL)
 249     return std::make_error_code(std::errc::bad_file_descriptor);
 250   return std::error_code();
 251 }
 252
 253 Expected<std::unique_ptr<raw_socket_stream>>
 254 ListeningSocket::accept(const std::chrono::milliseconds &Timeout) {
 255   auto getActiveFD = [this]() -> int { return FD; };
 256   std::error_code TimeoutErr = manageTimeout(Timeout, getActiveFD, PipeFD[0]);
 257   if (TimeoutErr)
 258     return llvm::make_error<StringError>(TimeoutErr, "Timeout error");
 259
 260   int AcceptFD;
 261 #ifdef _WIN32
 262   SOCKET WinAcceptSock = ::accept(_get_osfhandle(FD), NULL, NULL);
 263   AcceptFD = _open_osfhandle(WinAcceptSock, 0);
 264 #else
 265   AcceptFD = ::accept(FD, NULL, NULL);
 266 #endif
 267
 268   if (AcceptFD == -1)
 269     return llvm::make_error<StringError>(getLastSocketErrorCode(),
 270                                          "Socket accept failed");
 271   return std::make_unique<raw_socket_stream>(AcceptFD);
 272 }
 273
 274 void ListeningSocket::shutdown() {
 275   int ObservedFD = FD.load();
 276
 277   if (ObservedFD == -1)
 278     return;
 279
 280   // If FD equals ObservedFD set FD to -1; If FD doesn't equal ObservedFD then
 281   // another thread is responsible for shutdown so return
 282   if (!FD.compare_exchange_strong(ObservedFD, -1))
 283     return;
 284
 285   ::close(ObservedFD);
 286   ::unlink(SocketPath.c_str());
 287
 288   // Ensure ::poll returns if shutdown is called by a separate thread
 289   char Byte = 'A';
 290   ssize_t written = ::write(PipeFD[1], &Byte, 1);
 291
 292   // Ignore any write() error
 293   (void)written;
 294 }
 295
 296 ListeningSocket::~ListeningSocket() {
 297   shutdown();
 298
 299   // Close the pipe's FDs in the destructor instead of within
 300   // ListeningSocket::shutdown to avoid unnecessary synchronization issues that
 301   // would occur as PipeFD's values would have to be changed to -1
 302   //
 303   // The move constructor sets PipeFD to -1
 304   if (PipeFD[0] != -1)
 305     ::close(PipeFD[0]);
 306   if (PipeFD[1] != -1)
 307     ::close(PipeFD[1]);
 308 }
 309
 310 //===----------------------------------------------------------------------===//
 311 //  raw_socket_stream
 312 //===----------------------------------------------------------------------===//
 313
 314 raw_socket_stream::raw_socket_stream(int SocketFD)
 315     : raw_fd_stream(SocketFD, true) {}
 316
 317 raw_socket_stream::~raw_socket_stream() {}
 318
 319 Expected<std::unique_ptr<raw_socket_stream>>
 320 raw_socket_stream::createConnectedUnix(StringRef SocketPath) {
 321 #ifdef _WIN32
 322   WSABalancer _;
 323 #endif // _WIN32
 324   Expected<int> FD = getSocketFD(SocketPath);
 325   if (!FD)
 326     return FD.takeError();
 327   return std::make_unique<raw_socket_stream>(*FD);
 328 }
 329
 330 ssize_t raw_socket_stream::read(char *Ptr, size_t Size,
 331                                 const std::chrono::milliseconds &Timeout) {
 332   auto getActiveFD = [this]() -> int { return this->get_fd(); };
 333   std::error_code Err = manageTimeout(Timeout, getActiveFD);
 334   // Mimic raw_fd_stream::read error handling behavior
 335   if (Err) {
 336     raw_fd_stream::error_detected(Err);
 337     return -1;
 338   }
 339   return raw_fd_stream::read(Ptr, Size);
 340 }