[llvm] a5ffabc - [llvm][Support] Add UNIX socket support (#73603)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Dec 13 10:47:59 PST 2023
Author: criis
Date: 2023-12-13T10:47:55-08:00
New Revision: a5ffabce98a4b2e9d69009fa3e60f2b154100860
URL: https://github.com/llvm/llvm-project/commit/a5ffabce98a4b2e9d69009fa3e60f2b154100860
DIFF: https://github.com/llvm/llvm-project/commit/a5ffabce98a4b2e9d69009fa3e60f2b154100860.diff
LOG: [llvm][Support] Add UNIX socket support (#73603)
This adds support for UNIX socket communication to work similarly to
raw_stream.
---------
Patch by Christian Riis
Added:
llvm/unittests/Support/raw_socket_stream_test.cpp
Modified:
llvm/include/llvm/Support/raw_ostream.h
llvm/lib/Support/CMakeLists.txt
llvm/lib/Support/raw_ostream.cpp
llvm/unittests/Support/CMakeLists.txt
Removed:
################################################################################
diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 1e01eb9ea19c41..7c8d264afeff2f 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -16,6 +16,7 @@
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/DataTypes.h"
+#include "llvm/Support/Threading.h"
#include <cassert>
#include <cstddef>
#include <cstdint>
@@ -615,6 +616,8 @@ class raw_fd_stream : public raw_fd_ostream {
/// immediately destroyed.
raw_fd_stream(StringRef Filename, std::error_code &EC);
+ raw_fd_stream(int fd, bool shouldClose);
+
/// This reads the \p Size bytes into a buffer pointed by \p Ptr.
///
/// \param Ptr The start of the buffer to hold data to be read.
@@ -630,6 +633,54 @@ class raw_fd_stream : public raw_fd_ostream {
static bool classof(const raw_ostream *OS);
};
+//===----------------------------------------------------------------------===//
+// Socket Streams
+//===----------------------------------------------------------------------===//
+
+/// A raw stream for sockets reading/writing
+
+class raw_socket_stream;
+
+// Make sure that calls to WSAStartup and WSACleanup are balanced.
+#ifdef _WIN32
+class WSABalancer {
+public:
+ WSABalancer();
+ ~WSABalancer();
+};
+#endif // _WIN32
+
+class ListeningSocket {
+ int FD;
+ std::string SocketPath;
+ ListeningSocket(int SocketFD, StringRef SocketPath);
+#ifdef _WIN32
+ WSABalancer _;
+#endif // _WIN32
+
+public:
+ static Expected<ListeningSocket> createUnix(
+ StringRef SocketPath,
+ int MaxBacklog = llvm::hardware_concurrency().compute_thread_count());
+ Expected<std::unique_ptr<raw_socket_stream>> accept();
+ ListeningSocket(ListeningSocket &&LS);
+ ~ListeningSocket();
+};
+class raw_socket_stream : public raw_fd_stream {
+ uint64_t current_pos() const override { return 0; }
+#ifdef _WIN32
+ WSABalancer _;
+#endif // _WIN32
+
+public:
+ raw_socket_stream(int SocketFD);
+ /// Create a \p raw_socket_stream connected to the Unix domain socket at \p
+ /// SocketPath.
+ static Expected<std::unique_ptr<raw_socket_stream>>
+ createConnectedUnix(StringRef SocketPath);
+ ~raw_socket_stream();
+};
+
//===----------------------------------------------------------------------===//
// Output Stream Adaptors
//===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Support/CMakeLists.txt b/llvm/lib/Support/CMakeLists.txt
index b96d62c7a6224d..80854d1d09d985 100644
--- a/llvm/lib/Support/CMakeLists.txt
+++ b/llvm/lib/Support/CMakeLists.txt
@@ -40,7 +40,7 @@ endif()
if( MSVC OR MINGW )
# libuuid required for FOLDERID_Profile usage in lib/Support/Windows/Path.inc.
# advapi32 required for CryptAcquireContextW in lib/Support/Windows/Path.inc.
- set(system_libs ${system_libs} psapi shell32 ole32 uuid advapi32)
+ set(system_libs ${system_libs} psapi shell32 ole32 uuid advapi32 Ws2_32)
elseif( CMAKE_HOST_UNIX )
if( HAVE_LIBRT )
set(system_libs ${system_libs} rt)
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index d654ae450b3403..d7a1f91a27defc 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -16,6 +16,7 @@
#include "llvm/Support/AutoConvert.h"
#include "llvm/Support/Compiler.h"
#include "llvm/Support/Duration.h"
+#include "llvm/Support/Error.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Format.h"
@@ -24,11 +25,17 @@
#include "llvm/Support/NativeFormatting.h"
#include "llvm/Support/Process.h"
#include "llvm/Support/Program.h"
+#include "llvm/Support/Threading.h"
#include <algorithm>
#include <cerrno>
#include <cstdio>
#include <sys/stat.h>
+#ifndef _WIN32
+#include <sys/socket.h>
+#include <sys/un.h>
+#endif // _WIN32
+
// <fcntl.h> may provide O_BINARY.
#if defined(HAVE_FCNTL_H)
# include <fcntl.h>
@@ -59,6 +66,13 @@
#include "llvm/Support/ConvertUTF.h"
#include "llvm/Support/Signals.h"
#include "llvm/Support/Windows/WindowsSupport.h"
+// winsock2.h must be included before afunix.h. Briefly turn off clang-format to
+// avoid error.
+// clang-format off
+#include <winsock2.h>
+#include <afunix.h>
+// clang-format on
+#include <io.h>
#endif
using namespace llvm;
@@ -645,7 +659,7 @@ raw_fd_ostream::raw_fd_ostream(int fd, bool shouldClose, bool unbuffered,
// Check if this is a console device. This is not equivalent to isatty.
IsWindowsConsole =
::GetFileType((HANDLE)::_get_osfhandle(fd)) == FILE_TYPE_CHAR;
-#endif
+#endif // _WIN32
// Get the starting position.
off_t loc = ::lseek(FD, 0, SEEK_CUR);
@@ -937,6 +951,9 @@ raw_fd_stream::raw_fd_stream(StringRef Filename, std::error_code &EC)
EC = std::make_error_code(std::errc::invalid_argument);
}
+raw_fd_stream::raw_fd_stream(int fd, bool shouldClose)
+ : raw_fd_ostream(fd, shouldClose, false, OStreamKind::OK_FDStream) {}
+
ssize_t raw_fd_stream::read(char *Ptr, size_t Size) {
assert(get_fd() >= 0 && "File already closed.");
ssize_t Ret = ::read(get_fd(), (void *)Ptr, Size);
@@ -951,6 +968,145 @@ bool raw_fd_stream::classof(const raw_ostream *OS) {
return OS->get_kind() == OStreamKind::OK_FDStream;
}
+//===----------------------------------------------------------------------===//
+// raw_socket_stream
+//===----------------------------------------------------------------------===//
+
+#ifdef _WIN32
+WSABalancer::WSABalancer() {
+ WSADATA WsaData = {0};
+ if (WSAStartup(MAKEWORD(2, 2), &WsaData) != 0) {
+ llvm::report_fatal_error("WSAStartup failed");
+ }
+}
+
+WSABalancer::~WSABalancer() { WSACleanup(); }
+
+#endif // _WIN32
+
+static std::error_code getLastSocketErrorCode() {
+#ifdef _WIN32
+ return std::error_code(::WSAGetLastError(), std::system_category());
+#else
+ return std::error_code(errno, std::system_category());
+#endif
+}
+
+ListeningSocket::ListeningSocket(int SocketFD, StringRef SocketPath)
+ : FD(SocketFD), SocketPath(SocketPath) {}
+
+ListeningSocket::ListeningSocket(ListeningSocket &&LS)
+ : FD(LS.FD), SocketPath(LS.SocketPath) {
+ LS.FD = -1;
+}
+
+Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath,
+ int MaxBacklog) {
+
+#ifdef _WIN32
+ WSABalancer _;
+ SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+ if (MaybeWinsocket == INVALID_SOCKET) {
+#else
+ int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+ if (MaybeWinsocket == -1) {
+#endif
+ return llvm::make_error<StringError>(getLastSocketErrorCode(),
+ "socket create failed");
+ }
+
+ struct sockaddr_un Addr;
+ memset(&Addr, 0, sizeof(Addr));
+ Addr.sun_family = AF_UNIX;
+ strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
+
+ if (bind(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1) {
+ std::error_code Err = getLastSocketErrorCode();
+ if (Err == std::errc::address_in_use)
+ ::close(MaybeWinsocket);
+ return llvm::make_error<StringError>(Err, "Bind error");
+ }
+ if (listen(MaybeWinsocket, MaxBacklog) == -1) {
+ return llvm::make_error<StringError>(getLastSocketErrorCode(),
+ "Listen error");
+ }
+ int UnixSocket;
+#ifdef _WIN32
+ UnixSocket = _open_osfhandle(MaybeWinsocket, 0);
+#else
+ UnixSocket = MaybeWinsocket;
+#endif // _WIN32
+ ListeningSocket ListenSocket(UnixSocket, SocketPath);
+ return ListenSocket;
+}
+
+Expected<std::unique_ptr<raw_socket_stream>> ListeningSocket::accept() {
+ int AcceptFD;
+#ifdef _WIN32
+ SOCKET WinServerSock = _get_osfhandle(FD);
+ SOCKET WinAcceptSock = ::accept(WinServerSock, NULL, NULL);
+ AcceptFD = _open_osfhandle(WinAcceptSock, 0);
+#else
+ AcceptFD = ::accept(FD, NULL, NULL);
+#endif //_WIN32
+ if (AcceptFD == -1)
+ return llvm::make_error<StringError>(getLastSocketErrorCode(),
+ "Accept failed");
+ return std::make_unique<raw_socket_stream>(AcceptFD);
+}
+
+ListeningSocket::~ListeningSocket() {
+ if (FD == -1)
+ return;
+ ::close(FD);
+ unlink(SocketPath.c_str());
+}
+
+static Expected<int> GetSocketFD(StringRef SocketPath) {
+#ifdef _WIN32
+ SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+ if (MaybeWinsocket == INVALID_SOCKET) {
+#else
+ int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+ if (MaybeWinsocket == -1) {
+#endif // _WIN32
+ return llvm::make_error<StringError>(getLastSocketErrorCode(),
+ "Create socket failed");
+ }
+
+ struct sockaddr_un Addr;
+ memset(&Addr, 0, sizeof(Addr));
+ Addr.sun_family = AF_UNIX;
+ strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
+
+ int status = connect(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr));
+ if (status == -1) {
+ return llvm::make_error<StringError>(getLastSocketErrorCode(),
+ "Connect socket failed");
+ }
+#ifdef _WIN32
+ return _open_osfhandle(MaybeWinsocket, 0);
+#else
+ return MaybeWinsocket;
+#endif // _WIN32
+}
+
+raw_socket_stream::raw_socket_stream(int SocketFD)
+ : raw_fd_stream(SocketFD, true) {}
+
+Expected<std::unique_ptr<raw_socket_stream>>
+raw_socket_stream::createConnectedUnix(StringRef SocketPath) {
+#ifdef _WIN32
+ WSABalancer _;
+#endif // _WIN32
+ Expected<int> FD = GetSocketFD(SocketPath);
+ if (!FD)
+ return FD.takeError();
+ return std::make_unique<raw_socket_stream>(*FD);
+}
+
+raw_socket_stream::~raw_socket_stream() {}
+
//===----------------------------------------------------------------------===//
// raw_string_ostream
//===----------------------------------------------------------------------===//
diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
index e1bf793536b686..df35a7b7f3626a 100644
--- a/llvm/unittests/Support/CMakeLists.txt
+++ b/llvm/unittests/Support/CMakeLists.txt
@@ -103,6 +103,7 @@ add_llvm_unittest(SupportTests
raw_ostream_test.cpp
raw_pwrite_stream_test.cpp
raw_sha1_ostream_test.cpp
+ raw_socket_stream_test.cpp
xxhashTest.cpp
DEPENDS
diff --git a/llvm/unittests/Support/raw_socket_stream_test.cpp b/llvm/unittests/Support/raw_socket_stream_test.cpp
new file mode 100644
index 00000000000000..53eb86ae45d29d
--- /dev/null
+++ b/llvm/unittests/Support/raw_socket_stream_test.cpp
@@ -0,0 +1,52 @@
+#include "llvm/ADT/SmallString.h"
+#include "llvm/Config/llvm-config.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FileUtilities.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+#include <future>
+#include <iostream>
+#include <stdlib.h>
+
+using namespace llvm;
+
+namespace {
+
+TEST(raw_socket_streamTest, CLIENT_TO_SERVER_AND_SERVER_TO_CLIENT) {
+ SmallString<100> SocketPath;
+ llvm::sys::fs::createUniquePath("test_raw_socket_stream.sock", SocketPath,
+ true);
+
+ char Bytes[8];
+
+ Expected<ListeningSocket> MaybeServerListener =
+ ListeningSocket::createUnix(SocketPath);
+ ASSERT_THAT_EXPECTED(MaybeServerListener, llvm::Succeeded());
+
+ ListeningSocket ServerListener = std::move(*MaybeServerListener);
+
+ Expected<std::unique_ptr<raw_socket_stream>> MaybeClient =
+ raw_socket_stream::createConnectedUnix(SocketPath);
+ ASSERT_THAT_EXPECTED(MaybeClient, llvm::Succeeded());
+
+ raw_socket_stream &Client = **MaybeClient;
+
+ Expected<std::unique_ptr<raw_socket_stream>> MaybeServer =
+ ServerListener.accept();
+ ASSERT_THAT_EXPECTED(MaybeServer, llvm::Succeeded());
+
+ raw_socket_stream &Server = **MaybeServer;
+
+ Client << "01234567";
+ Client.flush();
+
+ ssize_t BytesRead = Server.read(Bytes, 8);
+
+ std::string string(Bytes, 8);
+
+ ASSERT_EQ(8, BytesRead);
+ ASSERT_EQ("01234567", string);
+}
+} // namespace
\ No newline at end of file
More information about the llvm-commits
mailing list