Skip to content

Commit d2c029f

Browse files
v-sathviksathvik vallivedu
andauthored
Add Unix domain socket support for gateway (documentdb#408)
* Add Unix domain socket support for gateway - Add unix_socket_enabled and unix_socket_path configuration options - Implement Unix socket listener alongside TCP in gateway - Add integration tests for the Unix socket - Add test helpers for Unix socket configuration Signed-off-by: sathvik vallivedu <tvsathvi@amazon.com> * Address comments on Unix socket implementation Changes per review feedback: - Simplify to single UnixSocketPath config option (remove separate enable flag) Empty path = disabled, specified path = enabled - Refactor Unix socket handling to gateway_stream.rs Signed-off-by: sathvik vallivedu <tvsathvi@amazon.com> * Resolve conflicts and add configurable Unix socket permissions Address review feedback and add UnixSocketFilePermissions parameter with octal notation support (e.g., UnixSocketFilePermissions: '0600') Signed-off-by: sathvik vallivedu <tvsathvi@amazon.com> * Address comment and clean up shutdown handler code Signed-off-by: sathvik vallivedu <tvsathvi@amazon.com> * Remove unused socket path Signed-off-by: sathvik vallivedu <tvsathvi@amazon.com> * fix: formatting Signed-off-by: sathvik vallivedu <tvsathvi@amazon.com> --------- Signed-off-by: sathvik vallivedu <tvsathvi@amazon.com> Co-authored-by: sathvik vallivedu <tvsathvi@amazon.com>
1 parent e4501f4 commit d2c029f

6 files changed

Lines changed: 376 additions & 5 deletions

File tree

‎pg_documentdb_gw/src/configuration/mod.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,12 +75,20 @@ pub trait SetupConfiguration: DynClone + Send + Sync + Debug {
7575
/// Returns the number of worker threads for the async runtime.
7676
fn async_runtime_worker_threads(&self) -> usize;
7777

78+
/// Returns the Unix domain socket path if Unix domain socket is enabled.
79+
/// Returns None if Unix socket is disabled.
80+
fn unix_socket_path(&self) -> Option<&str>;
81+
7882
/// Returns the timeout duration (in minutes) for PostgreSQL connections
7983
fn postgres_idle_connection_timeout_minutes(&self) -> u64;
8084

8185
/// Returns whether TLS should be enforced for all connections.
8286
fn enforce_tls(&self) -> bool;
8387

88+
/// Returns the file permissions for Unix socket files (octal format).
89+
/// Defaults to 0o660 (owner+group read/write) if not specified.
90+
fn unix_socket_file_permissions(&self) -> u32;
91+
8492
/// Provides a way to downcast the trait object to a concrete type.
8593
fn as_any(&self) -> &dyn std::any::Any;
8694
}

‎pg_documentdb_gw/src/configuration/setup.rs‎

Lines changed: 42 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,14 +48,43 @@ pub struct DocumentDBSetupConfiguration {
4848

4949
// Runtime configuration
5050
pub async_runtime_worker_threads: Option<usize>,
51+
52+
// Unix domain socket configuration
53+
// If specified with a non-empty path, Unix socket is enabled at that path.
54+
// If not specified (None), Unix socket is disabled.
55+
pub unix_socket_path: Option<String>,
56+
57+
// Unix socket file permissions (octal format string, e.g., "0660" for owner+group read/write)
58+
// If not specified, defaults to 0o660
59+
pub unix_socket_file_permissions: Option<String>,
5160
}
5261

5362
impl DocumentDBSetupConfiguration {
5463
pub fn new(config_path: &Path) -> Result<Self> {
5564
let config_file = File::open(config_path)?;
56-
serde_json::from_reader(config_file).map_err(|e| {
65+
let config: Self = serde_json::from_reader(config_file).map_err(|e| {
5766
DocumentDBError::internal_error(format!("Failed to parse configuration file: {e}"))
58-
})
67+
})?;
68+
69+
// Validate Unix socket path if provided
70+
if let Some(path) = &config.unix_socket_path {
71+
if path.trim().is_empty() {
72+
return Err(DocumentDBError::internal_error(
73+
"UnixSocketPath cannot be empty. Either provide a valid path or omit the field to disable Unix sockets.".to_string()
74+
));
75+
}
76+
}
77+
78+
// Validate Unix socket permissions if provided
79+
if let Some(perm_str) = &config.unix_socket_file_permissions {
80+
if u32::from_str_radix(perm_str, 8).is_err() {
81+
return Err(DocumentDBError::internal_error(
82+
format!("Invalid UnixSocketFilePermissions '{perm_str}'. Expected octal format like '0600', '0644'")
83+
));
84+
}
85+
}
86+
87+
Ok(config)
5988
}
6089
}
6190

@@ -138,11 +167,22 @@ impl SetupConfiguration for DocumentDBSetupConfiguration {
138167
})
139168
}
140169

