[llvm] Add raw_socket_stream (PR #73603)

via llvm-commits llvm-commits at lists.llvm.org
Wed Nov 29 10:09:20 PST 2023


https://github.com/criis updated https://github.com/llvm/llvm-project/pull/73603

>From abe3119e8cf0423974acbf2d2db5919ae2f00b9d Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Tue, 14 Nov 2023 14:01:17 -0800
Subject: [PATCH 1/3] Add raw_socket_stream

---
 llvm/include/llvm/Support/raw_ostream.h       |  26 +++
 llvm/lib/Support/raw_ostream.cpp              | 153 +++++++++++++++++-
 llvm/unittests/Support/CMakeLists.txt         |   1 +
 .../Support/raw_socket_stream_test.cpp        |  55 +++++++
 4 files changed, 233 insertions(+), 2 deletions(-)
 create mode 100644 llvm/unittests/Support/raw_socket_stream_test.cpp

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 1e01eb9ea19c418..f135ee5435379fc 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -630,6 +630,32 @@ 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 : public raw_fd_ostream {
+  StringRef SocketPath;
+  bool ShouldUnlink;
+
+  uint64_t current_pos() const override { return 0; }
+  
+public:
+  int get_socket() {
+    return get_fd();
+  }
+  
+  static int MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog, std::error_code &EC);
+
+  raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC);
+  raw_socket_stream(StringRef SockPath, std::error_code &EC);
+  ~raw_socket_stream();
+
+  Expected<std::string> read_impl();
+};
+
 //===----------------------------------------------------------------------===//
 // Output Stream Adaptors
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 8908e7b6a150cab..ff15b39ec69d8b1 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -1,4 +1,4 @@
-//===--- raw_ostream.cpp - Implement the raw_ostream classes --------------===//
+ //===--- raw_ostream.cpp - Implement the raw_ostream classes --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -23,11 +23,17 @@
 #include "llvm/Support/NativeFormatting.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/Program.h"
+#include "llvm/Support/Threading.h"
+#include "llvm/Support/Error.h"
 #include <algorithm>
 #include <cerrno>
 #include <cstdio>
 #include <sys/stat.h>
 
+#include <sys/socket.h>
+#include <sys/un.h>
+#include <iostream>
+
 // <fcntl.h> may provide O_BINARY.
 #if defined(HAVE_FCNTL_H)
 # include <fcntl.h>
@@ -58,6 +64,9 @@
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
+#include "raw_ostream.h"
+#include <afunix.h>
+#include <io.h>
 #endif
 
 using namespace llvm;
@@ -644,7 +653,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);
@@ -942,6 +951,146 @@ bool raw_fd_stream::classof(const raw_ostream *OS) {
   return OS->get_kind() == OStreamKind::OK_FDStream;
 }
 
