Skip to main content

sozu_lib/
tcp.rs

1use std::{
2    cell::RefCell,
3    collections::{BTreeMap, HashMap, hash_map::Entry},
4    io::ErrorKind,
5    net::{Shutdown, SocketAddr},
6    os::unix::io::AsRawFd,
7    rc::Rc,
8    time::{Duration, Instant},
9};
10
11use mio::{
12    Interest, Registry, Token,
13    net::{TcpListener as MioTcpListener, TcpStream as MioTcpStream},
14    unix::SourceFd,
15};
16use rusty_ulid::Ulid;
17use sozu_command::{
18    ObjectKind,
19    config::{
20        DEFAULT_SNI_PREREAD_MAX_BYTES, DEFAULT_SNI_PREREAD_TIMEOUT, MAX_LOOP_ITERATIONS,
21        MIN_SNI_PREREAD_MAX_BYTES, validate_sni_pattern,
22    },
23    logging::{EndpointRecord, LogContext, ansi_palette},
24    proto::command::request::RequestType,
25};
26
27use crate::metrics::names;
28use crate::router::pattern_trie::{InsertResult, TrieNode};
29use crate::{
30    AcceptError, BackendConnectAction, BackendConnectionError, BackendConnectionStatus, CachedTags,
31    ListenerError, ListenerHandler, Protocol, ProxyConfiguration, ProxyError, ProxySession,
32    Readiness, SessionIsToBeClosed, SessionMetrics, SessionResult, StateMachineBuilder,
33    backends::{Backend, BackendMap},
34    pool::{Checkout, Pool},
35    protocol::{
36        Pipe,
37        pipe::WebSocketContext,
38        proxy_protocol::{
39            expect::ExpectProxyProtocol, relay::RelayProxyProtocol, send::SendProxyProtocol,
40        },
41        tcp_preread::{AlpnMatcher, PrereadConfig, shell::SniPreread},
42    },
43    retry::RetryPolicy,
44    server::{CONN_RETRIES, ListenToken, SessionManager, push_event},
45    socket::{server_bind, stats::socket_rtt},
46    sozu_command::{
47        proto::command::{
48            Event, EventKind, ProxyProtocolConfig, RequestTcpFrontend, TcpListenerConfig,
49            UpdateTcpListenerConfig, WorkerRequest, WorkerResponse,
50        },
51        ready::Ready,
52        state::ClusterId,
53    },
54    timer::TimeoutContainer,
55};
56
57StateMachineBuilder! {
58    /// The various Stages of a TCP connection:
59    ///
60    /// 1. optional SniPreread (SNI-routed listeners only, sozu-proxy/sozu#1279)
61    /// 2. optional (ExpectProxyProtocol | SendProxyProtocol | RelayProxyProtocol)
62    /// 3. Pipe
63    enum TcpStateMachine {
64        Pipe(Pipe<MioTcpStream, TcpListener>),
65        SendProxyProtocol(SendProxyProtocol<MioTcpStream>),
66        RelayProxyProtocol(RelayProxyProtocol<MioTcpStream>),
67        ExpectProxyProtocol(ExpectProxyProtocol<MioTcpStream>),
68        SniPreread(SniPreread<MioTcpStream>),
69    }
70}
71
72/// This macro is defined uniquely in this module to help the tracking of kawa h1
73/// issues inside Sōzu. Colored output uses the unified log-context scheme:
74/// bold bright-white protocol label, light-grey `Session` keyword, gray keys
75/// and bright-white values.
76macro_rules! log_context {
77    ($self:expr) => {{
78        let (open, reset, grey, gray, white) = ansi_palette();
79        format!(
80            "{gray}{ctx}{reset}\t{open}TCP{reset}\t{grey}Session{reset}({gray}frontend{reset}={white}{frontend}{reset}, {gray}backend{reset}={white}{backend}{reset})\t >>>",
81            open = open,
82            reset = reset,
83            grey = grey,
84            gray = gray,
85            white = white,
86            ctx = $self.log_context(),
87            frontend = $self.frontend_token.0,
88            backend = $self
89                .backend_token
90                .map(|token| token.0.to_string())
91                .unwrap_or_else(|| "<none>".to_string()),
92        )
93    }};
94}
95
96/// Module-level prefix for log lines emitted from this file when no
97/// [`TcpSession`] is in scope. Produces a bold bright-white `TCP` label
98/// (uniform with the per-session `log_context!`) when the logger is in
99/// colored mode. Used by [`TcpProxy`] callbacks (notify, accept,
100/// create_session, soft_stop, hard_stop, status) and the `testing`
101/// helper module which own a listener/token map but have no
102/// `frontend_token` of their own.
103macro_rules! log_module_context {
104    () => {{
105        let (open, reset, _, _, _) = sozu_command::logging::ansi_palette();
106        format!("{open}TCP{reset}\t >>>", open = open, reset = reset)
107    }};
108}
109
110pub struct TcpSession {
111    backend_buffer: Option<Checkout>,
112    backend_connected: BackendConnectionStatus,
113    backend_id: Option<String>,
114    backend_token: Option<Token>,
115    backend: Option<Rc<RefCell<Backend>>>,
116    cluster_id: Option<String>,
117    configured_backend_timeout: Duration,
118    connection_attempt: u8,
119    container_backend_timeout: TimeoutContainer,
120    container_frontend_timeout: TimeoutContainer,
121    frontend_address: Option<SocketAddr>,
122    frontend_buffer: Option<Checkout>,
123    frontend_token: Token,
124    has_been_closed: SessionIsToBeClosed,
125    last_event: Instant,
126    listener: Rc<RefCell<TcpListener>>,
127    metrics: SessionMetrics,
128    proxy: Rc<RefCell<TcpProxy>>,
129    request_id: Ulid,
130    state: TcpStateMachine,
131    /// `true` once `connect_to_backend` has accounted this session
132    /// against the per-(cluster, source-IP) connection counter. Drives
133    /// the symmetric `untrack_all_cluster_ip` call in `close`. The flag
134    /// is per-session, not per-attempt: a TCP session has at most one
135    /// `(cluster, ip)` slot, so the SessionManager-side idempotency
136    /// already covers retries — this flag exists only to short-circuit
137    /// the close path's untrack when the feature is disabled or no
138    /// admit ever ran.
139    cluster_ip_tracked: bool,
140    /// SNI-preread routing result (sozu-proxy/sozu#1279), captured once by
141    /// `upgrade_sni_preread` for every `proxy_protocol` case and consumed
142    /// (`Option::take`) at the point the session actually reaches `Pipe`:
143    /// immediately in `build_pipe_from_preread` for
144    /// `Expect`/`Relay`/`None`, or one `ready()` cycle later in
145    /// `upgrade_send` for `SendHeader` (which transitions through
146    /// `SendProxyProtocol` first). `None` for every non-SNI-routed session.
147    routed_sni: Option<String>,
148    /// Paired with `routed_sni`: the client's first ALPN offer, mapped to a
149    /// known `&'static str` label (`"h2"` / `"http/1.1"`) for the access
150    /// log, or `None` if the client offered nothing recognized. Sōzu never
151    /// negotiates ALPN itself on the TCP passthrough path -- the backend
152    /// terminates TLS -- so this is informational (the client's
153    /// preference), not a negotiated value.
154    routed_alpn_label: Option<&'static str>,
155    /// Canonical access-log tags key for the MATCHED SNI/ALPN frontend
156    /// (`sni_tags_key`), rebuilt in `upgrade_sni_preread` from the route
157    /// decision's `matched_sni_pattern` + `matched_alpn`. `None` for every
158    /// non-SNI-routed session, whose tags stay keyed by the bare listener
159    /// address exactly as before SNI routing existed. Unlike `routed_sni`
160    /// it is never consumed: `log_request` reads it for the session-level
161    /// access log, and the `Pipe` receives a clone at upgrade time for the
162    /// post-upgrade log.
163    tags_key: Option<String>,
164}
165
166impl TcpSession {
167    #[allow(clippy::too_many_arguments)]
168    fn new(
169        backend_buffer: Checkout,
170        backend_id: Option<String>,
171        cluster_id: Option<String>,
172        configured_backend_timeout: Duration,
173        configured_connect_timeout: Duration,
174        configured_frontend_timeout: Duration,
175        frontend_buffer: Checkout,
176        frontend_token: Token,
177        listener: Rc<RefCell<TcpListener>>,
178        proxy_protocol: Option<ProxyProtocolConfig>,
179        proxy: Rc<RefCell<TcpProxy>>,
180        socket: MioTcpStream,
181        wait_time: Duration,
182    ) -> TcpSession {
183        let frontend_address = socket.peer_addr().ok();
184        let mut frontend_buffer_session = None;
185        let mut backend_buffer_session = None;
186
187        let request_id = Ulid::generate();
188
189        let container_frontend_timeout =
190            TimeoutContainer::new(configured_frontend_timeout, frontend_token);
191        let container_backend_timeout = TimeoutContainer::new_empty(configured_connect_timeout);
192
193        let state = match proxy_protocol {
194            Some(ProxyProtocolConfig::RelayHeader) => {
195                backend_buffer_session = Some(backend_buffer);
196                gauge_add!(names::protocol::PROXY_RELAY, 1);
197                TcpStateMachine::RelayProxyProtocol(RelayProxyProtocol::new(
198                    socket,
199                    frontend_token,
200                    request_id,
201                    None,
202                    frontend_buffer,
203                ))
204            }
205            Some(ProxyProtocolConfig::ExpectHeader) => {
206                frontend_buffer_session = Some(frontend_buffer);
207                backend_buffer_session = Some(backend_buffer);
208                gauge_add!(names::protocol::PROXY_EXPECT, 1);
209                TcpStateMachine::ExpectProxyProtocol(ExpectProxyProtocol::new(
210                    container_frontend_timeout.clone(),
211                    socket,
212                    frontend_token,
213                    request_id,
214                ))
215            }
216            Some(ProxyProtocolConfig::SendHeader) => {
217                frontend_buffer_session = Some(frontend_buffer);
218                backend_buffer_session = Some(backend_buffer);
219                gauge_add!(names::protocol::PROXY_SEND, 1);
220                TcpStateMachine::SendProxyProtocol(SendProxyProtocol::new(
221                    socket,
222                    frontend_token,
223                    request_id,
224                    None,
225                ))
226            }
227            None => {
228                gauge_add!(names::protocol::TCP, 1);
229                let mut pipe = Pipe::new(
230                    backend_buffer,
231                    backend_id.clone(),
232                    None,
233                    None,
234                    None,
235                    None,
236                    cluster_id.clone(),
237                    frontend_buffer,
238                    frontend_token,
239                    socket,
240                    listener.clone(),
241                    Protocol::TCP,
242                    request_id,
243                    request_id,
244                    frontend_address,
245                    WebSocketContext::Tcp,
246                );
247                pipe.set_cluster_id(cluster_id.clone());
248                TcpStateMachine::Pipe(pipe)
249            }
250        };
251
252        let metrics = SessionMetrics::new(Some(wait_time));
253        //FIXME: timeout usage
254
255        TcpSession {
256            backend_buffer: backend_buffer_session,
257            backend_connected: BackendConnectionStatus::NotConnected,
258            backend_id,
259            backend_token: None,
260            backend: None,
261            cluster_id,
262            configured_backend_timeout,
263            connection_attempt: 0,
264            container_backend_timeout,
265            container_frontend_timeout,
266            frontend_address,
267            frontend_buffer: frontend_buffer_session,
268            frontend_token,
269            has_been_closed: false,
270            last_event: Instant::now(),
271            listener,
272            metrics,
273            proxy,
274            request_id,
275            state,
276            cluster_ip_tracked: false,
277            routed_sni: None,
278            routed_alpn_label: None,
279            tags_key: None,
280        }
281    }
282
283    /// Construct a session that starts in [`TcpStateMachine::SniPreread`]
284    /// instead of resolving a `proxy_protocol` up front -- the cluster (and
285    /// therefore the per-cluster `proxy_protocol`) is only known once
286    /// [`crate::protocol::tcp_preread::SniPrereadCore`] decides a route.
287    /// Mirrors [`Self::new`]'s tail; kept as a separate constructor rather
288    /// than folding a synthetic sentinel into `proxy_protocol:
289    /// Option<ProxyProtocolConfig>` (a proto-generated enum this crate does
290    /// not own).
291    #[allow(clippy::too_many_arguments)]
292    fn new_sni_preread(
293        backend_buffer: Checkout,
294        configured_backend_timeout: Duration,
295        configured_connect_timeout: Duration,
296        frontend_buffer: Checkout,
297        frontend_token: Token,
298        listener: Rc<RefCell<TcpListener>>,
299        proxy: Rc<RefCell<TcpProxy>>,
300        socket: MioTcpStream,
301        wait_time: Duration,
302        preread_timeout: Duration,
303        effective_max_bytes: usize,
304    ) -> TcpSession {
305        let frontend_address = socket.peer_addr().ok();
306        let request_id = Ulid::generate();
307
308        // Armed with the SHORT preread timeout directly (not the listener's
309        // configured front_timeout) -- `upgrade_sni_preread` restores the
310        // configured duration on the SAME container once routed, so there is
311        // exactly one `TimeoutContainer` for the frontend token throughout,
312        // never a diverging clone (a clone independently rearmed to a
313        // shorter duration would strand `TcpSession::readable`'s own
314        // unconditional `reset()` on a since-cancelled timer-wheel entry).
315        let container_frontend_timeout = TimeoutContainer::new(preread_timeout, frontend_token);
316        let container_backend_timeout = TimeoutContainer::new_empty(configured_connect_timeout);
317
318        let state = TcpStateMachine::SniPreread(SniPreread::new(
319            socket,
320            frontend_token,
321            request_id,
322            frontend_buffer,
323            effective_max_bytes,
324        ));
325
326        // Enter the `SniPreread` state: +1 the active gauge exactly once, and
327        // unconditionally, so every one of the two `-1` decrements has a
328        // matching increment. The gauge is decremented on precisely one of the
329        // two mutually-exclusive exits: the "upgrade" exit in
330        // `upgrade_sni_preread` (which first transitions `self.state` away from
331        // `SniPreread`, so `close()` cannot re-decrement), and the
332        // "reject"/"teardown" exit in `close()`'s `StateMarker::SniPreread`
333        // arm. A session therefore nets to 0 and never underflows.
334        gauge_add!(names::tcp::sni_preread::ACTIVE, 1);
335
336        let metrics = SessionMetrics::new(Some(wait_time));
337
338        TcpSession {
339            backend_buffer: Some(backend_buffer),
340            backend_connected: BackendConnectionStatus::NotConnected,
341            backend_id: None,
342            backend_token: None,
343            backend: None,
344            cluster_id: None,
345            configured_backend_timeout,
346            connection_attempt: 0,
347            container_backend_timeout,
348            container_frontend_timeout,
349            frontend_address,
350            frontend_buffer: None,
351            frontend_token,
352            has_been_closed: false,
353            last_event: Instant::now(),
354            listener,
355            metrics,
356            proxy,
357            request_id,
358            state,
359            cluster_ip_tracked: false,
360            routed_sni: None,
361            routed_alpn_label: None,
362            tags_key: None,
363        }
364    }
365
366    /// Source-IP for per-(cluster, source-IP) accounting.
367    ///
368    /// Prefer the parsed PROXY-v2 source from whichever upgrade phase is
369    /// in flight, then the post-upgrade `Pipe.session_address`, finally
370    /// the raw TCP `peer_addr` captured at session creation. The
371    /// `Pipe::session_address` itself is already PROXY-v2-aware after
372    /// `expect.rs::into_pipe` and `relay.rs::into_pipe`.
373    fn effective_session_address(&self) -> Option<SocketAddr> {
374        match &self.state {
375            TcpStateMachine::Pipe(pipe) => pipe.get_session_address(),
376            TcpStateMachine::ExpectProxyProtocol(epp) => {
377                epp.addresses.as_ref().and_then(|pa| pa.source())
378            }
379            TcpStateMachine::RelayProxyProtocol(rpp) => {
380                rpp.addresses.as_ref().and_then(|pa| pa.source())
381            }
382            TcpStateMachine::SniPreread(preread) => preread.outcome().and_then(|o| o.proxy_source),
383            TcpStateMachine::SendProxyProtocol(_) | TcpStateMachine::FailedUpgrade(_) => None,
384        }
385        .or(self.frontend_address)
386    }
387
388    fn log_request(&self) {
389        let listener = self.listener.borrow();
390        let context = self.log_context();
391        self.metrics.register_end_of_session(&context);
392        // SNI-routed sessions carry the matched front's own tags key
393        // (`sni_tags_key`, stashed by `upgrade_sni_preread`); everything
394        // else keeps the historical bare-address key.
395        let address_key = TcpFrontendTagsKey::Address(*listener.get_addr()).to_string();
396        let tags_key = self.tags_key.as_deref().unwrap_or(&address_key);
397        info_access!(
398            on_failure: { incr!(names::access_logs::UNSENT) },
399            message: None,
400            context,
401            session_address: self.frontend_address,
402            backend_address: None,
403            protocol: "TCP",
404            endpoint: EndpointRecord::Tcp,
405            tags: listener.get_tags(tags_key),
406            client_rtt: socket_rtt(self.state.front_socket()),
407            server_rtt: None,
408            user_agent: None,
409            x_request_id: None,
410            // Sōzu never terminates TLS on the TCP path (the frontend is a
411            // raw `MioTcpStream`), so no negotiated version/cipher exists.
412            // A preread SNI/ALPN (SNI-routed listeners) is stamped on the
413            // `Pipe`'s own access log via `set_tls_metadata` at upgrade
414            // time; this pre-Pipe log site emits `None` for all four TLS
415            // fields and the parsed XFF chain.
416            tls_version: None,
417            tls_cipher: None,
418            tls_sni: None,
419            tls_alpn: None,
420            xff_chain: None,
421            service_time: self.metrics.service_time(),
422            response_time: self.metrics.backend_response_time(),
423            request_time: self.metrics.request_time(),
424            start_time_ns: self.metrics.start_wall_ns(),
425            bytes_in: self.metrics.bin,
426            bytes_out: self.metrics.bout,
427            otel: None,
428        );
429    }
430
431    fn front_hup(&mut self) -> SessionResult {
432        let listener = self.listener.borrow();
433        match &mut self.state {
434            TcpStateMachine::Pipe(pipe) => pipe.frontend_hup(&mut self.metrics),
435            // No access log here, mirroring `readable()`'s own error paths
436            // for the other pre-Pipe states (none of them call
437            // `log_request()` either): the shell itself decides silent vs.
438            // metered based on whether any bytes were ever received.
439            TcpStateMachine::SniPreread(preread) => {
440                let cfg = listener.preread_config(preread.effective_max_bytes());
441                preread.on_front_closed(&cfg);
442                SessionResult::Close
443            }
444            _ => {
445                self.log_request();
446                SessionResult::Close
447            }
448        }
449    }
450
451    fn back_hup(&mut self) -> SessionResult {
452        // `SniPreread` falls into the wildcard catch-all below (unconditional
453        // close + access log), same as Send/Relay/Expect: a backend HUP
454        // while still prereading is an ordinary connect-time failure with no
455        // preread-specific accounting to do (the core only ever reasons
456        // about frontend bytes).
457        match &mut self.state {
458            TcpStateMachine::Pipe(pipe) => pipe.backend_hup(&mut self.metrics),
459            _ => {
460                self.log_request();
461                SessionResult::Close
462            }
463        }
464    }
465
466    fn log_context(&self) -> LogContext<'_> {
467        LogContext {
468            session_id: self.request_id,
469            request_id: Some(self.request_id),
470            cluster_id: self.cluster_id.as_deref(),
471            backend_id: self.backend_id.as_deref(),
472        }
473    }
474
475    fn readable(&mut self) -> SessionResult {
476        // The absolute SNI-preread deadline (armed once, at session
477        // creation) must stand while undecided -- see
478        // `frontend_timeout_resets_on_readable`'s doc.
479        if frontend_timeout_resets_on_readable(&self.state)
480            && !self.container_frontend_timeout.reset()
481        {
482            error!(
483                "{} Could not reset frontend timeout on readable",
484                log_context!(self)
485            );
486        }
487        if self.backend_connected == BackendConnectionStatus::Connected
488            && !self.container_backend_timeout.reset()
489        {
490            error!(
491                "{} Could not reset backend timeout on readable",
492                log_context!(self)
493            );
494        }
495        let listener = self.listener.borrow();
496        let result = match &mut self.state {
497            TcpStateMachine::Pipe(pipe) => pipe.readable(&mut self.metrics),
498            TcpStateMachine::RelayProxyProtocol(pp) => pp.readable(&mut self.metrics),
499            TcpStateMachine::ExpectProxyProtocol(pp) => pp.readable(&mut self.metrics),
500            TcpStateMachine::SendProxyProtocol(_) => SessionResult::Continue,
501            TcpStateMachine::SniPreread(preread) => {
502                let cfg = listener.preread_config(preread.effective_max_bytes());
503                preread.readable(&mut self.metrics, &cfg)
504            }
505            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
506        };
507        drop(listener);
508
509        // Sync `cluster_id` the moment SNI preread lands a route, so
510        // `connect_to_backend`'s cluster source (`self.cluster_id.clone().or_else(...)`)
511        // sees it without waiting for a second dispatch.
512        if let TcpStateMachine::SniPreread(preread) = &self.state
513            && self.cluster_id.is_none()
514            && let Some(outcome) = preread.outcome()
515        {
516            self.cluster_id = Some(outcome.cluster.clone());
517            // Restore the listener's configured `front_timeout` THE MOMENT
518            // routing succeeds, not only once the backend connect completes
519            // (previously done only in `upgrade_sni_preread`, which can run
520            // one or more `ready()` cycles later): a slow-but-legitimate
521            // backend connect must be bounded by
522            // `front_timeout`/`connect_timeout`, never by the short
523            // `sni_preread_timeout` that only makes sense while a route
524            // decision is still pending (sozu-proxy/sozu#1290). This is
525            // also the point from which
526            // `frontend_timeout_resets_on_readable` starts resetting this
527            // container again on every future `readable()`.
528            self.container_frontend_timeout
529                .set_duration(Duration::from_secs(
530                    self.listener.borrow().config.front_timeout as u64,
531                ));
532        }
533
534        result
535    }
536
537    fn writable(&mut self) -> SessionResult {
538        match &mut self.state {
539            TcpStateMachine::Pipe(pipe) => pipe.writable(&mut self.metrics),
540            _ => SessionResult::Continue,
541        }
542    }
543
544    fn back_readable(&mut self) -> SessionResult {
545        if !self.container_frontend_timeout.reset() {
546            error!(
547                "{} Could not reset frontend timeout on back_readable",
548                log_context!(self)
549            );
550        }
551        if !self.container_backend_timeout.reset() {
552            error!(
553                "{} Could not reset backend timeout on back_readable",
554                log_context!(self)
555            );
556        }
557
558        match &mut self.state {
559            TcpStateMachine::Pipe(pipe) => pipe.backend_readable(&mut self.metrics),
560            _ => SessionResult::Continue,
561        }
562    }
563
564    fn back_writable(&mut self) -> SessionResult {
565        match &mut self.state {
566            TcpStateMachine::Pipe(pipe) => pipe.backend_writable(&mut self.metrics),
567            TcpStateMachine::RelayProxyProtocol(pp) => pp.back_writable(&mut self.metrics),
568            TcpStateMachine::SendProxyProtocol(pp) => pp.back_writable(&mut self.metrics),
569            // The FIRST backend-writable event while routed drives the
570            // upgrade out of `SniPreread` -- see
571            // `SniPreread::back_writable`'s doc and `upgrade_sni_preread`.
572            TcpStateMachine::SniPreread(preread) => preread.back_writable(),
573            TcpStateMachine::ExpectProxyProtocol(_) => SessionResult::Continue,
574            TcpStateMachine::FailedUpgrade(_) => {
575                unreachable!()
576            }
577        }
578    }
579
580    fn back_socket_mut(&mut self) -> Option<&mut MioTcpStream> {
581        match &mut self.state {
582            TcpStateMachine::Pipe(pipe) => pipe.back_socket_mut(),
583            TcpStateMachine::SendProxyProtocol(pp) => pp.back_socket_mut(),
584            TcpStateMachine::RelayProxyProtocol(pp) => pp.back_socket_mut(),
585            TcpStateMachine::SniPreread(preread) => preread.back_socket_mut(),
586            TcpStateMachine::ExpectProxyProtocol(_) => None,
587            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
588        }
589    }
590
591    pub fn upgrade(&mut self) -> SessionIsToBeClosed {
592        let new_state = match self.state.take() {
593            TcpStateMachine::SendProxyProtocol(spp) => self.upgrade_send(spp),
594            TcpStateMachine::RelayProxyProtocol(rpp) => self.upgrade_relay(rpp),
595            TcpStateMachine::ExpectProxyProtocol(epp) => self.upgrade_expect(epp),
596            TcpStateMachine::SniPreread(preread) => self.upgrade_sni_preread(preread),
597            TcpStateMachine::Pipe(_) => None,
598            TcpStateMachine::FailedUpgrade(_) => todo!(),
599        };
600
601        match new_state {
602            Some(state) => {
603                self.state = state;
604                false
605            } // The state stays FailedUpgrade, but the Session should be closed right after
606
607            None => true,
608        }
609    }
610
611    fn upgrade_send(
612        &mut self,
613        send_proxy_protocol: SendProxyProtocol<MioTcpStream>,
614    ) -> Option<TcpStateMachine> {
615        if self.backend_buffer.is_some() && self.frontend_buffer.is_some() {
616            let mut pipe = send_proxy_protocol.into_pipe(
617                self.frontend_buffer.take().unwrap(),
618                self.backend_buffer.take().unwrap(),
619                self.listener.clone(),
620            );
621
622            // `SendProxyProtocol::into_pipe` overwrites the whole readiness
623            // (clobbering the backend-writable interest `Pipe::new` armed for a
624            // non-empty inherited frontend buffer) and re-inserts only
625            // READABLE. That is harmless for a legacy `SendHeader` session
626            // (empty accumulator) but strands the coalesced payload tail
627            // carried in from `SniPreread`'s `SendHeader` branch -- the
628            // sozu-proxy/sozu#1279 close-before-flush truncation, where bytes
629            // already queued for the backend were dropped instead of flushed.
630            // Re-run the inherited-write arm by feeding the pipe's
631            // own current events back through `restore_readiness_events`; it is
632            // additive (never clears interest) and a no-op when the buffers are
633            // empty, so the legacy path is unaffected.
634            let frontend_event = pipe.frontend_readiness.event;
635            let backend_event = pipe.backend_readiness.event;
636            pipe.restore_readiness_events(frontend_event, backend_event);
637
638            pipe.set_cluster_id(self.cluster_id.clone());
639            // Only `Some` when this `SendProxyProtocol` was itself reached
640            // via `upgrade_sni_preread`'s `SendHeader` branch (sozu-proxy/sozu#1279)
641            // -- a legacy, non-SNI-routed `SendHeader` cluster never
642            // populates these fields, so this is a no-op for it.
643            if let Some(sni) = self.routed_sni.take() {
644                pipe.set_tls_metadata(None, None, Some(sni), self.routed_alpn_label.take());
645            }
646            // `None` for legacy sessions (keeps the bare-address tags
647            // lookup); the matched front's composed key for SNI-routed ones.
648            pipe.set_tags_key(self.tags_key.clone());
649            gauge_add!(names::protocol::PROXY_SEND, -1);
650            gauge_add!(names::protocol::TCP, 1);
651            return Some(TcpStateMachine::Pipe(pipe));
652        }
653
654        error!(
655            "{} Missing the frontend or backend buffer queue, we can't switch to a pipe",
656            log_context!(self)
657        );
658        None
659    }
660
661    fn upgrade_relay(&mut self, rpp: RelayProxyProtocol<MioTcpStream>) -> Option<TcpStateMachine> {
662        if self.backend_buffer.is_some() {
663            let mut pipe =
664                rpp.into_pipe(self.backend_buffer.take().unwrap(), self.listener.clone());
665            pipe.set_cluster_id(self.cluster_id.clone());
666            gauge_add!(names::protocol::PROXY_RELAY, -1);
667            gauge_add!(names::protocol::TCP, 1);
668            return Some(TcpStateMachine::Pipe(pipe));
669        }
670
671        error!(
672            "{} Missing the backend buffer queue, we can't switch to a pipe",
673            log_context!(self)
674        );
675        None
676    }
677
678    fn upgrade_expect(
679        &mut self,
680        epp: ExpectProxyProtocol<MioTcpStream>,
681    ) -> Option<TcpStateMachine> {
682        if self.frontend_buffer.is_some() && self.backend_buffer.is_some() {
683            let mut pipe = epp.into_pipe(
684                self.frontend_buffer.take().unwrap(),
685                self.backend_buffer.take().unwrap(),
686                None,
687                None,
688                self.listener.clone(),
689            );
690
691            pipe.set_cluster_id(self.cluster_id.clone());
692            gauge_add!(names::protocol::PROXY_EXPECT, -1);
693            gauge_add!(names::protocol::TCP, 1);
694            return Some(TcpStateMachine::Pipe(pipe));
695        }
696
697        error!(
698            "{} Missing the backend buffer queue, we can't switch to a pipe",
699            log_context!(self)
700        );
701        None
702    }
703
704    /// Dispatch out of [`TcpStateMachine::SniPreread`] once its backend has
705    /// connected, by the ROUTED cluster's `proxy_protocol` config:
706    ///
707    /// - `Some(SendHeader)` -> `SendProxyProtocol` synthesizes its OWN PPv2
708    ///   header for the backend, so any inbound PPv2 prefix this listener's
709    ///   `expect_proxy` preread already parsed (`content_offset` bytes) is
710    ///   dropped from the accumulator first: the wire order is `[synth
711    ///   PPv2][ClientHello...]`, never both headers back to back.
712    /// - `Some(ExpectHeader)` -> the inbound PPv2 prefix is consumed the
713    ///   same way (Sōzu terminates it locally; the backend never sees a
714    ///   PROXY header at all), then straight into `Pipe`.
715    /// - `Some(RelayHeader)` -> NO consume: the already-parsed inbound
716    ///   header bytes ARE the header this backend expects, replayed
717    ///   verbatim ahead of the ClientHello.
718    /// - `None` -> also consumed: a listener with `expect_proxy` but a
719    ///   `None`-proxy_protocol cluster still must not leak the stray
720    ///   inbound PPv2 prefix onto a backend that expects none -- the
721    ///   preread parsed those bytes for routing only, and a backend with
722    ///   no PROXY-protocol contract would read them as part of the TLS
723    ///   stream.
724    ///
725    /// `tcp.sni_preread.duration` is recorded on EVERY exit from this
726    /// function, including the four defensive early returns below: once
727    /// `preread.outcome()`/`preread`'s `SniPreread` value is consumed by the
728    /// `TcpStateMachine::FailedUpgrade`/`Pipe` transition its caller drives,
729    /// `close()`'s `StateMarker::SniPreread` arm can no longer reach a
730    /// `SniPreread` to read `started_at()` from, so recording later is not an
731    /// option. `tcp.sni_preread.active`, by contrast, is decremented exactly
732    /// once, on the "upgrade" exit named by the gauge's `-1 on every exit`
733    /// contract; the "reject"/"teardown" exits are each other's counterpart
734    /// in `SniPreread::handle_output` (metric only) and `TcpSession::close`'s
735    /// `StateMarker::SniPreread` arm (gauge).
736    /// Shared abort path for `upgrade_sni_preread`'s early-return guards:
737    /// logs `reason` through the same envelope as the rest of this module,
738    /// then records `tcp.sni_preread.duration` -- see the long comment on
739    /// `upgrade_sni_preread` for why that metric must fire on every exit --
740    /// before returning `None`.
741    fn abort_sni_preread_upgrade(
742        &self,
743        preread: &SniPreread<MioTcpStream>,
744        reason: &str,
745    ) -> Option<TcpStateMachine> {
746        error!("{} {}", log_context!(self), reason);
747        time!(
748            names::tcp::sni_preread::DURATION,
749            preread.started_at().elapsed().as_millis() as i64
750        );
751        None
752    }
753
754    fn upgrade_sni_preread(
755        &mut self,
756        mut preread: SniPreread<MioTcpStream>,
757    ) -> Option<TcpStateMachine> {
758        // Every early return below (a route decision missing, or the
759        // backend socket/token/buffer not yet wired) must happen BEFORE the
760        // `tcp.sni_preread.active` gauge is touched: `close()`'s
761        // `StateMarker::SniPreread` arm runs unconditionally whenever
762        // `self.state` is still (or, via `FailedUpgrade`, was last)
763        // `SniPreread` -- decrementing here AND there for the same session
764        // would underflow the gauge on this (defensive, should-never-happen)
765        // failure path. The gauge is deferred to `close()` on these paths,
766        // but the duration is NOT: it is recorded right before each `return
767        // None` below, symmetric with the success path's `time!` call.
768        let Some(outcome) = preread.outcome().cloned() else {
769            return self.abort_sni_preread_upgrade(
770                &preread,
771                "upgrade_sni_preread called before a route decision",
772            );
773        };
774        let Some(backend_socket) = preread.backend.take() else {
775            return self.abort_sni_preread_upgrade(
776                &preread,
777                "SNI preread upgrade with no backend socket set",
778            );
779        };
780        let Some(backend_token) = preread.backend_token else {
781            return self.abort_sni_preread_upgrade(
782                &preread,
783                "SNI preread upgrade with no backend token set",
784            );
785        };
786        let Some(back_buffer) = self.backend_buffer.take() else {
787            return self.abort_sni_preread_upgrade(
788                &preread,
789                "SNI preread upgrade with no backend buffer queued",
790            );
791        };
792
793        gauge_add!(names::tcp::sni_preread::ACTIVE, -1);
794        time!(
795            names::tcp::sni_preread::DURATION,
796            preread.started_at().elapsed().as_millis() as i64
797        );
798
799        self.cluster_id = Some(outcome.cluster.clone());
800        // `container_frontend_timeout` is NOT restored here anymore: by the
801        // time this runs, `TcpSession::readable`'s route-capture block has
802        // already restored it to the listener's configured `front_timeout`
803        // the moment the route decision first became visible (potentially
804        // one or more `ready()` cycles before this upgrade, while the
805        // backend was still connecting) -- see that block's doc
806        // (sozu-proxy/sozu#1290). Restoring it again here
807        // would just re-arm the same duration a second time.
808        // Access-log tagging: stash the routed SNI/ALPN for
809        // whichever of the four `proxy_protocol` branches below eventually
810        // reaches `Pipe` -- immediately via `build_pipe_from_preread` for
811        // `Expect`/`Relay`/`None`, or one `ready()` cycle later via
812        // `upgrade_send` for `SendHeader` (see that method and the
813        // `routed_sni` field doc).
814        self.routed_sni = Some(outcome.sni.clone());
815        self.routed_alpn_label = known_alpn_label(&outcome.alpn);
816        // Rebuild the MATCHED front's tags key from the route decision's
817        // identity (`matched_sni_pattern` is the trie key — the configured
818        // pattern, not the client's concrete SNI — and `matched_alpn` the
819        // winning matcher), so the access log emits the tags of the front
820        // that actually routed this session, not whichever front was added
821        // last. Must compose the same key `add_tcp_front` stored — see
822        // `sni_tags_key`'s canonical-form doc.
823        self.tags_key = Some(sni_tags_key(
824            self.listener.borrow().get_addr(),
825            &outcome.matched_sni_pattern,
826            &alpn_matcher_protocols(&outcome.matched_alpn),
827        ));
828
829        let proxy_protocol = self
830            .proxy
831            .borrow()
832            .configs
833            .get(&outcome.cluster)
834            .and_then(|c| c.proxy_protocol);
835
836        let frontend_event = preread.frontend_readiness.event;
837        let backend_event = preread.backend_readiness.event;
838        let mut frontend_buffer = preread.frontend_buffer;
839        let frontend = preread.frontend;
840        let frontend_token = preread.frontend_token;
841        let request_id = preread.request_id;
842
843        match proxy_protocol {
844            Some(ProxyProtocolConfig::SendHeader) => {
845                frontend_buffer.consume(outcome.content_offset);
846                self.frontend_buffer = Some(frontend_buffer);
847                self.backend_buffer = Some(back_buffer);
848                gauge_add!(names::protocol::PROXY_SEND, 1);
849                let mut spp = SendProxyProtocol::new(
850                    frontend,
851                    frontend_token,
852                    request_id,
853                    Some(backend_socket),
854                );
855                spp.frontend_readiness.event = frontend_event;
856                spp.backend_readiness.event = backend_event;
857                spp.set_back_token(backend_token);
858                spp.set_back_connected(BackendConnectionStatus::Connected);
859                Some(TcpStateMachine::SendProxyProtocol(spp))
860            }
861            Some(ProxyProtocolConfig::ExpectHeader) | None => {
862                frontend_buffer.consume(outcome.content_offset);
863                Some(self.build_pipe_from_preread(
864                    back_buffer,
865                    frontend_buffer,
866                    frontend,
867                    frontend_token,
868                    frontend_event,
869                    backend_event,
870                    backend_socket,
871                    backend_token,
872                    request_id,
873                    outcome.proxy_source,
874                    outcome.cluster,
875                ))
876            }
877            Some(ProxyProtocolConfig::RelayHeader) => Some(self.build_pipe_from_preread(
878                back_buffer,
879                frontend_buffer,
880                frontend,
881                frontend_token,
882                frontend_event,
883                backend_event,
884                backend_socket,
885                backend_token,
886                request_id,
887                outcome.proxy_source,
888                outcome.cluster,
889            )),
890        }
891    }
892
893    #[allow(clippy::too_many_arguments)]
894    fn build_pipe_from_preread(
895        &mut self,
896        back_buffer: Checkout,
897        frontend_buffer: Checkout,
898        frontend: MioTcpStream,
899        frontend_token: Token,
900        frontend_event: Ready,
901        backend_event: Ready,
902        backend_socket: MioTcpStream,
903        backend_token: Token,
904        request_id: Ulid,
905        proxy_source: Option<SocketAddr>,
906        cluster_id: ClusterId,
907    ) -> TcpStateMachine {
908        let addr = proxy_source.or(self.frontend_address);
909        let mut pipe = Pipe::new(
910            back_buffer,
911            self.backend_id.clone(),
912            Some(backend_socket),
913            None,
914            None,
915            None,
916            Some(cluster_id),
917            frontend_buffer,
918            frontend_token,
919            frontend,
920            self.listener.clone(),
921            Protocol::TCP,
922            request_id,
923            request_id,
924            addr,
925            WebSocketContext::Tcp,
926        );
927        // `Pipe::new` armed backend-writable for the inherited frontend
928        // accumulator (the ClientHello + any coalesced payload) via
929        // `arm_inherited_buffer_writes`. Restore the preread's readiness
930        // events through `restore_readiness_events` rather than a bare
931        // `pipe.frontend_readiness.event = …` / `pipe.backend_readiness.event
932        // = …` pair: it sets both `.event`s and THEN re-runs the inherited
933        // arm, so the synthetic backend-writable event survives even in the
934        // case where the restored `backend_event` does not itself carry
935        // WRITABLE (the byte-for-byte drain of the accumulator must not depend
936        // on that). The eventual flush-on-close of that accumulator is
937        // guaranteed by `Pipe::readable`'s half-close drain (sozu-proxy/sozu#1279).
938        pipe.restore_readiness_events(frontend_event, backend_event);
939        pipe.set_back_token(backend_token);
940        // Access-log tagging: reaching `Pipe` straight from
941        // `SniPreread` (Expect/Relay/None) -- unlike `SendHeader`, which
942        // detours through `SendProxyProtocol` first (see `upgrade_send`).
943        if let Some(sni) = self.routed_sni.take() {
944            pipe.set_tls_metadata(None, None, Some(sni), self.routed_alpn_label.take());
945        }
946        // The matched front's composed tags key (always `Some` here — this
947        // is only reachable from `upgrade_sni_preread`, which just set it).
948        pipe.set_tags_key(self.tags_key.clone());
949        gauge_add!(names::protocol::TCP, 1);
950        TcpStateMachine::Pipe(pipe)
951    }
952
953    fn front_readiness(&mut self) -> &mut Readiness {
954        match &mut self.state {
955            TcpStateMachine::Pipe(pipe) => &mut pipe.frontend_readiness,
956            TcpStateMachine::SendProxyProtocol(pp) => &mut pp.frontend_readiness,
957            TcpStateMachine::RelayProxyProtocol(pp) => &mut pp.frontend_readiness,
958            TcpStateMachine::ExpectProxyProtocol(pp) => &mut pp.frontend_readiness,
959            TcpStateMachine::SniPreread(preread) => &mut preread.frontend_readiness,
960            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
961        }
962    }
963
964    fn back_readiness(&mut self) -> Option<&mut Readiness> {
965        match &mut self.state {
966            TcpStateMachine::Pipe(pipe) => Some(&mut pipe.backend_readiness),
967            TcpStateMachine::SendProxyProtocol(pp) => Some(&mut pp.backend_readiness),
968            TcpStateMachine::RelayProxyProtocol(pp) => Some(&mut pp.backend_readiness),
969            TcpStateMachine::SniPreread(preread) => Some(&mut preread.backend_readiness),
970            TcpStateMachine::ExpectProxyProtocol(_) => None,
971            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
972        }
973    }
974
975    fn set_back_socket(&mut self, socket: MioTcpStream) {
976        match &mut self.state {
977            TcpStateMachine::Pipe(pipe) => pipe.set_back_socket(socket),
978            TcpStateMachine::SendProxyProtocol(pp) => pp.set_back_socket(socket),
979            TcpStateMachine::RelayProxyProtocol(pp) => pp.set_back_socket(socket),
980            TcpStateMachine::SniPreread(preread) => preread.set_back_socket(socket),
981            TcpStateMachine::ExpectProxyProtocol(_) => {
982                error!(
983                    "{} We should not set the back socket for the expect proxy protocol",
984                    log_context!(self)
985                );
986                panic!(
987                    "{} We should not set the back socket for the expect proxy protocol",
988                    log_context!(self)
989                );
990            }
991            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
992        }
993    }
994
995    fn set_back_token(&mut self, token: Token) {
996        // The frontend must own a token distinct from the backend's: the two
997        // index different slab slots, so wiring the same token to both would
998        // alias two sessions onto one slot.
999        debug_assert_ne!(
1000            token, self.frontend_token,
1001            "backend token must differ from the frontend token"
1002        );
1003        self.backend_token = Some(token);
1004
1005        match &mut self.state {
1006            TcpStateMachine::Pipe(pipe) => pipe.set_back_token(token),
1007            TcpStateMachine::SendProxyProtocol(pp) => pp.set_back_token(token),
1008            TcpStateMachine::SniPreread(preread) => preread.set_back_token(token),
1009            TcpStateMachine::RelayProxyProtocol(pp) => pp.set_back_token(token),
1010            TcpStateMachine::ExpectProxyProtocol(_) => self.backend_token = Some(token),
1011            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
1012        }
1013
1014        // Postcondition: the session now owns exactly the token it was asked
1015        // to register — every arm above (including the Expect arm, which only
1016        // stores the session-side token) leaves `backend_token == Some(token)`.
1017        debug_assert_eq!(
1018            self.backend_token,
1019            Some(token),
1020            "set_back_token must leave the session owning the registered token"
1021        );
1022    }
1023
1024    fn set_backend_id(&mut self, id: String) {
1025        self.backend_id = Some(id.clone());
1026        if let TcpStateMachine::Pipe(pipe) = &mut self.state {
1027            pipe.set_backend_id(Some(id));
1028        }
1029    }
1030
1031    fn back_connected(&self) -> BackendConnectionStatus {
1032        self.backend_connected
1033    }
1034
1035    fn set_back_connected(&mut self, status: BackendConnectionStatus) {
1036        let last = self.backend_connected;
1037        // Transitioning INTO `Connected` bumps the backend-connection gauge by
1038        // exactly +1. Doing so from an already-`Connected` state would
1039        // double-count (gauge drift that only `close_backend`'s single -1
1040        // would later reconcile, leaving the gauge permanently +1). The
1041        // promotion always comes from a `Connecting` (the normal handshake
1042        // completion in `ready_inner`) — never from `Connected` itself.
1043        debug_assert!(
1044            status != BackendConnectionStatus::Connected
1045                || last != BackendConnectionStatus::Connected,
1046            "set_back_connected(Connected) must not run on an already-Connected backend (gauge would double-count)"
1047        );
1048        self.backend_connected = status;
1049
1050        // Postcondition: the requested status is now in effect.
1051        debug_assert_eq!(
1052            self.backend_connected, status,
1053            "set_back_connected must record the requested status"
1054        );
1055
1056        if status == BackendConnectionStatus::Connected {
1057            gauge_add!(names::backend::CONNECTIONS, 1);
1058            gauge_add!(
1059                names::backend::CONNECTIONS_PER_BACKEND,
1060                1,
1061                self.cluster_id.as_deref(),
1062                self.metrics.backend_id.as_deref()
1063            );
1064
1065            // the back timeout was of connect_timeout duration before,
1066            // now that we're connected, move to backend_timeout duration
1067            self.container_backend_timeout
1068                .set_duration(self.configured_backend_timeout);
1069            self.container_frontend_timeout.reset();
1070
1071            if let TcpStateMachine::SendProxyProtocol(spp) = &mut self.state {
1072                spp.set_back_connected(BackendConnectionStatus::Connected);
1073            }
1074
1075            if let Some(backend) = self.backend.as_ref() {
1076                let mut backend = backend.borrow_mut();
1077
1078                if backend.retry_policy.is_down() {
1079                    incr!(
1080                        "backend.up",
1081                        self.cluster_id.as_deref(),
1082                        self.metrics.backend_id.as_deref()
1083                    );
1084                    gauge!(
1085                        names::backend::AVAILABLE,
1086                        1,
1087                        self.cluster_id.as_deref(),
1088                        self.metrics.backend_id.as_deref()
1089                    );
1090                    info!(
1091                        "{} backend server {} at {} is up",
1092                        log_context!(self),
1093                        backend.backend_id,
1094                        backend.address
1095                    );
1096                    push_event(Event {
1097                        kind: EventKind::BackendUp as i32,
1098                        backend_id: Some(backend.backend_id.to_owned()),
1099                        address: Some(backend.address.into()),
1100                        cluster_id: None,
1101                        metric_detail: None,
1102                    });
1103                }
1104
1105                if let BackendConnectionStatus::Connecting(start) = last {
1106                    backend.set_connection_time(Instant::now() - start);
1107                }
1108
1109                //successful connection, rest failure counter
1110                backend.failures = 0;
1111                backend.retry_policy.succeed();
1112            }
1113        }
1114    }
1115
1116    fn remove_backend(&mut self) {
1117        if let Some(backend) = self.backend.take() {
1118            (*backend.borrow_mut()).dec_connections();
1119        }
1120
1121        self.backend_token = None;
1122
1123        // Postcondition: the backend handle and its token are torn down
1124        // together — neither may outlive the other (a dangling token would
1125        // leave a stale slab reference; a dangling handle would over-count
1126        // backend connections).
1127        debug_assert!(
1128            self.backend.is_none(),
1129            "remove_backend must release the backend handle"
1130        );
1131        debug_assert!(
1132            self.backend_token.is_none(),
1133            "remove_backend must clear the backend token"
1134        );
1135    }
1136
1137    fn fail_backend_connection(&mut self) {
1138        if let Some(backend) = self.backend.as_ref() {
1139            let backend = &mut *backend.borrow_mut();
1140            backend.failures += 1;
1141
1142            let already_unavailable = backend.retry_policy.is_down();
1143            backend.retry_policy.fail();
1144            incr!(
1145                "backend.connections.error",
1146                self.cluster_id.as_deref(),
1147                self.metrics.backend_id.as_deref()
1148            );
1149            if !already_unavailable && backend.retry_policy.is_down() {
1150                error!(
1151                    "{} backend server {} at {} is down",
1152                    log_context!(self),
1153                    backend.backend_id,
1154                    backend.address
1155                );
1156                incr!(
1157                    "backend.down",
1158                    self.cluster_id.as_deref(),
1159                    self.metrics.backend_id.as_deref()
1160                );
1161                gauge!(
1162                    names::backend::AVAILABLE,
1163                    0,
1164                    self.cluster_id.as_deref(),
1165                    self.metrics.backend_id.as_deref()
1166                );
1167
1168                push_event(Event {
1169                    kind: EventKind::BackendDown as i32,
1170                    backend_id: Some(backend.backend_id.to_owned()),
1171                    address: Some(backend.address.into()),
1172                    cluster_id: None,
1173                    metric_detail: None,
1174                });
1175            }
1176        }
1177    }
1178
1179    pub fn test_back_socket(&mut self) -> SessionIsToBeClosed {
1180        match self.back_socket_mut() {
1181            Some(ref mut s) => {
1182                let mut tmp = [0u8; 1];
1183                let res = s.peek(&mut tmp[..]);
1184
1185                match res {
1186                    // if the socket is half open, it will report 0 bytes read (EOF)
1187                    Ok(0) => false,
1188                    Ok(_) => true,
1189                    Err(e) => matches!(e.kind(), std::io::ErrorKind::WouldBlock),
1190                }
1191            }
1192            None => false,
1193        }
1194    }
1195
1196    pub fn cancel_timeouts(&mut self) {
1197        self.container_frontend_timeout.cancel();
1198        self.container_backend_timeout.cancel();
1199    }
1200
1201    /// Full cross-field invariant sweep for the TCP session state machine.
1202    ///
1203    /// Run as a run-to-completion postcondition at the END of `ready()` (the
1204    /// only public entry point that drives the front/back token + readiness
1205    /// state machine). These are OUR-logic invariants — never reachable from
1206    /// hostile traffic — so a violation is a bug in Sōzu, not a malformed
1207    /// peer. Compiled out in release.
1208    #[cfg(debug_assertions)]
1209    fn check_invariants(&self) {
1210        // Connection-attempt budget: every retry path increments
1211        // `connection_attempt` and `connect_to_backend` refuses once the
1212        // counter reaches `CONN_RETRIES`, so the value can touch but never
1213        // exceed the configured ceiling (and resets to 0 on success).
1214        debug_assert!(
1215            self.connection_attempt <= CONN_RETRIES,
1216            "connection_attempt ({}) must never exceed CONN_RETRIES ({})",
1217            self.connection_attempt,
1218            CONN_RETRIES
1219        );
1220
1221        // Token ownership: a fully-connected backend always owns a backend
1222        // token (set by `set_back_token` during `connect_to_backend`, before
1223        // the status can ever flip to `Connected`). The `Connecting` phase is
1224        // deliberately excluded: there is a transient window inside
1225        // `connect_to_backend` where the status is `Connecting` but the token
1226        // has not been wired yet — that window never spans a `ready()`
1227        // boundary, so the postcondition still holds here.
1228        if self.backend_connected == BackendConnectionStatus::Connected {
1229            debug_assert!(
1230                self.backend_token.is_some(),
1231                "a Connected backend must own a backend token"
1232            );
1233        }
1234
1235        // A live backend handle implies the matching token is present: the
1236        // two are wired together in `connect_to_backend` and torn down
1237        // together in `remove_backend` (which clears the token) — they must
1238        // never drift apart. (For the pure-TCP proxy `backend` is currently
1239        // always `None`, so this is a guard against a future regression that
1240        // starts populating it without the token.)
1241        if self.backend.is_some() {
1242            debug_assert!(
1243                self.backend_token.is_some(),
1244                "a live backend handle must have a backend token"
1245            );
1246        }
1247
1248        // Once the session has been closed it is terminal: the backend has
1249        // been released and the per-(cluster, source-IP) slot untracked.
1250        if self.has_been_closed {
1251            debug_assert!(
1252                self.backend.is_none(),
1253                "a closed session must have released its backend handle"
1254            );
1255            debug_assert!(
1256                !self.cluster_ip_tracked,
1257                "a closed session must have untracked its (cluster, source-IP) slot"
1258            );
1259        }
1260    }
1261
1262    /// Attempt a fresh backend connect, exactly like `ready_inner`'s
1263    /// top-of-function gate -- but callable a second time from inside the
1264    /// dispatch loop. A `SniPreread` session's route decision can complete
1265    /// INSIDE `readable()`'s own dispatch (the SAME `ready_inner` call), and
1266    /// without a second attempt right after that dispatch the session would
1267    /// stall until an unrelated readiness event re-entered `ready_inner`.
1268    ///
1269    /// A no-op whenever `back_connected() != NotConnected` (already
1270    /// attempted, or backend already up) or the state is a NOT-YET-ROUTED
1271    /// `SniPreread` (the cluster -- and therefore the backend to dial -- is
1272    /// unknown until `SniPrereadCore` decides), so this changes nothing for
1273    /// any pre-existing state/path.
1274    fn attempt_backend_connect_if_needed(
1275        &mut self,
1276        session: &Rc<RefCell<dyn ProxySession>>,
1277    ) -> Option<SessionResult> {
1278        if self.back_connected() != BackendConnectionStatus::NotConnected {
1279            return None;
1280        }
1281        if matches!(&self.state, TcpStateMachine::SniPreread(preread) if !preread.is_routed()) {
1282            return None;
1283        }
1284
1285        let connection_result = self.connect_to_backend(session.clone());
1286        if let Err(err) = &connection_result {
1287            match err {
1288                // Already logged at warn! + metered at the retry-budget
1289                // gate in connect_to_backend; avoid double-emission.
1290                BackendConnectionError::MaxConnectionRetries(_) => trace!(
1291                    "{} Error connecting to backend: {}",
1292                    log_context!(self),
1293                    err
1294                ),
1295                _ => warn!(
1296                    "{} Error connecting to backend: {}",
1297                    log_context!(self),
1298                    err
1299                ),
1300            }
1301        }
1302        handle_connection_result(connection_result)
1303    }
1304
1305    fn ready_inner(&mut self, session: Rc<RefCell<dyn ProxySession>>) -> SessionResult {
1306        let mut counter = 0;
1307
1308        let back_connected = self.back_connected();
1309        if back_connected.is_connecting() {
1310            if self.back_readiness().unwrap().event.is_hup() && !self.test_back_socket() {
1311                //retry connecting the backend
1312                debug!(
1313                    "{} error connecting to backend, trying again",
1314                    log_context!(self)
1315                );
1316                self.connection_attempt += 1;
1317                self.fail_backend_connection();
1318
1319                // trigger a backend reconnection
1320                self.close_backend();
1321                let connection_result = self.connect_to_backend(session.clone());
1322                if let Err(err) = &connection_result {
1323                    match err {
1324                        // Already logged at warn! + metered at the retry-budget
1325                        // gate in connect_to_backend; avoid double-emission.
1326                        BackendConnectionError::MaxConnectionRetries(_) => trace!(
1327                            "{} Error connecting to backend: {}",
1328                            log_context!(self),
1329                            err
1330                        ),
1331                        _ => warn!(
1332                            "{} Error connecting to backend: {}",
1333                            log_context!(self),
1334                            err
1335                        ),
1336                    }
1337                }
1338
1339                if let Some(state_result) = handle_connection_result(connection_result) {
1340                    return state_result;
1341                }
1342            } else if self.back_readiness().unwrap().event != Ready::EMPTY {
1343                self.connection_attempt = 0;
1344                self.set_back_connected(BackendConnectionStatus::Connected);
1345            }
1346        } else if back_connected == BackendConnectionStatus::NotConnected
1347            && let Some(state_result) = self.attempt_backend_connect_if_needed(&session)
1348        {
1349            return state_result;
1350        }
1351
1352        if self.front_readiness().event.is_hup() {
1353            let session_result = self.front_hup();
1354            if session_result != SessionResult::Continue {
1355                return session_result;
1356            }
1357            // `front_hup` drained in-flight request bytes and wants the
1358            // session kept alive (`Pipe::frontend_hup`'s in-flight branch):
1359            // the client already sent FIN, so under edge-triggered epoll no
1360            // further frontend event will ever arrive -- returning here
1361            // would stall the session forever waiting for a wake-up that
1362            // never comes. Clear the now-consumed HUP bit and fall through
1363            // into the loop below so `readable` (drains the kernel tail to
1364            // EOF) and `back_writable` (flushes `frontend_buffer`) can run
1365            // synchronously in this same pass, exactly how a backend HUP is
1366            // already handled inside the loop.
1367            self.front_readiness().event.remove(Ready::HUP);
1368        }
1369
1370        while counter < MAX_LOOP_ITERATIONS {
1371            let front_interest = self.front_readiness().interest & self.front_readiness().event;
1372            let back_interest = self
1373                .back_readiness()
1374                .map(|r| r.interest & r.event)
1375                .unwrap_or(Ready::EMPTY);
1376
1377            trace!(
1378                "{} Frontend interest({:?}) and backend interest({:?})",
1379                log_context!(self),
1380                front_interest,
1381                back_interest
1382            );
1383
1384            if front_interest == Ready::EMPTY && back_interest == Ready::EMPTY {
1385                break;
1386            }
1387
1388            if self
1389                .back_readiness()
1390                .map(|r| r.event.is_hup())
1391                .unwrap_or(false)
1392                && self.front_readiness().interest.is_writable()
1393                && !self.front_readiness().event.is_writable()
1394            {
1395                break;
1396            }
1397
1398            if front_interest.is_readable() {
1399                let session_result = self.readable();
1400                if session_result != SessionResult::Continue {
1401                    return session_result;
1402                }
1403                // A `SniPreread` route decision can complete INSIDE this
1404                // very `readable()` call; without a second attempt here the
1405                // session would stall until an unrelated readiness event
1406                // re-entered `ready_inner` to reach the top-of-function
1407                // connect gate. A no-op for every other state/backend
1408                // status (see `attempt_backend_connect_if_needed`'s guard).
1409                if let Some(state_result) = self.attempt_backend_connect_if_needed(&session) {
1410                    return state_result;
1411                }
1412            }
1413
1414            if back_interest.is_writable() {
1415                let session_result = self.back_writable();
1416                if session_result != SessionResult::Continue {
1417                    return session_result;
1418                }
1419            }
1420
1421            if back_interest.is_readable() {
1422                let session_result = self.back_readable();
1423                if session_result != SessionResult::Continue {
1424                    return session_result;
1425                }
1426            }
1427
1428            if front_interest.is_writable() {
1429                let session_result = self.writable();
1430                if session_result != SessionResult::Continue {
1431                    return session_result;
1432                }
1433            }
1434
1435            if back_interest.is_hup() {
1436                let session_result = self.back_hup();
1437                if session_result != SessionResult::Continue {
1438                    return session_result;
1439                }
1440            }
1441
1442            if front_interest.is_error() {
1443                error!(
1444                    "{} Frontend socket error, disconnecting",
1445                    log_context!(self)
1446                );
1447                self.front_readiness().interest = Ready::EMPTY;
1448                if let Some(r) = self.back_readiness() {
1449                    r.interest = Ready::EMPTY;
1450                }
1451
1452                return SessionResult::Close;
1453            }
1454
1455            if back_interest.is_error() && self.back_hup() == SessionResult::Close {
1456                self.front_readiness().interest = Ready::EMPTY;
1457                if let Some(r) = self.back_readiness() {
1458                    r.interest = Ready::EMPTY;
1459                }
1460
1461                error!("{} backend socket error, disconnecting", log_context!(self));
1462                return SessionResult::Close;
1463            }
1464
1465            counter += 1;
1466        }
1467
1468        if counter >= MAX_LOOP_ITERATIONS {
1469            error!(
1470                "{} Handling session went through {} iterations, there's a probable infinite loop bug, closing the connection",
1471                log_context!(self),
1472                MAX_LOOP_ITERATIONS
1473            );
1474
1475            incr!(names::tcp::INFINITE_LOOP_ERROR);
1476
1477            let front_interest = self.front_readiness().interest & self.front_readiness().event;
1478            let back_interest = self
1479                .back_readiness()
1480                .map(|r| r.interest & r.event)
1481                .unwrap_or(Ready::EMPTY);
1482
1483            let back = self.back_readiness().cloned();
1484
1485            error!(
1486                "{} readiness: front {:?} / back {:?} | front: {:?} | back: {:?} ",
1487                log_context!(self),
1488                self.front_readiness(),
1489                back,
1490                front_interest,
1491                back_interest
1492            );
1493
1494            self.print_session();
1495
1496            return SessionResult::Close;
1497        }
1498
1499        SessionResult::Continue
1500    }
1501
1502    /// TCP session closes its backend on its own, without defering this task to the state
1503    fn close_backend(&mut self) {
1504        if let (Some(token), Some(fd)) = (
1505            self.backend_token,
1506            self.back_socket_mut().map(|s| s.as_raw_fd()),
1507        ) {
1508            let proxy = self.proxy.borrow();
1509            if let Err(e) = proxy.registry.deregister(&mut SourceFd(&fd)) {
1510                error!(
1511                    "{} Error deregistering socket({:?}): {:?}",
1512                    log_context!(self),
1513                    fd,
1514                    e
1515                );
1516            }
1517
1518            proxy.sessions.borrow_mut().slab.try_remove(token.0);
1519        }
1520        self.remove_backend();
1521
1522        let back_connected = self.back_connected();
1523        if back_connected != BackendConnectionStatus::NotConnected {
1524            if let Some(r) = self.back_readiness() {
1525                r.event = Ready::EMPTY;
1526            }
1527
1528            let log_context = log_context!(self);
1529            if let Some(sock) = self.back_socket_mut() {
1530                // TCP-only backend in the pure-TCP proxy: no outbound TLS
1531                // buffer to truncate, so `Shutdown::Both` is the right call.
1532                // If the TCP listener ever gains an inline TLS upgrade,
1533                // switch to `Shutdown::Write` here.
1534                if let Err(e) = sock.shutdown(Shutdown::Both)
1535                    && e.kind() != ErrorKind::NotConnected
1536                {
1537                    error!(
1538                        "{} Error closing back socket({:?}): {:?}",
1539                        log_context, sock, e
1540                    );
1541                }
1542            }
1543        }
1544
1545        // The -1 here pairs with the +1 in `set_back_connected(Connected)`:
1546        // we decrement the gauge exactly once, iff this session had actually
1547        // reached `Connected`. A `Connecting`/`NotConnected` backend never
1548        // bumped the gauge, so it must not decrement it either — that
1549        // asymmetry would underflow the gauge (a correctness bug, never a
1550        // rounding issue).
1551        if back_connected == BackendConnectionStatus::Connected {
1552            gauge_add!(names::backend::CONNECTIONS, -1);
1553            gauge_add!(
1554                names::backend::CONNECTIONS_PER_BACKEND,
1555                -1,
1556                self.cluster_id.as_deref(),
1557                self.metrics.backend_id.as_deref()
1558            );
1559        }
1560
1561        self.set_back_connected(BackendConnectionStatus::NotConnected);
1562
1563        // Postcondition: the backend is fully torn down — `remove_backend`
1564        // cleared the token/handle above and the status is now `NotConnected`,
1565        // so a subsequent `connect_to_backend` starts from a clean slate.
1566        debug_assert_eq!(
1567            self.backend_connected,
1568            BackendConnectionStatus::NotConnected,
1569            "close_backend must leave the backend NotConnected"
1570        );
1571        debug_assert!(
1572            self.backend_token.is_none(),
1573            "close_backend must clear the backend token"
1574        );
1575    }
1576
1577    fn connect_to_backend(
1578        &mut self,
1579        session_rc: Rc<RefCell<dyn ProxySession>>,
1580    ) -> Result<BackendConnectAction, BackendConnectionError> {
1581        // Precondition: the retry budget can sit AT the ceiling (the gate
1582        // below converts that into `MaxConnectionRetries`) but the increment
1583        // in `ready_inner` must never have pushed it past `CONN_RETRIES`.
1584        debug_assert!(
1585            self.connection_attempt <= CONN_RETRIES,
1586            "connection_attempt ({}) overflowed CONN_RETRIES ({}) before the retry gate",
1587            self.connection_attempt,
1588            CONN_RETRIES
1589        );
1590
1591        // Prefer the SNI-routed cluster (set by `TcpSession::readable` once
1592        // `SniPrereadCore` decides) over the listener's legacy no-SNI
1593        // catch-all -- a listener never configures both (sozu-proxy/sozu#1279),
1594        // but this order is also simply correct for the routed case, where
1595        // `listener.cluster_id` is `None`.
1596        let cluster_id = self
1597            .cluster_id
1598            .clone()
1599            .or_else(|| self.listener.borrow().cluster_id.clone())
1600            .ok_or(BackendConnectionError::NotFound(ObjectKind::TcpCluster))?;
1601
1602        self.cluster_id = Some(cluster_id.clone());
1603
1604        if self.connection_attempt >= CONN_RETRIES {
1605            incr!(
1606                "backend.connect.retries_exhausted",
1607                self.cluster_id.as_deref(),
1608                self.metrics.backend_id.as_deref()
1609            );
1610            warn!(
1611                "{} Max connection attempt reached ({})",
1612                log_context!(self),
1613                self.connection_attempt
1614            );
1615            return Err(BackendConnectionError::MaxConnectionRetries(Some(
1616                cluster_id,
1617            )));
1618        }
1619
1620        if self.proxy.borrow().sessions.borrow().at_capacity() {
1621            return Err(BackendConnectionError::MaxSessionsMemory);
1622        }
1623
1624        // Per-(cluster, source-IP) connection limit gate (TCP). The
1625        // source IP comes from `effective_session_address`, which folds
1626        // a parsed PROXY-v2 source over the raw `peer_addr`. The mux's
1627        // Router does the same gate for HTTP/HTTPS sessions; here it
1628        // runs for raw TCP. Rejection produces a graceful TCP FIN via
1629        // `BackendConnectionError::TooManyConnectionsPerIp` →
1630        // `handle_connection_result` → `SessionResult::Close` — TCP has
1631        // no HTTP envelope to carry a 429 / `Retry-After`.
1632        let cluster_max_connections_per_ip = self
1633            .proxy
1634            .borrow()
1635            .configs
1636            .get(&cluster_id)
1637            .and_then(|c| c.max_connections_per_ip);
1638        if let Some(ip) = self.effective_session_address().map(|sa| sa.ip()) {
1639            let sessions_rc = self.proxy.borrow().sessions.clone();
1640            let at_limit = sessions_rc.borrow().cluster_ip_at_limit(
1641                self.frontend_token,
1642                &cluster_id,
1643                &ip,
1644                cluster_max_connections_per_ip,
1645            );
1646            if at_limit {
1647                debug!(
1648                    "{} per-(cluster, source-IP) limit hit for cluster {} from {}",
1649                    log_context!(self),
1650                    cluster_id,
1651                    ip
1652                );
1653                return Err(BackendConnectionError::TooManyConnectionsPerIp { cluster_id });
1654            }
1655            sessions_rc
1656                .borrow_mut()
1657                .track_cluster_ip(self.frontend_token, cluster_id.clone(), ip);
1658            self.cluster_ip_tracked = true;
1659        }
1660
1661        let (backend, mut stream) = self
1662            .proxy
1663            .borrow()
1664            .backends
1665            .borrow_mut()
1666            .backend_from_cluster_id(&cluster_id)
1667            .map_err(BackendConnectionError::Backend)?;
1668
1669        if let Err(e) = stream.set_nodelay(true) {
1670            error!(
1671                "{} Error setting nodelay on back socket({:?}): {:?}",
1672                log_context!(self),
1673                stream,
1674                e
1675            );
1676        }
1677        self.backend_connected = BackendConnectionStatus::Connecting(Instant::now());
1678
1679        let back_token = {
1680            let proxy = self.proxy.borrow();
1681            let mut s = proxy.sessions.borrow_mut();
1682            let entry = s.slab.vacant_entry();
1683            let back_token = Token(entry.key());
1684            let _entry = entry.insert(session_rc.clone());
1685            back_token
1686        };
1687
1688        if let Err(e) = self.proxy.borrow().registry.register(
1689            &mut stream,
1690            back_token,
1691            Interest::READABLE | Interest::WRITABLE,
1692        ) {
1693            error!(
1694                "{} Error registering back socket({:?}): {:?}",
1695                log_context!(self),
1696                stream,
1697                e
1698            );
1699        }
1700
1701        self.container_backend_timeout.set(back_token);
1702
1703        self.set_back_token(back_token);
1704        self.set_back_socket(stream);
1705
1706        self.metrics.backend_id = Some(backend.borrow().backend_id.clone());
1707        self.metrics.backend_start();
1708        self.set_backend_id(backend.borrow().backend_id.clone());
1709
1710        // Postcondition of a successful New connect: the session is wired to
1711        // its freshly-registered backend token and the status reflects an
1712        // in-flight handshake (`Connecting`). The promotion to `Connected`
1713        // happens later in `ready_inner` once the socket signals writable.
1714        debug_assert!(
1715            self.backend_token.is_some(),
1716            "a New backend connection must own its backend token"
1717        );
1718        debug_assert!(
1719            self.backend_connected.is_connecting(),
1720            "a New backend connection must be in the Connecting state"
1721        );
1722
1723        Ok(BackendConnectAction::New)
1724    }
1725}
1726
1727impl ProxySession for TcpSession {
1728    fn close(&mut self) {
1729        if self.has_been_closed {
1730            return;
1731        }
1732
1733        // Past the idempotency guard the session is closing for the first
1734        // time: every gauge-restore / untrack below must run exactly once, so
1735        // re-entry on an already-closed session would double-decrement.
1736        debug_assert!(
1737            !self.has_been_closed,
1738            "close() body must only run on a not-yet-closed session"
1739        );
1740
1741        // TODO: the state should handle the timeouts
1742        trace!("{} Closing TCP session", log_context!(self));
1743        self.metrics.service_stop();
1744
1745        // Drain the per-(cluster, source-IP) accounting before any
1746        // early-return path below. The fail / non-fail close branches
1747        // both count, and the SessionManager-side untrack is idempotent
1748        // (no-op when the slot was never tracked) so this is safe even
1749        // when `cluster_ip_tracked` is false.
1750        if self.cluster_ip_tracked {
1751            self.proxy
1752                .borrow()
1753                .sessions
1754                .borrow_mut()
1755                .untrack_all_cluster_ip(self.frontend_token);
1756            self.cluster_ip_tracked = false;
1757        }
1758
1759        // Restore gauges. `SniPreread` is the "reject"/"teardown" half of
1760        // `tcp.sni_preread.active`'s "-1 on every exit" contract -- the
1761        // "upgrade" exit already decremented it in `upgrade_sni_preread`,
1762        // which also transitions `self.state` away from `SniPreread` before
1763        // `close()` can ever observe that marker again. A session that
1764        // reaches `close()` still marked `SniPreread` (directly, or via
1765        // `FailedUpgrade(SniPreread)` if `upgrade_sni_preread` itself failed)
1766        // therefore never had its gauge/duration accounted for yet.
1767        match self.state.marker() {
1768            StateMarker::Pipe => gauge_add!(names::protocol::TCP, -1),
1769            StateMarker::SendProxyProtocol => gauge_add!(names::protocol::PROXY_SEND, -1),
1770            StateMarker::RelayProxyProtocol => gauge_add!(names::protocol::PROXY_RELAY, -1),
1771            StateMarker::ExpectProxyProtocol => gauge_add!(names::protocol::PROXY_EXPECT, -1),
1772            StateMarker::SniPreread => {
1773                gauge_add!(names::tcp::sni_preread::ACTIVE, -1);
1774                if let TcpStateMachine::SniPreread(preread) = &self.state {
1775                    time!(
1776                        names::tcp::sni_preread::DURATION,
1777                        preread.started_at().elapsed().as_millis() as i64
1778                    );
1779                }
1780            }
1781        }
1782
1783        if self.state.failed() {
1784            match self.state.marker() {
1785                StateMarker::Pipe => incr!(names::tcp::UPGRADE_PIPE_FAILED),
1786                StateMarker::SendProxyProtocol => incr!(names::tcp::UPGRADE_SEND_FAILED),
1787                StateMarker::RelayProxyProtocol => incr!(names::tcp::UPGRADE_RELAY_FAILED),
1788                StateMarker::ExpectProxyProtocol => incr!(names::tcp::UPGRADE_EXPECT_FAILED),
1789                StateMarker::SniPreread => incr!(names::tcp::UPGRADE_SNI_PREREAD_FAILED),
1790            }
1791            return;
1792        }
1793
1794        self.cancel_timeouts();
1795
1796        let front_socket = self.state.front_socket();
1797        // TCP listener is plaintext at this layer — `Shutdown::Both` does not
1798        // truncate any TLS write buffer, so the canonical anti-pattern
1799        // (forces a TCP RST on the read direction, dropping in-flight bytes)
1800        // does not apply. Move to `Shutdown::Write` if a TLS upgrade ever
1801        // wraps this listener.
1802        if let Err(e) = front_socket.shutdown(Shutdown::Both) {
1803            // error 107 NotConnected can happen when was never fully connected, or was already disconnected due to error
1804            if e.kind() != ErrorKind::NotConnected {
1805                error!(
1806                    "{} Error shutting down front socket({:?}): {:?}",
1807                    log_context!(self),
1808                    front_socket,
1809                    e
1810                );
1811            }
1812        }
1813
1814        // deregister the frontend and remove it, in a separate scope to drop proxy when done
1815        {
1816            let proxy = self.proxy.borrow();
1817            let fd = front_socket.as_raw_fd();
1818            if let Err(e) = proxy.registry.deregister(&mut SourceFd(&fd)) {
1819                error!(
1820                    "{} Error deregistering front socket({:?}) while closing TCP session: {:?}",
1821                    log_context!(self),
1822                    fd,
1823                    e
1824                );
1825            }
1826            proxy
1827                .sessions
1828                .borrow_mut()
1829                .slab
1830                .try_remove(self.frontend_token.0);
1831        }
1832
1833        self.close_backend();
1834        self.has_been_closed = true;
1835
1836        // Postcondition of the normal close path: the session is terminal and
1837        // every accounting slot has been released — `close_backend` cleared
1838        // the backend token, and the per-(cluster, source-IP) untrack above
1839        // reset the flag. The idempotency guard now short-circuits any repeat.
1840        debug_assert!(self.has_been_closed, "close() must mark the session closed");
1841        debug_assert!(
1842            self.backend_token.is_none(),
1843            "close() must leave no dangling backend token"
1844        );
1845        debug_assert!(
1846            !self.cluster_ip_tracked,
1847            "close() must untrack the (cluster, source-IP) slot"
1848        );
1849    }
1850
1851    fn timeout(&mut self, token: Token) -> SessionIsToBeClosed {
1852        // The frontend and backend slots are distinct tokens, so the two
1853        // dispatch arms below are mutually exclusive — a single token can
1854        // never match both. (Obsolete tokens matching neither are tolerated
1855        // and fall through to the `false` arm.)
1856        debug_assert!(
1857            self.backend_token != Some(self.frontend_token),
1858            "frontend and backend tokens must never collide"
1859        );
1860        if self.frontend_token == token {
1861            self.container_frontend_timeout.triggered();
1862            // The preread deadline firing always closes the session either
1863            // way (matches every other state's front-timeout behavior).
1864            // Route-aware: only feed `Input::Timeout` into the core while
1865            // still UNDECIDED, for its `tcp.sni_preread.rejected.fragmented`
1866            // metric + log side effect. Once a route has already latched
1867            // (backend connect still pending), this same front-timeout
1868            // firing is a plain "connect/upgrade took too long" close, NOT a
1869            // fresh preread verdict -- re-feeding `Input::Timeout` into an
1870            // already-decided core would just replay the SAME latched
1871            // `Output::Routed` through `SniPreread::handle_output`'s
1872            // `Routed` arm a SECOND time: double-incrementing
1873            // `tcp.sni_preread.routed` in release, and tripping its
1874            // `debug_assert!(self.outcome.is_none(), ...)` in debug
1875            // (sozu-proxy/sozu#1290).
1876            if let TcpStateMachine::SniPreread(preread) = &mut self.state {
1877                if preread.is_routed() {
1878                    debug!(
1879                        "{} frontend timeout while a routed SNI-preread session was still \
1880                         waiting on its backend connect",
1881                        log_context!(self)
1882                    );
1883                } else {
1884                    let listener = self.listener.borrow();
1885                    let cfg = listener.preread_config(preread.effective_max_bytes());
1886                    preread.on_timeout(&cfg);
1887                }
1888            }
1889            return true;
1890        }
1891        if self.backend_token == Some(token) {
1892            self.container_backend_timeout.triggered();
1893            return true;
1894        }
1895        // invalid token, obsolete timeout triggered
1896        false
1897    }
1898
1899    fn protocol(&self) -> Protocol {
1900        Protocol::TCP
1901    }
1902
1903    fn update_readiness(&mut self, token: Token, events: Ready) {
1904        trace!(
1905            "{} token {:?} got event {}",
1906            log_context!(self),
1907            token,
1908            super::ready_to_string(events)
1909        );
1910
1911        self.last_event = Instant::now();
1912        self.metrics.wait_start();
1913
1914        if self.frontend_token == token {
1915            self.front_readiness().event = self.front_readiness().event | events;
1916        } else if self.backend_token == Some(token)
1917            && let Some(r) = self.back_readiness()
1918        {
1919            r.event |= events;
1920        }
1921    }
1922
1923    fn ready(&mut self, session: Rc<RefCell<dyn ProxySession>>) -> SessionIsToBeClosed {
1924        self.metrics.service_start();
1925
1926        let session_result = self.ready_inner(session.clone());
1927
1928        let to_bo_closed = match session_result {
1929            SessionResult::Close => true,
1930            SessionResult::Continue => false,
1931            SessionResult::Upgrade => match self.upgrade() {
1932                false => self.ready(session),
1933                true => true,
1934            },
1935        };
1936
1937        self.metrics.service_stop();
1938
1939        // Run-to-completion postcondition: the front/back token + readiness
1940        // state machine must satisfy its cross-field invariants after every
1941        // `ready()` pass. Cfg-guarded so the call (and `check_invariants`
1942        // itself) is absent from release builds.
1943        #[cfg(debug_assertions)]
1944        self.check_invariants();
1945
1946        to_bo_closed
1947    }
1948
1949    fn shutting_down(&mut self) -> SessionIsToBeClosed {
1950        true
1951    }
1952
1953    fn last_event(&self) -> Instant {
1954        self.last_event
1955    }
1956
1957    fn print_session(&self) {
1958        let state: String = match &self.state {
1959            TcpStateMachine::ExpectProxyProtocol(_) => String::from("Expect"),
1960            TcpStateMachine::SendProxyProtocol(_) => String::from("Send"),
1961            TcpStateMachine::RelayProxyProtocol(_) => String::from("Relay"),
1962            TcpStateMachine::Pipe(_) => String::from("TCP"),
1963            TcpStateMachine::SniPreread(_) => String::from("SniPreread"),
1964            TcpStateMachine::FailedUpgrade(marker) => format!("FailedUpgrade({marker:?})"),
1965        };
1966
1967        let front_readiness = match &self.state {
1968            TcpStateMachine::ExpectProxyProtocol(expect) => Some(&expect.frontend_readiness),
1969            TcpStateMachine::SendProxyProtocol(send) => Some(&send.frontend_readiness),
1970            TcpStateMachine::RelayProxyProtocol(relay) => Some(&relay.frontend_readiness),
1971            TcpStateMachine::Pipe(pipe) => Some(&pipe.frontend_readiness),
1972            TcpStateMachine::SniPreread(preread) => Some(&preread.frontend_readiness),
1973            TcpStateMachine::FailedUpgrade(_) => None,
1974        };
1975
1976        let back_readiness = match &self.state {
1977            TcpStateMachine::SendProxyProtocol(send) => Some(&send.backend_readiness),
1978            TcpStateMachine::RelayProxyProtocol(relay) => Some(&relay.backend_readiness),
1979            TcpStateMachine::Pipe(pipe) => Some(&pipe.backend_readiness),
1980            TcpStateMachine::SniPreread(preread) => Some(&preread.backend_readiness),
1981            TcpStateMachine::ExpectProxyProtocol(_) => None,
1982            TcpStateMachine::FailedUpgrade(_) => None,
1983        };
1984
1985        error!(
1986            "\
1987{} Session ({:?})
1988\tFrontend:
1989\t\ttoken: {:?}\treadiness: {:?}
1990\tBackend:
1991\t\ttoken: {:?}\treadiness: {:?}\tstatus: {:?}\tcluster id: {:?}",
1992            log_context!(self),
1993            state,
1994            self.frontend_token,
1995            front_readiness,
1996            self.backend_token,
1997            back_readiness,
1998            self.backend_connected,
1999            self.cluster_id
2000        );
2001        error!("Metrics: {:?}", self.metrics);
2002    }
2003
2004    fn frontend_token(&self) -> Token {
2005        self.frontend_token
2006    }
2007}
2008
2009pub struct TcpListener {
2010    active: SessionIsToBeClosed,
2011    address: SocketAddr,
2012    cluster_id: Option<String>,
2013    config: TcpListenerConfig,
2014    listener: Option<MioTcpListener>,
2015    /// SNI -> `(AlpnMatcher, ClusterId)` route table (sozu-proxy/sozu#1279).
2016    /// Populated by `add_tcp_front`/`remove_tcp_front` from
2017    /// `RequestTcpFrontend.sni`/`.alpn`; empty for a listener whose fronts
2018    /// are all no-SNI (the legacy `cluster_id` catch-all). A listener never
2019    /// mixes both (enforced at config load, `command/src/config.rs`), but
2020    /// `create_session`'s routing gate stays defensive and checks both.
2021    sni_routes: TrieNode<Vec<(AlpnMatcher, ClusterId)>>,
2022    tags: BTreeMap<String, CachedTags>,
2023    token: Token,
2024}
2025
2026impl ListenerHandler for TcpListener {
2027    fn get_addr(&self) -> &SocketAddr {
2028        &self.address
2029    }
2030
2031    fn get_tags(&self, key: &str) -> Option<&CachedTags> {
2032        self.tags.get(key)
2033    }
2034
2035    fn set_tags(&mut self, key: String, tags: Option<BTreeMap<String, String>>) {
2036        match tags {
2037            Some(tags) => self.tags.insert(key, CachedTags::new(tags)),
2038            None => self.tags.remove(&key),
2039        };
2040    }
2041
2042    fn protocol(&self) -> Protocol {
2043        Protocol::TCP
2044    }
2045
2046    fn public_address(&self) -> SocketAddr {
2047        self.config
2048            .public_address
2049            .map(|addr| addr.into())
2050            .unwrap_or(self.address)
2051    }
2052}
2053
2054impl TcpListener {
2055    fn new(config: TcpListenerConfig, token: Token) -> Result<TcpListener, ListenerError> {
2056        Ok(TcpListener {
2057            cluster_id: None,
2058            listener: None,
2059            token,
2060            address: config.address.into(),
2061            config,
2062            active: false,
2063            sni_routes: TrieNode::root(),
2064            tags: BTreeMap::new(),
2065        })
2066    }
2067
2068    /// Validate that a worker can build this TCP listener configuration WITHOUT
2069    /// constructing the full listener or binding a socket. TCP listener
2070    /// construction has no fallible config today (no rustls context, no answer
2071    /// templates; a bad bind surfaces later as an `ActivateListener` failure),
2072    /// so this currently always succeeds. It exists for surface parity with the
2073    /// HTTP/HTTPS validators the main process calls before committing an
2074    /// `Add*Listener` to `ConfigState` and fanning it out (sozu#1301), and is
2075    /// the hook for any future TCP-config validation.
2076    pub fn validate_config(_config: &TcpListenerConfig) -> Result<(), ListenerError> {
2077        Ok(())
2078    }
2079
2080    /// Build the [`PrereadConfig`] this listener's `SniPreread` sessions
2081    /// feed to [`crate::protocol::tcp_preread::SniPrereadCore::handle_input`].
2082    /// `routes`/`inbound_proxy`/`timeout` come straight from the listener
2083    /// config; `effective_max_bytes` is session-specific (already clamped to
2084    /// that session's buffer capacity at construction, see `create_session`),
2085    /// so it is passed in rather than re-derived here.
2086    fn preread_config(&self, effective_max_bytes: usize) -> PrereadConfig<'_> {
2087        PrereadConfig {
2088            routes: &self.sni_routes,
2089            inbound_proxy: self.config.expect_proxy,
2090            max_bytes: effective_max_bytes,
2091            timeout: Duration::from_secs(u64::from(
2092                self.config
2093                    .sni_preread_timeout
2094                    .unwrap_or(DEFAULT_SNI_PREREAD_TIMEOUT),
2095            )),
2096            accept_wildcard: true,
2097        }
2098    }
2099
2100    /// Validate an incoming `AddTcpFrontend` against this listener's
2101    /// CURRENT routing state before any mutation, mirroring
2102    /// `command/src/config.rs`'s TOML config-load TCP SNI/ALPN invariants
2103    /// (sozu-proxy/sozu#1279):
2104    ///
2105    /// - `alpn` set with no `sni`: the worker's no-SNI catch-all path never
2106    ///   consults `alpn`, so the protocol list would silently never be
2107    ///   enforced (mirrors `ConfigError::AlpnWithoutSni`).
2108    /// - a no-SNI frontend added to a listener that already has SNI-scoped
2109    ///   routes, or an SNI-scoped frontend added to a listener that already
2110    ///   has a no-SNI catch-all cluster (mirrors
2111    ///   `ConfigError::TcpListenerMixesSniAndNoSni`).
2112    /// - an ALPN protocol, or a catch-all (empty `alpn`), that overlaps an
2113    ///   existing route already registered for the same `(address, sni)`
2114    ///   (mirrors `ConfigError::TcpFrontendAlpnOverlap` /
2115    ///   `TcpFrontendMultipleAlpnCatchAll`).
2116    ///
2117    /// Config-load already rejects all of these shapes for requests built
2118    /// from a TOML file, but `AddTcpFrontend` can also arrive directly over
2119    /// the command socket, or via `LoadState` replay of a hand-edited or
2120    /// stale state file, bypassing config.rs entirely -- the worker must
2121    /// not silently corrupt its own routing table when that happens.
2122    fn validate_new_tcp_front(&self, front: &RequestTcpFrontend) -> Result<(), ProxyError> {
2123        let reject = |reason: String| {
2124            Err(ProxyError::InvalidTcpFrontend {
2125                address: self.address,
2126                reason,
2127            })
2128        };
2129
2130        match &front.sni {
2131            None => {
2132                if !front.alpn.is_empty() {
2133                    return reject(format!(
2134                        "alpn = {:?} set without sni: alpn only matches within an SNI-scoped \
2135                         preread, so a frontend without sni would silently ignore its alpn list",
2136                        front.alpn
2137                    ));
2138                }
2139                if !self.sni_routes.is_empty() {
2140                    return reject(
2141                        "a no-SNI frontend cannot be added to a listener that already has \
2142                         SNI-scoped routes"
2143                            .to_string(),
2144                    );
2145                }
2146            }
2147            Some(sni) => {
2148                if self.cluster_id.is_some() {
2149                    return reject(
2150                        "an SNI-scoped frontend cannot be added to a listener that already has \
2151                         a no-SNI catch-all cluster"
2152                            .to_string(),
2153                    );
2154                }
2155
2156                // SNI SHAPE check: delegate to the SAME validator config-load
2157                // uses (`sozu_command::config::validate_sni_pattern`) rather
2158                // than a hand-rolled partial check. A direct `AddTcpFrontend`
2159                // over the command socket, or a `LoadState` replay, bypasses
2160                // config.rs entirely, so the worker boundary must enforce
2161                // the identical rule -- including the checks a bare '/'/'*'
2162                // scan used to miss: empty string, non-ASCII, and any empty
2163                // label (leading/trailing/consecutive dots). An unvalidated
2164                // leading-empty-label pattern like `.example.com` would
2165                // otherwise reach `insert_sni_route` ->
2166                // `pattern_trie::insert_recursive`'s RELEASE-mode
2167                // `assert_ne!(partial_key, &b""[..])` and crash the worker.
2168                let normalized_sni = match validate_sni_pattern(sni) {
2169                    Ok(normalized) => normalized,
2170                    Err(config_error) => {
2171                        return reject(format!(
2172                            "sni {sni:?} failed SNI shape validation: {config_error}"
2173                        ));
2174                    }
2175                };
2176
2177                // Same key as `insert_sni_route`'s own lookup, so this checks
2178                // against exactly the entries the new route would be appended
2179                // alongside. This is bookkeeping over the key's OWN node, not
2180                // the routing lookup (`preread_config` keeps `true` for
2181                // that), so `accept_wildcard` must make the lookup EXACT for
2182                // both key shapes:
2183                //
2184                // - literal key -> `false`. With `true`, a literal key with
2185                //   no child yet (e.g. `a.example.com` when only
2186                //   `*.example.com` exists) falls back to the sibling
2187                //   wildcard's entries (`pattern_trie.rs`'s `lookup`
2188                //   wildcard-fallback branch), misattributing the wildcard's
2189                //   catch-all to the exact key (falsely rejecting a
2190                //   legitimate exact catch-all) — and symmetrically,
2191                //   `insert_sni_route`/`remove_sni_route` would corrupt the
2192                //   WILDCARD's `Vec` instead of touching a distinct
2193                //   exact-key node.
2194                // - wildcard key -> `true`. Unlike `lookup_mut` (which
2195                //   short-circuits `partial_key == b"*"` before consulting
2196                //   `accept_wildcard`, so insert/remove stay on `false`),
2197                //   the immutable `lookup` reaches a wildcard entry ONLY
2198                //   through the fallback branch; with `false` an existing
2199                //   `*.example.com` entry is invisible here and a duplicate
2200                //   catch-all / overlapping-ALPN wildcard front bypasses
2201                //   validation. For a wildcard key, `true` IS the exact
2202                //   self-lookup: the traversal descends the key's own
2203                //   literal ancestry, and `insert` never creates a literal
2204                //   `*` child that could shadow the node's `wildcard` slot.
2205                //
2206                // `starts_with(b"*.")` matches every wildcard shape
2207                // config-load can emit: `command/src/config.rs`'s
2208                // `validate_sni_pattern` only admits a single leading `*.`
2209                // label (any other `*` placement is rejected), and a plain
2210                // hostname key traverses literal children in both `lookup`
2211                // and `insert_recursive`. A bare `*` key (config-load
2212                // rejects it; only a bypassing IPC request can carry one)
2213                // is the known gap: `lookup_mut`'s short-circuit maps it to
2214                // the wildcard slot while this immutable self-lookup (flag
2215                // `false`) cannot see an existing entry there, so duplicate
2216                // bare-`*` fronts are not detected -- sibling keys stay
2217                // untouched either way.
2218                //
2219                // `normalized_sni` (not a fresh `sni.to_ascii_lowercase()`)
2220                // is used here: `validate_sni_pattern` already returned the
2221                // canonical lowercased form above, and reusing it keeps this
2222                // function's notion of "the key" identical to what
2223                // `insert_sni_route`/`remove_sni_route` compute from the
2224                // same input.
2225                let key = normalized_sni.into_bytes();
2226                let accept_wildcard_for_self_lookup = key.starts_with(b"*.");
2227                if let Some((_, existing)) = self
2228                    .sni_routes
2229                    .domain_lookup(&key, accept_wildcard_for_self_lookup)
2230                {
2231                    let new_is_catch_all = front.alpn.is_empty();
2232                    for (matcher, _cluster_id) in existing {
2233                        match matcher {
2234                            AlpnMatcher::Any if new_is_catch_all => {
2235                                return reject(format!(
2236                                    "sni {sni:?} already has a catch-all (empty alpn) \
2237                                     frontend: at most one frontend per (address, sni) may \
2238                                     omit alpn"
2239                                ));
2240                            }
2241                            AlpnMatcher::OneOf(protocols) => {
2242                                if let Some(overlap) = front
2243                                    .alpn
2244                                    .iter()
2245                                    .find(|protocol| protocols.contains(protocol.as_bytes()))
2246                                {
2247                                    return reject(format!(
2248                                        "sni {sni:?} already has a frontend matching ALPN \
2249                                         protocol {overlap:?}: ALPN matchers for the same \
2250                                         (address, sni) must not overlap"
2251                                    ));
2252                                }
2253                            }
2254                            AlpnMatcher::Any => {}
2255                        }
2256                    }
2257                }
2258            }
2259        }
2260
2261        Ok(())
2262    }
2263
2264    /// Add one `(AlpnMatcher, ClusterId)` entry to this listener's SNI route
2265    /// table, appending to the SNI key's existing `Vec` if one is already
2266    /// present rather than clobbering it (multiple ALPN-scoped fronts can
2267    /// share the same SNI). `sni` is defensively lowercased: the
2268    /// SNI-preread core normalizes (lowercase, no trailing dot) before ever
2269    /// looking a route up, so the table key must already be in that form.
2270    ///
2271    /// Returns `Err` when the underlying trie insert reports
2272    /// [`InsertResult::Failed`] (a malformed key, e.g. an empty label) —
2273    /// the caller (`TcpProxy::add_tcp_front`) must propagate this rather
2274    /// than silently keep going with an add that reports success while the
2275    /// route table never gained a usable entry. With `validate_new_tcp_front`
2276    /// enforcing the shared SNI shape validator before this ever runs, a
2277    /// `Failed` result should be unreachable here — the `debug_assert_ne!`
2278    /// below asserts OUR OWN validated invariant, not raw wire/IPC input,
2279    /// and is a loud Sōzu-internal-bug canary in debug builds; the `Err`
2280    /// path is the release-mode graceful fallback if that invariant is ever
2281    /// violated by a future bug.
2282    fn insert_sni_route(
2283        &mut self,
2284        sni: String,
2285        alpn: Vec<String>,
2286        cluster_id: ClusterId,
2287    ) -> Result<(), ProxyError> {
2288        let (key, matcher) = route_key_and_matcher(&sni, alpn);
2289        // `accept_wildcard: false` — see `validate_new_tcp_front`'s comment:
2290        // an exact key must never fall back to a sibling wildcard's entry,
2291        // or this push would corrupt the WILDCARD's route `Vec` instead of
2292        // creating this key's own node.
2293        match self.sni_routes.domain_lookup_mut(&key, false) {
2294            Some((_, entries)) => {
2295                entries.push((matcher, cluster_id));
2296                Ok(())
2297            }
2298            None => {
2299                let insert_result = self
2300                    .sni_routes
2301                    .domain_insert(key, vec![(matcher, cluster_id)]);
2302                debug_assert_ne!(
2303                    insert_result,
2304                    InsertResult::Failed,
2305                    "insert_sni_route's key must already be validated by validate_new_tcp_front"
2306                );
2307                if insert_result == InsertResult::Failed {
2308                    error!(
2309                        "{} SNI route insert failed for {:?} despite passing shape \
2310                         validation -- rejecting to avoid a silently dead route",
2311                        log_module_context!(),
2312                        sni
2313                    );
2314                    return Err(ProxyError::InvalidTcpFrontend {
2315                        address: self.address,
2316                        reason: format!(
2317                            "internal error: the route table rejected sni {sni:?} despite \
2318                             passing shape validation"
2319                        ),
2320                    });
2321                }
2322                Ok(())
2323            }
2324        }
2325    }
2326
2327    /// Symmetric counterpart to [`Self::insert_sni_route`]: removes the
2328    /// matching `(AlpnMatcher, ClusterId)` entry, then `domain_remove`s the
2329    /// SNI key itself once its `Vec` is empty (no stranded empty entries in
2330    /// the trie).
2331    fn remove_sni_route(&mut self, sni: String, alpn: Vec<String>, cluster_id: &ClusterId) {
2332        let (key, matcher) = route_key_and_matcher(&sni, alpn);
2333        // `accept_wildcard: false` — same reasoning as `insert_sni_route`:
2334        // removing the exact key must never reach into and strip a sibling
2335        // wildcard's entries.
2336        if let Some((_, entries)) = self.sni_routes.domain_lookup_mut(&key, false) {
2337            entries.retain(|(m, c)| !(*m == matcher && c == cluster_id));
2338            if entries.is_empty() {
2339                self.sni_routes.domain_remove(&key);
2340            }
2341        }
2342    }
2343
2344    pub fn activate(
2345        &mut self,
2346        registry: &Registry,
2347        tcp_listener: Option<MioTcpListener>,
2348    ) -> Result<Token, ProxyError> {
2349        if self.active {
2350            return Ok(self.token);
2351        }
2352
2353        let mut listener = match tcp_listener {
2354            Some(listener) => listener,
2355            None => {
2356                let address = self.config.address.into();
2357                server_bind(address).map_err(|e| ProxyError::BindToSocket(address, e))?
2358            }
2359        };
2360
2361        registry
2362            .register(&mut listener, self.token, Interest::READABLE)
2363            .map_err(ProxyError::RegisterListener)?;
2364
2365        self.listener = Some(listener);
2366        self.active = true;
2367        Ok(self.token)
2368    }
2369
2370    /// Apply a partial-update patch to this TCP listener's live configuration.
2371    ///
2372    /// Fields absent in the patch (i.e. `None`) are preserved unchanged.
2373    pub fn update_config(&mut self, patch: &UpdateTcpListenerConfig) -> Result<(), ListenerError> {
2374        if let Some(v) = patch.public_address {
2375            self.config.public_address = Some(v);
2376        }
2377        if let Some(v) = patch.expect_proxy {
2378            self.config.expect_proxy = v;
2379        }
2380        if let Some(v) = patch.front_timeout {
2381            self.config.front_timeout = v;
2382        }
2383        if let Some(v) = patch.back_timeout {
2384            self.config.back_timeout = v;
2385        }
2386        if let Some(v) = patch.connect_timeout {
2387            self.config.connect_timeout = v;
2388        }
2389        Ok(())
2390    }
2391}
2392
2393fn handle_connection_result(
2394    connection_result: Result<BackendConnectAction, BackendConnectionError>,
2395) -> Option<SessionResult> {
2396    match connection_result {
2397        // reuse connection or send a default answer, we can continue
2398        Ok(BackendConnectAction::Reuse) => None,
2399        Ok(BackendConnectAction::New) | Ok(BackendConnectAction::Replace) => {
2400            // we must wait for an event
2401            Some(SessionResult::Continue)
2402        }
2403        Err(_) => {
2404            // in case of BackendConnectionError::Backend(BackendError::ConnectionFailures(..))
2405            // we may want to retry instead of closing
2406            Some(SessionResult::Close)
2407        }
2408    }
2409}
2410
2411/// `min(listener.config.sni_preread_max_bytes, frontend_buffer.capacity())`,
2412/// floored at [`MIN_SNI_PREREAD_MAX_BYTES`] -- the SNI-preread core has no
2413/// independent backstop of its own (`SniPrereadCore::handle_input` trusts
2414/// `PrereadConfig::max_bytes` entirely), so the shell must never hand it a
2415/// cap the checked-out buffer cannot actually hold (the `min`), NOR a cap so
2416/// small the preread read is zero-length and spins until the loop guard (the
2417/// `max`). Config-load rejects a sub-floor `sni_preread_max_bytes` loudly
2418/// (`ConfigError::SniPrereadMaxBytesTooSmall`), but a `0` knob from a direct
2419/// `sozu listener tcp add` CLI/IPC request, or a stale `LoadState`
2420/// replay, bypasses that check and reaches the worker -- this floor degrades
2421/// it to the 5-byte TLS-record-header minimum instead, killing the spin for
2422/// EVERY config source at the single point of use. `buffer_capacity` is
2423/// always `>= MIN_SNI_PREREAD_MAX_BYTES` in practice (buffers are KB-sized),
2424/// so the `min` never fights the `max`.
2425fn effective_sni_preread_max_bytes(configured: Option<u32>, buffer_capacity: usize) -> usize {
2426    (configured.unwrap_or(DEFAULT_SNI_PREREAD_MAX_BYTES) as usize)
2427        .min(buffer_capacity)
2428        .max(MIN_SNI_PREREAD_MAX_BYTES as usize)
2429}
2430
2431/// Whether `TcpSession::readable`'s frontend-timeout reset should fire for
2432/// the CURRENT state. `false` only while the session is an UNDECIDED
2433/// `SniPreread`: the preread deadline armed at session creation
2434/// (`new_sni_preread`) is an ABSOLUTE budget, not a per-fragment idle timer
2435/// -- resetting it on every `readable()` event would let a client
2436/// trickling one byte just before each expiry hold the session (and both
2437/// its checked-out buffers) open far past the configured
2438/// `sni_preread_timeout` (sozu-proxy/sozu#1290). Every
2439/// other state -- a `SniPreread` that has already routed, or any state
2440/// reached after it -- resets normally: once routed, the frontend timeout
2441/// reverts to being a genuine per-read idle timer (see the `front_timeout`
2442/// restore in `readable()`'s route-capture block).
2443fn frontend_timeout_resets_on_readable(state: &TcpStateMachine) -> bool {
2444    !matches!(
2445        state,
2446        TcpStateMachine::SniPreread(preread) if preread.outcome().is_none()
2447    )
2448}
2449
2450/// Access-log ALPN tag for an SNI-routed TCP session: the client's FIRST
2451/// offered protocol (client preference order, matching
2452/// `SniPrereadCore::route`'s own routing precedence), mapped to a known
2453/// `&'static str` label -- the same two labels `https.rs`'s own ALPN
2454/// negotiation records (`"h2"` / `"http/1.1"`) -- so `tcp.sni_preread`
2455/// sessions and terminated-TLS sessions chart under the same values.
2456/// `None` for an empty offer or an unrecognized protocol: Sōzu never
2457/// terminates TLS on this path, so this is the client's stated preference,
2458/// not a negotiated outcome.
2459fn known_alpn_label(offered: &[Vec<u8>]) -> Option<&'static str> {
2460    match offered.first().map(Vec::as_slice) {
2461        Some(b"h2") => Some("h2"),
2462        Some(b"http/1.1") => Some("http/1.1"),
2463        _ => None,
2464    }
2465}
2466
2467/// Worker-internal, collision-free identity for a TCP frontend's access-log
2468/// tags entry (sozu-proxy/sozu#1290). ALPN protocol
2469/// identifiers are opaque RFC 7301 byte strings -- nothing forbids a `,` or
2470/// `|` inside one -- so a naive `sorted_alpn.join(",")` string key let two
2471/// legal, DISJOINT fronts collide: a single protocol `"a,b"` and the pair
2472/// `["a", "b"]` both joined to the identical string `"a,b"`, so the second
2473/// `add_tcp_front` silently clobbered the first front's tags entry.
2474///
2475/// `TcpListener::tags: BTreeMap<String, CachedTags>` and
2476/// `Pipe::set_tags_key(Option<String>)` (`lib/src/protocol/pipe.rs`) both
2477/// key on a plain `String` -- crossing either boundary still needs one --
2478/// so this type does not replace that storage; it is the SINGLE place that
2479/// composes the string, via [`Self::fmt`]'s length-prefixed per-protocol
2480/// encoding, so no choice of in-band separator can be confused with a
2481/// protocol-name boundary the way a bare `join(",")` could.
2482#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
2483enum TcpFrontendTagsKey {
2484    /// Legacy no-SNI catch-all front: keyed by the bare listener address,
2485    /// exactly as before SNI routing existed.
2486    Address(SocketAddr),
2487    /// An SNI-scoped front. `alpn` is sorted (and deduped, since a plain
2488    /// `Vec<String>` cannot enforce uniqueness the way
2489    /// [`AlpnMatcher::OneOf`]'s `BTreeSet` does) so two constructions from
2490    /// the same logical protocol set always compare equal regardless of
2491    /// the caller's original order.
2492    Sni {
2493        address: SocketAddr,
2494        sni: String,
2495        alpn: Vec<String>,
2496    },
2497}
2498
2499impl TcpFrontendTagsKey {
2500    /// The identity triple `command/src/state.rs`'s `add_tcp_frontend`
2501    /// deduplicates on: `sni` is lowercased to match `insert_sni_route`'s
2502    /// trie key, `alpn` sorted + deduped so `add_tcp_front`'s operator
2503    /// order and the route-time rebuild's [`AlpnMatcher::OneOf`] `BTreeSet`
2504    /// order (`alpn_matcher_protocols`) always agree.
2505    fn sni(address: SocketAddr, sni: &str, alpn: &[String]) -> Self {
2506        let mut alpn: Vec<String> = alpn.to_vec();
2507        alpn.sort_unstable();
2508        alpn.dedup();
2509        TcpFrontendTagsKey::Sni {
2510            address,
2511            sni: sni.to_ascii_lowercase(),
2512            alpn,
2513        }
2514    }
2515}
2516
2517impl std::fmt::Display for TcpFrontendTagsKey {
2518    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2519        match self {
2520            TcpFrontendTagsKey::Address(address) => write!(f, "{address}"),
2521            TcpFrontendTagsKey::Sni { address, sni, alpn } => {
2522                write!(f, "{address}|{sni}|")?;
2523                for protocol in alpn {
2524                    // `<len>:<bytes>` per protocol: the length prefix is
2525                    // CHECKED against the following bytes, never scanned
2526                    // for, so a `,`/`:` embedded in one protocol name can
2527                    // never be misread as a boundary between two --
2528                    // unlike the old bare `join(",")`.
2529                    write!(f, "{}:{protocol}", protocol.len())?;
2530                }
2531                Ok(())
2532            }
2533        }
2534    }
2535}
2536
2537/// Canonical access-log tags key for an SNI-scoped TCP frontend: one
2538/// listener can carry many SNI/ALPN fronts, each with its own `tags`, so
2539/// keying tags by the bare listener address (the pre-SNI behavior, kept
2540/// verbatim for no-SNI fronts) would let the LAST added front clobber every
2541/// sibling and let removing ANY front clear tags for the whole address. See
2542/// [`TcpFrontendTagsKey`] for the collision-free composition.
2543fn sni_tags_key(address: &SocketAddr, sni: &str, alpn: &[String]) -> String {
2544    TcpFrontendTagsKey::sni(*address, sni, alpn).to_string()
2545}
2546
2547/// The matched [`AlpnMatcher`]'s protocols as strings, for rebuilding the
2548/// [`sni_tags_key`] a route decision maps to: `Any` is the empty-`alpn`
2549/// catch-all front, `OneOf` yields its (BTreeSet-sorted) protocol list.
2550/// `from_utf8_lossy` is defensive — matcher protocols originate from
2551/// config/IPC `String`s, never raw network bytes.
2552fn alpn_matcher_protocols(matcher: &AlpnMatcher) -> Vec<String> {
2553    match matcher {
2554        AlpnMatcher::Any => Vec::new(),
2555        AlpnMatcher::OneOf(set) => set
2556            .iter()
2557            .map(|protocol| String::from_utf8_lossy(protocol).into_owned())
2558            .collect(),
2559    }
2560}
2561
2562/// Shared `(trie key, ALPN matcher)` construction for
2563/// [`TcpListener::insert_sni_route`] and [`TcpListener::remove_sni_route`]:
2564/// `sni` lowercased into the trie key, `alpn` collapsed to
2565/// [`AlpnMatcher::Any`] when empty (catch-all) or [`AlpnMatcher::OneOf`]
2566/// otherwise. `alpn` is taken by value since neither caller needs it again
2567/// afterward.
2568fn route_key_and_matcher(sni: &str, alpn: Vec<String>) -> (Vec<u8>, AlpnMatcher) {
2569    let key = sni.to_ascii_lowercase().into_bytes();
2570    let matcher = if alpn.is_empty() {
2571        AlpnMatcher::Any
2572    } else {
2573        AlpnMatcher::OneOf(alpn.into_iter().map(String::into_bytes).collect())
2574    };
2575    (key, matcher)
2576}
2577
2578#[derive(Debug)]
2579pub struct ClusterConfiguration {
2580    proxy_protocol: Option<ProxyProtocolConfig>,
2581    // Uncomment this when implementing new load balancing algorithms
2582    // load_balancing: LoadBalancingAlgorithms,
2583    /// Per-cluster override of the global per-(cluster, source-IP)
2584    /// connection limit. `None` inherits the global default,
2585    /// `Some(0)` is explicit "unlimited", `Some(n > 0)` overrides.
2586    /// Resolved against `SessionManager::effective_max_connections_per_ip`
2587    /// at admit time in `connect_to_backend`.
2588    pub max_connections_per_ip: Option<u64>,
2589}
2590
2591pub struct TcpProxy {
2592    fronts: HashMap<String, Token>,
2593    backends: Rc<RefCell<BackendMap>>,
2594    listeners: HashMap<Token, Rc<RefCell<TcpListener>>>,
2595    configs: HashMap<ClusterId, ClusterConfiguration>,
2596    registry: Registry,
2597    sessions: Rc<RefCell<SessionManager>>,
2598    pool: Rc<RefCell<Pool>>,
2599}
2600
2601impl TcpProxy {
2602    pub fn new(
2603        registry: Registry,
2604        sessions: Rc<RefCell<SessionManager>>,
2605        pool: Rc<RefCell<Pool>>,
2606        backends: Rc<RefCell<BackendMap>>,
2607    ) -> TcpProxy {
2608        TcpProxy {
2609            backends,
2610            listeners: HashMap::new(),
2611            configs: HashMap::new(),
2612            fronts: HashMap::new(),
2613            registry,
2614            sessions,
2615            pool,
2616        }
2617    }
2618
2619    pub fn add_listener(
2620        &mut self,
2621        config: TcpListenerConfig,
2622        token: Token,
2623    ) -> Result<Token, ProxyError> {
2624        match self.listeners.entry(token) {
2625            Entry::Vacant(entry) => {
2626                let tcp_listener =
2627                    TcpListener::new(config, token).map_err(ProxyError::AddListener)?;
2628                entry.insert(Rc::new(RefCell::new(tcp_listener)));
2629                Ok(token)
2630            }
2631            _ => Err(ProxyError::ListenerAlreadyPresent),
2632        }
2633    }
2634
2635    pub fn remove_listener(&mut self, address: SocketAddr) -> SessionIsToBeClosed {
2636        let len = self.listeners.len();
2637
2638        self.listeners.retain(|_, l| l.borrow().address != address);
2639        self.listeners.len() < len
2640    }
2641
2642    pub fn activate_listener(
2643        &self,
2644        addr: &SocketAddr,
2645        tcp_listener: Option<MioTcpListener>,
2646    ) -> Result<Token, ProxyError> {
2647        let listener = self
2648            .listeners
2649            .values()
2650            .find(|listener| listener.borrow().address == *addr)
2651            .ok_or(ProxyError::NoListenerFound(*addr))?;
2652
2653        listener.borrow_mut().activate(&self.registry, tcp_listener)
2654    }
2655
2656    pub fn give_back_listeners(&mut self) -> Vec<(SocketAddr, MioTcpListener)> {
2657        self.listeners
2658            .values()
2659            .filter_map(|listener| {
2660                let mut owned = listener.borrow_mut();
2661                if let Some(listener) = owned.listener.take() {
2662                    // Reset `active` so a subsequent `activate()` re-binds
2663                    // instead of short-circuiting on the stale flag.
2664                    owned.active = false;
2665                    return Some((owned.address, listener));
2666                }
2667
2668                None
2669            })
2670            .collect()
2671    }
2672
2673    pub fn give_back_listener(
2674        &mut self,
2675        address: SocketAddr,
2676    ) -> Result<(Token, MioTcpListener), ProxyError> {
2677        let listener = self
2678            .listeners
2679            .values()
2680            .find(|listener| listener.borrow().address == address)
2681            .ok_or(ProxyError::NoListenerFound(address))?;
2682
2683        let mut owned = listener.borrow_mut();
2684
2685        let taken_listener = owned
2686            .listener
2687            .take()
2688            .ok_or(ProxyError::UnactivatedListener)?;
2689
2690        // Reset `active` so a subsequent `activate()` re-binds instead of
2691        // short-circuiting on the stale flag.
2692        owned.active = false;
2693
2694        Ok((owned.token, taken_listener))
2695    }
2696
2697    /// Apply a partial-update patch to the identified TCP listener.
2698    pub fn update_listener(&mut self, patch: UpdateTcpListenerConfig) -> Result<(), ProxyError> {
2699        let address: SocketAddr = patch.address.into();
2700        let listener = self
2701            .listeners
2702            .values()
2703            .find(|l| l.borrow().address == address)
2704            .ok_or(ProxyError::NoListenerFound(address))?;
2705        listener
2706            .borrow_mut()
2707            .update_config(&patch)
2708            .map_err(|listener_error| ProxyError::ListenerActivation {
2709                address,
2710                listener_error,
2711            })
2712    }
2713
2714    pub fn add_tcp_front(&mut self, front: RequestTcpFrontend) -> Result<(), ProxyError> {
2715        let address = front.address.into();
2716
2717        let mut listener = self
2718            .listeners
2719            .values()
2720            .find(|l| l.borrow().address == address)
2721            .ok_or(ProxyError::NoListenerFound(address))?
2722            .borrow_mut();
2723
2724        // Hard-reject a request that would corrupt this listener's SNI/ALPN
2725        // routing invariants BEFORE any mutation below. Config-load
2726        // (`command/src/config.rs`, sozu-proxy/sozu#1279) already rejects
2727        // the same shapes for TOML-sourced requests, but `AddTcpFrontend`
2728        // can also arrive directly over the command socket, or via
2729        // `LoadState` replay of a hand-edited/stale state file, bypassing
2730        // config.rs entirely.
2731        listener.validate_new_tcp_front(&front)?;
2732
2733        self.fronts
2734            .insert(front.cluster_id.to_string(), listener.token);
2735
2736        match front.sni {
2737            Some(sni) => {
2738                // Per-frontend tags key: many SNI/ALPN fronts share one
2739                // listener, so the bare-address key (kept for no-SNI fronts
2740                // below) would clobber siblings — see `sni_tags_key`.
2741                listener.set_tags(sni_tags_key(&address, &sni, &front.alpn), Some(front.tags));
2742                listener.insert_sni_route(sni, front.alpn, front.cluster_id)?;
2743            }
2744            None => {
2745                listener.set_tags(
2746                    TcpFrontendTagsKey::Address(address).to_string(),
2747                    Some(front.tags),
2748                );
2749                listener.cluster_id = Some(front.cluster_id);
2750            }
2751        }
2752
2753        // POST: the mixing invariant must hold after every successful add —
2754        // `validate_new_tcp_front` is the enforcement point above, this is
2755        // the cheap live re-check that it actually held.
2756        debug_assert!(
2757            listener.cluster_id.is_none() || listener.sni_routes.is_empty(),
2758            "a TCP listener must never mix a no-SNI catch-all cluster with SNI-scoped routes"
2759        );
2760
2761        Ok(())
2762    }
2763
2764    pub fn remove_tcp_front(&mut self, front: RequestTcpFrontend) -> Result<(), ProxyError> {
2765        let address = front.address.into();
2766
2767        let mut listener = match self
2768            .listeners
2769            .values()
2770            .find(|l| l.borrow().address == address)
2771        {
2772            Some(l) => l.borrow_mut(),
2773            None => return Err(ProxyError::NoListenerFound(address)),
2774        };
2775
2776        match front.sni {
2777            Some(sni) => {
2778                // Clear ONLY this front's own tags entry (`sni_tags_key`) —
2779                // the pre-SNI bare-address removal here used to strip tags
2780                // for every sibling front on the listener.
2781                listener.set_tags(sni_tags_key(&address, &sni, &front.alpn), None);
2782                listener.remove_sni_route(sni, front.alpn, &front.cluster_id);
2783                self.fronts.remove(&front.cluster_id);
2784            }
2785            None => {
2786                listener.set_tags(TcpFrontendTagsKey::Address(address).to_string(), None);
2787                if let Some(cluster_id) = listener.cluster_id.take() {
2788                    self.fronts.remove(&cluster_id);
2789                }
2790            }
2791        }
2792
2793        Ok(())
2794    }
2795}
2796
2797impl ProxyConfiguration for TcpProxy {
2798    fn notify(&mut self, message: WorkerRequest) -> WorkerResponse {
2799        let request_type = match message.content.request_type {
2800            Some(t) => t,
2801            None => return WorkerResponse::error(message.id, "Empty request"),
2802        };
2803        match request_type {
2804            RequestType::AddTcpFrontend(front) => {
2805                if let Err(err) = self.add_tcp_front(front) {
2806                    return WorkerResponse::error(message.id, err);
2807                }
2808
2809                WorkerResponse::ok(message.id)
2810            }
2811            RequestType::RemoveTcpFrontend(front) => {
2812                if let Err(err) = self.remove_tcp_front(front) {
2813                    return WorkerResponse::error(message.id, err);
2814                }
2815
2816                WorkerResponse::ok(message.id)
2817            }
2818            RequestType::SoftStop(_) => {
2819                info!(
2820                    "{} {} processing soft shutdown",
2821                    log_module_context!(),
2822                    message.id
2823                );
2824                let listeners: HashMap<_, _> = self.listeners.drain().collect();
2825                for l in listeners.values() {
2826                    l.borrow_mut()
2827                        .listener
2828                        .take()
2829                        .map(|mut sock| self.registry.deregister(&mut sock));
2830                }
2831                WorkerResponse::processing(message.id)
2832            }
2833            RequestType::HardStop(_) => {
2834                info!("{} {} hard shutdown", log_module_context!(), message.id);
2835                let mut listeners: HashMap<_, _> = self.listeners.drain().collect();
2836                for (_, l) in listeners.drain() {
2837                    l.borrow_mut()
2838                        .listener
2839                        .take()
2840                        .map(|mut sock| self.registry.deregister(&mut sock));
2841                }
2842                WorkerResponse::ok(message.id)
2843            }
2844            RequestType::Status(_) => {
2845                info!("{} {} status", log_module_context!(), message.id);
2846                WorkerResponse::ok(message.id)
2847            }
2848            RequestType::AddCluster(cluster) => {
2849                let config = ClusterConfiguration {
2850                    proxy_protocol: cluster
2851                        .proxy_protocol
2852                        .and_then(|n| ProxyProtocolConfig::try_from(n).ok()),
2853                    //load_balancing: cluster.load_balancing,
2854                    max_connections_per_ip: cluster.max_connections_per_ip,
2855                };
2856                self.configs.insert(cluster.cluster_id, config);
2857                WorkerResponse::ok(message.id)
2858            }
2859            RequestType::RemoveCluster(cluster_id) => {
2860                self.configs.remove(&cluster_id);
2861                WorkerResponse::ok(message.id)
2862            }
2863            RequestType::RemoveListener(remove) => {
2864                if !self.remove_listener(remove.address.into()) {
2865                    WorkerResponse::error(
2866                        message.id,
2867                        format!("no TCP listener to remove at address {:?}", remove.address),
2868                    )
2869                } else {
2870                    WorkerResponse::ok(message.id)
2871                }
2872            }
2873            command => {
2874                debug!(
2875                    "{} {} unsupported message for TCP proxy, ignoring {:?}",
2876                    log_module_context!(),
2877                    message.id,
2878                    command
2879                );
2880                WorkerResponse::error(message.id, "unsupported message")
2881            }
2882        }
2883    }
2884
2885    fn accept(&mut self, token: ListenToken) -> Result<MioTcpStream, AcceptError> {
2886        let internal_token = Token(token.0);
2887        if let Some(listener) = self.listeners.get(&internal_token) {
2888            if let Some(tcp_listener) = &listener.borrow().listener {
2889                tcp_listener
2890                    .accept()
2891                    .map(|(frontend_sock, _)| frontend_sock)
2892                    .map_err(|e| match e.kind() {
2893                        ErrorKind::WouldBlock => AcceptError::WouldBlock,
2894                        _ => {
2895                            error!("{} accept() IO error: {:?}", log_module_context!(), e);
2896                            AcceptError::IoError
2897                        }
2898                    })
2899            } else {
2900                Err(AcceptError::IoError)
2901            }
2902        } else {
2903            Err(AcceptError::IoError)
2904        }
2905    }
2906
2907    fn create_session(
2908        &mut self,
2909        mut frontend_sock: MioTcpStream,
2910        token: ListenToken,
2911        wait_time: Duration,
2912        proxy: Rc<RefCell<Self>>,
2913    ) -> Result<(), AcceptError> {
2914        let listener_token = Token(token.0);
2915
2916        let listener = self
2917            .listeners
2918            .get(&listener_token)
2919            .ok_or(AcceptError::IoError)?;
2920
2921        let owned = listener.borrow();
2922        let mut pool = self.pool.borrow_mut();
2923
2924        let (front_buffer, back_buffer) = match (pool.checkout(), pool.checkout()) {
2925            (Some(fb), Some(bb)) => (fb, bb),
2926            _ => {
2927                error!("{} could not get buffers from pool", log_module_context!());
2928                error!(
2929                    "{} Buffer capacity has been reached, stopping to accept new connections for now",
2930                    log_module_context!()
2931                );
2932                gauge!(names::accept_queue::BACKPRESSURE, 1);
2933                self.sessions.borrow_mut().can_accept = false;
2934
2935                return Err(AcceptError::BufferCapacityReached);
2936            }
2937        };
2938
2939        // A listener may route either by a legacy no-SNI catch-all cluster
2940        // OR by SNI-scoped routes (never both -- enforced at config load,
2941        // sozu-proxy/sozu#1279); reject only when NEITHER is configured.
2942        if owned.cluster_id.is_none() && owned.sni_routes.is_empty() {
2943            error!(
2944                "{} listener at address {:?} has no linked cluster",
2945                log_module_context!(),
2946                owned.address
2947            );
2948            return Err(AcceptError::IoError);
2949        }
2950
2951        if let Err(e) = frontend_sock.set_nodelay(true) {
2952            error!(
2953                "{} error setting nodelay on front socket({:?}): {:?}",
2954                log_module_context!(),
2955                frontend_sock,
2956                e
2957            );
2958        }
2959
2960        let mut session_manager = self.sessions.borrow_mut();
2961        let entry = session_manager.slab.vacant_entry();
2962        let frontend_token = Token(entry.key());
2963
2964        if let Err(register_error) = self.registry.register(
2965            &mut frontend_sock,
2966            frontend_token,
2967            Interest::READABLE | Interest::WRITABLE,
2968        ) {
2969            error!(
2970                "{} error registering front socket({:?}): {:?}",
2971                log_module_context!(),
2972                frontend_sock,
2973                register_error
2974            );
2975            return Err(AcceptError::RegisterError);
2976        }
2977
2978        let session = if !owned.sni_routes.is_empty() {
2979            // Routing decides the cluster post-accept; the effective
2980            // preread cap can never exceed what the checked-out buffer can
2981            // actually hold, regardless of the configured knob.
2982            let effective_max_bytes = effective_sni_preread_max_bytes(
2983                owned.config.sni_preread_max_bytes,
2984                front_buffer.capacity(),
2985            );
2986            let preread_timeout = Duration::from_secs(u64::from(
2987                owned
2988                    .config
2989                    .sni_preread_timeout
2990                    .unwrap_or(DEFAULT_SNI_PREREAD_TIMEOUT),
2991            ));
2992            TcpSession::new_sni_preread(
2993                back_buffer,
2994                Duration::from_secs(owned.config.back_timeout as u64),
2995                Duration::from_secs(owned.config.connect_timeout as u64),
2996                front_buffer,
2997                frontend_token,
2998                listener.clone(),
2999                proxy,
3000                frontend_sock,
3001                wait_time,
3002                preread_timeout,
3003                effective_max_bytes,
3004            )
3005        } else {
3006            let proxy_protocol = self
3007                .configs
3008                .get(owned.cluster_id.as_ref().unwrap())
3009                .and_then(|c| c.proxy_protocol);
3010            TcpSession::new(
3011                back_buffer,
3012                None,
3013                owned.cluster_id.clone(),
3014                Duration::from_secs(owned.config.back_timeout as u64),
3015                Duration::from_secs(owned.config.connect_timeout as u64),
3016                Duration::from_secs(owned.config.front_timeout as u64),
3017                front_buffer,
3018                frontend_token,
3019                listener.clone(),
3020                proxy_protocol,
3021                proxy,
3022                frontend_sock,
3023                wait_time,
3024            )
3025        };
3026        incr!(names::tcp::REQUESTS);
3027
3028        let session = Rc::new(RefCell::new(session));
3029        entry.insert(session);
3030
3031        Ok(())
3032    }
3033}
3034
3035pub mod testing {
3036    use crate::testing::*;
3037
3038    /// This is not directly used by Sōzu but is available for example and testing purposes
3039    pub fn start_tcp_worker(
3040        config: TcpListenerConfig,
3041        max_buffers: usize,
3042        buffer_size: usize,
3043        channel: ProxyChannel,
3044    ) -> anyhow::Result<()> {
3045        let address = config.address.into();
3046
3047        let ServerParts {
3048            event_loop,
3049            registry,
3050            sessions,
3051            pool,
3052            backends,
3053            client_scm_socket: _,
3054            server_scm_socket,
3055            server_config,
3056        } = prebuild_server(max_buffers, buffer_size, true)?;
3057
3058        let token = {
3059            let mut sessions = sessions.borrow_mut();
3060            let entry = sessions.slab.vacant_entry();
3061            let key = entry.key();
3062            let _ = entry.insert(Rc::new(RefCell::new(ListenSession {
3063                protocol: Protocol::TCPListen,
3064            })));
3065            Token(key)
3066        };
3067
3068        let mut proxy = TcpProxy::new(registry, sessions.clone(), pool.clone(), backends.clone());
3069        proxy
3070            .add_listener(config, token)
3071            .with_context(|| "Failed at creating adding the listener")?;
3072        proxy
3073            .activate_listener(&address, None)
3074            .with_context(|| "Failed at creating activating the listener")?;
3075
3076        let mut server = Server::new(
3077            event_loop,
3078            channel,
3079            server_scm_socket,
3080            sessions,
3081            pool,
3082            backends,
3083            None,
3084            None,
3085            Some(proxy),
3086            server_config,
3087            None,
3088            false,
3089        )
3090        .with_context(|| "Failed at creating server")?;
3091
3092        debug!("{} starting event loop", log_module_context!());
3093        server.run();
3094        debug!("{} ending event loop", log_module_context!());
3095        Ok(())
3096    }
3097}
3098
3099#[cfg(test)]
3100mod tests {
3101    use std::{
3102        io::{Read, Write},
3103        net::{Shutdown, TcpListener, TcpStream},
3104        str,
3105        sync::{
3106            Arc, Barrier,
3107            atomic::{AtomicBool, Ordering},
3108        },
3109        thread,
3110        time::Duration,
3111    };
3112
3113    use sozu_command::{
3114        channel::Channel,
3115        config::ListenerBuilder,
3116        proto::command::{
3117            LoadBalancingParams, RequestTcpFrontend, SocketAddress, SoftStop, WorkerRequest,
3118            WorkerResponse, request::RequestType,
3119        },
3120    };
3121
3122    use super::testing::start_tcp_worker;
3123    use crate::testing::*;
3124
3125    /*
3126    #[test]
3127    #[cfg(target_pointer_width = "64")]
3128    fn size_test() {
3129      assert_size!(Pipe<mio::net::TcpStream>, 224);
3130      assert_size!(SendProxyProtocol<mio::net::TcpStream>, 144);
3131      assert_size!(RelayProxyProtocol<mio::net::TcpStream>, 152);
3132      assert_size!(ExpectProxyProtocol<mio::net::TcpStream>, 520);
3133      assert_size!(State, 528);
3134      // fails depending on the platform?
3135      //assert_size!(Session, 808);
3136    }*/
3137
3138    #[test]
3139    fn round_trip() {
3140        setup_test_logger!();
3141        let barrier = Arc::new(Barrier::new(2));
3142        let test_finished = Arc::new(AtomicBool::new(false));
3143
3144        let front_port1 = provide_port();
3145        let front_port2 = provide_port();
3146
3147        let backend_port = start_server(barrier.clone(), test_finished.clone());
3148        let mut command =
3149            start_proxy(backend_port, front_port1, front_port2).expect("Could not start proxy");
3150        barrier.wait();
3151
3152        thread::scope(|_s| {
3153            let front_addr = format!("127.0.0.1:{front_port1}");
3154
3155            let mut s1 = TcpStream::connect(&front_addr).expect("could not connect");
3156            s1.set_read_timeout(Some(Duration::from_secs(5)))
3157                .expect("could not set read timeout on s1");
3158
3159            let s3 = TcpStream::connect(&front_addr).expect("could not connect");
3160
3161            let mut s2 = TcpStream::connect(&front_addr).expect("could not connect");
3162            s2.set_read_timeout(Some(Duration::from_secs(5)))
3163                .expect("could not set read timeout on s2");
3164
3165            s1.write_all(b"hello ").expect("could not write to s1");
3166            println!("s1 sent");
3167
3168            s2.write_all(b"pouet pouet").expect("could not write to s2");
3169            println!("s2 sent");
3170
3171            let mut res = [0; 128];
3172            s1.write_all(b"coucou").expect("could not write to s1");
3173
3174            s3.shutdown(Shutdown::Both).expect("could not shutdown s3");
3175
3176            let sz2 = s2
3177                .read(&mut res[..])
3178                .expect("could not read from socket s2");
3179            println!("s2 received {:?}", str::from_utf8(&res[..sz2]));
3180            assert_eq!(&res[..sz2], &b"pouet pouet"[..]);
3181
3182            // Read in a loop: a single read() on a TCP stream is not
3183            // guaranteed to return all echoed data if the second write's
3184            // round trip (client → proxy → backend → proxy → client) is
3185            // still in flight when we poll.
3186            let expected = b"hello coucou";
3187            let mut total = 0;
3188            while total < expected.len() {
3189                let sz = s1
3190                    .read(&mut res[total..])
3191                    .expect("could not read from socket s1");
3192                assert!(sz > 0, "connection closed before receiving all data");
3193                total += sz;
3194            }
3195            println!(
3196                "s1 received again({}): {:?}",
3197                total,
3198                str::from_utf8(&res[..total])
3199            );
3200            assert_eq!(&res[..total], &expected[..]);
3201
3202            // Signal the echo server to stop
3203            test_finished.store(true, Ordering::Relaxed);
3204
3205            // Send SoftStop to the sozu worker so server.run() exits cleanly
3206            command
3207                .write_message(&WorkerRequest {
3208                    id: "ID_SOFTSTOP".to_owned(),
3209                    content: RequestType::SoftStop(SoftStop {}).into(),
3210                })
3211                .expect("could not send SoftStop to sozu worker");
3212        });
3213    }
3214
3215    /// Start an echo server on an ephemeral port.
3216    /// Returns the port the server is listening on.
3217    fn start_server(barrier: Arc<Barrier>, test_finished: Arc<AtomicBool>) -> u16 {
3218        let listener =
3219            TcpListener::bind("127.0.0.1:0").expect("could not bind echo server listener");
3220        let port = listener
3221            .local_addr()
3222            .expect("could not get echo server local address")
3223            .port();
3224
3225        listener
3226            .set_nonblocking(true)
3227            .expect("could not set echo server listener to non-blocking");
3228
3229        thread::spawn(move || {
3230            barrier.wait();
3231            let mut count: u8 = 0;
3232            loop {
3233                match listener.accept() {
3234                    Ok((mut stream, _)) => {
3235                        let finished = test_finished.clone();
3236                        thread::spawn(move || {
3237                            println!("got a new client: {count}");
3238                            stream
3239                                .set_read_timeout(Some(Duration::from_secs(2)))
3240                                .expect("could not set read timeout on echo client");
3241                            let mut buf = [0; 128];
3242                            loop {
3243                                match stream.read(&mut buf[..]) {
3244                                    Ok(0) => break,
3245                                    Ok(sz) => {
3246                                        println!(
3247                                            "ECHO[{count}] got \"{:?}\"",
3248                                            str::from_utf8(&buf[..sz])
3249                                        );
3250                                        stream
3251                                            .write_all(&buf[..sz])
3252                                            .expect("could not echo data back");
3253                                    }
3254                                    Err(ref e)
3255                                        if e.kind() == std::io::ErrorKind::WouldBlock
3256                                            || e.kind() == std::io::ErrorKind::TimedOut =>
3257                                    {
3258                                        if finished.load(Ordering::Relaxed) {
3259                                            println!("backend server stopping (client handler)");
3260                                            break;
3261                                        }
3262                                    }
3263                                    Err(_) => break,
3264                                }
3265                            }
3266                        });
3267                        count = count.wrapping_add(1);
3268                    }
3269                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
3270                        if test_finished.load(Ordering::Relaxed) {
3271                            println!("backend server stopping (accept loop)");
3272                            break;
3273                        }
3274                        thread::sleep(Duration::from_millis(50));
3275                    }
3276                    Err(e) => {
3277                        println!("connection failed: {e:?}");
3278                    }
3279                }
3280            }
3281        });
3282
3283        port
3284    }
3285
3286    /// Start a sozu TCP proxy worker with the given backend and frontend ports.
3287    fn start_proxy(
3288        backend_port: u16,
3289        front_port1: u16,
3290        front_port2: u16,
3291    ) -> anyhow::Result<Channel<WorkerRequest, WorkerResponse>> {
3292        let config = ListenerBuilder::new_tcp(SocketAddress::new_v4(127, 0, 0, 1, front_port1))
3293            .to_tcp(None)
3294            .expect("could not create listener config");
3295
3296        let (mut command, channel) =
3297            Channel::generate(1000, 10000).with_context(|| "should create a channel")?;
3298        let _jg = thread::spawn(move || {
3299            setup_test_logger!();
3300            start_tcp_worker(config, 100, 16384, channel).expect("could not start the tcp server");
3301        });
3302
3303        command
3304            .blocking()
3305            .expect("could not set command channel to blocking");
3306        {
3307            let front = RequestTcpFrontend {
3308                cluster_id: "yolo".to_owned(),
3309                address: SocketAddress::new_v4(127, 0, 0, 1, front_port1),
3310                ..Default::default()
3311            };
3312            let backend = sozu_command_lib::response::Backend {
3313                cluster_id: "yolo".to_owned(),
3314                backend_id: "yolo-0".to_owned(),
3315                address: SocketAddress::new_v4(127, 0, 0, 1, backend_port).into(),
3316                load_balancing_parameters: Some(LoadBalancingParams::default()),
3317                sticky_id: None,
3318                backup: None,
3319            };
3320
3321            command
3322                .write_message(&WorkerRequest {
3323                    id: "ID_YOLO1".to_owned(),
3324                    content: RequestType::AddTcpFrontend(front).into(),
3325                })
3326                .expect("could not send AddTcpFrontend for front1");
3327            command
3328                .write_message(&WorkerRequest {
3329                    id: "ID_YOLO2".to_owned(),
3330                    content: RequestType::AddBackend(backend.to_add_backend()).into(),
3331                })
3332                .expect("could not send AddBackend for front1");
3333        }
3334        {
3335            let front = RequestTcpFrontend {
3336                cluster_id: "yolo".to_owned(),
3337                address: SocketAddress::new_v4(127, 0, 0, 1, front_port2),
3338                ..Default::default()
3339            };
3340            let backend = sozu_command::response::Backend {
3341                cluster_id: "yolo".to_owned(),
3342                backend_id: "yolo-0".to_owned(),
3343                address: SocketAddress::new_v4(127, 0, 0, 1, backend_port).into(),
3344                load_balancing_parameters: Some(LoadBalancingParams::default()),
3345                sticky_id: None,
3346                backup: None,
3347            };
3348            command
3349                .write_message(&WorkerRequest {
3350                    id: "ID_YOLO3".to_owned(),
3351                    content: RequestType::AddTcpFrontend(front).into(),
3352                })
3353                .expect("could not send AddTcpFrontend for front2");
3354            command
3355                .write_message(&WorkerRequest {
3356                    id: "ID_YOLO4".to_owned(),
3357                    content: RequestType::AddBackend(backend.to_add_backend()).into(),
3358                })
3359                .expect("could not send AddBackend for front2");
3360        }
3361
3362        for _ in 0..4 {
3363            println!(
3364                "read_message: {:?}",
3365                command
3366                    .read_message()
3367                    .with_context(|| "could not read message")?
3368            );
3369        }
3370
3371        Ok(command)
3372    }
3373}
3374
3375/// Unit coverage for the SNI-preread routing shell added for
3376/// sozu-proxy/sozu#1279: the route-table mutations (`add_tcp_front` /
3377/// `remove_tcp_front`), the `AlpnMatcher` mapping, the effective preread
3378/// cap, and the routing gate's data invariants. None of this needs a live
3379/// socket or event loop -- it is a separate module (rather than nested in
3380/// the `tests` module above) purely to avoid that module's `use
3381/// std::net::TcpListener` import shadowing `super::TcpListener` (this
3382/// crate's listener struct).
3383#[cfg(test)]
3384mod sni_routing_tests {
3385    use sozu_command::{config::ListenerBuilder, proto::command::SocketAddress};
3386
3387    use super::*;
3388    use crate::testing::{ServerParts, prebuild_server, provide_port};
3389
3390    fn test_listener() -> TcpListener {
3391        let config = ListenerBuilder::new_tcp(SocketAddress::new_v4(127, 0, 0, 1, provide_port()))
3392            .to_tcp(None)
3393            .expect("could not build a TcpListenerConfig for the test");
3394        TcpListener::new(config, Token(0)).expect("could not build a bare TcpListener for the test")
3395    }
3396
3397    fn frontend(cluster_id: &str, sni: Option<&str>, alpn: &[&str]) -> RequestTcpFrontend {
3398        RequestTcpFrontend {
3399            cluster_id: cluster_id.to_owned(),
3400            address: SocketAddress::new_v4(127, 0, 0, 1, provide_port()),
3401            sni: sni.map(str::to_owned),
3402            alpn: alpn.iter().map(|p| p.to_string()).collect(),
3403            ..Default::default()
3404        }
3405    }
3406
3407    // ---- effective_sni_preread_max_bytes ------------------------------
3408
3409    #[test]
3410    fn effective_max_bytes_falls_back_to_default_when_unconfigured() {
3411        assert_eq!(
3412            effective_sni_preread_max_bytes(None, 65536),
3413            DEFAULT_SNI_PREREAD_MAX_BYTES as usize
3414        );
3415    }
3416
3417    #[test]
3418    fn effective_max_bytes_is_the_min_of_knob_and_capacity() {
3419        assert_eq!(effective_sni_preread_max_bytes(Some(8192), 16384), 8192);
3420        assert_eq!(effective_sni_preread_max_bytes(Some(32768), 16384), 16384);
3421        assert_eq!(effective_sni_preread_max_bytes(Some(16384), 16384), 16384);
3422    }
3423
3424    #[test]
3425    fn effective_max_bytes_never_below_the_floor() {
3426        // A `sni_preread_max_bytes = 0` knob reaching the worker from a
3427        // direct `sozu listener tcp add`/`update` CLI/IPC request (or a stale
3428        // LoadState replay) bypasses config.rs's loud MIN_SNI_PREREAD_MAX_BYTES
3429        // load-time reject. Without the floor the shell would issue
3430        // zero-length preread reads and spin until the loop guard closes each
3431        // session; the floor degrades a sub-minimum knob to the 5-byte
3432        // TLS-record-header minimum instead.
3433        assert_eq!(
3434            effective_sni_preread_max_bytes(Some(0), 16384),
3435            MIN_SNI_PREREAD_MAX_BYTES as usize,
3436            "a 0 knob must degrade to the floor, never 0 (would spin the preread)"
3437        );
3438        assert_eq!(
3439            effective_sni_preread_max_bytes(Some(3), 16384),
3440            MIN_SNI_PREREAD_MAX_BYTES as usize,
3441            "any sub-floor knob must be raised to the floor"
3442        );
3443        // The floor itself, and anything above it, are respected unchanged.
3444        assert_eq!(
3445            effective_sni_preread_max_bytes(Some(MIN_SNI_PREREAD_MAX_BYTES), 16384),
3446            MIN_SNI_PREREAD_MAX_BYTES as usize
3447        );
3448        assert_eq!(
3449            effective_sni_preread_max_bytes(Some(MIN_SNI_PREREAD_MAX_BYTES + 1), 16384),
3450            (MIN_SNI_PREREAD_MAX_BYTES + 1) as usize
3451        );
3452    }
3453
3454    // ---- known_alpn_label (access-log tagging) -------------------------
3455
3456    #[test]
3457    fn known_alpn_label_picks_the_clients_first_offer() {
3458        assert_eq!(
3459            known_alpn_label(&[b"h2".to_vec(), b"http/1.1".to_vec()]),
3460            Some("h2")
3461        );
3462        assert_eq!(
3463            known_alpn_label(&[b"http/1.1".to_vec(), b"h2".to_vec()]),
3464            Some("http/1.1"),
3465            "client preference order must win, not a fixed h2-first priority"
3466        );
3467    }
3468
3469    #[test]
3470    fn known_alpn_label_is_none_for_empty_or_unrecognized_offers() {
3471        assert_eq!(known_alpn_label(&[]), None);
3472        assert_eq!(known_alpn_label(&[b"spdy/1".to_vec()]), None);
3473    }
3474
3475    // ---- AlpnMatcher mapping + route-table add/remove symmetry --------
3476
3477    #[test]
3478    fn empty_alpn_maps_to_any_non_empty_maps_to_one_of() {
3479        let mut listener = test_listener();
3480        listener
3481            .insert_sni_route("example.com".to_owned(), vec![], "cluster-any".to_owned())
3482            .expect("insert_sni_route must succeed for a valid test SNI");
3483        listener
3484            .insert_sni_route(
3485                "h2.example.com".to_owned(),
3486                vec!["h2".to_owned(), "http/1.1".to_owned()],
3487                "cluster-h2".to_owned(),
3488            )
3489            .expect("insert_sni_route must succeed for a valid test SNI");
3490
3491        let (_, any_entries) = listener
3492            .sni_routes
3493            .domain_lookup(b"example.com", true)
3494            .expect("example.com must be routable");
3495        assert_eq!(
3496            any_entries,
3497            &vec![(AlpnMatcher::Any, "cluster-any".to_owned())]
3498        );
3499
3500        let (_, h2_entries) = listener
3501            .sni_routes
3502            .domain_lookup(b"h2.example.com", true)
3503            .expect("h2.example.com must be routable");
3504        assert_eq!(
3505            h2_entries,
3506            &vec![(
3507                AlpnMatcher::OneOf([b"h2".to_vec(), b"http/1.1".to_vec()].into_iter().collect()),
3508                "cluster-h2".to_owned()
3509            )]
3510        );
3511    }
3512
3513    #[test]
3514    fn insert_sni_route_appends_under_the_same_sni() {
3515        let mut listener = test_listener();
3516        listener
3517            .insert_sni_route(
3518                "example.com".to_owned(),
3519                vec!["h2".to_owned()],
3520                "cluster-h2".to_owned(),
3521            )
3522            .expect("insert_sni_route must succeed for a valid test SNI");
3523        listener
3524            .insert_sni_route(
3525                "example.com".to_owned(),
3526                vec![],
3527                "cluster-default".to_owned(),
3528            )
3529            .expect("insert_sni_route must succeed for a valid test SNI");
3530
3531        let (_, entries) = listener
3532            .sni_routes
3533            .domain_lookup(b"example.com", true)
3534            .expect("example.com must be routable");
3535        assert_eq!(entries.len(), 2, "both fronts must share the SNI's Vec");
3536    }
3537
3538    #[test]
3539    fn remove_sni_route_drops_only_the_matching_entry() {
3540        let mut listener = test_listener();
3541        listener
3542            .insert_sni_route(
3543                "example.com".to_owned(),
3544                vec!["h2".to_owned()],
3545                "cluster-h2".to_owned(),
3546            )
3547            .expect("insert_sni_route must succeed for a valid test SNI");
3548        listener
3549            .insert_sni_route(
3550                "example.com".to_owned(),
3551                vec![],
3552                "cluster-default".to_owned(),
3553            )
3554            .expect("insert_sni_route must succeed for a valid test SNI");
3555
3556        listener.remove_sni_route(
3557            "example.com".to_owned(),
3558            vec!["h2".to_owned()],
3559            &"cluster-h2".to_owned(),
3560        );
3561
3562        let (_, entries) = listener
3563            .sni_routes
3564            .domain_lookup(b"example.com", true)
3565            .expect("example.com must still be routable via the remaining entry");
3566        assert_eq!(
3567            entries,
3568            &vec![(AlpnMatcher::Any, "cluster-default".to_owned())],
3569            "removing one entry must not disturb the other"
3570        );
3571    }
3572
3573    #[test]
3574    fn remove_sni_route_empties_the_trie_key_when_the_last_entry_goes() {
3575        let mut listener = test_listener();
3576        listener
3577            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".to_owned())
3578            .expect("insert_sni_route must succeed for a valid test SNI");
3579        assert!(!listener.sni_routes.is_empty());
3580
3581        listener.remove_sni_route("example.com".to_owned(), vec![], &"cluster-a".to_owned());
3582
3583        assert!(
3584            listener.sni_routes.is_empty(),
3585            "domain_remove must run once the SNI's Vec empties, leaving no stranded key"
3586        );
3587        assert!(
3588            listener
3589                .sni_routes
3590                .domain_lookup(b"example.com", true)
3591                .is_none()
3592        );
3593    }
3594
3595    #[test]
3596    fn remove_sni_route_on_an_absent_sni_is_a_harmless_no_op() {
3597        let mut listener = test_listener();
3598        listener
3599            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".to_owned())
3600            .expect("insert_sni_route must succeed for a valid test SNI");
3601
3602        // Removing a route for a SNI that was never inserted must not panic
3603        // and must not disturb the existing route.
3604        listener.remove_sni_route(
3605            "other.example.net".to_owned(),
3606            vec![],
3607            &"cluster-a".to_owned(),
3608        );
3609
3610        assert!(
3611            listener
3612                .sni_routes
3613                .domain_lookup(b"example.com", true)
3614                .is_some()
3615        );
3616    }
3617
3618    // ---- exact-key bookkeeping must not fall back to a sibling wildcard
3619    // (route-table corruption caught in sozu-proxy/sozu#1290 review) ----
3620
3621    #[test]
3622    fn insert_sni_route_creates_a_distinct_node_for_an_exact_key_over_a_sibling_wildcard() {
3623        let mut listener = test_listener();
3624        // Wildcard catch-all first.
3625        listener
3626            .insert_sni_route(
3627                "*.example.com".to_owned(),
3628                vec![],
3629                "cluster-wildcard".to_owned(),
3630            )
3631            .expect("insert_sni_route must succeed for a valid test SNI");
3632        // Exact ALPN-scoped route for one specific subdomain.
3633        listener
3634            .insert_sni_route(
3635                "a.example.com".to_owned(),
3636                vec!["h2".to_owned()],
3637                "cluster-a-h2".to_owned(),
3638            )
3639            .expect("insert_sni_route must succeed for a valid test SNI");
3640
3641        // The exact key must have gotten its OWN trie node -- with
3642        // `accept_wildcard: true` this lookup would instead fall back to
3643        // (and the insert above would have corrupted) the wildcard's node,
3644        // since no literal `a` child existed yet at insert time.
3645        let (_, a_entries) = listener
3646            .sni_routes
3647            .domain_lookup(b"a.example.com", false)
3648            .expect("a.example.com must have a distinct exact-key node");
3649        assert_eq!(
3650            a_entries,
3651            &vec![(
3652                AlpnMatcher::OneOf([b"h2".to_vec()].into_iter().collect()),
3653                "cluster-a-h2".to_owned()
3654            )],
3655            "the exact key's own Vec must hold only its own entry, not the wildcard's"
3656        );
3657
3658        // Any OTHER subdomain must still resolve to ONLY the wildcard, via
3659        // the same `accept_wildcard: true` lookup the routing path uses
3660        // (`preread_config`) -- it must never see `a.example.com`'s h2 route.
3661        let (_, b_entries) = listener
3662            .sni_routes
3663            .domain_lookup(b"b.example.com", true)
3664            .expect("b.example.com must fall back to the wildcard catch-all");
3665        assert_eq!(
3666            b_entries,
3667            &vec![(AlpnMatcher::Any, "cluster-wildcard".to_owned())],
3668            "an unrelated subdomain must see ONLY the wildcard's entry"
3669        );
3670    }
3671
3672    #[test]
3673    fn validate_new_tcp_front_accepts_an_exact_catch_all_sibling_of_a_wildcard_catch_all() {
3674        let mut listener = test_listener();
3675        listener
3676            .insert_sni_route(
3677                "*.example.com".to_owned(),
3678                vec![],
3679                "cluster-wildcard".to_owned(),
3680            )
3681            .expect("insert_sni_route must succeed for a valid test SNI");
3682
3683        // An exact catch-all for one subdomain must be accepted: it is a
3684        // SIBLING of the wildcard's catch-all, not a duplicate of it. With
3685        // `accept_wildcard: true` this lookup would wrongly find the
3686        // wildcard's own `AlpnMatcher::Any` entry and reject it as "already
3687        // has a catch-all".
3688        let front = frontend("cluster-a", Some("a.example.com"), &[]);
3689        assert!(
3690            listener.validate_new_tcp_front(&front).is_ok(),
3691            "an exact catch-all must be accepted when only a SIBLING wildcard has a catch-all"
3692        );
3693    }
3694
3695    #[test]
3696    fn validate_new_tcp_front_rejects_a_duplicate_wildcard_catch_all() {
3697        let mut listener = test_listener();
3698        listener
3699            .insert_sni_route(
3700                "*.example.com".to_owned(),
3701                vec![],
3702                "cluster-wildcard".to_owned(),
3703            )
3704            .expect("insert_sni_route must succeed for a valid test SNI");
3705
3706        // A SECOND catch-all for the SAME wildcard key is an ambiguous
3707        // duplicate and must be rejected. The immutable trie `lookup` has no
3708        // literal-`*` short-circuit (only `lookup_mut` does), so a plain
3709        // `accept_wildcard: false` self-lookup never sees the existing
3710        // wildcard entry and waves the duplicate through -- which
3711        // `insert_sni_route` (lookup_mut, short-circuit present) would then
3712        // happily append.
3713        let front = frontend("cluster-dup", Some("*.example.com"), &[]);
3714        assert!(
3715            listener.validate_new_tcp_front(&front).is_err(),
3716            "a second catch-all on the same wildcard SNI must be rejected as a duplicate"
3717        );
3718    }
3719
3720    #[test]
3721    fn validate_new_tcp_front_rejects_overlapping_alpn_on_the_same_wildcard() {
3722        let mut listener = test_listener();
3723        listener
3724            .insert_sni_route(
3725                "*.example.com".to_owned(),
3726                vec!["h2".to_owned()],
3727                "cluster-wildcard-h2".to_owned(),
3728            )
3729            .expect("insert_sni_route must succeed for a valid test SNI");
3730
3731        // Same wildcard key, overlapping ALPN protocol: ambiguous, must be
3732        // rejected (same bypass as the duplicate catch-all above).
3733        let front = frontend("cluster-dup", Some("*.example.com"), &["h2"]);
3734        assert!(
3735            listener.validate_new_tcp_front(&front).is_err(),
3736            "an overlapping ALPN matcher on the same wildcard SNI must be rejected"
3737        );
3738    }
3739
3740    #[test]
3741    fn validate_new_tcp_front_accepts_a_disjoint_alpn_addition_on_the_same_wildcard() {
3742        let mut listener = test_listener();
3743        listener
3744            .insert_sni_route(
3745                "*.example.com".to_owned(),
3746                vec![],
3747                "cluster-wildcard".to_owned(),
3748            )
3749            .expect("insert_sni_route must succeed for a valid test SNI");
3750
3751        // A non-overlapping ALPN-scoped addition alongside the wildcard's
3752        // catch-all stays legal -- the wildcard-aware self-lookup must not
3753        // over-reject.
3754        let front = frontend("cluster-h2", Some("*.example.com"), &["h2"]);
3755        assert!(
3756            listener.validate_new_tcp_front(&front).is_ok(),
3757            "a disjoint ALPN addition on the same wildcard SNI must be accepted"
3758        );
3759    }
3760
3761    /// The worker boundary must reject malformed SNI SHAPES that a direct
3762    /// `AddTcpFrontend` (command socket) or `LoadState` replay could carry
3763    /// past config.rs's `validate_sni_pattern`: a `/.../` label would be
3764    /// inserted into the `pattern_trie` as a REGEX route and a misplaced
3765    /// `*` as an unintended wildcard — silently widening routing.
3766    #[test]
3767    fn validate_new_tcp_front_rejects_malformed_sni_shapes() {
3768        let listener = test_listener();
3769
3770        for bad in [
3771            "/[a-z]+/.example.com",
3772            "foo/bar.example.com",
3773            "a.*.example.com",
3774            "*.*.example.com",
3775            "*",
3776        ] {
3777            let front = frontend("cluster-a", Some(bad), &[]);
3778            assert!(
3779                listener.validate_new_tcp_front(&front).is_err(),
3780                "malformed SNI shape {bad:?} must be rejected at the worker boundary"
3781            );
3782        }
3783
3784        // The two documented-legal shapes must still pass.
3785        for good in ["a.example.com", "*.example.com"] {
3786            let front = frontend("cluster-a", Some(good), &[]);
3787            assert!(
3788                listener.validate_new_tcp_front(&front).is_ok(),
3789                "legal SNI shape {good:?} must still be accepted"
3790            );
3791        }
3792    }
3793
3794    /// `validate_new_tcp_front`'s OLD hand-rolled shape check only rejected
3795    /// `/` and a misplaced `*`, letting an empty label (leading dot,
3796    /// trailing dot, consecutive dots), an empty string, or a non-ASCII
3797    /// pattern reach `insert_sni_route` -> `pattern_trie::insert_recursive`,
3798    /// whose RELEASE-mode `assert_ne!(partial_key, &b""[..])` panics the
3799    /// worker on a leading-empty-label key like `.example.com`. This is the
3800    /// worker-boundary counterpart to `config.rs`'s `validate_sni_pattern`
3801    /// tests: every shape the shared validator rejects at config-load must
3802    /// also be rejected here, for `AddTcpFrontend`/`LoadState` requests that
3803    /// bypass config.rs entirely.
3804    #[test]
3805    fn add_tcp_front_enforces_the_shared_sni_validator_at_the_worker_boundary() {
3806        let mut proxy = test_proxy();
3807        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
3808        let config = ListenerBuilder::new_tcp(address)
3809            .to_tcp(None)
3810            .expect("could not build listener config");
3811        let token = Token(0);
3812        proxy
3813            .add_listener(config, token)
3814            .expect("could not add listener");
3815
3816        for bad in [
3817            ".example.com",     // leading empty label -- the pattern_trie-crashing shape
3818            "example.com.",     // trailing empty label
3819            "a..b.example.com", // consecutive dots -- empty middle label
3820            "",                 // empty pattern
3821            "exämple.com",      // non-ASCII
3822        ] {
3823            let front = frontend("cluster-a", Some(bad), &[]);
3824            let front = RequestTcpFrontend { address, ..front };
3825            match proxy.add_tcp_front(front) {
3826                Err(ProxyError::InvalidTcpFrontend { .. }) => {}
3827                other => panic!("malformed sni {bad:?} must be rejected, got {other:?}"),
3828            }
3829
3830            let listener = proxy
3831                .listeners
3832                .get(&token)
3833                .expect("listener must be present")
3834                .borrow();
3835            assert!(
3836                listener.sni_routes.is_empty(),
3837                "a rejected sni {bad:?} must leave sni_routes empty"
3838            );
3839            assert!(
3840                listener.cluster_id.is_none(),
3841                "a rejected sni {bad:?} must leave cluster_id unset"
3842            );
3843        }
3844
3845        // A space IS valid ASCII, is not '*'/'/', and splits into non-empty
3846        // labels ("exa", "mple.com") -- none of the shared validator's rules
3847        // (empty pattern, non-ASCII, misplaced '*', empty label) cover "not a
3848        // valid hostname character" in general, so this shape is ACCEPTED,
3849        // not rejected. Documenting the validator's actual behavior rather
3850        // than assuming a hostname-shaped string with a space would be
3851        // caught too.
3852        let space_front = frontend("cluster-space", Some("exa mple.com"), &[]);
3853        let space_front = RequestTcpFrontend {
3854            address,
3855            ..space_front
3856        };
3857        assert!(
3858            proxy.add_tcp_front(space_front).is_ok(),
3859            "the shared SNI validator does not reject an embedded space -- \
3860             it is valid ASCII with no empty label"
3861        );
3862    }
3863
3864    // ---- per-frontend access-log tags keying (sozu-proxy/sozu#1290) ----
3865
3866    #[test]
3867    fn sni_fronts_keep_distinct_tags_under_distinct_keys() {
3868        let mut proxy = test_proxy();
3869        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
3870        let config = ListenerBuilder::new_tcp(address)
3871            .to_tcp(None)
3872            .expect("could not build listener config");
3873        let token = Token(0);
3874        proxy
3875            .add_listener(config, token)
3876            .expect("could not add listener");
3877
3878        let front_a = RequestTcpFrontend {
3879            cluster_id: "cluster-a".to_owned(),
3880            address,
3881            sni: Some("a.example.com".to_owned()),
3882            alpn: vec![],
3883            tags: std::collections::BTreeMap::from([("team".to_owned(), "alpha".to_owned())]),
3884        };
3885        let front_b = RequestTcpFrontend {
3886            cluster_id: "cluster-b".to_owned(),
3887            address,
3888            sni: Some("b.example.com".to_owned()),
3889            alpn: vec!["h2".to_owned()],
3890            tags: std::collections::BTreeMap::from([("team".to_owned(), "beta".to_owned())]),
3891        };
3892        proxy
3893            .add_tcp_front(front_a)
3894            .expect("add_tcp_front A must succeed");
3895        proxy
3896            .add_tcp_front(front_b)
3897            .expect("add_tcp_front B must succeed");
3898
3899        let std_address: SocketAddr = address.into();
3900        let key_a = sni_tags_key(&std_address, "a.example.com", &[]);
3901        let key_b = sni_tags_key(&std_address, "b.example.com", &["h2".to_owned()]);
3902
3903        let listener = proxy
3904            .listeners
3905            .get(&token)
3906            .expect("listener must be present")
3907            .borrow();
3908        let tags_a = listener
3909            .get_tags(&key_a)
3910            .expect("front A's tags must live under its own composed key");
3911        assert_eq!(
3912            tags_a.tags.get("team").map(String::as_str),
3913            Some("alpha"),
3914            "front A's tags must survive front B's add, not be clobbered by it"
3915        );
3916        let tags_b = listener
3917            .get_tags(&key_b)
3918            .expect("front B's tags must live under its own composed key");
3919        assert_eq!(tags_b.tags.get("team").map(String::as_str), Some("beta"));
3920    }
3921
3922    #[test]
3923    fn removing_one_sni_front_clears_only_its_own_tags() {
3924        let mut proxy = test_proxy();
3925        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
3926        let config = ListenerBuilder::new_tcp(address)
3927            .to_tcp(None)
3928            .expect("could not build listener config");
3929        let token = Token(0);
3930        proxy
3931            .add_listener(config, token)
3932            .expect("could not add listener");
3933
3934        let front_a = RequestTcpFrontend {
3935            cluster_id: "cluster-a".to_owned(),
3936            address,
3937            sni: Some("a.example.com".to_owned()),
3938            alpn: vec![],
3939            tags: std::collections::BTreeMap::from([("team".to_owned(), "alpha".to_owned())]),
3940        };
3941        let front_b = RequestTcpFrontend {
3942            cluster_id: "cluster-b".to_owned(),
3943            address,
3944            sni: Some("b.example.com".to_owned()),
3945            alpn: vec!["h2".to_owned()],
3946            tags: std::collections::BTreeMap::from([("team".to_owned(), "beta".to_owned())]),
3947        };
3948        proxy
3949            .add_tcp_front(front_a.clone())
3950            .expect("add_tcp_front A must succeed");
3951        proxy
3952            .add_tcp_front(front_b)
3953            .expect("add_tcp_front B must succeed");
3954
3955        proxy
3956            .remove_tcp_front(front_a)
3957            .expect("remove_tcp_front A must succeed");
3958
3959        let std_address: SocketAddr = address.into();
3960        let key_a = sni_tags_key(&std_address, "a.example.com", &[]);
3961        let key_b = sni_tags_key(&std_address, "b.example.com", &["h2".to_owned()]);
3962
3963        let listener = proxy
3964            .listeners
3965            .get(&token)
3966            .expect("listener must be present")
3967            .borrow();
3968        assert!(
3969            listener.get_tags(&key_a).is_none(),
3970            "removing front A must clear its own tags entry"
3971        );
3972        assert!(
3973            listener.get_tags(&key_b).is_some(),
3974            "removing front A must NOT clear sibling front B's tags"
3975        );
3976    }
3977
3978    /// A session routed to front A must look tags up under front A's key:
3979    /// the key `add_tcp_front` stores and the key `upgrade_sni_preread`
3980    /// rebuilds from the route decision (`matched_sni_pattern` +
3981    /// `matched_alpn`) must be identical, regardless of the operator's
3982    /// original ALPN order or SNI casing, and must never collide with the
3983    /// bare-address key used by no-SNI fronts.
3984    #[test]
3985    fn route_time_tags_key_rebuild_matches_the_add_time_key() {
3986        let address: SocketAddr = "127.0.0.1:9000".parse().expect("test address");
3987
3988        // Wildcard front, operator wrote mixed case + reverse ALPN order.
3989        let add_time = sni_tags_key(
3990            &address,
3991            "*.Example.COM",
3992            &["http/1.1".to_owned(), "h2".to_owned()],
3993        );
3994        let matcher =
3995            AlpnMatcher::OneOf([b"h2".to_vec(), b"http/1.1".to_vec()].into_iter().collect());
3996        let route_time = sni_tags_key(
3997            &address,
3998            "*.example.com", // matched_sni_pattern: the lowercased trie key
3999            &alpn_matcher_protocols(&matcher),
4000        );
4001        assert_eq!(
4002            add_time, route_time,
4003            "add-time and route-time keys must agree for the same front"
4004        );
4005
4006        // Catch-all front: empty alpn at add time <-> AlpnMatcher::Any.
4007        assert_eq!(
4008            sni_tags_key(&address, "a.example.com", &[]),
4009            sni_tags_key(
4010                &address,
4011                "a.example.com",
4012                &alpn_matcher_protocols(&AlpnMatcher::Any)
4013            )
4014        );
4015
4016        // A composed key never collides with the bare-address key.
4017        assert_ne!(add_time, address.to_string());
4018    }
4019
4020    /// Regression from the sozu-proxy/sozu#1290 review: ALPN protocol
4021    /// identifiers are opaque byte strings -- nothing forbids a `,` inside
4022    /// one -- so a SINGLE protocol `"a,b"` and the DISJOINT pair `["a",
4023    /// "b"]` are two legal, distinct `AlpnMatcher`s on the same
4024    /// `(address, sni)`, yet the naive `sorted_alpn.join(",")` key collapses
4025    /// both to the literal string `"a,b"`.
4026    #[test]
4027    fn alpn_sets_differing_only_by_an_embedded_separator_get_distinct_tags_keys() {
4028        let address: SocketAddr = "127.0.0.1:9001".parse().expect("test address");
4029        let key_joined = sni_tags_key(&address, "example.com", &["a,b".to_owned()]);
4030        let key_split = sni_tags_key(&address, "example.com", &["a".to_owned(), "b".to_owned()]);
4031        assert_ne!(
4032            key_joined, key_split,
4033            "a single \"a,b\" protocol must not collide with the disjoint [\"a\", \"b\"] pair"
4034        );
4035    }
4036
4037    /// End-to-end through `TcpProxy`: both fronts are accepted (they are
4038    /// genuinely disjoint `AlpnMatcher`s, so `validate_new_tcp_front`'s
4039    /// overlap check does not reject either), each keeps its OWN tags under
4040    /// its own composed key, and removing one leaves the other's tags
4041    /// intact.
4042    #[test]
4043    fn alpn_sets_differing_only_by_an_embedded_separator_keep_distinct_tags() {
4044        let mut proxy = test_proxy();
4045        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4046        let config = ListenerBuilder::new_tcp(address)
4047            .to_tcp(None)
4048            .expect("could not build listener config");
4049        let token = Token(0);
4050        proxy
4051            .add_listener(config, token)
4052            .expect("could not add listener");
4053
4054        let front_joined = RequestTcpFrontend {
4055            cluster_id: "cluster-joined".to_owned(),
4056            address,
4057            sni: Some("example.com".to_owned()),
4058            alpn: vec!["a,b".to_owned()],
4059            tags: std::collections::BTreeMap::from([("variant".to_owned(), "joined".to_owned())]),
4060        };
4061        let front_split = RequestTcpFrontend {
4062            cluster_id: "cluster-split".to_owned(),
4063            address,
4064            sni: Some("example.com".to_owned()),
4065            alpn: vec!["a".to_owned(), "b".to_owned()],
4066            tags: std::collections::BTreeMap::from([("variant".to_owned(), "split".to_owned())]),
4067        };
4068
4069        proxy
4070            .add_tcp_front(front_joined.clone())
4071            .expect("the single \"a,b\" protocol front must be accepted");
4072        proxy.add_tcp_front(front_split.clone()).expect(
4073            "the disjoint [\"a\", \"b\"] front must be accepted -- it is NOT the same ALPN \
4074                 set as [\"a,b\"]",
4075        );
4076
4077        let std_address: SocketAddr = address.into();
4078        let key_joined = sni_tags_key(&std_address, "example.com", &["a,b".to_owned()]);
4079        let key_split = sni_tags_key(
4080            &std_address,
4081            "example.com",
4082            &["a".to_owned(), "b".to_owned()],
4083        );
4084        assert_ne!(key_joined, key_split);
4085
4086        {
4087            let listener = proxy
4088                .listeners
4089                .get(&token)
4090                .expect("listener must be present")
4091                .borrow();
4092            assert_eq!(
4093                listener
4094                    .get_tags(&key_joined)
4095                    .and_then(|t| t.tags.get("variant"))
4096                    .map(String::as_str),
4097                Some("joined"),
4098                "the joined front's tags must live under its own composed key"
4099            );
4100            assert_eq!(
4101                listener
4102                    .get_tags(&key_split)
4103                    .and_then(|t| t.tags.get("variant"))
4104                    .map(String::as_str),
4105                Some("split"),
4106                "the split front's tags must live under its own DISTINCT composed key"
4107            );
4108        }
4109
4110        proxy
4111            .remove_tcp_front(front_joined)
4112            .expect("remove the joined front");
4113
4114        let listener = proxy
4115            .listeners
4116            .get(&token)
4117            .expect("listener must be present")
4118            .borrow();
4119        assert!(
4120            listener.get_tags(&key_joined).is_none(),
4121            "removing the joined front must clear its own tags entry"
4122        );
4123        assert_eq!(
4124            listener
4125                .get_tags(&key_split)
4126                .and_then(|t| t.tags.get("variant"))
4127                .map(String::as_str),
4128            Some("split"),
4129            "removing the joined front must not disturb its sibling's tags"
4130        );
4131    }
4132
4133    #[test]
4134    fn remove_sni_route_for_an_absent_exact_key_does_not_strip_a_sibling_wildcards_entry() {
4135        let mut listener = test_listener();
4136        listener
4137            .insert_sni_route(
4138                "*.example.com".to_owned(),
4139                vec![],
4140                "cluster-wildcard".to_owned(),
4141            )
4142            .expect("insert_sni_route must succeed for a valid test SNI");
4143
4144        // "a.example.com" was never inserted as its own route -- only the
4145        // wildcard catch-all exists. A remove targeting the exact host
4146        // (e.g. a stale `RemoveTcpFrontend` replayed from a hand-edited
4147        // `LoadState`) must be a no-op here, not reach into and strip the
4148        // WILDCARD's own catch-all entry.
4149        listener.remove_sni_route(
4150            "a.example.com".to_owned(),
4151            vec![],
4152            &"cluster-wildcard".to_owned(),
4153        );
4154
4155        let (_, wildcard_entries) = listener
4156            .sni_routes
4157            .domain_lookup(b"b.example.com", true)
4158            .expect(
4159                "the wildcard catch-all must survive a remove targeting an unrelated exact key",
4160            );
4161        assert_eq!(
4162            wildcard_entries,
4163            &vec![(AlpnMatcher::Any, "cluster-wildcard".to_owned())]
4164        );
4165    }
4166
4167    // ---- routing gate data invariant: never both cluster_id AND routes ----
4168
4169    #[test]
4170    fn a_no_sni_front_leaves_the_route_table_empty() {
4171        let mut listener = test_listener();
4172        listener.cluster_id = Some("legacy-catch-all".to_owned());
4173        assert!(
4174            listener.sni_routes.is_empty(),
4175            "a listener with only a no-SNI front must never populate sni_routes"
4176        );
4177    }
4178
4179    #[test]
4180    fn an_sni_scoped_front_leaves_cluster_id_unset() {
4181        let mut listener = test_listener();
4182        listener
4183            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".to_owned())
4184            .expect("insert_sni_route must succeed for a valid test SNI");
4185        assert!(
4186            listener.cluster_id.is_none(),
4187            "a listener with only SNI-scoped fronts must never populate the legacy cluster_id"
4188        );
4189    }
4190
4191    // ---- end-to-end through TcpProxy::add_tcp_front / remove_tcp_front ----
4192
4193    fn test_proxy() -> TcpProxy {
4194        let ServerParts {
4195            registry,
4196            sessions,
4197            pool,
4198            backends,
4199            ..
4200        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
4201        TcpProxy::new(registry, sessions, pool, backends)
4202    }
4203
4204    #[test]
4205    fn add_then_remove_sni_front_round_trips_through_tcp_proxy() {
4206        let mut proxy = test_proxy();
4207        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4208        let config = ListenerBuilder::new_tcp(address)
4209            .to_tcp(None)
4210            .expect("could not build listener config");
4211        let token = Token(0);
4212        proxy
4213            .add_listener(config, token)
4214            .expect("could not add listener");
4215
4216        let front = RequestTcpFrontend {
4217            cluster_id: "cluster-a".to_owned(),
4218            address,
4219            sni: Some("Example.COM".to_owned()),
4220            alpn: vec![],
4221            ..Default::default()
4222        };
4223        proxy
4224            .add_tcp_front(front.clone())
4225            .expect("add_tcp_front must succeed");
4226
4227        {
4228            let listener = proxy
4229                .listeners
4230                .get(&token)
4231                .expect("listener must be present")
4232                .borrow();
4233            assert!(listener.cluster_id.is_none());
4234            let (_, entries) = listener
4235                .sni_routes
4236                // Lowercased at insert time regardless of wire-form casing.
4237                .domain_lookup(b"example.com", true)
4238                .expect("example.com must be routable after add_tcp_front");
4239            assert_eq!(entries, &vec![(AlpnMatcher::Any, "cluster-a".to_owned())]);
4240        }
4241        assert_eq!(proxy.fronts.get("cluster-a"), Some(&token));
4242
4243        proxy
4244            .remove_tcp_front(front)
4245            .expect("remove_tcp_front must succeed");
4246
4247        {
4248            let listener = proxy
4249                .listeners
4250                .get(&token)
4251                .expect("listener must be present")
4252                .borrow();
4253            assert!(
4254                listener.sni_routes.is_empty(),
4255                "remove_tcp_front must leave no stranded route"
4256            );
4257        }
4258        assert_eq!(
4259            proxy.fronts.get("cluster-a"),
4260            None,
4261            "remove_tcp_front must undo add_tcp_front's self.fronts bookkeeping"
4262        );
4263    }
4264
4265    #[test]
4266    fn add_then_remove_legacy_no_sni_front_round_trips() {
4267        let mut proxy = test_proxy();
4268        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4269        let config = ListenerBuilder::new_tcp(address)
4270            .to_tcp(None)
4271            .expect("could not build listener config");
4272        let token = Token(0);
4273        proxy
4274            .add_listener(config, token)
4275            .expect("could not add listener");
4276
4277        let front = frontend("cluster-legacy", None, &[]);
4278        let front = RequestTcpFrontend { address, ..front };
4279        proxy
4280            .add_tcp_front(front.clone())
4281            .expect("add_tcp_front must succeed");
4282
4283        assert_eq!(
4284            proxy
4285                .listeners
4286                .get(&token)
4287                .expect("listener must be present")
4288                .borrow()
4289                .cluster_id,
4290            Some("cluster-legacy".to_owned())
4291        );
4292
4293        proxy
4294            .remove_tcp_front(front)
4295            .expect("remove_tcp_front must succeed");
4296
4297        assert_eq!(
4298            proxy
4299                .listeners
4300                .get(&token)
4301                .expect("listener must be present")
4302                .borrow()
4303                .cluster_id,
4304            None
4305        );
4306        assert_eq!(proxy.fronts.get("cluster-legacy"), None);
4307    }
4308
4309    // ---- add_tcp_front hard-rejects routing-corrupting requests --------
4310    //
4311    // Worker-side mirror of `command/src/config.rs`'s TOML config-load
4312    // invariants (sozu-proxy/sozu#1279 hardening): `AddTcpFrontend` can
4313    // reach the worker directly over the command socket, or via `LoadState`
4314    // replay, bypassing config.rs entirely, so `add_tcp_front` must defend
4315    // itself rather than rely on a debug-only assertion.
4316
4317    #[test]
4318    fn add_tcp_front_rejects_alpn_without_sni() {
4319        let mut proxy = test_proxy();
4320        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4321        let config = ListenerBuilder::new_tcp(address)
4322            .to_tcp(None)
4323            .expect("could not build listener config");
4324        proxy
4325            .add_listener(config, Token(0))
4326            .expect("could not add listener");
4327
4328        let front = frontend("cluster-a", None, &["h2"]);
4329        let front = RequestTcpFrontend { address, ..front };
4330        match proxy.add_tcp_front(front) {
4331            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4332            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4333        }
4334    }
4335
4336    #[test]
4337    fn add_tcp_front_rejects_no_sni_front_on_listener_with_sni_routes() {
4338        let mut proxy = test_proxy();
4339        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4340        let config = ListenerBuilder::new_tcp(address)
4341            .to_tcp(None)
4342            .expect("could not build listener config");
4343        proxy
4344            .add_listener(config, Token(0))
4345            .expect("could not add listener");
4346
4347        let sni_front = frontend("cluster-a", Some("example.com"), &[]);
4348        let sni_front = RequestTcpFrontend {
4349            address,
4350            ..sni_front
4351        };
4352        proxy
4353            .add_tcp_front(sni_front)
4354            .expect("the first, SNI-scoped frontend must be accepted");
4355
4356        let no_sni_front = frontend("cluster-b", None, &[]);
4357        let no_sni_front = RequestTcpFrontend {
4358            address,
4359            ..no_sni_front
4360        };
4361        match proxy.add_tcp_front(no_sni_front) {
4362            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4363            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4364        }
4365    }
4366
4367    #[test]
4368    fn add_tcp_front_rejects_sni_front_on_listener_with_no_sni_cluster() {
4369        let mut proxy = test_proxy();
4370        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4371        let config = ListenerBuilder::new_tcp(address)
4372            .to_tcp(None)
4373            .expect("could not build listener config");
4374        proxy
4375            .add_listener(config, Token(0))
4376            .expect("could not add listener");
4377
4378        let no_sni_front = frontend("cluster-a", None, &[]);
4379        let no_sni_front = RequestTcpFrontend {
4380            address,
4381            ..no_sni_front
4382        };
4383        proxy
4384            .add_tcp_front(no_sni_front)
4385            .expect("the first, no-SNI frontend must be accepted");
4386
4387        let sni_front = frontend("cluster-b", Some("example.com"), &[]);
4388        let sni_front = RequestTcpFrontend {
4389            address,
4390            ..sni_front
4391        };
4392        match proxy.add_tcp_front(sni_front) {
4393            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4394            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4395        }
4396    }
4397
4398    #[test]
4399    fn add_tcp_front_rejects_alpn_overlap_on_same_sni() {
4400        let mut proxy = test_proxy();
4401        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4402        let config = ListenerBuilder::new_tcp(address)
4403            .to_tcp(None)
4404            .expect("could not build listener config");
4405        proxy
4406            .add_listener(config, Token(0))
4407            .expect("could not add listener");
4408
4409        let first = frontend("cluster-a", Some("example.com"), &["h2"]);
4410        let first = RequestTcpFrontend { address, ..first };
4411        proxy
4412            .add_tcp_front(first)
4413            .expect("the first frontend must be accepted");
4414
4415        // Second frontend shares "h2" with the first on the same sni.
4416        let second = frontend("cluster-b", Some("example.com"), &["h2", "http/1.1"]);
4417        let second = RequestTcpFrontend { address, ..second };
4418        match proxy.add_tcp_front(second) {
4419            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4420            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4421        }
4422    }
4423
4424    #[test]
4425    fn add_tcp_front_rejects_duplicate_catch_all_on_same_sni() {
4426        let mut proxy = test_proxy();
4427        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4428        let config = ListenerBuilder::new_tcp(address)
4429            .to_tcp(None)
4430            .expect("could not build listener config");
4431        proxy
4432            .add_listener(config, Token(0))
4433            .expect("could not add listener");
4434
4435        let first = frontend("cluster-a", Some("example.com"), &[]);
4436        let first = RequestTcpFrontend { address, ..first };
4437        proxy
4438            .add_tcp_front(first)
4439            .expect("the first catch-all frontend must be accepted");
4440
4441        let second = frontend("cluster-b", Some("example.com"), &[]);
4442        let second = RequestTcpFrontend { address, ..second };
4443        match proxy.add_tcp_front(second) {
4444            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4445            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4446        }
4447    }
4448
4449    /// The valid, intended shape (sozu-proxy/sozu#1279's whole reason for
4450    /// existing) must still be accepted: disjoint, non-empty `alpn` lists
4451    /// on the same `(address, sni)`, and a catch-all alongside a
4452    /// specific-protocol entry.
4453    #[test]
4454    fn add_tcp_front_accepts_disjoint_alpn_and_catch_all_on_same_sni() {
4455        let mut proxy = test_proxy();
4456        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4457        let config = ListenerBuilder::new_tcp(address)
4458            .to_tcp(None)
4459            .expect("could not build listener config");
4460        proxy
4461            .add_listener(config, Token(0))
4462            .expect("could not add listener");
4463
4464        let h2 = frontend("cluster-h2", Some("example.com"), &["h2"]);
4465        let h2 = RequestTcpFrontend { address, ..h2 };
4466        proxy
4467            .add_tcp_front(h2)
4468            .expect("disjoint alpn frontend must be accepted");
4469
4470        let http11 = frontend("cluster-http11", Some("example.com"), &["http/1.1"]);
4471        let http11 = RequestTcpFrontend { address, ..http11 };
4472        proxy
4473            .add_tcp_front(http11)
4474            .expect("second disjoint alpn frontend must be accepted");
4475
4476        let catch_all = frontend("cluster-default", Some("example.com"), &[]);
4477        let catch_all = RequestTcpFrontend {
4478            address,
4479            ..catch_all
4480        };
4481        proxy
4482            .add_tcp_front(catch_all)
4483            .expect("a catch-all alongside specific-protocol entries must be accepted");
4484    }
4485
4486    // ---- tcp.sni_preread.active gauge accounting ----------------------
4487
4488    /// Read the current process-local `tcp.sni_preread.active` gauge,
4489    /// treating an absent key as 0. `dump_local_proxy_metrics` is a
4490    /// non-draining filter over the proxy `MetricsMap`, so repeated reads are
4491    /// side-effect free and the key is the raw metric name.
4492    fn sni_preread_active_gauge() -> i64 {
4493        use sozu_command::proto::command::filtered_metrics::Inner;
4494        crate::metrics::METRICS.with(|metrics| {
4495            metrics
4496                .borrow_mut()
4497                .dump_local_proxy_metrics()
4498                .get(names::tcp::sni_preread::ACTIVE)
4499                .and_then(|fm| fm.inner.as_ref())
4500                .and_then(|inner| match inner {
4501                    Inner::Gauge(v) => Some(*v as i64),
4502                    _ => None,
4503                })
4504                .unwrap_or(0)
4505        })
4506    }
4507
4508    #[test]
4509    fn entering_sni_preread_increments_the_active_gauge() {
4510        // Regression guard for the missing-`+1` gauge bug
4511        // (sozu-proxy/sozu#1279): `new_sni_preread` must bump
4512        // `tcp.sni_preread.active` by exactly one when a session ENTERS the
4513        // state, so each of the two `-1` decrements -- the "upgrade" exit in
4514        // `upgrade_sni_preread` and the "reject"/"teardown" exit in `close()`'s
4515        // `StateMarker::SniPreread` arm -- has a matching increment. Without
4516        // this `+1` the first `-1` underflows a fresh-zero gauge (clamped to 0,
4517        // ERROR-logged), pinning the gauge at 0 and rendering the e2e gauge
4518        // assertion vacuous.
4519        //
4520        // `METRICS` is a thread-local shared across unit tests on the same
4521        // worker thread, so this asserts the DELTA around one constructor call
4522        // (robust to any starting value), not an absolute reading. The
4523        // net-zero-per-session contract spans the full lifecycle (accept ->
4524        // live backend connect -> upgrade/teardown) and is the behavioural job
4525        // of the e2e gauge assertion
4526        // (`test_tcp_sni_reject_then_valid_connection_not_limited` in
4527        // `e2e/src/tests/tcp_sni_tests.rs`), not reproducible at this unit
4528        // level.
4529        let ServerParts {
4530            registry,
4531            sessions,
4532            pool,
4533            backends,
4534            ..
4535        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
4536
4537        let proxy = Rc::new(RefCell::new(TcpProxy::new(
4538            registry,
4539            sessions,
4540            pool.clone(),
4541            backends,
4542        )));
4543        let listener = Rc::new(RefCell::new(test_listener()));
4544
4545        let (front_buffer, back_buffer) = {
4546            let mut pool = pool.borrow_mut();
4547            (
4548                pool.checkout().expect("front buffer checkout must succeed"),
4549                pool.checkout().expect("back buffer checkout must succeed"),
4550            )
4551        };
4552
4553        // A non-blocking connect to a (likely unused) loopback port returns a
4554        // real `MioTcpStream` handle immediately, regardless of whether the
4555        // connection completes; `new_sni_preread` only reads `peer_addr()`.
4556        let socket = MioTcpStream::connect(
4557            format!("127.0.0.1:{}", provide_port())
4558                .parse()
4559                .expect("loopback address must parse"),
4560        )
4561        .expect("mio connect must return a socket handle");
4562
4563        let before = sni_preread_active_gauge();
4564        let session = TcpSession::new_sni_preread(
4565            back_buffer,
4566            Duration::from_secs(30),
4567            Duration::from_secs(30),
4568            front_buffer,
4569            Token(0),
4570            listener,
4571            proxy,
4572            socket,
4573            Duration::from_millis(0),
4574            Duration::from_secs(3),
4575            16384,
4576        );
4577        let after = sni_preread_active_gauge();
4578
4579        // The session is measured while still in `SniPreread`; hold it across
4580        // the read so no future `Drop` side effect could race the measurement.
4581        assert!(matches!(session.state, TcpStateMachine::SniPreread(_)));
4582
4583        assert_eq!(
4584            after - before,
4585            1,
4586            "entering the SniPreread state must increment tcp.sni_preread.active by exactly one"
4587        );
4588    }
4589
4590    // ---- absolute preread deadline + route-aware timeout (sozu-proxy/sozu#1290) ----
4591
4592    /// `frontend_timeout_resets_on_readable` must gate the reset on exactly
4593    /// one thing: whether the CURRENT state is an undecided `SniPreread`.
4594    /// Constructed directly (no socket I/O needed) since the predicate is a
4595    /// pure function of `&TcpStateMachine`.
4596    #[test]
4597    fn frontend_timeout_reset_is_gated_on_sni_preread_decision() {
4598        let mut pool = crate::pool::Pool::with_capacity(1, 1, 16 * 1024);
4599        let frontend_buffer = pool.checkout().expect("frontend buffer");
4600        let socket = MioTcpStream::connect(
4601            format!("127.0.0.1:{}", provide_port())
4602                .parse()
4603                .expect("loopback address must parse"),
4604        )
4605        .expect("mio connect must return a socket handle");
4606
4607        let undecided = TcpStateMachine::SniPreread(SniPreread::new(
4608            socket,
4609            Token(0),
4610            Ulid::generate(),
4611            frontend_buffer,
4612            16384,
4613        ));
4614        assert!(
4615            !frontend_timeout_resets_on_readable(&undecided),
4616            "an undecided SniPreread must not have its absolute deadline reset"
4617        );
4618
4619        // Every OTHER state resets normally -- represented here by
4620        // `ExpectProxyProtocol`, cheaply constructible without driving a
4621        // real route decision (unlike a ROUTED `SniPreread`, whose
4622        // `outcome` field has no test-only setter and can only become
4623        // `Some` by parsing a real ClientHello -- see
4624        // `frontend_timeout_restored_and_timeout_after_route_is_not_double_counted`
4625        // below for that scenario end-to-end).
4626        let socket2 = MioTcpStream::connect(
4627            format!("127.0.0.1:{}", provide_port())
4628                .parse()
4629                .expect("loopback address must parse"),
4630        )
4631        .expect("mio connect must return a socket handle");
4632        let container = crate::timer::TimeoutContainer::new_empty(Duration::from_secs(5));
4633        let other = TcpStateMachine::ExpectProxyProtocol(ExpectProxyProtocol::new(
4634            container,
4635            socket2,
4636            Token(1),
4637            Ulid::generate(),
4638        ));
4639        assert!(
4640            frontend_timeout_resets_on_readable(&other),
4641            "every non-preread-undecided state must keep resetting its frontend timeout"
4642        );
4643    }
4644
4645    /// Minimal single-record TLS ClientHello wire carrying only a
4646    /// `server_name` extension for `host` -- hand-built rather than reusing
4647    /// `tcp_preread::parser`'s test-only wire-building helpers (`mod
4648    /// parser` is private to `tcp_preread`, unreachable from this sibling
4649    /// module) to drive a REAL route decision through
4650    /// `TcpSession::readable()` for the regression test below.
4651    fn minimal_client_hello_wire(host: &str) -> Vec<u8> {
4652        let mut name_list = vec![0u8]; // name_type = host_name
4653        name_list.extend_from_slice(&(host.len() as u16).to_be_bytes());
4654        name_list.extend_from_slice(host.as_bytes());
4655        let mut sni_ext_data = Vec::new();
4656        sni_ext_data.extend_from_slice(&(name_list.len() as u16).to_be_bytes());
4657        sni_ext_data.extend_from_slice(&name_list);
4658        let mut sni_ext = Vec::new();
4659        sni_ext.extend_from_slice(&0x0000u16.to_be_bytes()); // server_name extension type
4660        sni_ext.extend_from_slice(&(sni_ext_data.len() as u16).to_be_bytes());
4661        sni_ext.extend_from_slice(&sni_ext_data);
4662
4663        let mut body = Vec::new();
4664        body.extend_from_slice(&[0x03, 0x03]); // legacy_version
4665        body.extend_from_slice(&[0u8; 32]); // random
4666        body.push(0); // session_id: empty
4667        body.extend_from_slice(&[0x00, 0x02, 0x13, 0x01]); // cipher_suites
4668        body.push(1); // compression_methods length
4669        body.push(0); // compression_method: null
4670        body.extend_from_slice(&(sni_ext.len() as u16).to_be_bytes()); // extensions block length
4671        body.extend_from_slice(&sni_ext);
4672
4673        let mut handshake = Vec::new();
4674        handshake.push(1u8); // msg_type = client_hello
4675        let hs_len = body.len() as u32;
4676        handshake.extend_from_slice(&hs_len.to_be_bytes()[1..4]);
4677        handshake.extend_from_slice(&body);
4678
4679        let mut record = Vec::new();
4680        record.push(22u8); // ContentType::handshake
4681        record.extend_from_slice(&[0x03, 0x03]); // legacy record version
4682        record.extend_from_slice(&(handshake.len() as u16).to_be_bytes());
4683        record.extend_from_slice(&handshake);
4684        record
4685    }
4686
4687    /// Read the current process-local `tcp.sni_preread.routed` counter.
4688    /// Same non-draining-read pattern as `sni_preread_active_gauge`, but
4689    /// for a `Count` metric instead of a `Gauge`.
4690    fn sni_preread_routed_count() -> i64 {
4691        use sozu_command::proto::command::filtered_metrics::Inner;
4692        crate::metrics::METRICS.with(|metrics| {
4693            metrics
4694                .borrow_mut()
4695                .dump_local_proxy_metrics()
4696                .get(names::tcp::sni_preread::ROUTED)
4697                .and_then(|fm| fm.inner.as_ref())
4698                .and_then(|inner| match inner {
4699                    Inner::Count(v) => Some(*v),
4700                    _ => None,
4701                })
4702                .unwrap_or(0)
4703        })
4704    }
4705
4706    /// End-to-end regression from the sozu-proxy/sozu#1290 review:
4707    ///
4708    /// (a) the moment a real ClientHello routes, `container_frontend_timeout`
4709    ///     must already carry the listener's configured `front_timeout` --
4710    ///     not just once `upgrade_sni_preread` eventually runs (which can be
4711    ///     one or more `ready()` cycles later, after the backend connects);
4712    /// (b) a front-timeout firing AFTER that route decision (backend connect
4713    ///     still pending) must not re-feed `Input::Timeout` into the
4714    ///     already-decided core: pre-fix, doing so replayed the SAME latched
4715    ///     `Output::Routed` through `SniPreread::handle_output`'s `Routed`
4716    ///     arm a second time, double-incrementing `tcp.sni_preread.routed`
4717    ///     (release) and tripping `debug_assert!(self.outcome.is_none(),
4718    ///     ...)` (debug -- this test runs in a debug build, so pre-fix it
4719    ///     panics here).
4720    #[test]
4721    fn frontend_timeout_restored_and_timeout_after_route_is_not_double_counted() {
4722        let ServerParts {
4723            registry,
4724            sessions,
4725            pool,
4726            backends,
4727            ..
4728        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
4729        let proxy = Rc::new(RefCell::new(TcpProxy::new(
4730            registry,
4731            sessions,
4732            pool.clone(),
4733            backends,
4734        )));
4735
4736        let mut bare_listener = test_listener();
4737        bare_listener
4738            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".to_owned())
4739            .expect("insert_sni_route must succeed for a valid test SNI");
4740        let configured_front_timeout = bare_listener.config.front_timeout;
4741        let listener = Rc::new(RefCell::new(bare_listener));
4742
4743        let std_listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind test listener");
4744        let addr = std_listener.local_addr().expect("listener local addr");
4745        let mut client = std::net::TcpStream::connect(addr).expect("connect test client");
4746        let (server, _) = std_listener.accept().expect("accept test server");
4747        server.set_nonblocking(true).expect("server nonblocking");
4748
4749        let (front_buffer, back_buffer) = {
4750            let mut pool = pool.borrow_mut();
4751            (
4752                pool.checkout().expect("front buffer checkout must succeed"),
4753                pool.checkout().expect("back buffer checkout must succeed"),
4754            )
4755        };
4756
4757        let mut session = TcpSession::new_sni_preread(
4758            back_buffer,
4759            Duration::from_secs(30),
4760            Duration::from_secs(30),
4761            front_buffer,
4762            Token(0),
4763            listener,
4764            proxy,
4765            MioTcpStream::from_std(server),
4766            Duration::from_millis(0),
4767            Duration::from_secs(3),
4768            16384,
4769        );
4770
4771        {
4772            use std::io::Write as _;
4773            client
4774                .write_all(&minimal_client_hello_wire("example.com"))
4775                .expect("write ClientHello");
4776            client.flush().ok();
4777        }
4778
4779        let routed_before = sni_preread_routed_count();
4780        for _ in 0..10 {
4781            if session.cluster_id.is_some() {
4782                break;
4783            }
4784            let _ = session.readable();
4785        }
4786        assert_eq!(
4787            session.cluster_id.as_deref(),
4788            Some("cluster-a"),
4789            "the session must have routed on a valid ClientHello for a configured SNI"
4790        );
4791        assert_eq!(
4792            sni_preread_routed_count() - routed_before,
4793            1,
4794            "routing must count tcp.sni_preread.routed exactly once"
4795        );
4796
4797        // (a) front_timeout is restored the moment routing succeeds.
4798        assert_eq!(
4799            session.container_frontend_timeout.duration(),
4800            Duration::from_secs(configured_front_timeout as u64),
4801            "the frontend timeout must already carry the configured front_timeout right after \
4802             routing, not only once the backend connects"
4803        );
4804
4805        // (b) a timeout firing after the route already latched must not
4806        // double-count tcp.sni_preread.routed, and (running in a debug
4807        // build) must not panic on SniPreread::handle_output's
4808        // `debug_assert!(self.outcome.is_none(), ...)`.
4809        let _ = session.timeout(Token(0));
4810        assert_eq!(
4811            sni_preread_routed_count() - routed_before,
4812            1,
4813            "a timeout firing after the route already latched must not double-count \
4814             tcp.sni_preread.routed"
4815        );
4816
4817        drop(client);
4818    }
4819}