170+
fn unix_socket_path(&self) -> Option<&str> {
171+
self.unix_socket_path.as_deref()
172+
}
173+
141174
fn postgres_idle_connection_timeout_minutes(&self) -> u64 {
142175
self.postgres_idle_connection_timeout_minutes.unwrap_or(5)
143176
}
144177

145178
fn enforce_tls(&self) -> bool {
146179
self.enforce_tls.unwrap_or(true)
147180
}
181+
182+
fn unix_socket_file_permissions(&self) -> u32 {
183+
match &self.unix_socket_file_permissions {
184+
None => 0o660, // Default when not provided
185+
Some(perm_str) => u32::from_str_radix(perm_str, 8).unwrap(),
186+
}
187+
}
148188
}

‎pg_documentdb_gw/src/lib.rs‎

Lines changed: 163 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ use std::net::IpAddr;
3131
use std::{pin::Pin, sync::Arc, time::Duration};
3232
use tokio::{
3333
io::{AsyncRead, AsyncWrite, BufStream},
34-
net::{TcpListener, TcpStream},
34+
net::{TcpListener, TcpStream, UnixListener, UnixStream},
3535
};
3636
use tokio_openssl::SslStream;
3737
use tokio_util::sync::CancellationToken;
@@ -61,6 +61,71 @@ const STREAM_WRITE_BUFFER_SIZE: usize = 8 * 1024;
6161
// TLS detection timeout
6262
const TLS_PEEK_TIMEOUT_SECS: u64 = 5;
6363