+//===----------------------------------------------------------------------===//
+//  raw_socket_stream
+//===----------------------------------------------------------------------===//
+
+int raw_socket_stream::MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog, std::error_code &EC) {
+
+#ifdef _WIN32
+  SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+#else
+  int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+#endif // defined(_WIN32)
+
+#ifdef _WIN32
+  if (MaybeWinsocket == INVALID_SOCKET) {
+#else
+  if (MaybeWinsocket == -1) {
+#endif // _WIN32
+    std::string Msg = "socket create error" + std::string(strerror(errno));
+    std::perror(Msg.c_str());
+    std::cout << Msg << std::endl;
+    EC = std::make_error_code(std::errc::connection_aborted);
+    return -1;  
+  }
+
+  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) {
+    if (errno == EADDRINUSE) {
+      ::close(MaybeWinsocket);
+      EC = std::make_error_code(std::errc::address_in_use);
+    } else {
+      EC = std::make_error_code(std::errc::inappropriate_io_control_operation);
+    }
+    return -1;
+  }
+
+  if (listen(MaybeWinsocket, MaxBacklog) == -1) {
+    EC = std::make_error_code(std::errc::address_not_available);
+    return -1;
+  }
+#ifdef _WIN32
+  return _open_osfhandle(MaybeWinsocket, 0); // flags? https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/open-osfhandle?view=msvc-170
+#else
+  return MaybeWinsocket;
+#endif // _WIN32
+}
+
+int GetSocketFD(StringRef SocketPath, std::error_code &EC) {
+#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
+    std::string Msg = "socket create error" + std::string(strerror(errno));
+    std::perror(Msg.c_str());
+    EC = std::make_error_code(std::errc::connection_aborted);
+    return -1;
+  }
+
+  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) {
+    std::string Msg = "socket connect error" + std::string(strerror(errno));
+    std::perror(Msg.c_str());
+    EC = std::make_error_code(std::errc::connection_aborted);
+    return -1;
+  }
+#ifdef _WIN32
+  return _open_osfhandle(MaybeWinsocket, 0);
+#else
+  return MaybeWinsocket;
+#endif // _WIN32
+}
+
+static int ServerAccept(int FD) {
+  int AcceptFD;
+#ifdef _WIN32
+  SOCKET WinServerSock = _get_osfhandle(FD);
+  SOCKET WinAcceptSock = ::accept(WinServerSock, NULL, NULL);
+  AcceptFD = _open_osfhandle(WinAcceptSock, 0); // flags?
+#else
+  AcceptFD = ::accept(FD, NULL, NULL);
+#endif //_WIN32
+  return AcceptFD;
+}
+
+// Server
+// Call raw_fd_ostream with ShouldClose=false
+raw_socket_stream::raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC) : raw_fd_ostream(ServerAccept(SocketFD), true) {
+  SocketPath = SockPath;
+  ShouldUnlink = true;
+}
+
+// Client
+raw_socket_stream::raw_socket_stream(StringRef SockPath, std::error_code &EC) : raw_fd_ostream(GetSocketFD(SockPath, EC), true, true, OStreamKind::OK_OStream ) {
+  SocketPath = SockPath;
+  ShouldUnlink = false;
+}
+
+raw_socket_stream::~raw_socket_stream() {
+  if (ShouldUnlink) {
+    unlink(SocketPath.str().c_str());
+  }
+}
+
+Expected<std::string> raw_socket_stream::read_impl() {
+  const size_t BUFFER_SIZE = 4096;
+  std::vector<char> Buffer(BUFFER_SIZE);
+
+  int Socket = get_socket();
+  assert(Socket >= 0 && "Socket not found.");
+
+  ssize_t n;
+#ifdef _WIN32
+  SOCKET MaybeWinsocket = _get_osfhandle(Socket);
+#else
+  int MaybeWinsocket = Socket;
+#endif // _WIN32
+  n = ::read(MaybeWinsocket, Buffer.data(), Buffer.size());
+
+  if (n < 0) {
+      std::string Msg = "Buffer read error: " + std::string(strerror(errno));
+      return llvm::make_error<StringError>(Msg, inconvertibleErrorCode());
+  }
+
+  if (n == 0) {
+      return llvm::make_error<StringError>("EOF", inconvertibleErrorCode());
+  }
+  return std::string(Buffer.data());
+}
+
 //===----------------------------------------------------------------------===//
 //  raw_string_ostream
 //===----------------------------------------------------------------------===//
diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
index e1bf793536b6862..df35a7b7f3626ac 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 000000000000000..b2f756b792289b1
--- /dev/null
+++ b/llvm/unittests/Support/raw_socket_stream_test.cpp
@@ -0,0 +1,55 @@
+#include <stdlib.h>
+#include <iostream>
+#include <future>
+#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 "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+TEST(raw_socket_streamTest, CLIENT_TO_SERVER_AND_SERVER_TO_CLIENT) {
+
+  SmallString<100> SocketPath("/tmp/test_raw_socket_stream.sock");
+  std::error_code ECServer, ECClient;
+
+  int ServerFD = raw_socket_stream::MakeServerSocket(SocketPath, 3, ECServer);
+
+  raw_socket_stream Client(SocketPath, ECClient);
+  EXPECT_TRUE(!ECClient);
+
+  raw_socket_stream Client2(SocketPath, ECClient);
+
+  raw_socket_stream Server(ServerFD, SocketPath, ECServer);
+  EXPECT_TRUE(!ECServer);
+
+  Client << "01234567";
+  Client.flush();
+
+  Client2 << "abcdefgh";
+  Client2.flush();
+
+  Expected<std::string> from_client = Server.read_impl();
+
+  if (auto E = from_client.takeError()) {
+    return; // FIXME: Do something.
+  }
+  EXPECT_EQ("01234567", (*from_client));
+
+  Server << "76543210";
+  Server.flush();
+
+  Expected<std::string> from_server = Client.read_impl();
+    if (auto E = from_server.takeError()) {
+    return;
+    // YIKES! 😩
+  }
+  EXPECT_EQ("76543210", (*from_server));
+
+}
+} // namespace
\ No newline at end of file

