diff --git a/Cargo.lock b/Cargo.lock index ab93968..1ab38d7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -69,6 +69,17 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] + [[package]] name = "futures-sink" version = "0.3.34" @@ -88,6 +99,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ "futures-core", + "futures-macro", "futures-task", "pin-project-lite", "slab", @@ -160,6 +172,7 @@ version = "0.1.0" dependencies = [ "argh", "camino", + "futures-util", "ldap3_proto", "serde", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 686e0c7..b44a4b7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ edition = "2024" [dependencies] argh = "0.1.19" camino = "1.2.5" +futures-util = "0.3.34" ldap3_proto = "0.8.1" serde = { version = "1.0.229", features = ["derive"] } tokio = { version = "1.53.1", features = ["macros", "net", "rt", "time", "sync"] } diff --git a/src/ldap/handler.rs b/src/ldap/handler.rs new file mode 100644 index 0000000..f13a992 --- /dev/null +++ b/src/ldap/handler.rs @@ -0,0 +1,37 @@ +use ldap3_proto::LdapMsg; + +use crate::ldap::{LdapStream, LdapStreamError}; + +#[tracing::instrument(name = "ldap", skip(s), fields(session = %s.session))] +pub async fn ldap_handler(mut s: LdapStream) { + tracing::info! { + remote_addr = ?s.remote_addr, + "New client connection" + }; + + loop { + match s.next().await { + Ok(msg) => { + if let Err(e) = ldap_handler_inner(msg).await { + tracing::debug!( + reason = ?e, + "Failed to respond" + ); + return; + } + } + Err(e) => { + tracing::debug!( + reason = ?e, + "Closing connection" + ); + return; + } + } + } +} + +pub async fn ldap_handler_inner(msg: LdapMsg) -> Result<(), LdapStreamError> { + tracing::debug!(msg = ?msg, "Received LDAP message"); + Ok(()) +} diff --git a/src/ldap/mod.rs b/src/ldap/mod.rs new file mode 100644 index 0000000..a829da5 --- /dev/null +++ b/src/ldap/mod.rs @@ -0,0 +1,4 @@ +mod handler; +pub use handler::ldap_handler; +mod stream; +pub use stream::{LdapStream, LdapStreamError}; diff --git a/src/ldap/stream.rs b/src/ldap/stream.rs new file mode 100644 index 0000000..3975e81 --- /dev/null +++ b/src/ldap/stream.rs @@ -0,0 +1,69 @@ +use futures_util::StreamExt; +use ldap3_proto::{LdapCodec, LdapMsg}; +use tokio::time::{Duration, timeout}; +use tokio_util::codec::Framed; +use uuid::Uuid; + +use std::fmt; + +use crate::stream::{AbstractSocketAddr, AbstractStreamKind}; + +#[derive(Debug)] +pub enum LdapStreamError { + ClientClosed, + Timeout(Duration), + LdapError(std::io::Error), +} + +impl fmt::Display for LdapStreamError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::ClientClosed => write!(f, "Client closed the connection"), + Self::Timeout(d) => write!(f, "Client timeout after {}ms", d.as_millis()), + Self::LdapError(e) => write!(f, "Failed to parse LDAP message: {e:?}"), + } + } +} + +impl std::error::Error for LdapStreamError {} + +pub struct LdapStream { + pub remote_addr: AbstractSocketAddr, + pub inner: Framed, + /// An arbitrary ID to correlate logs with a specific TCP session + pub session: Uuid, + pub timeout: Duration, +} + +impl LdapStream { + pub fn new( + remote_addr: AbstractSocketAddr, + kind: AbstractStreamKind, + timeout: tokio::time::Duration, + ) -> Self { + Self { + remote_addr, + inner: Framed::new(kind, LdapCodec::new(None, None)), + session: Uuid::new_v4(), + timeout, + } + } + + pub async fn next(&mut self) -> Result { + // Check for timeout + let Ok(msg) = timeout(self.timeout, self.inner.next()).await else { + return Err(LdapStreamError::Timeout(self.timeout)); + }; + + // Check for closed client connection + let Some(msg) = msg else { + return Err(LdapStreamError::ClientClosed); + }; + + // So far the message may or may not be a valid LDAP message, + // but using the codec some other conditions (eg. OOM) may trigger + // an error. We abort whether the message failed because it was malformed, + // because of transient IO error, or any other error condition. + msg.map_err(LdapStreamError::LdapError) + } +} diff --git a/src/listener.rs b/src/listener.rs index 0b3b068..d1dda07 100644 --- a/src/listener.rs +++ b/src/listener.rs @@ -1,12 +1,13 @@ use camino::Utf8PathBuf; use tokio::net::{TcpListener, UnixListener}; +use tokio::time::Duration; use std::fmt; use std::net::SocketAddr; use std::str::FromStr; use crate::error::GlobalError; -use crate::stream::AbstractStream; +use crate::ldap::LdapStream; #[derive(Debug)] pub struct AcceptError { @@ -98,10 +99,17 @@ impl Listener { } } - pub async fn accept(&self) -> Result, AcceptError> { + pub async fn accept_ldap(&self) -> Result, AcceptError> { + // TODO: configurable timeout + let timeout = Duration::from_millis(500); + let res = match &self.kind { - ListenerKind::Tcp(l) => l.accept().await.map(AbstractStream::from), - ListenerKind::Uds(l) => l.accept().await.map(AbstractStream::from), + ListenerKind::Tcp(l) => l.accept().await.map(|(stream, remote_addr)| { + LdapStream::new(remote_addr.into(), stream.into(), timeout) + }), + ListenerKind::Uds(l) => l.accept().await.map(|(stream, remote_addr)| { + LdapStream::new(remote_addr.into(), stream.into(), timeout) + }), }; match res { diff --git a/src/main.rs b/src/main.rs index 3982ec5..2c5430e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,21 +2,14 @@ mod cli; mod error; +mod ldap; mod listener; mod stream; use cli::CliArgs; use error::GlobalError; +use ldap::ldap_handler; use listener::ListenerPath; -use stream::AbstractStream; - -#[tracing::instrument(name = "client", skip(s), fields(session = %s.session))] -async fn client(s: AbstractStream) { - tracing::info! { - remote_addr = ?s.remote_addr, - "New client connection" - }; -} #[tokio::main(flavor = "current_thread")] async fn main() -> Result<(), GlobalError> { @@ -34,8 +27,8 @@ async fn main() -> Result<(), GlobalError> { // If the connection is None, it's because the client aborted early // so there's nothing to do about it. - while let Some(stream) = listener.accept().await? { - client(stream).await; + while let Some(stream) = listener.accept_ldap().await? { + tokio::spawn(ldap_handler(stream)); } Ok(()) diff --git a/src/stream.rs b/src/stream.rs index 204b036..0278e43 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1,9 +1,8 @@ use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use tokio::net::unix::SocketAddr as UnixSocketAddr; use tokio::net::{TcpStream, UnixStream}; -use uuid::Uuid; -use std::io::Result; +use std::io::Result as IOResult; use std::marker::Unpin; use std::net::SocketAddr as TcpSocketAddr; use std::pin::{Pin, pin}; @@ -47,70 +46,42 @@ impl From for AbstractStreamKind { } } -#[derive(Debug)] -pub struct AbstractStream { - pub remote_addr: AbstractSocketAddr, - pub kind: AbstractStreamKind, - /// An arbitrary ID to correlate logs with a specific TCP session - pub session: Uuid, -} +impl Unpin for AbstractStreamKind {} -impl From<(UnixStream, UnixSocketAddr)> for AbstractStream { - fn from(res: (UnixStream, UnixSocketAddr)) -> Self { - Self { - remote_addr: res.1.into(), - kind: res.0.into(), - session: Uuid::new_v4(), - } - } -} - -impl From<(TcpStream, TcpSocketAddr)> for AbstractStream { - fn from(res: (TcpStream, std::net::SocketAddr)) -> Self { - Self { - remote_addr: res.1.into(), - kind: res.0.into(), - session: Uuid::new_v4(), - } - } -} - -impl Unpin for AbstractStream {} - -impl AsyncRead for AbstractStream { +impl AsyncRead for AbstractStreamKind { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, - ) -> Poll> { - match &mut self.kind { + ) -> Poll> { + match &mut *self { AbstractStreamKind::Tcp(stream) => pin!(stream).poll_read(cx, buf), AbstractStreamKind::Uds(stream) => pin!(stream).poll_read(cx, buf), } } } -impl AsyncWrite for AbstractStream { +impl AsyncWrite for AbstractStreamKind { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], - ) -> Poll> { - match &mut self.kind { + ) -> Poll> { + match &mut *self { AbstractStreamKind::Tcp(stream) => pin!(stream).poll_write(cx, buf), AbstractStreamKind::Uds(stream) => pin!(stream).poll_write(cx, buf), } } - fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match &mut self.kind { + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match &mut *self { AbstractStreamKind::Tcp(stream) => pin!(stream).poll_flush(cx), AbstractStreamKind::Uds(stream) => pin!(stream).poll_flush(cx), } } - fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match &mut self.kind { + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match &mut *self { AbstractStreamKind::Tcp(stream) => pin!(stream).poll_shutdown(cx), AbstractStreamKind::Uds(stream) => pin!(stream).poll_shutdown(cx), }