[llvm][Support] Lift raw_socket_stream implementation out into own files (#75653)

Move the implementation of raw_socket_stream from raw_ostream.h/cpp to
raw_socket_stream.h/cpp as requested in #73603.
This commit is contained in:
criis 2023-12-18 17:13:47 -08:00 committed by GitHub
parent 3effc191e3
commit c0d5d36dda
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 248 additions and 212 deletions

View File

@ -16,7 +16,6 @@
#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>
@ -633,54 +632,6 @@ public:
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
//===----------------------------------------------------------------------===//

View File

@ -0,0 +1,66 @@
//===-- llvm/Support/raw_socket_stream.h - Socket streams --*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file contains raw_ostream implementations for streams to communicate
// via UNIX sockets
//
//===----------------------------------------------------------------------===//
#ifndef LLVM_SUPPORT_RAW_SOCKET_STREAM_H
#define LLVM_SUPPORT_RAW_SOCKET_STREAM_H
#include "llvm/Support/Threading.h"
#include "llvm/Support/raw_ostream.h"
namespace llvm {
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();
};
} // end namespace llvm
#endif

View File

@ -252,6 +252,7 @@ add_llvm_component_library(LLVMSupport
YAMLTraits.cpp
raw_os_ostream.cpp
raw_ostream.cpp
raw_socket_stream.cpp
regcomp.c
regerror.c
regexec.c

View File

@ -16,7 +16,6 @@
#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"
@ -25,17 +24,11 @@
#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>
@ -66,13 +59,6 @@
#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;
@ -659,7 +645,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 // _WIN32
#endif
// Get the starting position.
off_t loc = ::lseek(FD, 0, SEEK_CUR);
@ -968,153 +954,6 @@ 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;
::memset(&WsaData, 0, sizeof(WsaData));
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
return ListeningSocket{UnixSocket, SocketPath};
}
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
//===----------------------------------------------------------------------===//
void raw_string_ostream::write_impl(const char *Ptr, size_t Size) {
OS.append(Ptr, Size);
}
//===----------------------------------------------------------------------===//
// raw_svector_ostream
//===----------------------------------------------------------------------===//

View File

@ -0,0 +1,179 @@
//===-- llvm/Support/raw_socket_stream.cpp - Socket streams --*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file contains raw_ostream implementations for streams to communicate
// via UNIX sockets
//
//===----------------------------------------------------------------------===//
#include "llvm/Support/raw_socket_stream.h"
#include "llvm/Config/config.h"
#include "llvm/Support/Error.h"
#ifndef _WIN32
#include <sys/socket.h>
#include <sys/un.h>
#else
#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 // _WIN32
#if defined(HAVE_UNISTD_H)
#include <unistd.h>
#endif
using namespace llvm;
#ifdef _WIN32
WSABalancer::WSABalancer() {
WSADATA WsaData;
::memset(&WsaData, 0, sizeof(WsaData));
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
return ListeningSocket{UnixSocket, SocketPath};
}
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
//===----------------------------------------------------------------------===//
void raw_string_ostream::write_impl(const char *Ptr, size_t Size) {
OS.append(Ptr, Size);
}

View File

@ -3,7 +3,7 @@
#include "llvm/Support/Casting.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/FileUtilities.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/raw_socket_stream.h"
#include "llvm/Testing/Support/Error.h"
#include "gtest/gtest.h"
#include <future>