>From 67a0ed599be8843e027419881c2eab4456f988ef Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Mon, 27 Nov 2023 18:12:18 -0800
Subject: [PATCH 2/3] clang-format

---
 llvm/include/llvm/Support/raw_ostream.h       | 11 +++---
 llvm/lib/Support/raw_ostream.cpp              | 35 ++++++++++++-------
 .../Support/raw_socket_stream_test.cpp        |  9 +++--
 3 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index f135ee5435379fc..c126abace7c25f2 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -641,13 +641,12 @@ class raw_socket_stream : public raw_fd_ostream {
   bool ShouldUnlink;
 
   uint64_t current_pos() const override { return 0; }
-  
+
 public:
-  int get_socket() {
-    return get_fd();
-  }
-  
-  static int MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog, std::error_code &EC);
+  int get_socket() { return get_fd(); }
+
+  static int MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog,
+                              std::error_code &EC);
 
   raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC);
   raw_socket_stream(StringRef SockPath, std::error_code &EC);
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index ff15b39ec69d8b1..505b62e33b3888d 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -1,4 +1,4 @@
- //===--- raw_ostream.cpp - Implement the raw_ostream classes --------------===//
+//===--- raw_ostream.cpp - Implement the raw_ostream classes --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -15,6 +15,7 @@
 #include "llvm/Config/config.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,15 +25,14 @@
 #include "llvm/Support/Process.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Threading.h"
-#include "llvm/Support/Error.h"
 #include <algorithm>
 #include <cerrno>
 #include <cstdio>
 #include <sys/stat.h>
 
+#include <iostream>
 #include <sys/socket.h>
 #include <sys/un.h>
-#include <iostream>
 
 // <fcntl.h> may provide O_BINARY.
 #if defined(HAVE_FCNTL_H)
@@ -61,10 +61,10 @@
 #endif
 
 #ifdef _WIN32
+#include "raw_ostream.h"
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
-#include "raw_ostream.h"
 #include <afunix.h>
 #include <io.h>
 #endif
