xref: /llvm-project/llvm/lib/Support/raw_socket_stream.cpp (revision d44ea7186befe38eb2b3804b15cd1ee1777458ed)
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 }
341