64+
/// Applies configurable permissions to Unix domain socket file.
65+
///
66+
/// On Unix systems: Sets permissions to the specified octal value
67+
/// On other platforms: No-op (permissions handled by OS defaults)
68+
#[cfg(unix)]
69+
fn apply_socket_permissions(path: &str, permissions: u32) -> std::io::Result<()> {
70+
use std::os::unix::fs::PermissionsExt;
71+
let file_permissions = std::fs::Permissions::from_mode(permissions);
72+
std::fs::set_permissions(path, file_permissions)
73+
}
74+
75+
#[cfg(not(unix))]
76+
fn apply_socket_permissions(_path: &str, _permissions: u32) -> std::io::Result<()> {
77+
Ok(())
78+
}
79+
80+
/// Creates and configures a Unix domain socket listener.
81+
///
82+
/// This function handles the complete Unix socket setup:
83+
/// 1. Removes stale socket file from previous run (crash recovery)
84+
/// 2. Binds to the socket path
85+
/// 3. Sets appropriate permissions via platform-specific logic
86+
///
87+
/// # Arguments
88+
///
89+
/// * `socket_path` - Path where the Unix socket should be created
90+
/// * `permissions` - Octal file permissions (e.g., 0o660)
91+
///
92+
/// # Returns
93+
///
94+
/// Returns the configured `UnixListener` or an error if setup fails.
95+
///
96+
/// # Errors
97+
///
98+
/// This function will return an error if:
99+
/// * Failed to bind to the socket path
100+
/// * Failed to set socket file permissions
101+
fn create_unix_socket_listener(socket_path: &str, permissions: u32) -> Result<UnixListener> {
102+
// Attempt to remove stale socket file from previous run (e.g., after crash).
103+
// This is standard Unix socket practice - PostgreSQL, MySQL, Redis all use this pattern.
104+
// Socket files cannot be cleaned up during crash (SIGKILL, segfault, power loss, etc.),
105+
// so cleanup at startup is the only reliable approach for crash recovery.
106+
// If the file couldn't be removed, bind() will fail with clear error "Address already in use".
107+
if let Err(e) = std::fs::remove_file(socket_path) {
108+
if e.kind() != std::io::ErrorKind::NotFound {
109+
tracing::warn!(
110+
"Could not remove existing socket file {}: {}.",
111+
socket_path,
112+
e
113+
);
114+
}
115+
}
116+
117+
let listener = UnixListener::bind(socket_path)?;
118+
119+
apply_socket_permissions(socket_path, permissions)?;
120+
121+
tracing::info!(
122+
"Unix socket listener bound to {} with permissions {:o}",
123+
socket_path,
124+
permissions
125+
);
126+
Ok(listener)
127+
}
128+
64129
/// Runs the DocumentDB gateway server, accepting and handling incoming connections.
65130
///
66131
/// This function sets up a TCP listener and SSL context, then continuously accepts
@@ -104,7 +169,24 @@ where
104169
))
105170
.await?;
106171

107-
// Listen for new tcp connections until token is not cancelled
172+
tracing::info!(
173+
"TCP listener bound to port {}",
174+
service_context.setup_configuration().gateway_listen_port()
175+
);
176+
177+
let unix_listener =
178+
if let Some(unix_socket_path) = service_context.setup_configuration().unix_socket_path() {
179+
let permissions = service_context
180+
.setup_configuration()
181+
.unix_socket_file_permissions();
182+
let unix_listener = create_unix_socket_listener(unix_socket_path, permissions)?;
183+
Some(unix_listener)
184+
} else {
185+
tracing::info!("Unix socket disabled (not configured)");
186+
None
187+
};
188+
189+
// Listen for new tcp and unix socket connections
108190
loop {
109191
tokio::select! {
110192
stream_and_address = tcp_listener.accept() => {
@@ -119,7 +201,28 @@ where
119201
.await;
120202

121203
if let Err(conn_err) = conn_res {
122-
tracing::error!("Failed to accept a connection: {conn_err:?}.");
204+
tracing::error!("Failed to accept a TCP connection: {conn_err:?}.");
205+
}
206+
});
207+
}
208+
stream_result = async {
209+
match &unix_listener {
210+
Some(listener) => listener.accept().await,
211+
None => std::future::pending().await,
212+
}
213+
}, if unix_listener.is_some() => {
214+
let conn_service_context = service_context.clone();
215+
let conn_telemetry = telemetry.clone();
216+
tokio::spawn(async move {
217+
let conn_res = handle_unix_connection::<T>(
218+
stream_result,
219+
conn_service_context,
220+
conn_telemetry,
221+
)
222+
.await;
223+
224+
if let Err(conn_err) = conn_res {
225+
tracing::error!("Failed to accept a Unix socket connection: {conn_err:?}.");
123226
}
124227
});
125228
}
@@ -337,6 +440,63 @@ where
337440
Ok(())
338441
}
339442

