praxis_filter/
tcp_filter.rs

1// SPDX-License-Identifier: MIT
2// Copyright (c) 2024 Praxis Contributors
3
4//! The [`TcpFilter`] trait and per-connection [`TcpFilterContext`].
5
6use std::{borrow::Cow, sync::Arc, time::Instant};
7
8use async_trait::async_trait;
9use praxis_core::{health::HealthRegistry, kv::KvStoreRegistry};
10
11use crate::{actions::FilterAction, filter::FilterError};
12
13// -----------------------------------------------------------------------------
14// TcpFilter Trait
15// -----------------------------------------------------------------------------
16
17/// A filter that participates in TCP connection processing.
18///
19/// ```
20/// use std::{borrow::Cow, time::Instant};
21///
22/// use async_trait::async_trait;
23/// use praxis_filter::{FilterAction, FilterError, TcpFilter, TcpFilterContext};
24///
25/// struct LogFilter;
26///
27/// #[async_trait]
28/// impl TcpFilter for LogFilter {
29///     fn name(&self) -> &'static str {
30///         "log"
31///     }
32///
33///     async fn on_connect(
34///         &self,
35///         ctx: &mut TcpFilterContext<'_>,
36///     ) -> Result<FilterAction, FilterError> {
37///         println!("connection from {}", ctx.remote_addr);
38///         Ok(FilterAction::Continue)
39///     }
40/// }
41///
42/// # fn example() {
43/// let mut ctx = TcpFilterContext {
44///     remote_addr: "127.0.0.1:1234",
45///     local_addr: "0.0.0.0:8080",
46///     sni: None,
47///     upstream_addr: Some(Cow::Borrowed("10.0.0.1:80")),
48///     cluster: None,
49///     health_registry: None,
50///     kv_stores: None,
51///     connect_time: Instant::now(),
52///     bytes_in: 0,
53///     bytes_out: 0,
54/// };
55/// # }
56/// ```
57#[async_trait]
58pub trait TcpFilter: Send + Sync {
59    /// Unique name identifying this filter type.
60    fn name(&self) -> &'static str;
61
62    /// Called when a new TCP connection is accepted.
63    async fn on_connect(&self, ctx: &mut TcpFilterContext<'_>) -> Result<FilterAction, FilterError> {
64        let _ = ctx;
65        Ok(FilterAction::Continue)
66    }
67
68    /// Called when a TCP connection is closed.
69    async fn on_disconnect(&self, ctx: &mut TcpFilterContext<'_>) -> Result<(), FilterError> {
70        let _ = ctx;
71        Ok(())
72    }
73}
74
75// -----------------------------------------------------------------------------
76// TcpFilterContext
77// -----------------------------------------------------------------------------
78
79/// Per-connection state for TCP filters.
80pub struct TcpFilterContext<'a> {
81    /// Remote client address.
82    pub remote_addr: &'a str,
83
84    /// Local listener address.
85    pub local_addr: &'a str,
86
87    /// SNI hostname extracted from the TLS `ClientHello`, if present.
88    pub sni: Option<&'a str>,
89
90    /// Upstream address being proxied to.
91    ///
92    /// `None` until a static upstream or a filter (e.g. `sni_router`)
93    /// provides one.
94    pub upstream_addr: Option<Cow<'a, str>>,
95
96    /// Cluster name selected for this connection.
97    ///
98    /// Set by the listener config when `cluster` is configured.
99    /// Read by `tcp_load_balancer` to look up the strategy.
100    pub cluster: Option<Arc<str>>,
101
102    /// Shared health registry for endpoint health lookups.
103    pub health_registry: Option<&'a HealthRegistry>,
104
105    /// Named key-value stores for runtime mappings.
106    pub kv_stores: Option<&'a KvStoreRegistry>,
107
108    /// When the connection was accepted.
109    pub connect_time: Instant,
110
111    /// Bytes received from client (populated after forwarding completes).
112    pub bytes_in: u64,
113
114    /// Bytes sent to client (populated after forwarding completes).
115    pub bytes_out: u64,
116}
117
118// -----------------------------------------------------------------------------
119// Tests
120// -----------------------------------------------------------------------------
121
122#[cfg(test)]
123#[expect(clippy::allow_attributes, reason = "blanket test suppressions")]
124#[allow(
125    clippy::unwrap_used,
126    clippy::expect_used,
127    clippy::indexing_slicing,
128    clippy::panic,
129    reason = "tests"
130)]
131mod tests {
132    use super::*;
133
134    #[tokio::test]
135    async fn default_on_connect_returns_continue() {
136        let filter = NoopTcpFilter;
137        let mut ctx = TcpFilterContext {
138            remote_addr: "127.0.0.1:12345",
139            local_addr: "0.0.0.0:5432",
140            sni: None,
141            upstream_addr: Some(Cow::Borrowed("10.0.0.1:5432")),
142            cluster: None,
143            health_registry: None,
144            kv_stores: None,
145            connect_time: Instant::now(),
146            bytes_in: 0,
147            bytes_out: 0,
148        };
149        let action = filter.on_connect(&mut ctx).await.unwrap();
150        assert!(matches!(action, FilterAction::Continue));
151    }
152
153    #[tokio::test]
154    async fn default_on_disconnect_succeeds() {
155        let filter = NoopTcpFilter;
156        let mut ctx = TcpFilterContext {
157            remote_addr: "127.0.0.1:12345",
158            local_addr: "0.0.0.0:5432",
159            sni: None,
160            upstream_addr: Some(Cow::Borrowed("10.0.0.1:5432")),
161            cluster: None,
162            health_registry: None,
163            kv_stores: None,
164            connect_time: Instant::now(),
165            bytes_in: 0,
166            bytes_out: 0,
167        };
168        filter.on_disconnect(&mut ctx).await.unwrap();
169    }
170
171    // -------------------------------------------------------------------------
172    // Test Utilities
173    // -------------------------------------------------------------------------
174
175    /// Minimal TCP filter that uses all trait defaults.
176    struct NoopTcpFilter;
177
178    #[async_trait]
179    impl TcpFilter for NoopTcpFilter {
180        fn name(&self) -> &'static str {
181            "noop_tcp"
182        }
183    }
184}
praxis_filter/tcp_filter.rs

praxis_filter/
tcp_filter.rs