@@ -955,7 +955,9 @@ bool raw_fd_stream::classof(const raw_ostream *OS) {
 //  raw_socket_stream
 //===----------------------------------------------------------------------===//
 
-int raw_socket_stream::MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog, std::error_code &EC) {
+int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
+                                        unsigned int MaxBacklog,
+                                        std::error_code &EC) {
 
 #ifdef _WIN32
   SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
@@ -972,14 +974,14 @@ int raw_socket_stream::MakeServerSocket(StringRef SocketPath, unsigned int MaxBa
     std::perror(Msg.c_str());
     std::cout << Msg << std::endl;
     EC = std::make_error_code(std::errc::connection_aborted);
-    return -1;  
+    return -1;
   }
 
   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) {
     if (errno == EADDRINUSE) {
       ::close(MaybeWinsocket);
@@ -995,7 +997,10 @@ int raw_socket_stream::MakeServerSocket(StringRef SocketPath, unsigned int MaxBa
     return -1;
   }
 #ifdef _WIN32
-  return _open_osfhandle(MaybeWinsocket, 0); // flags? https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/open-osfhandle?view=msvc-170
+  return _open_osfhandle(
+      MaybeWinsocket,
+      0); // flags?
+          // https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/open-osfhandle?view=msvc-170
 #else
   return MaybeWinsocket;
 #endif // _WIN32
@@ -1048,13 +1053,17 @@ static int ServerAccept(int FD) {
 
 // Server
 // Call raw_fd_ostream with ShouldClose=false
-raw_socket_stream::raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC) : raw_fd_ostream(ServerAccept(SocketFD), true) {
+raw_socket_stream::raw_socket_stream(int SocketFD, StringRef SockPath,
+                                     std::error_code &EC)
+    : raw_fd_ostream(ServerAccept(SocketFD), true) {
   SocketPath = SockPath;
   ShouldUnlink = true;
 }
 
 // Client
-raw_socket_stream::raw_socket_stream(StringRef SockPath, std::error_code &EC) : raw_fd_ostream(GetSocketFD(SockPath, EC), true, true, OStreamKind::OK_OStream ) {
+raw_socket_stream::raw_socket_stream(StringRef SockPath, std::error_code &EC)
+    : raw_fd_ostream(GetSocketFD(SockPath, EC), true, true,
+                     OStreamKind::OK_OStream) {
   SocketPath = SockPath;
   ShouldUnlink = false;
 }
@@ -1081,12 +1090,12 @@ Expected<std::string> raw_socket_stream::read_impl() {
   n = ::read(MaybeWinsocket, Buffer.data(), Buffer.size());
 
   if (n < 0) {
-      std::string Msg = "Buffer read error: " + std::string(strerror(errno));
-      return llvm::make_error<StringError>(Msg, inconvertibleErrorCode());
+    std::string Msg = "Buffer read error: " + std::string(strerror(errno));
+    return llvm::make_error<StringError>(Msg, inconvertibleErrorCode());
   }
 
   if (n == 0) {
-      return llvm::make_error<StringError>("EOF", inconvertibleErrorCode());
+    return llvm::make_error<StringError>("EOF", inconvertibleErrorCode());
   }
   return std::string(Buffer.data());
 }
diff --git a/llvm/unittests/Support/raw_socket_stream_test.cpp b/llvm/unittests/Support/raw_socket_stream_test.cpp
index b2f756b792289b1..c2bc9c4b3e1b2b8 100644
--- a/llvm/unittests/Support/raw_socket_stream_test.cpp
+++ b/llvm/unittests/Support/raw_socket_stream_test.cpp
@@ -1,6 +1,3 @@
-#include <stdlib.h>
-#include <iostream>
-#include <future>
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Casting.h"
@@ -8,6 +5,9 @@
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/raw_ostream.h"
 #include "gtest/gtest.h"
+#include <future>
+#include <iostream>
+#include <stdlib.h>
 
 using namespace llvm;
 
@@ -45,11 +45,10 @@ TEST(raw_socket_streamTest, CLIENT_TO_SERVER_AND_SERVER_TO_CLIENT) {
   Server.flush();
 
   Expected<std::string> from_server = Client.read_impl();
-    if (auto E = from_server.takeError()) {
+  if (auto E = from_server.takeError()) {
     return;
     // YIKES! 😩
   }
   EXPECT_EQ("76543210", (*from_server));
-
 }
 } // namespace
\ No newline at end of file

>From 91a587464093ca8e2452ab6954711d7314434404 Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Tue, 28 Nov 2023 09:58:38 -0800
Subject: [PATCH 3/3] Windows

---
 llvm/lib/Support/raw_ostream.cpp | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 505b62e33b3888d..3ea7c2cc4c974c5 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -31,8 +31,13 @@
 #include <sys/stat.h>
 
 #include <iostream>
+#ifdef _WIN32
+#include <afunix.h>
+#include <winsock2.h>
+#else
 #include <sys/socket.h>
 #include <sys/un.h>
+#endif // _WIN32
 
 // <fcntl.h> may provide O_BINARY.
 #if defined(HAVE_FCNTL_H)
@@ -61,7 +66,6 @@
 #endif
 
 #ifdef _WIN32
-#include "raw_ostream.h"
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
@@ -979,7 +983,9 @@ int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
 
   struct sockaddr_un Addr;
   memset(&Addr, 0, sizeof(Addr));
+#ifndef _WIN32
   Addr.sun_family = AF_UNIX;
+#endif // _WIN32
   strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
 
   if (bind(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1) {
@@ -997,10 +1003,7 @@ int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
     return -1;
   }
 #ifdef _WIN32
-  return _open_osfhandle(
-      MaybeWinsocket,
-      0); // flags?
-          // https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/open-osfhandle?view=msvc-170
+  return _open_osfhandle(MaybeWinsocket, 0);
 #else
   return MaybeWinsocket;
 #endif // _WIN32
@@ -1044,7 +1047,7 @@ static int ServerAccept(int FD) {
 #ifdef _WIN32
   SOCKET WinServerSock = _get_osfhandle(FD);
   SOCKET WinAcceptSock = ::accept(WinServerSock, NULL, NULL);
-  AcceptFD = _open_osfhandle(WinAcceptSock, 0); // flags?
+  AcceptFD = _open_osfhandle(WinAcceptSock, 0);
 #else
   AcceptFD = ::accept(FD, NULL, NULL);
 #endif //_WIN32



More information about the llvm-commits mailing list