| 1 | //===- ConnectionUtils.h - Connection helpers for llvm-jitlink tools -----===// |
| 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 | // Connection-establishment helpers shared between llvm-jitlink and |
| 10 | // llvm-jitlink-executor. Header-only: the two tools are separate |
| 11 | // executables and never link against each other, so there is no shared |
| 12 | // object to put these definitions in. |
| 13 | // |
| 14 | //===----------------------------------------------------------------------===// |
| 15 | |
| 16 | #ifndef LLVM_TOOLS_LLVM_JITLINK_CONNECTIONUTILS_H |
| 17 | #define LLVM_TOOLS_LLVM_JITLINK_CONNECTIONUTILS_H |
| 18 | |
| 19 | #include "llvm/ADT/ScopeExit.h" |
| 20 | #include "llvm/ADT/StringRef.h" |
| 21 | #include "llvm/ADT/Twine.h" |
| 22 | #include "llvm/Config/llvm-config.h" // for LLVM_ON_UNIX |
| 23 | #include "llvm/Support/Error.h" |
| 24 | #include "llvm/Support/ErrorHandling.h" |
| 25 | #include "llvm/Support/ExponentialBackoff.h" |
| 26 | |
| 27 | #include <cstring> |
| 28 | #include <memory> |
| 29 | #include <string> |
| 30 | |
| 31 | #ifdef LLVM_ON_UNIX |
| 32 | #include <netdb.h> |
| 33 | #include <netinet/in.h> |
| 34 | #include <sys/socket.h> |
| 35 | #include <unistd.h> |
| 36 | #endif // LLVM_ON_UNIX |
| 37 | |
| 38 | namespace llvm { |
| 39 | |
| 40 | /// ConnectFn must make one attempt to connect to the executor, returning an |
| 41 | /// Expected<T>, where T is the connection handle type (e.g. T=int for a |
| 42 | /// socket). |
| 43 | /// If Retry is non-null this function will retry with exponential backoff |
| 44 | /// until Retry's timeout elapses. |
| 45 | /// If Retry is null this function will make a single attempt to connect and |
| 46 | /// return the result. |
| 47 | template <typename ConnectorFn> |
| 48 | decltype(auto) |
| 49 | connectWithRetry(ConnectorFn &&Connect, |
| 50 | std::unique_ptr<ExponentialBackoff> Retry = nullptr) { |
| 51 | if (!Retry) |
| 52 | return Connect(); |
| 53 | |
| 54 | while (true) { |
| 55 | if (auto Handle = Connect()) |
| 56 | return Handle; // Return success. |
| 57 | else if (Retry->waitForNextAttempt()) |
| 58 | consumeError(Handle.takeError()); // Ignore error and retry. |
| 59 | else |
| 60 | return Handle; // Returns error. |
| 61 | } |
| 62 | |
| 63 | llvm_unreachable("should exit from loop above" ); |
| 64 | } |
| 65 | |
| 66 | #ifdef LLVM_ON_UNIX |
| 67 | |
| 68 | /// Connects to Host:PortStr over TCP. Returns the connected socket |
| 69 | /// descriptor. |
| 70 | inline Expected<int> connectTCPSocket(StringRef Host, StringRef PortStr) { |
| 71 | addrinfo Hints{}; |
| 72 | Hints.ai_family = AF_INET; |
| 73 | Hints.ai_socktype = SOCK_STREAM; |
| 74 | Hints.ai_flags = AI_NUMERICSERV; |
| 75 | |
| 76 | addrinfo *AI; |
| 77 | if (int EC = |
| 78 | getaddrinfo(name: Host.str().c_str(), service: PortStr.str().c_str(), req: &Hints, pai: &AI)) |
| 79 | return make_error<StringError>(Args: Twine("Address resolution failed for '" ) + |
| 80 | Host + ":" + PortStr + |
| 81 | "': " + gai_strerror(ecode: EC), |
| 82 | Args: inconvertibleErrorCode()); |
| 83 | auto FreeAI = scope_exit([&]() { freeaddrinfo(ai: AI); }); |
| 84 | |
| 85 | // Cycle through the returned addrinfo structures and connect to the first |
| 86 | // reachable endpoint. |
| 87 | int SockFD = -1; |
| 88 | addrinfo *Server; |
| 89 | for (Server = AI; Server != nullptr; Server = Server->ai_next) { |
| 90 | // socket might fail, e.g. if the address family is not supported. Skip |
| 91 | // to the next addrinfo structure in such a case. |
| 92 | if ((SockFD = socket(domain: Server->ai_family, type: Server->ai_socktype, |
| 93 | protocol: Server->ai_protocol)) < 0) |
| 94 | continue; |
| 95 | |
| 96 | // If connect returns 0 we exit the loop with a working socket. |
| 97 | if (connect(fd: SockFD, addr: Server->ai_addr, len: Server->ai_addrlen) == 0) |
| 98 | break; |
| 99 | |
| 100 | close(fd: SockFD); |
| 101 | } |
| 102 | |
| 103 | // If we reached the end of the loop without connecting to a valid |
| 104 | // endpoint, report the last error logged by socket() or connect(). |
| 105 | if (Server == nullptr) |
| 106 | return make_error<StringError>(Args: Twine("Failed to connect to '" ) + Host + |
| 107 | ":" + PortStr + |
| 108 | "': " + std::strerror(errno), |
| 109 | Args: inconvertibleErrorCode()); |
| 110 | |
| 111 | return SockFD; |
| 112 | } |
| 113 | |
| 114 | /// Binds and listens for a single incoming TCP connection on Host:PortStr. |
| 115 | /// Host may be empty to bind the wildcard address; PortStr may be "0" to |
| 116 | /// request an OS-assigned ephemeral port. Returns the listening socket |
| 117 | /// descriptor -- pass it to acceptTCPConnection to accept the connection. |
| 118 | /// If ResolvedPortStr is non-null it is set to the concrete bound port, |
| 119 | /// which is required to learn the real port when PortStr was "0". |
| 120 | inline Expected<int> listenTCPSocket(StringRef Host, StringRef PortStr, |
| 121 | std::string *ResolvedPortStr = nullptr) { |
| 122 | addrinfo Hints{}; |
| 123 | Hints.ai_family = AF_INET; |
| 124 | Hints.ai_socktype = SOCK_STREAM; |
| 125 | Hints.ai_flags = AI_PASSIVE; |
| 126 | |
| 127 | std::string HostStr = Host.str(); |
| 128 | const char *Node = Host.empty() ? nullptr : HostStr.c_str(); |
| 129 | |
| 130 | addrinfo *AI; |
| 131 | if (int EC = getaddrinfo(name: Node, service: PortStr.str().c_str(), req: &Hints, pai: &AI)) |
| 132 | return make_error<StringError>(Args: Twine("Address resolution failed for '" ) + |
| 133 | Host + ":" + PortStr + |
| 134 | "': " + gai_strerror(ecode: EC), |
| 135 | Args: inconvertibleErrorCode()); |
| 136 | auto FreeAI = scope_exit([&]() { freeaddrinfo(ai: AI); }); |
| 137 | |
| 138 | int SockFD = socket(domain: AI->ai_family, type: AI->ai_socktype, protocol: AI->ai_protocol); |
| 139 | if (SockFD < 0) |
| 140 | return make_error<StringError>(Args: Twine("Error creating socket: " ) + |
| 141 | std::strerror(errno), |
| 142 | Args: inconvertibleErrorCode()); |
| 143 | auto CloseSockFD = scope_exit([&]() { close(fd: SockFD); }); |
| 144 | |
| 145 | // Avoid "Address already in use" errors. |
| 146 | const int Yes = 1; |
| 147 | if (setsockopt(fd: SockFD, SOL_SOCKET, SO_REUSEADDR, optval: &Yes, optlen: sizeof(int)) == -1) |
| 148 | return make_error<StringError>(Args: Twine("Error calling setsockopt: " ) + |
| 149 | std::strerror(errno), |
| 150 | Args: inconvertibleErrorCode()); |
| 151 | |
| 152 | if (bind(fd: SockFD, addr: AI->ai_addr, len: AI->ai_addrlen) < 0) |
| 153 | return make_error<StringError>(Args: Twine("Error binding to port '" ) + PortStr + |
| 154 | "': " + std::strerror(errno), |
| 155 | Args: inconvertibleErrorCode()); |
| 156 | |
| 157 | static constexpr int ConnectionQueueLen = 1; |
| 158 | if (listen(fd: SockFD, n: ConnectionQueueLen) < 0) |
| 159 | return make_error<StringError>(Args: Twine("Error listening on port '" ) + |
| 160 | PortStr + "': " + std::strerror(errno), |
| 161 | Args: inconvertibleErrorCode()); |
| 162 | |
| 163 | if (ResolvedPortStr) { |
| 164 | sockaddr_in BoundAddr{}; |
| 165 | socklen_t BoundAddrLen = sizeof(BoundAddr); |
| 166 | if (getsockname(fd: SockFD, addr: reinterpret_cast<sockaddr *>(&BoundAddr), |
| 167 | len: &BoundAddrLen) < 0) |
| 168 | return make_error<StringError>(Args: Twine("Error resolving bound port: " ) + |
| 169 | std::strerror(errno), |
| 170 | Args: inconvertibleErrorCode()); |
| 171 | *ResolvedPortStr = std::to_string(ntohs(BoundAddr.sin_port)); |
| 172 | } |
| 173 | |
| 174 | CloseSockFD.release(); |
| 175 | return SockFD; |
| 176 | } |
| 177 | |
| 178 | /// Accepts a single incoming connection on ListeningSockFD (as returned by |
| 179 | /// listenTCPSocket) and closes the listening socket, whether or not the |
| 180 | /// accept succeeds -- listenTCPSocket only ever queues one connection. |
| 181 | /// Returns the accepted connection's socket descriptor. |
| 182 | inline Expected<int> acceptTCPConnection(int ListeningSockFD) { |
| 183 | auto CloseListeningSockFD = scope_exit([&]() { close(fd: ListeningSockFD); }); |
| 184 | |
| 185 | int FD = accept(fd: ListeningSockFD, addr: nullptr, addr_len: nullptr); |
| 186 | if (FD < 0) |
| 187 | return make_error<StringError>(Args: Twine("Error accepting connection: " ) + |
| 188 | std::strerror(errno), |
| 189 | Args: inconvertibleErrorCode()); |
| 190 | |
| 191 | return FD; |
| 192 | } |
| 193 | |
| 194 | #endif // LLVM_ON_UNIX |
| 195 | |
| 196 | } // namespace llvm |
| 197 | |
| 198 | #endif // LLVM_TOOLS_LLVM_JITLINK_CONNECTIONUTILS_H |
| 199 | |