443+
/// Handles a single Unix socket connection without TLS.
444+
///
445+
/// Unix socket connections are local-only and don't require TLS encryption.
446+
/// This function creates a connection context and processes the connection.
447+
///
448+
/// # Arguments
449+
///
450+
/// * `stream_result` - Result containing the Unix stream from accept()
451+
/// * `service_context` - Service configuration and shared state
452+
/// * `telemetry` - Optional telemetry provider for metrics collection
453+
///
454+
/// # Returns
455+
///
456+
/// Returns `Ok(())` on successful connection handling, or an error if connection
457+
/// setup fails.
458+
async fn handle_unix_connection<T>(
459+
stream_result: std::result::Result<(UnixStream, tokio::net::unix::SocketAddr), std::io::Error>,
460+
service_context: ServiceContext,
461+
telemetry: Option<Box<dyn TelemetryProvider>>,
462+
) -> Result<()>
463+
where
464+
T: PgDataClient,
465+
{
466+
let (unix_stream, _socket_addr) = stream_result?;
467+
468+
let connection_id = Uuid::new_v4();
469+
tracing::info!(
470+
activity_id = connection_id.to_string().as_str(),
471+
"New Unix socket connection established"
472+
);
473+
474+
// For Unix sockets, use localhost as the address since they don't have IP addresses
475+
476+
let connection_context = ConnectionContext::new(
477+
service_context,
478+
telemetry,
479+
"localhost".to_string(),
480+
None, // No TLS for Unix sockets
481+
connection_id,
482+
"UnixSocket".to_string(),
483+
);
484+
485+
let buffered_stream = BufStream::with_capacity(
486+
STREAM_READ_BUFFER_SIZE,
487+
STREAM_WRITE_BUFFER_SIZE,
488+
unix_stream,
489+
);
490+
491+
tracing::info!(
492+
activity_id = connection_id.to_string().as_str(),
493+
"Unix socket connection established - Connection Id {connection_id}"
494+
);
495+
496+
handle_stream::<T, _>(buffered_stream, connection_context).await;
497+
Ok(())
498+
}
499+
340500
async fn handle_stream<T, S>(mut stream: S, mut connection_context: ConnectionContext)
341501
where
342502
T: PgDataClient,

‎pg_documentdb_gw/tests/common/mod.rs‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,16 @@ pub fn setup_configuration() -> DocumentDBSetupConfiguration {
133133
}
134134
}
135135

136+
pub fn setup_configuration_with_unix_socket_custom(
137+
path: Option<String>,
138+
permissions: Option<String>,
139+
) -> DocumentDBSetupConfiguration {
140+
let mut config = setup_configuration();
141+
config.unix_socket_path = path;
142+
config.unix_socket_file_permissions = permissions;
143+
config
144+
}
145+
136146
pub fn get_client() -> Client {
137147
let credential = Credential::builder()
138148
.username("test".to_string())
@@ -179,6 +189,35 @@ pub async fn initialize_with_config(config: DocumentDBSetupConfiguration) -> Cli
179189
get_client()
180190
}
181191

192+
#[allow(dead_code)]
193+
pub fn get_unix_socket_client_custom(path: &str) -> Client {
194+
use std::time::Duration;
195+
let credential = Credential::builder()
196+
.username("test".to_string())
197+
.password("test".to_string())
198+
.mechanism(AuthMechanism::ScramSha256)
199+
.build();
200+
201+
let client_options = ClientOptions::builder()
202+
.credential(credential)
203+
.hosts(vec![ServerAddress::parse(path).unwrap()])
204+
.connect_timeout(Duration::from_millis(100))
205+
.server_selection_timeout(Duration::from_millis(100))
206+
.build();
207+
Client::with_options(client_options).unwrap()
208+
}
209+
210+
#[allow(dead_code)]
211+
pub async fn initialize_with_config_and_unix(path: Option<String>) -> (Client, Option<Client>) {
212+
let config = setup_configuration_with_unix_socket_custom(path.clone(), None);
213+
initialize_full(config).await;
214+
215+
let tcp_client = get_client();
216+
let unix_client = path.map(|socket_path| get_unix_socket_client_custom(&socket_path));
217+
218+
(tcp_client, unix_client)
219+
}
220+
182221
#[allow(dead_code)]
183222
pub async fn initialize() -> Client {
184223
initialize_full(setup_configuration()).await;

0 commit comments

Comments
 (0)