diff --git a/Cargo.toml b/Cargo.toml index 9ac7fd31..4b4366f5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,7 +83,7 @@ getrandom.workspace = true lazy_static.workspace = true unftp-core.workspace = true moka = { version = "0.12.13", default-features = false, features = ["sync"] } -nix = { version = "0.30.1", default-features = false, features = ["fs"] } +nix = { version = "0.30.1", default-features = false, features = ["fs", "net", "socket"] } prometheus = { version = "0.14.0", default-features = false, optional = true } proxy-protocol = { version = "0.5.0", optional = true } rustls = { version = "0.23.36", default-features = false } diff --git a/src/lib.rs b/src/lib.rs index fea2fd7d..bf4ad1e1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -50,6 +50,8 @@ pub mod notification; pub(crate) mod server; pub mod storage; +#[cfg(unix)] +pub use crate::server::RETR_SOCKETS; pub use crate::server::ftpserver::{Server, ServerBuilder, error::ServerError, options}; type BoxError = Box; diff --git a/src/server/datachan.rs b/src/server/datachan.rs index 2bedb22e..fc97a613 100644 --- a/src/server/datachan.rs +++ b/src/server/datachan.rs @@ -11,6 +11,12 @@ use crate::{ }; use crate::server::chancomms::DataChanCmd; +#[cfg(unix)] +use std::{ + net::SocketAddr, + os::fd::{AsRawFd, BorrowedFd, RawFd}, + sync::atomic::{AtomicU64, Ordering}, +}; use std::{path::PathBuf, sync::Arc}; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf}; use tokio::net::TcpStream; @@ -42,6 +48,61 @@ use std::pin::Pin; use std::task::{Context, Poll}; use std::time::Instant; +/// Holds information about a socket processing a RETR command +#[cfg(unix)] +#[derive(Debug)] +pub struct RetrSocket { + bytes: AtomicU64, + fd: RawFd, + peer: SocketAddr, +} + +#[cfg(unix)] +impl RetrSocket { + /// How many bytes have been written to the socket so far? + /// + /// Note that this tracks bytes written to the socket, not sent on the wire. + pub fn bytes(&self) -> u64 { + self.bytes.load(Ordering::Relaxed) + } + + pub fn fd(&self) -> BorrowedFd<'_> { + // Safe because we always destroy the RetrSocket when the MeasuringWriter drops + #[allow(unsafe_code)] + unsafe { + BorrowedFd::borrow_raw(self.fd) + } + } + + fn new(w: &W) -> nix::Result { + let fd = w.as_raw_fd(); + let ss: nix::sys::socket::SockaddrStorage = nix::sys::socket::getpeername(fd)?; + let peer = if let Some(sin) = ss.as_sockaddr_in() { + SocketAddr::V4((*sin).into()) + } else if let Some(sin6) = ss.as_sockaddr_in6() { + SocketAddr::V6((*sin6).into()) + } else { + return Err(nix::errno::Errno::EINVAL); + }; + let bytes = Default::default(); + Ok(RetrSocket { bytes, fd, peer }) + } + + pub fn peer(&self) -> &SocketAddr { + &self.peer + } +} + +/// Collection of all sockets currently serving RETR commands +#[cfg(unix)] +pub static RETR_SOCKETS: std::sync::RwLock> = std::sync::RwLock::new(std::collections::BTreeMap::new()); + +#[cfg(unix)] +struct MeasuringWriter { + writer: W, + command: &'static str, +} +#[cfg(not(unix))] struct MeasuringWriter { writer: W, command: &'static str, @@ -52,12 +113,46 @@ struct MeasuringReader { command: &'static str, } +#[cfg(unix)] +impl AsyncWrite for MeasuringWriter { + fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> std::task::Poll> { + let this = self.get_mut(); + + let result = Pin::new(&mut this.writer).poll_write(cx, buf); + if let Poll::Ready(Ok(bytes_written)) = &result { + let bw = *bytes_written as u64; + RETR_SOCKETS + .read() + .unwrap() + .get(&this.writer.as_raw_fd()) + .expect("TODO: better error handling") + .bytes + .fetch_add(bw, Ordering::Relaxed); + metrics::inc_sent_bytes(*bytes_written, this.command); + } + + result + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.writer).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.writer).poll_shutdown(cx) + } +} + +#[cfg(not(unix))] impl AsyncWrite for MeasuringWriter { fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> std::task::Poll> { let this = self.get_mut(); let result = Pin::new(&mut this.writer).poll_write(cx, buf); if let Poll::Ready(Ok(bytes_written)) = &result { + let bw = *bytes_written as u64; metrics::inc_sent_bytes(*bytes_written, this.command); } @@ -87,12 +182,30 @@ impl AsyncRead for MeasuringReader { } } +#[cfg(unix)] +impl MeasuringWriter { + fn new(writer: W, command: &'static str) -> MeasuringWriter { + let retr_socket = RetrSocket::new(&writer).expect("TODO: better error handling"); + RETR_SOCKETS.write().unwrap().insert(retr_socket.fd, retr_socket); + Self { writer, command } + } +} +#[cfg(not(unix))] impl MeasuringWriter { fn new(writer: W, command: &'static str) -> MeasuringWriter { Self { writer, command } } } +#[cfg(unix)] +impl Drop for MeasuringWriter { + fn drop(&mut self) { + if let Ok(mut guard) = RETR_SOCKETS.write() { + guard.remove(&self.writer.as_raw_fd()); + } + } +} + impl MeasuringReader { fn new(reader: R, command: &'static str) -> MeasuringReader { Self { reader, command } diff --git a/src/server/mod.rs b/src/server/mod.rs index 3cac7202..103abc2f 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -19,5 +19,6 @@ pub(crate) use controlchan::Event; pub(crate) use controlchan::command::Command; pub(crate) use controlchan::reply::{Reply, ReplyCode}; pub(crate) use controlchan::{ControlChanError, ControlChanErrorKind}; - +#[cfg(unix)] +pub use datachan::RETR_SOCKETS; use session::{Session, SessionState};