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    /// Build the [`PrereadConfig`] this listener's `SniPreread` sessions
2069    /// feed to [`crate::protocol::tcp_preread::SniPrereadCore::handle_input`].
2070    /// `routes`/`inbound_proxy`/`timeout` come straight from the listener
2071    /// config; `effective_max_bytes` is session-specific (already clamped to
2072    /// that session's buffer capacity at construction, see `create_session`),
2073    /// so it is passed in rather than re-derived here.
2074    fn preread_config(&self, effective_max_bytes: usize) -> PrereadConfig<'_> {
2075        PrereadConfig {
2076            routes: &self.sni_routes,
2077            inbound_proxy: self.config.expect_proxy,
2078            max_bytes: effective_max_bytes,
2079            timeout: Duration::from_secs(u64::from(
2080                self.config
2081                    .sni_preread_timeout
2082                    .unwrap_or(DEFAULT_SNI_PREREAD_TIMEOUT),
2083            )),
2084            accept_wildcard: true,
2085        }
2086    }
2087
2088    /// Validate an incoming `AddTcpFrontend` against this listener's
2089    /// CURRENT routing state before any mutation, mirroring
2090    /// `command/src/config.rs`'s TOML config-load TCP SNI/ALPN invariants
2091    /// (sozu-proxy/sozu#1279):
2092    ///
2093    /// - `alpn` set with no `sni`: the worker's no-SNI catch-all path never
2094    ///   consults `alpn`, so the protocol list would silently never be
2095    ///   enforced (mirrors `ConfigError::AlpnWithoutSni`).
2096    /// - a no-SNI frontend added to a listener that already has SNI-scoped
2097    ///   routes, or an SNI-scoped frontend added to a listener that already
2098    ///   has a no-SNI catch-all cluster (mirrors
2099    ///   `ConfigError::TcpListenerMixesSniAndNoSni`).
2100    /// - an ALPN protocol, or a catch-all (empty `alpn`), that overlaps an
2101    ///   existing route already registered for the same `(address, sni)`
2102    ///   (mirrors `ConfigError::TcpFrontendAlpnOverlap` /
2103    ///   `TcpFrontendMultipleAlpnCatchAll`).
2104    ///
2105    /// Config-load already rejects all of these shapes for requests built
2106    /// from a TOML file, but `AddTcpFrontend` can also arrive directly over
2107    /// the command socket, or via `LoadState` replay of a hand-edited or
2108    /// stale state file, bypassing config.rs entirely -- the worker must
2109    /// not silently corrupt its own routing table when that happens.
2110    fn validate_new_tcp_front(&self, front: &RequestTcpFrontend) -> Result<(), ProxyError> {
2111        let reject = |reason: String| {
2112            Err(ProxyError::InvalidTcpFrontend {
2113                address: self.address,
2114                reason,
2115            })
2116        };
2117
2118        match &front.sni {
2119            None => {
2120                if !front.alpn.is_empty() {
2121                    return reject(format!(
2122                        "alpn = {:?} set without sni: alpn only matches within an SNI-scoped \
2123                         preread, so a frontend without sni would silently ignore its alpn list",
2124                        front.alpn
2125                    ));
2126                }
2127                if !self.sni_routes.is_empty() {
2128                    return reject(
2129                        "a no-SNI frontend cannot be added to a listener that already has \
2130                         SNI-scoped routes"
2131                            .to_string(),
2132                    );
2133                }
2134            }
2135            Some(sni) => {
2136                if self.cluster_id.is_some() {
2137                    return reject(
2138                        "an SNI-scoped frontend cannot be added to a listener that already has \
2139                         a no-SNI catch-all cluster"
2140                            .to_string(),
2141                    );
2142                }
2143
2144                // SNI SHAPE check: delegate to the SAME validator config-load
2145                // uses (`sozu_command::config::validate_sni_pattern`) rather
2146                // than a hand-rolled partial check. A direct `AddTcpFrontend`
2147                // over the command socket, or a `LoadState` replay, bypasses
2148                // config.rs entirely, so the worker boundary must enforce
2149                // the identical rule -- including the checks a bare '/'/'*'
2150                // scan used to miss: empty string, non-ASCII, and any empty
2151                // label (leading/trailing/consecutive dots). An unvalidated
2152                // leading-empty-label pattern like `.example.com` would
2153                // otherwise reach `insert_sni_route` ->
2154                // `pattern_trie::insert_recursive`'s RELEASE-mode
2155                // `assert_ne!(partial_key, &b""[..])` and crash the worker.
2156                let normalized_sni = match validate_sni_pattern(sni) {
2157                    Ok(normalized) => normalized,
2158                    Err(config_error) => {
2159                        return reject(format!(
2160                            "sni {sni:?} failed SNI shape validation: {config_error}"
2161                        ));
2162                    }
2163                };
2164
2165                // Same key as `insert_sni_route`'s own lookup, so this checks
2166                // against exactly the entries the new route would be appended
2167                // alongside. This is bookkeeping over the key's OWN node, not
2168                // the routing lookup (`preread_config` keeps `true` for
2169                // that), so `accept_wildcard` must make the lookup EXACT for
2170                // both key shapes:
2171                //
2172                // - literal key -> `false`. With `true`, a literal key with
2173                //   no child yet (e.g. `a.example.com` when only
2174                //   `*.example.com` exists) falls back to the sibling
2175                //   wildcard's entries (`pattern_trie.rs`'s `lookup`
2176                //   wildcard-fallback branch), misattributing the wildcard's
2177                //   catch-all to the exact key (falsely rejecting a
2178                //   legitimate exact catch-all) — and symmetrically,
2179                //   `insert_sni_route`/`remove_sni_route` would corrupt the
2180                //   WILDCARD's `Vec` instead of touching a distinct
2181                //   exact-key node.
2182                // - wildcard key -> `true`. Unlike `lookup_mut` (which
2183                //   short-circuits `partial_key == b"*"` before consulting
2184                //   `accept_wildcard`, so insert/remove stay on `false`),
2185                //   the immutable `lookup` reaches a wildcard entry ONLY
2186                //   through the fallback branch; with `false` an existing
2187                //   `*.example.com` entry is invisible here and a duplicate
2188                //   catch-all / overlapping-ALPN wildcard front bypasses
2189                //   validation. For a wildcard key, `true` IS the exact
2190                //   self-lookup: the traversal descends the key's own
2191                //   literal ancestry, and `insert` never creates a literal
2192                //   `*` child that could shadow the node's `wildcard` slot.
2193                //
2194                // `starts_with(b"*.")` matches every wildcard shape
2195                // config-load can emit: `command/src/config.rs`'s
2196                // `validate_sni_pattern` only admits a single leading `*.`
2197                // label (any other `*` placement is rejected), and a plain
2198                // hostname key traverses literal children in both `lookup`
2199                // and `insert_recursive`. A bare `*` key (config-load
2200                // rejects it; only a bypassing IPC request can carry one)
2201                // is the known gap: `lookup_mut`'s short-circuit maps it to
2202                // the wildcard slot while this immutable self-lookup (flag
2203                // `false`) cannot see an existing entry there, so duplicate
2204                // bare-`*` fronts are not detected -- sibling keys stay
2205                // untouched either way.
2206                //
2207                // `normalized_sni` (not a fresh `sni.to_ascii_lowercase()`)
2208                // is used here: `validate_sni_pattern` already returned the
2209                // canonical lowercased form above, and reusing it keeps this
2210                // function's notion of "the key" identical to what
2211                // `insert_sni_route`/`remove_sni_route` compute from the
2212                // same input.
2213                let key = normalized_sni.into_bytes();
2214                let accept_wildcard_for_self_lookup = key.starts_with(b"*.");
2215                if let Some((_, existing)) = self
2216                    .sni_routes
2217                    .domain_lookup(&key, accept_wildcard_for_self_lookup)
2218                {
2219                    let new_is_catch_all = front.alpn.is_empty();
2220                    for (matcher, _cluster_id) in existing {
2221                        match matcher {
2222                            AlpnMatcher::Any if new_is_catch_all => {
2223                                return reject(format!(
2224                                    "sni {sni:?} already has a catch-all (empty alpn) \
2225                                     frontend: at most one frontend per (address, sni) may \
2226                                     omit alpn"
2227                                ));
2228                            }
2229                            AlpnMatcher::OneOf(protocols) => {
2230                                if let Some(overlap) = front
2231                                    .alpn
2232                                    .iter()
2233                                    .find(|protocol| protocols.contains(protocol.as_bytes()))
2234                                {
2235                                    return reject(format!(
2236                                        "sni {sni:?} already has a frontend matching ALPN \
2237                                         protocol {overlap:?}: ALPN matchers for the same \
2238                                         (address, sni) must not overlap"
2239                                    ));
2240                                }
2241                            }
2242                            AlpnMatcher::Any => {}
2243                        }
2244                    }
2245                }
2246            }
2247        }
2248
2249        Ok(())
2250    }
2251
2252    /// Add one `(AlpnMatcher, ClusterId)` entry to this listener's SNI route
2253    /// table, appending to the SNI key's existing `Vec` if one is already
2254    /// present rather than clobbering it (multiple ALPN-scoped fronts can
2255    /// share the same SNI). `sni` is defensively lowercased: the
2256    /// SNI-preread core normalizes (lowercase, no trailing dot) before ever
2257    /// looking a route up, so the table key must already be in that form.
2258    ///
2259    /// Returns `Err` when the underlying trie insert reports
2260    /// [`InsertResult::Failed`] (a malformed key, e.g. an empty label) —
2261    /// the caller (`TcpProxy::add_tcp_front`) must propagate this rather
2262    /// than silently keep going with an add that reports success while the
2263    /// route table never gained a usable entry. With `validate_new_tcp_front`
2264    /// enforcing the shared SNI shape validator before this ever runs, a
2265    /// `Failed` result should be unreachable here — the `debug_assert_ne!`
2266    /// below asserts OUR OWN validated invariant, not raw wire/IPC input,
2267    /// and is a loud Sōzu-internal-bug canary in debug builds; the `Err`
2268    /// path is the release-mode graceful fallback if that invariant is ever
2269    /// violated by a future bug.
2270    fn insert_sni_route(
2271        &mut self,
2272        sni: String,
2273        alpn: Vec<String>,
2274        cluster_id: ClusterId,
2275    ) -> Result<(), ProxyError> {
2276        let (key, matcher) = route_key_and_matcher(&sni, alpn);
2277        // `accept_wildcard: false` — see `validate_new_tcp_front`'s comment:
2278        // an exact key must never fall back to a sibling wildcard's entry,
2279        // or this push would corrupt the WILDCARD's route `Vec` instead of
2280        // creating this key's own node.
2281        match self.sni_routes.domain_lookup_mut(&key, false) {
2282            Some((_, entries)) => {
2283                entries.push((matcher, cluster_id));
2284                Ok(())
2285            }
2286            None => {
2287                let insert_result = self
2288                    .sni_routes
2289                    .domain_insert(key, vec![(matcher, cluster_id)]);
2290                debug_assert_ne!(
2291                    insert_result,
2292                    InsertResult::Failed,
2293                    "insert_sni_route's key must already be validated by validate_new_tcp_front"
2294                );
2295                if insert_result == InsertResult::Failed {
2296                    error!(
2297                        "{} SNI route insert failed for {:?} despite passing shape \
2298                         validation -- rejecting to avoid a silently dead route",
2299                        log_module_context!(),
2300                        sni
2301                    );
2302                    return Err(ProxyError::InvalidTcpFrontend {
2303                        address: self.address,
2304                        reason: format!(
2305                            "internal error: the route table rejected sni {sni:?} despite \
2306                             passing shape validation"
2307                        ),
2308                    });
2309                }
2310                Ok(())
2311            }
2312        }
2313    }
2314
2315    /// Symmetric counterpart to [`Self::insert_sni_route`]: removes the
2316    /// matching `(AlpnMatcher, ClusterId)` entry, then `domain_remove`s the
2317    /// SNI key itself once its `Vec` is empty (no stranded empty entries in
2318    /// the trie).
2319    fn remove_sni_route(&mut self, sni: String, alpn: Vec<String>, cluster_id: &ClusterId) {
2320        let (key, matcher) = route_key_and_matcher(&sni, alpn);
2321        // `accept_wildcard: false` — same reasoning as `insert_sni_route`:
2322        // removing the exact key must never reach into and strip a sibling
2323        // wildcard's entries.
2324        if let Some((_, entries)) = self.sni_routes.domain_lookup_mut(&key, false) {
2325            entries.retain(|(m, c)| !(*m == matcher && c == cluster_id));
2326            if entries.is_empty() {
2327                self.sni_routes.domain_remove(&key);
2328            }
2329        }
2330    }
2331
2332    pub fn activate(
2333        &mut self,
2334        registry: &Registry,
2335        tcp_listener: Option<MioTcpListener>,
2336    ) -> Result<Token, ProxyError> {
2337        if self.active {
2338            return Ok(self.token);
2339        }
2340
2341        let mut listener = match tcp_listener {
2342            Some(listener) => listener,
2343            None => {
2344                let address = self.config.address.into();
2345                server_bind(address).map_err(|e| ProxyError::BindToSocket(address, e))?
2346            }
2347        };
2348
2349        registry
2350            .register(&mut listener, self.token, Interest::READABLE)
2351            .map_err(ProxyError::RegisterListener)?;
2352
2353        self.listener = Some(listener);
2354        self.active = true;
2355        Ok(self.token)
2356    }
2357
2358    /// Apply a partial-update patch to this TCP listener's live configuration.
2359    ///
2360    /// Fields absent in the patch (i.e. `None`) are preserved unchanged.
2361    pub fn update_config(&mut self, patch: &UpdateTcpListenerConfig) -> Result<(), ListenerError> {
2362        if let Some(v) = patch.public_address {
2363            self.config.public_address = Some(v);
2364        }
2365        if let Some(v) = patch.expect_proxy {
2366            self.config.expect_proxy = v;
2367        }
2368        if let Some(v) = patch.front_timeout {
2369            self.config.front_timeout = v;
2370        }
2371        if let Some(v) = patch.back_timeout {
2372            self.config.back_timeout = v;
2373        }
2374        if let Some(v) = patch.connect_timeout {
2375            self.config.connect_timeout = v;
2376        }
2377        Ok(())
2378    }
2379}
2380
2381fn handle_connection_result(
2382    connection_result: Result<BackendConnectAction, BackendConnectionError>,
2383) -> Option<SessionResult> {
2384    match connection_result {
2385        // reuse connection or send a default answer, we can continue
2386        Ok(BackendConnectAction::Reuse) => None,
2387        Ok(BackendConnectAction::New) | Ok(BackendConnectAction::Replace) => {
2388            // we must wait for an event
2389            Some(SessionResult::Continue)
2390        }
2391        Err(_) => {
2392            // in case of BackendConnectionError::Backend(BackendError::ConnectionFailures(..))
2393            // we may want to retry instead of closing
2394            Some(SessionResult::Close)
2395        }
2396    }
2397}
2398
2399/// `min(listener.config.sni_preread_max_bytes, frontend_buffer.capacity())`,
2400/// floored at [`MIN_SNI_PREREAD_MAX_BYTES`] -- the SNI-preread core has no
2401/// independent backstop of its own (`SniPrereadCore::handle_input` trusts
2402/// `PrereadConfig::max_bytes` entirely), so the shell must never hand it a
2403/// cap the checked-out buffer cannot actually hold (the `min`), NOR a cap so
2404/// small the preread read is zero-length and spins until the loop guard (the
2405/// `max`). Config-load rejects a sub-floor `sni_preread_max_bytes` loudly
2406/// (`ConfigError::SniPrereadMaxBytesTooSmall`), but a `0` knob from a direct
2407/// `sozu listener tcp add` CLI/IPC request, or a stale `LoadState`
2408/// replay, bypasses that check and reaches the worker -- this floor degrades
2409/// it to the 5-byte TLS-record-header minimum instead, killing the spin for
2410/// EVERY config source at the single point of use. `buffer_capacity` is
2411/// always `>= MIN_SNI_PREREAD_MAX_BYTES` in practice (buffers are KB-sized),
2412/// so the `min` never fights the `max`.
2413fn effective_sni_preread_max_bytes(configured: Option<u32>, buffer_capacity: usize) -> usize {
2414    (configured.unwrap_or(DEFAULT_SNI_PREREAD_MAX_BYTES) as usize)
2415        .min(buffer_capacity)
2416        .max(MIN_SNI_PREREAD_MAX_BYTES as usize)
2417}
2418
2419/// Whether `TcpSession::readable`'s frontend-timeout reset should fire for
2420/// the CURRENT state. `false` only while the session is an UNDECIDED
2421/// `SniPreread`: the preread deadline armed at session creation
2422/// (`new_sni_preread`) is an ABSOLUTE budget, not a per-fragment idle timer
2423/// -- resetting it on every `readable()` event would let a client
2424/// trickling one byte just before each expiry hold the session (and both
2425/// its checked-out buffers) open far past the configured
2426/// `sni_preread_timeout` (sozu-proxy/sozu#1290). Every
2427/// other state -- a `SniPreread` that has already routed, or any state
2428/// reached after it -- resets normally: once routed, the frontend timeout
2429/// reverts to being a genuine per-read idle timer (see the `front_timeout`
2430/// restore in `readable()`'s route-capture block).
2431fn frontend_timeout_resets_on_readable(state: &TcpStateMachine) -> bool {
2432    !matches!(
2433        state,
2434        TcpStateMachine::SniPreread(preread) if preread.outcome().is_none()
2435    )
2436}
2437
2438/// Access-log ALPN tag for an SNI-routed TCP session: the client's FIRST
2439/// offered protocol (client preference order, matching
2440/// `SniPrereadCore::route`'s own routing precedence), mapped to a known
2441/// `&'static str` label -- the same two labels `https.rs`'s own ALPN
2442/// negotiation records (`"h2"` / `"http/1.1"`) -- so `tcp.sni_preread`
2443/// sessions and terminated-TLS sessions chart under the same values.
2444/// `None` for an empty offer or an unrecognized protocol: Sōzu never
2445/// terminates TLS on this path, so this is the client's stated preference,
2446/// not a negotiated outcome.
2447fn known_alpn_label(offered: &[Vec<u8>]) -> Option<&'static str> {
2448    match offered.first().map(Vec::as_slice) {
2449        Some(b"h2") => Some("h2"),
2450        Some(b"http/1.1") => Some("http/1.1"),
2451        _ => None,
2452    }
2453}
2454
2455/// Worker-internal, collision-free identity for a TCP frontend's access-log
2456/// tags entry (sozu-proxy/sozu#1290). ALPN protocol
2457/// identifiers are opaque RFC 7301 byte strings -- nothing forbids a `,` or
2458/// `|` inside one -- so a naive `sorted_alpn.join(",")` string key let two
2459/// legal, DISJOINT fronts collide: a single protocol `"a,b"` and the pair
2460/// `["a", "b"]` both joined to the identical string `"a,b"`, so the second
2461/// `add_tcp_front` silently clobbered the first front's tags entry.
2462///
2463/// `TcpListener::tags: BTreeMap<String, CachedTags>` and
2464/// `Pipe::set_tags_key(Option<String>)` (`lib/src/protocol/pipe.rs`) both
2465/// key on a plain `String` -- crossing either boundary still needs one --
2466/// so this type does not replace that storage; it is the SINGLE place that
2467/// composes the string, via [`Self::fmt`]'s length-prefixed per-protocol
2468/// encoding, so no choice of in-band separator can be confused with a
2469/// protocol-name boundary the way a bare `join(",")` could.
2470#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
2471enum TcpFrontendTagsKey {
2472    /// Legacy no-SNI catch-all front: keyed by the bare listener address,
2473    /// exactly as before SNI routing existed.
2474    Address(SocketAddr),
2475    /// An SNI-scoped front. `alpn` is sorted (and deduped, since a plain
2476    /// `Vec<String>` cannot enforce uniqueness the way
2477    /// [`AlpnMatcher::OneOf`]'s `BTreeSet` does) so two constructions from
2478    /// the same logical protocol set always compare equal regardless of
2479    /// the caller's original order.
2480    Sni {
2481        address: SocketAddr,
2482        sni: String,
2483        alpn: Vec<String>,
2484    },
2485}
2486
2487impl TcpFrontendTagsKey {
2488    /// The identity triple `command/src/state.rs`'s `add_tcp_frontend`
2489    /// deduplicates on: `sni` is lowercased to match `insert_sni_route`'s
2490    /// trie key, `alpn` sorted + deduped so `add_tcp_front`'s operator
2491    /// order and the route-time rebuild's [`AlpnMatcher::OneOf`] `BTreeSet`
2492    /// order (`alpn_matcher_protocols`) always agree.
2493    fn sni(address: SocketAddr, sni: &str, alpn: &[String]) -> Self {
2494        let mut alpn: Vec<String> = alpn.to_vec();
2495        alpn.sort_unstable();
2496        alpn.dedup();
2497        TcpFrontendTagsKey::Sni {
2498            address,
2499            sni: sni.to_ascii_lowercase(),
2500            alpn,
2501        }
2502    }
2503}
2504
2505impl std::fmt::Display for TcpFrontendTagsKey {
2506    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2507        match self {
2508            TcpFrontendTagsKey::Address(address) => write!(f, "{address}"),
2509            TcpFrontendTagsKey::Sni { address, sni, alpn } => {
2510                write!(f, "{address}|{sni}|")?;
2511                for protocol in alpn {
2512                    // `<len>:<bytes>` per protocol: the length prefix is
2513                    // CHECKED against the following bytes, never scanned
2514                    // for, so a `,`/`:` embedded in one protocol name can
2515                    // never be misread as a boundary between two --
2516                    // unlike the old bare `join(",")`.
2517                    write!(f, "{}:{protocol}", protocol.len())?;
2518                }
2519                Ok(())
2520            }
2521        }
2522    }
2523}
2524
2525/// Canonical access-log tags key for an SNI-scoped TCP frontend: one
2526/// listener can carry many SNI/ALPN fronts, each with its own `tags`, so
2527/// keying tags by the bare listener address (the pre-SNI behavior, kept
2528/// verbatim for no-SNI fronts) would let the LAST added front clobber every
2529/// sibling and let removing ANY front clear tags for the whole address. See
2530/// [`TcpFrontendTagsKey`] for the collision-free composition.
2531fn sni_tags_key(address: &SocketAddr, sni: &str, alpn: &[String]) -> String {
2532    TcpFrontendTagsKey::sni(*address, sni, alpn).to_string()
2533}
2534
2535/// The matched [`AlpnMatcher`]'s protocols as strings, for rebuilding the
2536/// [`sni_tags_key`] a route decision maps to: `Any` is the empty-`alpn`
2537/// catch-all front, `OneOf` yields its (BTreeSet-sorted) protocol list.
2538/// `from_utf8_lossy` is defensive — matcher protocols originate from
2539/// config/IPC `String`s, never raw network bytes.
2540fn alpn_matcher_protocols(matcher: &AlpnMatcher) -> Vec<String> {
2541    match matcher {
2542        AlpnMatcher::Any => Vec::new(),
2543        AlpnMatcher::OneOf(set) => set
2544            .iter()
2545            .map(|protocol| String::from_utf8_lossy(protocol).into_owned())
2546            .collect(),
2547    }
2548}
2549
2550/// Shared `(trie key, ALPN matcher)` construction for
2551/// [`TcpListener::insert_sni_route`] and [`TcpListener::remove_sni_route`]:
2552/// `sni` lowercased into the trie key, `alpn` collapsed to
2553/// [`AlpnMatcher::Any`] when empty (catch-all) or [`AlpnMatcher::OneOf`]
2554/// otherwise. `alpn` is taken by value since neither caller needs it again
2555/// afterward.
2556fn route_key_and_matcher(sni: &str, alpn: Vec<String>) -> (Vec<u8>, AlpnMatcher) {
2557    let key = sni.to_ascii_lowercase().into_bytes();
2558    let matcher = if alpn.is_empty() {
2559        AlpnMatcher::Any
2560    } else {
2561        AlpnMatcher::OneOf(alpn.into_iter().map(String::into_bytes).collect())
2562    };
2563    (key, matcher)
2564}
2565
2566#[derive(Debug)]
2567pub struct ClusterConfiguration {
2568    proxy_protocol: Option<ProxyProtocolConfig>,
2569    // Uncomment this when implementing new load balancing algorithms
2570    // load_balancing: LoadBalancingAlgorithms,
2571    /// Per-cluster override of the global per-(cluster, source-IP)
2572    /// connection limit. `None` inherits the global default,
2573    /// `Some(0)` is explicit "unlimited", `Some(n > 0)` overrides.
2574    /// Resolved against `SessionManager::effective_max_connections_per_ip`
2575    /// at admit time in `connect_to_backend`.
2576    pub max_connections_per_ip: Option<u64>,
2577}
2578
2579pub struct TcpProxy {
2580    fronts: HashMap<String, Token>,
2581    backends: Rc<RefCell<BackendMap>>,
2582    listeners: HashMap<Token, Rc<RefCell<TcpListener>>>,
2583    configs: HashMap<ClusterId, ClusterConfiguration>,
2584    registry: Registry,
2585    sessions: Rc<RefCell<SessionManager>>,
2586    pool: Rc<RefCell<Pool>>,
2587}
2588
2589impl TcpProxy {
2590    pub fn new(
2591        registry: Registry,
2592        sessions: Rc<RefCell<SessionManager>>,
2593        pool: Rc<RefCell<Pool>>,
2594        backends: Rc<RefCell<BackendMap>>,
2595    ) -> TcpProxy {
2596        TcpProxy {
2597            backends,
2598            listeners: HashMap::new(),
2599            configs: HashMap::new(),
2600            fronts: HashMap::new(),
2601            registry,
2602            sessions,
2603            pool,
2604        }
2605    }
2606
2607    pub fn add_listener(
2608        &mut self,
2609        config: TcpListenerConfig,
2610        token: Token,
2611    ) -> Result<Token, ProxyError> {
2612        match self.listeners.entry(token) {
2613            Entry::Vacant(entry) => {
2614                let tcp_listener =
2615                    TcpListener::new(config, token).map_err(ProxyError::AddListener)?;
2616                entry.insert(Rc::new(RefCell::new(tcp_listener)));
2617                Ok(token)
2618            }
2619            _ => Err(ProxyError::ListenerAlreadyPresent),
2620        }
2621    }
2622
2623    pub fn remove_listener(&mut self, address: SocketAddr) -> SessionIsToBeClosed {
2624        let len = self.listeners.len();
2625
2626        self.listeners.retain(|_, l| l.borrow().address != address);
2627        self.listeners.len() < len
2628    }
2629
2630    pub fn activate_listener(
2631        &self,
2632        addr: &SocketAddr,
2633        tcp_listener: Option<MioTcpListener>,
2634    ) -> Result<Token, ProxyError> {
2635        let listener = self
2636            .listeners
2637            .values()
2638            .find(|listener| listener.borrow().address == *addr)
2639            .ok_or(ProxyError::NoListenerFound(*addr))?;
2640
2641        listener.borrow_mut().activate(&self.registry, tcp_listener)
2642    }
2643
2644    pub fn give_back_listeners(&mut self) -> Vec<(SocketAddr, MioTcpListener)> {
2645        self.listeners
2646            .values()
2647            .filter_map(|listener| {
2648                let mut owned = listener.borrow_mut();
2649                if let Some(listener) = owned.listener.take() {
2650                    // Reset `active` so a subsequent `activate()` re-binds
2651                    // instead of short-circuiting on the stale flag.
2652                    owned.active = false;
2653                    return Some((owned.address, listener));
2654                }
2655
2656                None
2657            })
2658            .collect()
2659    }
2660
2661    pub fn give_back_listener(
2662        &mut self,
2663        address: SocketAddr,
2664    ) -> Result<(Token, MioTcpListener), ProxyError> {
2665        let listener = self
2666            .listeners
2667            .values()
2668            .find(|listener| listener.borrow().address == address)
2669            .ok_or(ProxyError::NoListenerFound(address))?;
2670
2671        let mut owned = listener.borrow_mut();
2672
2673        let taken_listener = owned
2674            .listener
2675            .take()
2676            .ok_or(ProxyError::UnactivatedListener)?;
2677
2678        // Reset `active` so a subsequent `activate()` re-binds instead of
2679        // short-circuiting on the stale flag.
2680        owned.active = false;
2681
2682        Ok((owned.token, taken_listener))
2683    }
2684
2685    /// Apply a partial-update patch to the identified TCP listener.
2686    pub fn update_listener(&mut self, patch: UpdateTcpListenerConfig) -> Result<(), ProxyError> {
2687        let address: SocketAddr = patch.address.into();
2688        let listener = self
2689            .listeners
2690            .values()
2691            .find(|l| l.borrow().address == address)
2692            .ok_or(ProxyError::NoListenerFound(address))?;
2693        listener
2694            .borrow_mut()
2695            .update_config(&patch)
2696            .map_err(|listener_error| ProxyError::ListenerActivation {
2697                address,
2698                listener_error,
2699            })
2700    }
2701
2702    pub fn add_tcp_front(&mut self, front: RequestTcpFrontend) -> Result<(), ProxyError> {
2703        let address = front.address.into();
2704
2705        let mut listener = self
2706            .listeners
2707            .values()
2708            .find(|l| l.borrow().address == address)
2709            .ok_or(ProxyError::NoListenerFound(address))?
2710            .borrow_mut();
2711
2712        // Hard-reject a request that would corrupt this listener's SNI/ALPN
2713        // routing invariants BEFORE any mutation below. Config-load
2714        // (`command/src/config.rs`, sozu-proxy/sozu#1279) already rejects
2715        // the same shapes for TOML-sourced requests, but `AddTcpFrontend`
2716        // can also arrive directly over the command socket, or via
2717        // `LoadState` replay of a hand-edited/stale state file, bypassing
2718        // config.rs entirely.
2719        listener.validate_new_tcp_front(&front)?;
2720
2721        self.fronts
2722            .insert(front.cluster_id.to_string(), listener.token);
2723
2724        match front.sni {
2725            Some(sni) => {
2726                // Per-frontend tags key: many SNI/ALPN fronts share one
2727                // listener, so the bare-address key (kept for no-SNI fronts
2728                // below) would clobber siblings — see `sni_tags_key`.
2729                listener.set_tags(sni_tags_key(&address, &sni, &front.alpn), Some(front.tags));
2730                listener.insert_sni_route(sni, front.alpn, front.cluster_id)?;
2731            }
2732            None => {
2733                listener.set_tags(
2734                    TcpFrontendTagsKey::Address(address).to_string(),
2735                    Some(front.tags),
2736                );
2737                listener.cluster_id = Some(front.cluster_id);
2738            }
2739        }
2740
2741        // POST: the mixing invariant must hold after every successful add —
2742        // `validate_new_tcp_front` is the enforcement point above, this is
2743        // the cheap live re-check that it actually held.
2744        debug_assert!(
2745            listener.cluster_id.is_none() || listener.sni_routes.is_empty(),
2746            "a TCP listener must never mix a no-SNI catch-all cluster with SNI-scoped routes"
2747        );
2748
2749        Ok(())
2750    }
2751
2752    pub fn remove_tcp_front(&mut self, front: RequestTcpFrontend) -> Result<(), ProxyError> {
2753        let address = front.address.into();
2754
2755        let mut listener = match self
2756            .listeners
2757            .values()
2758            .find(|l| l.borrow().address == address)
2759        {
2760            Some(l) => l.borrow_mut(),
2761            None => return Err(ProxyError::NoListenerFound(address)),
2762        };
2763
2764        match front.sni {
2765            Some(sni) => {
2766                // Clear ONLY this front's own tags entry (`sni_tags_key`) —
2767                // the pre-SNI bare-address removal here used to strip tags
2768                // for every sibling front on the listener.
2769                listener.set_tags(sni_tags_key(&address, &sni, &front.alpn), None);
2770                listener.remove_sni_route(sni, front.alpn, &front.cluster_id);
2771                self.fronts.remove(&front.cluster_id);
2772            }
2773            None => {
2774                listener.set_tags(TcpFrontendTagsKey::Address(address).to_string(), None);
2775                if let Some(cluster_id) = listener.cluster_id.take() {
2776                    self.fronts.remove(&cluster_id);
2777                }
2778            }
2779        }
2780
2781        Ok(())
2782    }
2783}
2784
2785impl ProxyConfiguration for TcpProxy {
2786    fn notify(&mut self, message: WorkerRequest) -> WorkerResponse {
2787        let request_type = match message.content.request_type {
2788            Some(t) => t,
2789            None => return WorkerResponse::error(message.id, "Empty request"),
2790        };
2791        match request_type {
2792            RequestType::AddTcpFrontend(front) => {
2793                if let Err(err) = self.add_tcp_front(front) {
2794                    return WorkerResponse::error(message.id, err);
2795                }
2796
2797                WorkerResponse::ok(message.id)
2798            }
2799            RequestType::RemoveTcpFrontend(front) => {
2800                if let Err(err) = self.remove_tcp_front(front) {
2801                    return WorkerResponse::error(message.id, err);
2802                }
2803
2804                WorkerResponse::ok(message.id)
2805            }
2806            RequestType::SoftStop(_) => {
2807                info!(
2808                    "{} {} processing soft shutdown",
2809                    log_module_context!(),
2810                    message.id
2811                );
2812                let listeners: HashMap<_, _> = self.listeners.drain().collect();
2813                for (_, l) in listeners.iter() {
2814                    l.borrow_mut()
2815                        .listener
2816                        .take()
2817                        .map(|mut sock| self.registry.deregister(&mut sock));
2818                }
2819                WorkerResponse::processing(message.id)
2820            }
2821            RequestType::HardStop(_) => {
2822                info!("{} {} hard shutdown", log_module_context!(), message.id);
2823                let mut listeners: HashMap<_, _> = self.listeners.drain().collect();
2824                for (_, l) in listeners.drain() {
2825                    l.borrow_mut()
2826                        .listener
2827                        .take()
2828                        .map(|mut sock| self.registry.deregister(&mut sock));
2829                }
2830                WorkerResponse::ok(message.id)
2831            }
2832            RequestType::Status(_) => {
2833                info!("{} {} status", log_module_context!(), message.id);
2834                WorkerResponse::ok(message.id)
2835            }
2836            RequestType::AddCluster(cluster) => {
2837                let config = ClusterConfiguration {
2838                    proxy_protocol: cluster
2839                        .proxy_protocol
2840                        .and_then(|n| ProxyProtocolConfig::try_from(n).ok()),
2841                    //load_balancing: cluster.load_balancing,
2842                    max_connections_per_ip: cluster.max_connections_per_ip,
2843                };
2844                self.configs.insert(cluster.cluster_id, config);
2845                WorkerResponse::ok(message.id)
2846            }
2847            RequestType::RemoveCluster(cluster_id) => {
2848                self.configs.remove(&cluster_id);
2849                WorkerResponse::ok(message.id)
2850            }
2851            RequestType::RemoveListener(remove) => {
2852                if !self.remove_listener(remove.address.into()) {
2853                    WorkerResponse::error(
2854                        message.id,
2855                        format!("no TCP listener to remove at address {:?}", remove.address),
2856                    )
2857                } else {
2858                    WorkerResponse::ok(message.id)
2859                }
2860            }
2861            command => {
2862                debug!(
2863                    "{} {} unsupported message for TCP proxy, ignoring {:?}",
2864                    log_module_context!(),
2865                    message.id,
2866                    command
2867                );
2868                WorkerResponse::error(message.id, "unsupported message")
2869            }
2870        }
2871    }
2872
2873    fn accept(&mut self, token: ListenToken) -> Result<MioTcpStream, AcceptError> {
2874        let internal_token = Token(token.0);
2875        if let Some(listener) = self.listeners.get(&internal_token) {
2876            if let Some(tcp_listener) = &listener.borrow().listener {
2877                tcp_listener
2878                    .accept()
2879                    .map(|(frontend_sock, _)| frontend_sock)
2880                    .map_err(|e| match e.kind() {
2881                        ErrorKind::WouldBlock => AcceptError::WouldBlock,
2882                        _ => {
2883                            error!("{} accept() IO error: {:?}", log_module_context!(), e);
2884                            AcceptError::IoError
2885                        }
2886                    })
2887            } else {
2888                Err(AcceptError::IoError)
2889            }
2890        } else {
2891            Err(AcceptError::IoError)
2892        }
2893    }
2894
2895    fn create_session(
2896        &mut self,
2897        mut frontend_sock: MioTcpStream,
2898        token: ListenToken,
2899        wait_time: Duration,
2900        proxy: Rc<RefCell<Self>>,
2901    ) -> Result<(), AcceptError> {
2902        let listener_token = Token(token.0);
2903
2904        let listener = self
2905            .listeners
2906            .get(&listener_token)
2907            .ok_or(AcceptError::IoError)?;
2908
2909        let owned = listener.borrow();
2910        let mut pool = self.pool.borrow_mut();
2911
2912        let (front_buffer, back_buffer) = match (pool.checkout(), pool.checkout()) {
2913            (Some(fb), Some(bb)) => (fb, bb),
2914            _ => {
2915                error!("{} could not get buffers from pool", log_module_context!());
2916                error!(
2917                    "{} Buffer capacity has been reached, stopping to accept new connections for now",
2918                    log_module_context!()
2919                );
2920                gauge!(names::accept_queue::BACKPRESSURE, 1);
2921                self.sessions.borrow_mut().can_accept = false;
2922
2923                return Err(AcceptError::BufferCapacityReached);
2924            }
2925        };
2926
2927        // A listener may route either by a legacy no-SNI catch-all cluster
2928        // OR by SNI-scoped routes (never both -- enforced at config load,
2929        // sozu-proxy/sozu#1279); reject only when NEITHER is configured.
2930        if owned.cluster_id.is_none() && owned.sni_routes.is_empty() {
2931            error!(
2932                "{} listener at address {:?} has no linked cluster",
2933                log_module_context!(),
2934                owned.address
2935            );
2936            return Err(AcceptError::IoError);
2937        }
2938
2939        if let Err(e) = frontend_sock.set_nodelay(true) {
2940            error!(
2941                "{} error setting nodelay on front socket({:?}): {:?}",
2942                log_module_context!(),
2943                frontend_sock,
2944                e
2945            );
2946        }
2947
2948        let mut session_manager = self.sessions.borrow_mut();
2949        let entry = session_manager.slab.vacant_entry();
2950        let frontend_token = Token(entry.key());
2951
2952        if let Err(register_error) = self.registry.register(
2953            &mut frontend_sock,
2954            frontend_token,
2955            Interest::READABLE | Interest::WRITABLE,
2956        ) {
2957            error!(
2958                "{} error registering front socket({:?}): {:?}",
2959                log_module_context!(),
2960                frontend_sock,
2961                register_error
2962            );
2963            return Err(AcceptError::RegisterError);
2964        }
2965
2966        let session = if !owned.sni_routes.is_empty() {
2967            // Routing decides the cluster post-accept; the effective
2968            // preread cap can never exceed what the checked-out buffer can
2969            // actually hold, regardless of the configured knob.
2970            let effective_max_bytes = effective_sni_preread_max_bytes(
2971                owned.config.sni_preread_max_bytes,
2972                front_buffer.capacity(),
2973            );
2974            let preread_timeout = Duration::from_secs(u64::from(
2975                owned
2976                    .config
2977                    .sni_preread_timeout
2978                    .unwrap_or(DEFAULT_SNI_PREREAD_TIMEOUT),
2979            ));
2980            TcpSession::new_sni_preread(
2981                back_buffer,
2982                Duration::from_secs(owned.config.back_timeout as u64),
2983                Duration::from_secs(owned.config.connect_timeout as u64),
2984                front_buffer,
2985                frontend_token,
2986                listener.clone(),
2987                proxy,
2988                frontend_sock,
2989                wait_time,
2990                preread_timeout,
2991                effective_max_bytes,
2992            )
2993        } else {
2994            let proxy_protocol = self
2995                .configs
2996                .get(owned.cluster_id.as_ref().unwrap())
2997                .and_then(|c| c.proxy_protocol);
2998            TcpSession::new(
2999                back_buffer,
3000                None,
3001                owned.cluster_id.clone(),
3002                Duration::from_secs(owned.config.back_timeout as u64),
3003                Duration::from_secs(owned.config.connect_timeout as u64),
3004                Duration::from_secs(owned.config.front_timeout as u64),
3005                front_buffer,
3006                frontend_token,
3007                listener.clone(),
3008                proxy_protocol,
3009                proxy,
3010                frontend_sock,
3011                wait_time,
3012            )
3013        };
3014        incr!(names::tcp::REQUESTS);
3015
3016        let session = Rc::new(RefCell::new(session));
3017        entry.insert(session);
3018
3019        Ok(())
3020    }
3021}
3022
3023pub mod testing {
3024    use crate::testing::*;
3025
3026    /// This is not directly used by Sōzu but is available for example and testing purposes
3027    pub fn start_tcp_worker(
3028        config: TcpListenerConfig,
3029        max_buffers: usize,
3030        buffer_size: usize,
3031        channel: ProxyChannel,
3032    ) -> anyhow::Result<()> {
3033        let address = config.address.into();
3034
3035        let ServerParts {
3036            event_loop,
3037            registry,
3038            sessions,
3039            pool,
3040            backends,
3041            client_scm_socket: _,
3042            server_scm_socket,
3043            server_config,
3044        } = prebuild_server(max_buffers, buffer_size, true)?;
3045
3046        let token = {
3047            let mut sessions = sessions.borrow_mut();
3048            let entry = sessions.slab.vacant_entry();
3049            let key = entry.key();
3050            let _ = entry.insert(Rc::new(RefCell::new(ListenSession {
3051                protocol: Protocol::TCPListen,
3052            })));
3053            Token(key)
3054        };
3055
3056        let mut proxy = TcpProxy::new(registry, sessions.clone(), pool.clone(), backends.clone());
3057        proxy
3058            .add_listener(config, token)
3059            .with_context(|| "Failed at creating adding the listener")?;
3060        proxy
3061            .activate_listener(&address, None)
3062            .with_context(|| "Failed at creating activating the listener")?;
3063
3064        let mut server = Server::new(
3065            event_loop,
3066            channel,
3067            server_scm_socket,
3068            sessions,
3069            pool,
3070            backends,
3071            None,
3072            None,
3073            Some(proxy),
3074            server_config,
3075            None,
3076            false,
3077        )
3078        .with_context(|| "Failed at creating server")?;
3079
3080        debug!("{} starting event loop", log_module_context!());
3081        server.run();
3082        debug!("{} ending event loop", log_module_context!());
3083        Ok(())
3084    }
3085}
3086
3087#[cfg(test)]
3088mod tests {
3089    use std::{
3090        io::{Read, Write},
3091        net::{Shutdown, TcpListener, TcpStream},
3092        str,
3093        sync::{
3094            Arc, Barrier,
3095            atomic::{AtomicBool, Ordering},
3096        },
3097        thread,
3098        time::Duration,
3099    };
3100
3101    use sozu_command::{
3102        channel::Channel,
3103        config::ListenerBuilder,
3104        proto::command::{
3105            LoadBalancingParams, RequestTcpFrontend, SocketAddress, SoftStop, WorkerRequest,
3106            WorkerResponse, request::RequestType,
3107        },
3108    };
3109
3110    use super::testing::start_tcp_worker;
3111    use crate::testing::*;
3112
3113    /*
3114    #[test]
3115    #[cfg(target_pointer_width = "64")]
3116    fn size_test() {
3117      assert_size!(Pipe<mio::net::TcpStream>, 224);
3118      assert_size!(SendProxyProtocol<mio::net::TcpStream>, 144);
3119      assert_size!(RelayProxyProtocol<mio::net::TcpStream>, 152);
3120      assert_size!(ExpectProxyProtocol<mio::net::TcpStream>, 520);
3121      assert_size!(State, 528);
3122      // fails depending on the platform?
3123      //assert_size!(Session, 808);
3124    }*/
3125
3126    #[test]
3127    fn round_trip() {
3128        setup_test_logger!();
3129        let barrier = Arc::new(Barrier::new(2));
3130        let test_finished = Arc::new(AtomicBool::new(false));
3131
3132        let front_port1 = provide_port();
3133        let front_port2 = provide_port();
3134
3135        let backend_port = start_server(barrier.clone(), test_finished.clone());
3136        let mut command =
3137            start_proxy(backend_port, front_port1, front_port2).expect("Could not start proxy");
3138        barrier.wait();
3139
3140        thread::scope(|_s| {
3141            let front_addr = format!("127.0.0.1:{front_port1}");
3142
3143            let mut s1 = TcpStream::connect(&front_addr).expect("could not connect");
3144            s1.set_read_timeout(Some(Duration::from_secs(5)))
3145                .expect("could not set read timeout on s1");
3146
3147            let s3 = TcpStream::connect(&front_addr).expect("could not connect");
3148
3149            let mut s2 = TcpStream::connect(&front_addr).expect("could not connect");
3150            s2.set_read_timeout(Some(Duration::from_secs(5)))
3151                .expect("could not set read timeout on s2");
3152
3153            s1.write_all(b"hello ").expect("could not write to s1");
3154            println!("s1 sent");
3155
3156            s2.write_all(b"pouet pouet").expect("could not write to s2");
3157            println!("s2 sent");
3158
3159            let mut res = [0; 128];
3160            s1.write_all(b"coucou").expect("could not write to s1");
3161
3162            s3.shutdown(Shutdown::Both).expect("could not shutdown s3");
3163
3164            let sz2 = s2
3165                .read(&mut res[..])
3166                .expect("could not read from socket s2");
3167            println!("s2 received {:?}", str::from_utf8(&res[..sz2]));
3168            assert_eq!(&res[..sz2], &b"pouet pouet"[..]);
3169
3170            // Read in a loop: a single read() on a TCP stream is not
3171            // guaranteed to return all echoed data if the second write's
3172            // round trip (client → proxy → backend → proxy → client) is
3173            // still in flight when we poll.
3174            let expected = b"hello coucou";
3175            let mut total = 0;
3176            while total < expected.len() {
3177                let sz = s1
3178                    .read(&mut res[total..])
3179                    .expect("could not read from socket s1");
3180                assert!(sz > 0, "connection closed before receiving all data");
3181                total += sz;
3182            }
3183            println!(
3184                "s1 received again({}): {:?}",
3185                total,
3186                str::from_utf8(&res[..total])
3187            );
3188            assert_eq!(&res[..total], &expected[..]);
3189
3190            // Signal the echo server to stop
3191            test_finished.store(true, Ordering::Relaxed);
3192
3193            // Send SoftStop to the sozu worker so server.run() exits cleanly
3194            command
3195                .write_message(&WorkerRequest {
3196                    id: "ID_SOFTSTOP".to_owned(),
3197                    content: RequestType::SoftStop(SoftStop {}).into(),
3198                })
3199                .expect("could not send SoftStop to sozu worker");
3200        });
3201    }
3202
3203    /// Start an echo server on an ephemeral port.
3204    /// Returns the port the server is listening on.
3205    fn start_server(barrier: Arc<Barrier>, test_finished: Arc<AtomicBool>) -> u16 {
3206        let listener =
3207            TcpListener::bind("127.0.0.1:0").expect("could not bind echo server listener");
3208        let port = listener
3209            .local_addr()
3210            .expect("could not get echo server local address")
3211            .port();
3212
3213        listener
3214            .set_nonblocking(true)
3215            .expect("could not set echo server listener to non-blocking");
3216
3217        thread::spawn(move || {
3218            barrier.wait();
3219            let mut count: u8 = 0;
3220            loop {
3221                match listener.accept() {
3222                    Ok((mut stream, _)) => {
3223                        let finished = test_finished.clone();
3224                        thread::spawn(move || {
3225                            println!("got a new client: {count}");
3226                            stream
3227                                .set_read_timeout(Some(Duration::from_secs(2)))
3228                                .expect("could not set read timeout on echo client");
3229                            let mut buf = [0; 128];
3230                            loop {
3231                                match stream.read(&mut buf[..]) {
3232                                    Ok(0) => break,
3233                                    Ok(sz) => {
3234                                        println!(
3235                                            "ECHO[{count}] got \"{:?}\"",
3236                                            str::from_utf8(&buf[..sz])
3237                                        );
3238                                        stream
3239                                            .write_all(&buf[..sz])
3240                                            .expect("could not echo data back");
3241                                    }
3242                                    Err(ref e)
3243                                        if e.kind() == std::io::ErrorKind::WouldBlock
3244                                            || e.kind() == std::io::ErrorKind::TimedOut =>
3245                                    {
3246                                        if finished.load(Ordering::Relaxed) {
3247                                            println!("backend server stopping (client handler)");
3248                                            break;
3249                                        }
3250                                    }
3251                                    Err(_) => break,
3252                                }
3253                            }
3254                        });
3255                        count = count.wrapping_add(1);
3256                    }
3257                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
3258                        if test_finished.load(Ordering::Relaxed) {
3259                            println!("backend server stopping (accept loop)");
3260                            break;
3261                        }
3262                        thread::sleep(Duration::from_millis(50));
3263                    }
3264                    Err(e) => {
3265                        println!("connection failed: {e:?}");
3266                    }
3267                }
3268            }
3269        });
3270
3271        port
3272    }
3273
3274    /// Start a sozu TCP proxy worker with the given backend and frontend ports.
3275    fn start_proxy(
3276        backend_port: u16,
3277        front_port1: u16,
3278        front_port2: u16,
3279    ) -> anyhow::Result<Channel<WorkerRequest, WorkerResponse>> {
3280        let config = ListenerBuilder::new_tcp(SocketAddress::new_v4(127, 0, 0, 1, front_port1))
3281            .to_tcp(None)
3282            .expect("could not create listener config");
3283
3284        let (mut command, channel) =
3285            Channel::generate(1000, 10000).with_context(|| "should create a channel")?;
3286        let _jg = thread::spawn(move || {
3287            setup_test_logger!();
3288            start_tcp_worker(config, 100, 16384, channel).expect("could not start the tcp server");
3289        });
3290
3291        command
3292            .blocking()
3293            .expect("could not set command channel to blocking");
3294        {
3295            let front = RequestTcpFrontend {
3296                cluster_id: "yolo".to_owned(),
3297                address: SocketAddress::new_v4(127, 0, 0, 1, front_port1),
3298                ..Default::default()
3299            };
3300            let backend = sozu_command_lib::response::Backend {
3301                cluster_id: "yolo".to_owned(),
3302                backend_id: "yolo-0".to_owned(),
3303                address: SocketAddress::new_v4(127, 0, 0, 1, backend_port).into(),
3304                load_balancing_parameters: Some(LoadBalancingParams::default()),
3305                sticky_id: None,
3306                backup: None,
3307            };
3308
3309            command
3310                .write_message(&WorkerRequest {
3311                    id: "ID_YOLO1".to_owned(),
3312                    content: RequestType::AddTcpFrontend(front).into(),
3313                })
3314                .expect("could not send AddTcpFrontend for front1");
3315            command
3316                .write_message(&WorkerRequest {
3317                    id: "ID_YOLO2".to_owned(),
3318                    content: RequestType::AddBackend(backend.to_add_backend()).into(),
3319                })
3320                .expect("could not send AddBackend for front1");
3321        }
3322        {
3323            let front = RequestTcpFrontend {
3324                cluster_id: "yolo".to_owned(),
3325                address: SocketAddress::new_v4(127, 0, 0, 1, front_port2),
3326                ..Default::default()
3327            };
3328            let backend = sozu_command::response::Backend {
3329                cluster_id: "yolo".to_owned(),
3330                backend_id: "yolo-0".to_owned(),
3331                address: SocketAddress::new_v4(127, 0, 0, 1, backend_port).into(),
3332                load_balancing_parameters: Some(LoadBalancingParams::default()),
3333                sticky_id: None,
3334                backup: None,
3335            };
3336            command
3337                .write_message(&WorkerRequest {
3338                    id: "ID_YOLO3".to_owned(),
3339                    content: RequestType::AddTcpFrontend(front).into(),
3340                })
3341                .expect("could not send AddTcpFrontend for front2");
3342            command
3343                .write_message(&WorkerRequest {
3344                    id: "ID_YOLO4".to_owned(),
3345                    content: RequestType::AddBackend(backend.to_add_backend()).into(),
3346                })
3347                .expect("could not send AddBackend for front2");
3348        }
3349
3350        for _ in 0..4 {
3351            println!(
3352                "read_message: {:?}",
3353                command
3354                    .read_message()
3355                    .with_context(|| "could not read message")?
3356            );
3357        }
3358
3359        Ok(command)
3360    }
3361}
3362
3363/// Unit coverage for the SNI-preread routing shell added for
3364/// sozu-proxy/sozu#1279: the route-table mutations (`add_tcp_front` /
3365/// `remove_tcp_front`), the `AlpnMatcher` mapping, the effective preread
3366/// cap, and the routing gate's data invariants. None of this needs a live
3367/// socket or event loop -- it is a separate module (rather than nested in
3368/// the `tests` module above) purely to avoid that module's `use
3369/// std::net::TcpListener` import shadowing `super::TcpListener` (this
3370/// crate's listener struct).
3371#[cfg(test)]
3372mod sni_routing_tests {
3373    use sozu_command::{config::ListenerBuilder, proto::command::SocketAddress};
3374
3375    use super::*;
3376    use crate::testing::{ServerParts, prebuild_server, provide_port};
3377
3378    fn test_listener() -> TcpListener {
3379        let config = ListenerBuilder::new_tcp(SocketAddress::new_v4(127, 0, 0, 1, provide_port()))
3380            .to_tcp(None)
3381            .expect("could not build a TcpListenerConfig for the test");
3382        TcpListener::new(config, Token(0)).expect("could not build a bare TcpListener for the test")
3383    }
3384
3385    fn frontend(cluster_id: &str, sni: Option<&str>, alpn: &[&str]) -> RequestTcpFrontend {
3386        RequestTcpFrontend {
3387            cluster_id: cluster_id.to_owned(),
3388            address: SocketAddress::new_v4(127, 0, 0, 1, provide_port()),
3389            sni: sni.map(str::to_owned),
3390            alpn: alpn.iter().map(|p| p.to_string()).collect(),
3391            ..Default::default()
3392        }
3393    }
3394
3395    // ---- effective_sni_preread_max_bytes ------------------------------
3396
3397    #[test]
3398    fn effective_max_bytes_falls_back_to_default_when_unconfigured() {
3399        assert_eq!(
3400            effective_sni_preread_max_bytes(None, 65536),
3401            DEFAULT_SNI_PREREAD_MAX_BYTES as usize
3402        );
3403    }
3404
3405    #[test]
3406    fn effective_max_bytes_is_the_min_of_knob_and_capacity() {
3407        assert_eq!(effective_sni_preread_max_bytes(Some(8192), 16384), 8192);
3408        assert_eq!(effective_sni_preread_max_bytes(Some(32768), 16384), 16384);
3409        assert_eq!(effective_sni_preread_max_bytes(Some(16384), 16384), 16384);
3410    }
3411
3412    #[test]
3413    fn effective_max_bytes_never_below_the_floor() {
3414        // A `sni_preread_max_bytes = 0` knob reaching the worker from a
3415        // direct `sozu listener tcp add`/`update` CLI/IPC request (or a stale
3416        // LoadState replay) bypasses config.rs's loud MIN_SNI_PREREAD_MAX_BYTES
3417        // load-time reject. Without the floor the shell would issue
3418        // zero-length preread reads and spin until the loop guard closes each
3419        // session; the floor degrades a sub-minimum knob to the 5-byte
3420        // TLS-record-header minimum instead.
3421        assert_eq!(
3422            effective_sni_preread_max_bytes(Some(0), 16384),
3423            MIN_SNI_PREREAD_MAX_BYTES as usize,
3424            "a 0 knob must degrade to the floor, never 0 (would spin the preread)"
3425        );
3426        assert_eq!(
3427            effective_sni_preread_max_bytes(Some(3), 16384),
3428            MIN_SNI_PREREAD_MAX_BYTES as usize,
3429            "any sub-floor knob must be raised to the floor"
3430        );
3431        // The floor itself, and anything above it, are respected unchanged.
3432        assert_eq!(
3433            effective_sni_preread_max_bytes(Some(MIN_SNI_PREREAD_MAX_BYTES), 16384),
3434            MIN_SNI_PREREAD_MAX_BYTES as usize
3435        );
3436        assert_eq!(
3437            effective_sni_preread_max_bytes(Some(MIN_SNI_PREREAD_MAX_BYTES + 1), 16384),
3438            (MIN_SNI_PREREAD_MAX_BYTES + 1) as usize
3439        );
3440    }
3441
3442    // ---- known_alpn_label (access-log tagging) -------------------------
3443
3444    #[test]
3445    fn known_alpn_label_picks_the_clients_first_offer() {
3446        assert_eq!(
3447            known_alpn_label(&[b"h2".to_vec(), b"http/1.1".to_vec()]),
3448            Some("h2")
3449        );
3450        assert_eq!(
3451            known_alpn_label(&[b"http/1.1".to_vec(), b"h2".to_vec()]),
3452            Some("http/1.1"),
3453            "client preference order must win, not a fixed h2-first priority"
3454        );
3455    }
3456
3457    #[test]
3458    fn known_alpn_label_is_none_for_empty_or_unrecognized_offers() {
3459        assert_eq!(known_alpn_label(&[]), None);
3460        assert_eq!(known_alpn_label(&[b"spdy/1".to_vec()]), None);
3461    }
3462
3463    // ---- AlpnMatcher mapping + route-table add/remove symmetry --------
3464
3465    #[test]
3466    fn empty_alpn_maps_to_any_non_empty_maps_to_one_of() {
3467        let mut listener = test_listener();
3468        listener
3469            .insert_sni_route("example.com".to_owned(), vec![], "cluster-any".to_owned())
3470            .expect("insert_sni_route must succeed for a valid test SNI");
3471        listener
3472            .insert_sni_route(
3473                "h2.example.com".to_owned(),
3474                vec!["h2".to_owned(), "http/1.1".to_owned()],
3475                "cluster-h2".to_owned(),
3476            )
3477            .expect("insert_sni_route must succeed for a valid test SNI");
3478
3479        let (_, any_entries) = listener
3480            .sni_routes
3481            .domain_lookup(b"example.com", true)
3482            .expect("example.com must be routable");
3483        assert_eq!(
3484            any_entries,
3485            &vec![(AlpnMatcher::Any, "cluster-any".to_owned())]
3486        );
3487
3488        let (_, h2_entries) = listener
3489            .sni_routes
3490            .domain_lookup(b"h2.example.com", true)
3491            .expect("h2.example.com must be routable");
3492        assert_eq!(
3493            h2_entries,
3494            &vec![(
3495                AlpnMatcher::OneOf([b"h2".to_vec(), b"http/1.1".to_vec()].into_iter().collect()),
3496                "cluster-h2".to_owned()
3497            )]
3498        );
3499    }
3500
3501    #[test]
3502    fn insert_sni_route_appends_under_the_same_sni() {
3503        let mut listener = test_listener();
3504        listener
3505            .insert_sni_route(
3506                "example.com".to_owned(),
3507                vec!["h2".to_owned()],
3508                "cluster-h2".to_owned(),
3509            )
3510            .expect("insert_sni_route must succeed for a valid test SNI");
3511        listener
3512            .insert_sni_route(
3513                "example.com".to_owned(),
3514                vec![],
3515                "cluster-default".to_owned(),
3516            )
3517            .expect("insert_sni_route must succeed for a valid test SNI");
3518
3519        let (_, entries) = listener
3520            .sni_routes
3521            .domain_lookup(b"example.com", true)
3522            .expect("example.com must be routable");
3523        assert_eq!(entries.len(), 2, "both fronts must share the SNI's Vec");
3524    }
3525
3526    #[test]
3527    fn remove_sni_route_drops_only_the_matching_entry() {
3528        let mut listener = test_listener();
3529        listener
3530            .insert_sni_route(
3531                "example.com".to_owned(),
3532                vec!["h2".to_owned()],
3533                "cluster-h2".to_owned(),
3534            )
3535            .expect("insert_sni_route must succeed for a valid test SNI");
3536        listener
3537            .insert_sni_route(
3538                "example.com".to_owned(),
3539                vec![],
3540                "cluster-default".to_owned(),
3541            )
3542            .expect("insert_sni_route must succeed for a valid test SNI");
3543
3544        listener.remove_sni_route(
3545            "example.com".to_owned(),
3546            vec!["h2".to_owned()],
3547            &"cluster-h2".to_owned(),
3548        );
3549
3550        let (_, entries) = listener
3551            .sni_routes
3552            .domain_lookup(b"example.com", true)
3553            .expect("example.com must still be routable via the remaining entry");
3554        assert_eq!(
3555            entries,
3556            &vec![(AlpnMatcher::Any, "cluster-default".to_owned())],
3557            "removing one entry must not disturb the other"
3558        );
3559    }
3560
3561    #[test]
3562    fn remove_sni_route_empties_the_trie_key_when_the_last_entry_goes() {
3563        let mut listener = test_listener();
3564        listener
3565            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".to_owned())
3566            .expect("insert_sni_route must succeed for a valid test SNI");
3567        assert!(!listener.sni_routes.is_empty());
3568
3569        listener.remove_sni_route("example.com".to_owned(), vec![], &"cluster-a".to_owned());
3570
3571        assert!(
3572            listener.sni_routes.is_empty(),
3573            "domain_remove must run once the SNI's Vec empties, leaving no stranded key"
3574        );
3575        assert!(
3576            listener
3577                .sni_routes
3578                .domain_lookup(b"example.com", true)
3579                .is_none()
3580        );
3581    }
3582
3583    #[test]
3584    fn remove_sni_route_on_an_absent_sni_is_a_harmless_no_op() {
3585        let mut listener = test_listener();
3586        listener
3587            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".to_owned())
3588            .expect("insert_sni_route must succeed for a valid test SNI");
3589
3590        // Removing a route for a SNI that was never inserted must not panic
3591        // and must not disturb the existing route.
3592        listener.remove_sni_route(
3593            "other.example.net".to_owned(),
3594            vec![],
3595            &"cluster-a".to_owned(),
3596        );
3597
3598        assert!(
3599            listener
3600                .sni_routes
3601                .domain_lookup(b"example.com", true)
3602                .is_some()
3603        );
3604    }
3605
3606    // ---- exact-key bookkeeping must not fall back to a sibling wildcard
3607    // (route-table corruption caught in sozu-proxy/sozu#1290 review) ----
3608
3609    #[test]
3610    fn insert_sni_route_creates_a_distinct_node_for_an_exact_key_over_a_sibling_wildcard() {
3611        let mut listener = test_listener();
3612        // Wildcard catch-all first.
3613        listener
3614            .insert_sni_route(
3615                "*.example.com".to_owned(),
3616                vec![],
3617                "cluster-wildcard".to_owned(),
3618            )
3619            .expect("insert_sni_route must succeed for a valid test SNI");
3620        // Exact ALPN-scoped route for one specific subdomain.
3621        listener
3622            .insert_sni_route(
3623                "a.example.com".to_owned(),
3624                vec!["h2".to_owned()],
3625                "cluster-a-h2".to_owned(),
3626            )
3627            .expect("insert_sni_route must succeed for a valid test SNI");
3628
3629        // The exact key must have gotten its OWN trie node -- with
3630        // `accept_wildcard: true` this lookup would instead fall back to
3631        // (and the insert above would have corrupted) the wildcard's node,
3632        // since no literal `a` child existed yet at insert time.
3633        let (_, a_entries) = listener
3634            .sni_routes
3635            .domain_lookup(b"a.example.com", false)
3636            .expect("a.example.com must have a distinct exact-key node");
3637        assert_eq!(
3638            a_entries,
3639            &vec![(
3640                AlpnMatcher::OneOf([b"h2".to_vec()].into_iter().collect()),
3641                "cluster-a-h2".to_owned()
3642            )],
3643            "the exact key's own Vec must hold only its own entry, not the wildcard's"
3644        );
3645
3646        // Any OTHER subdomain must still resolve to ONLY the wildcard, via
3647        // the same `accept_wildcard: true` lookup the routing path uses
3648        // (`preread_config`) -- it must never see `a.example.com`'s h2 route.
3649        let (_, b_entries) = listener
3650            .sni_routes
3651            .domain_lookup(b"b.example.com", true)
3652            .expect("b.example.com must fall back to the wildcard catch-all");
3653        assert_eq!(
3654            b_entries,
3655            &vec![(AlpnMatcher::Any, "cluster-wildcard".to_owned())],
3656            "an unrelated subdomain must see ONLY the wildcard's entry"
3657        );
3658    }
3659
3660    #[test]
3661    fn validate_new_tcp_front_accepts_an_exact_catch_all_sibling_of_a_wildcard_catch_all() {
3662        let mut listener = test_listener();
3663        listener
3664            .insert_sni_route(
3665                "*.example.com".to_owned(),
3666                vec![],
3667                "cluster-wildcard".to_owned(),
3668            )
3669            .expect("insert_sni_route must succeed for a valid test SNI");
3670
3671        // An exact catch-all for one subdomain must be accepted: it is a
3672        // SIBLING of the wildcard's catch-all, not a duplicate of it. With
3673        // `accept_wildcard: true` this lookup would wrongly find the
3674        // wildcard's own `AlpnMatcher::Any` entry and reject it as "already
3675        // has a catch-all".
3676        let front = frontend("cluster-a", Some("a.example.com"), &[]);
3677        assert!(
3678            listener.validate_new_tcp_front(&front).is_ok(),
3679            "an exact catch-all must be accepted when only a SIBLING wildcard has a catch-all"
3680        );
3681    }
3682
3683    #[test]
3684    fn validate_new_tcp_front_rejects_a_duplicate_wildcard_catch_all() {
3685        let mut listener = test_listener();
3686        listener
3687            .insert_sni_route(
3688                "*.example.com".to_owned(),
3689                vec![],
3690                "cluster-wildcard".to_owned(),
3691            )
3692            .expect("insert_sni_route must succeed for a valid test SNI");
3693
3694        // A SECOND catch-all for the SAME wildcard key is an ambiguous
3695        // duplicate and must be rejected. The immutable trie `lookup` has no
3696        // literal-`*` short-circuit (only `lookup_mut` does), so a plain
3697        // `accept_wildcard: false` self-lookup never sees the existing
3698        // wildcard entry and waves the duplicate through -- which
3699        // `insert_sni_route` (lookup_mut, short-circuit present) would then
3700        // happily append.
3701        let front = frontend("cluster-dup", Some("*.example.com"), &[]);
3702        assert!(
3703            listener.validate_new_tcp_front(&front).is_err(),
3704            "a second catch-all on the same wildcard SNI must be rejected as a duplicate"
3705        );
3706    }
3707
3708    #[test]
3709    fn validate_new_tcp_front_rejects_overlapping_alpn_on_the_same_wildcard() {
3710        let mut listener = test_listener();
3711        listener
3712            .insert_sni_route(
3713                "*.example.com".to_owned(),
3714                vec!["h2".to_owned()],
3715                "cluster-wildcard-h2".to_owned(),
3716            )
3717            .expect("insert_sni_route must succeed for a valid test SNI");
3718
3719        // Same wildcard key, overlapping ALPN protocol: ambiguous, must be
3720        // rejected (same bypass as the duplicate catch-all above).
3721        let front = frontend("cluster-dup", Some("*.example.com"), &["h2"]);
3722        assert!(
3723            listener.validate_new_tcp_front(&front).is_err(),
3724            "an overlapping ALPN matcher on the same wildcard SNI must be rejected"
3725        );
3726    }
3727
3728    #[test]
3729    fn validate_new_tcp_front_accepts_a_disjoint_alpn_addition_on_the_same_wildcard() {
3730        let mut listener = test_listener();
3731        listener
3732            .insert_sni_route(
3733                "*.example.com".to_owned(),
3734                vec![],
3735                "cluster-wildcard".to_owned(),
3736            )
3737            .expect("insert_sni_route must succeed for a valid test SNI");
3738
3739        // A non-overlapping ALPN-scoped addition alongside the wildcard's
3740        // catch-all stays legal -- the wildcard-aware self-lookup must not
3741        // over-reject.
3742        let front = frontend("cluster-h2", Some("*.example.com"), &["h2"]);
3743        assert!(
3744            listener.validate_new_tcp_front(&front).is_ok(),
3745            "a disjoint ALPN addition on the same wildcard SNI must be accepted"
3746        );
3747    }
3748
3749    /// The worker boundary must reject malformed SNI SHAPES that a direct
3750    /// `AddTcpFrontend` (command socket) or `LoadState` replay could carry
3751    /// past config.rs's `validate_sni_pattern`: a `/.../` label would be
3752    /// inserted into the `pattern_trie` as a REGEX route and a misplaced
3753    /// `*` as an unintended wildcard — silently widening routing.
3754    #[test]
3755    fn validate_new_tcp_front_rejects_malformed_sni_shapes() {
3756        let listener = test_listener();
3757
3758        for bad in [
3759            "/[a-z]+/.example.com",
3760            "foo/bar.example.com",
3761            "a.*.example.com",
3762            "*.*.example.com",
3763            "*",
3764        ] {
3765            let front = frontend("cluster-a", Some(bad), &[]);
3766            assert!(
3767                listener.validate_new_tcp_front(&front).is_err(),
3768                "malformed SNI shape {bad:?} must be rejected at the worker boundary"
3769            );
3770        }
3771
3772        // The two documented-legal shapes must still pass.
3773        for good in ["a.example.com", "*.example.com"] {
3774            let front = frontend("cluster-a", Some(good), &[]);
3775            assert!(
3776                listener.validate_new_tcp_front(&front).is_ok(),
3777                "legal SNI shape {good:?} must still be accepted"
3778            );
3779        }
3780    }
3781
3782    /// `validate_new_tcp_front`'s OLD hand-rolled shape check only rejected
3783    /// `/` and a misplaced `*`, letting an empty label (leading dot,
3784    /// trailing dot, consecutive dots), an empty string, or a non-ASCII
3785    /// pattern reach `insert_sni_route` -> `pattern_trie::insert_recursive`,
3786    /// whose RELEASE-mode `assert_ne!(partial_key, &b""[..])` panics the
3787    /// worker on a leading-empty-label key like `.example.com`. This is the
3788    /// worker-boundary counterpart to `config.rs`'s `validate_sni_pattern`
3789    /// tests: every shape the shared validator rejects at config-load must
3790    /// also be rejected here, for `AddTcpFrontend`/`LoadState` requests that
3791    /// bypass config.rs entirely.
3792    #[test]
3793    fn add_tcp_front_enforces_the_shared_sni_validator_at_the_worker_boundary() {
3794        let mut proxy = test_proxy();
3795        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
3796        let config = ListenerBuilder::new_tcp(address)
3797            .to_tcp(None)
3798            .expect("could not build listener config");
3799        let token = Token(0);
3800        proxy
3801            .add_listener(config, token)
3802            .expect("could not add listener");
3803
3804        for bad in [
3805            ".example.com",     // leading empty label -- the pattern_trie-crashing shape
3806            "example.com.",     // trailing empty label
3807            "a..b.example.com", // consecutive dots -- empty middle label
3808            "",                 // empty pattern
3809            "exämple.com",      // non-ASCII
3810        ] {
3811            let front = frontend("cluster-a", Some(bad), &[]);
3812            let front = RequestTcpFrontend { address, ..front };
3813            match proxy.add_tcp_front(front) {
3814                Err(ProxyError::InvalidTcpFrontend { .. }) => {}
3815                other => panic!("malformed sni {bad:?} must be rejected, got {other:?}"),
3816            }
3817
3818            let listener = proxy
3819                .listeners
3820                .get(&token)
3821                .expect("listener must be present")
3822                .borrow();
3823            assert!(
3824                listener.sni_routes.is_empty(),
3825                "a rejected sni {bad:?} must leave sni_routes empty"
3826            );
3827            assert!(
3828                listener.cluster_id.is_none(),
3829                "a rejected sni {bad:?} must leave cluster_id unset"
3830            );
3831        }
3832
3833        // A space IS valid ASCII, is not '*'/'/', and splits into non-empty
3834        // labels ("exa", "mple.com") -- none of the shared validator's rules
3835        // (empty pattern, non-ASCII, misplaced '*', empty label) cover "not a
3836        // valid hostname character" in general, so this shape is ACCEPTED,
3837        // not rejected. Documenting the validator's actual behavior rather
3838        // than assuming a hostname-shaped string with a space would be
3839        // caught too.
3840        let space_front = frontend("cluster-space", Some("exa mple.com"), &[]);
3841        let space_front = RequestTcpFrontend {
3842            address,
3843            ..space_front
3844        };
3845        assert!(
3846            proxy.add_tcp_front(space_front).is_ok(),
3847            "the shared SNI validator does not reject an embedded space -- \
3848             it is valid ASCII with no empty label"
3849        );
3850    }
3851
3852    // ---- per-frontend access-log tags keying (sozu-proxy/sozu#1290) ----
3853
3854    #[test]
3855    fn sni_fronts_keep_distinct_tags_under_distinct_keys() {
3856        let mut proxy = test_proxy();
3857        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
3858        let config = ListenerBuilder::new_tcp(address)
3859            .to_tcp(None)
3860            .expect("could not build listener config");
3861        let token = Token(0);
3862        proxy
3863            .add_listener(config, token)
3864            .expect("could not add listener");
3865
3866        let front_a = RequestTcpFrontend {
3867            cluster_id: "cluster-a".to_owned(),
3868            address,
3869            sni: Some("a.example.com".to_owned()),
3870            alpn: vec![],
3871            tags: std::collections::BTreeMap::from([("team".to_owned(), "alpha".to_owned())]),
3872        };
3873        let front_b = RequestTcpFrontend {
3874            cluster_id: "cluster-b".to_owned(),
3875            address,
3876            sni: Some("b.example.com".to_owned()),
3877            alpn: vec!["h2".to_owned()],
3878            tags: std::collections::BTreeMap::from([("team".to_owned(), "beta".to_owned())]),
3879        };
3880        proxy
3881            .add_tcp_front(front_a)
3882            .expect("add_tcp_front A must succeed");
3883        proxy
3884            .add_tcp_front(front_b)
3885            .expect("add_tcp_front B must succeed");
3886
3887        let std_address: SocketAddr = address.into();
3888        let key_a = sni_tags_key(&std_address, "a.example.com", &[]);
3889        let key_b = sni_tags_key(&std_address, "b.example.com", &["h2".to_owned()]);
3890
3891        let listener = proxy
3892            .listeners
3893            .get(&token)
3894            .expect("listener must be present")
3895            .borrow();
3896        let tags_a = listener
3897            .get_tags(&key_a)
3898            .expect("front A's tags must live under its own composed key");
3899        assert_eq!(
3900            tags_a.tags.get("team").map(String::as_str),
3901            Some("alpha"),
3902            "front A's tags must survive front B's add, not be clobbered by it"
3903        );
3904        let tags_b = listener
3905            .get_tags(&key_b)
3906            .expect("front B's tags must live under its own composed key");
3907        assert_eq!(tags_b.tags.get("team").map(String::as_str), Some("beta"));
3908    }
3909
3910    #[test]
3911    fn removing_one_sni_front_clears_only_its_own_tags() {
3912        let mut proxy = test_proxy();
3913        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
3914        let config = ListenerBuilder::new_tcp(address)
3915            .to_tcp(None)
3916            .expect("could not build listener config");
3917        let token = Token(0);
3918        proxy
3919            .add_listener(config, token)
3920            .expect("could not add listener");
3921
3922        let front_a = RequestTcpFrontend {
3923            cluster_id: "cluster-a".to_owned(),
3924            address,
3925            sni: Some("a.example.com".to_owned()),
3926            alpn: vec![],
3927            tags: std::collections::BTreeMap::from([("team".to_owned(), "alpha".to_owned())]),
3928        };
3929        let front_b = RequestTcpFrontend {
3930            cluster_id: "cluster-b".to_owned(),
3931            address,
3932            sni: Some("b.example.com".to_owned()),
3933            alpn: vec!["h2".to_owned()],
3934            tags: std::collections::BTreeMap::from([("team".to_owned(), "beta".to_owned())]),
3935        };
3936        proxy
3937            .add_tcp_front(front_a.clone())
3938            .expect("add_tcp_front A must succeed");
3939        proxy
3940            .add_tcp_front(front_b)
3941            .expect("add_tcp_front B must succeed");
3942
3943        proxy
3944            .remove_tcp_front(front_a)
3945            .expect("remove_tcp_front A must succeed");
3946
3947        let std_address: SocketAddr = address.into();
3948        let key_a = sni_tags_key(&std_address, "a.example.com", &[]);
3949        let key_b = sni_tags_key(&std_address, "b.example.com", &["h2".to_owned()]);
3950
3951        let listener = proxy
3952            .listeners
3953            .get(&token)
3954            .expect("listener must be present")
3955            .borrow();
3956        assert!(
3957            listener.get_tags(&key_a).is_none(),
3958            "removing front A must clear its own tags entry"
3959        );
3960        assert!(
3961            listener.get_tags(&key_b).is_some(),
3962            "removing front A must NOT clear sibling front B's tags"
3963        );
3964    }
3965
3966    /// A session routed to front A must look tags up under front A's key:
3967    /// the key `add_tcp_front` stores and the key `upgrade_sni_preread`
3968    /// rebuilds from the route decision (`matched_sni_pattern` +
3969    /// `matched_alpn`) must be identical, regardless of the operator's
3970    /// original ALPN order or SNI casing, and must never collide with the
3971    /// bare-address key used by no-SNI fronts.
3972    #[test]
3973    fn route_time_tags_key_rebuild_matches_the_add_time_key() {
3974        let address: SocketAddr = "127.0.0.1:9000".parse().expect("test address");
3975
3976        // Wildcard front, operator wrote mixed case + reverse ALPN order.
3977        let add_time = sni_tags_key(
3978            &address,
3979            "*.Example.COM",
3980            &["http/1.1".to_owned(), "h2".to_owned()],
3981        );
3982        let matcher =
3983            AlpnMatcher::OneOf([b"h2".to_vec(), b"http/1.1".to_vec()].into_iter().collect());
3984        let route_time = sni_tags_key(
3985            &address,
3986            "*.example.com", // matched_sni_pattern: the lowercased trie key
3987            &alpn_matcher_protocols(&matcher),
3988        );
3989        assert_eq!(
3990            add_time, route_time,
3991            "add-time and route-time keys must agree for the same front"
3992        );
3993
3994        // Catch-all front: empty alpn at add time <-> AlpnMatcher::Any.
3995        assert_eq!(
3996            sni_tags_key(&address, "a.example.com", &[]),
3997            sni_tags_key(
3998                &address,
3999                "a.example.com",
4000                &alpn_matcher_protocols(&AlpnMatcher::Any)
4001            )
4002        );
4003
4004        // A composed key never collides with the bare-address key.
4005        assert_ne!(add_time, address.to_string());
4006    }
4007
4008    /// Regression from the sozu-proxy/sozu#1290 review: ALPN protocol
4009    /// identifiers are opaque byte strings -- nothing forbids a `,` inside
4010    /// one -- so a SINGLE protocol `"a,b"` and the DISJOINT pair `["a",
4011    /// "b"]` are two legal, distinct `AlpnMatcher`s on the same
4012    /// `(address, sni)`, yet the naive `sorted_alpn.join(",")` key collapses
4013    /// both to the literal string `"a,b"`.
4014    #[test]
4015    fn alpn_sets_differing_only_by_an_embedded_separator_get_distinct_tags_keys() {
4016        let address: SocketAddr = "127.0.0.1:9001".parse().expect("test address");
4017        let key_joined = sni_tags_key(&address, "example.com", &["a,b".to_owned()]);
4018        let key_split = sni_tags_key(&address, "example.com", &["a".to_owned(), "b".to_owned()]);
4019        assert_ne!(
4020            key_joined, key_split,
4021            "a single \"a,b\" protocol must not collide with the disjoint [\"a\", \"b\"] pair"
4022        );
4023    }
4024
4025    /// End-to-end through `TcpProxy`: both fronts are accepted (they are
4026    /// genuinely disjoint `AlpnMatcher`s, so `validate_new_tcp_front`'s
4027    /// overlap check does not reject either), each keeps its OWN tags under
4028    /// its own composed key, and removing one leaves the other's tags
4029    /// intact.
4030    #[test]
4031    fn alpn_sets_differing_only_by_an_embedded_separator_keep_distinct_tags() {
4032        let mut proxy = test_proxy();
4033        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4034        let config = ListenerBuilder::new_tcp(address)
4035            .to_tcp(None)
4036            .expect("could not build listener config");
4037        let token = Token(0);
4038        proxy
4039            .add_listener(config, token)
4040            .expect("could not add listener");
4041
4042        let front_joined = RequestTcpFrontend {
4043            cluster_id: "cluster-joined".to_owned(),
4044            address,
4045            sni: Some("example.com".to_owned()),
4046            alpn: vec!["a,b".to_owned()],
4047            tags: std::collections::BTreeMap::from([("variant".to_owned(), "joined".to_owned())]),
4048        };
4049        let front_split = RequestTcpFrontend {
4050            cluster_id: "cluster-split".to_owned(),
4051            address,
4052            sni: Some("example.com".to_owned()),
4053            alpn: vec!["a".to_owned(), "b".to_owned()],
4054            tags: std::collections::BTreeMap::from([("variant".to_owned(), "split".to_owned())]),
4055        };
4056
4057        proxy
4058            .add_tcp_front(front_joined.clone())
4059            .expect("the single \"a,b\" protocol front must be accepted");
4060        proxy.add_tcp_front(front_split.clone()).expect(
4061            "the disjoint [\"a\", \"b\"] front must be accepted -- it is NOT the same ALPN \
4062                 set as [\"a,b\"]",
4063        );
4064
4065        let std_address: SocketAddr = address.into();
4066        let key_joined = sni_tags_key(&std_address, "example.com", &["a,b".to_owned()]);
4067        let key_split = sni_tags_key(
4068            &std_address,
4069            "example.com",
4070            &["a".to_owned(), "b".to_owned()],
4071        );
4072        assert_ne!(key_joined, key_split);
4073
4074        {
4075            let listener = proxy
4076                .listeners
4077                .get(&token)
4078                .expect("listener must be present")
4079                .borrow();
4080            assert_eq!(
4081                listener
4082                    .get_tags(&key_joined)
4083                    .and_then(|t| t.tags.get("variant"))
4084                    .map(String::as_str),
4085                Some("joined"),
4086                "the joined front's tags must live under its own composed key"
4087            );
4088            assert_eq!(
4089                listener
4090                    .get_tags(&key_split)
4091                    .and_then(|t| t.tags.get("variant"))
4092                    .map(String::as_str),
4093                Some("split"),
4094                "the split front's tags must live under its own DISTINCT composed key"
4095            );
4096        }
4097
4098        proxy
4099            .remove_tcp_front(front_joined)
4100            .expect("remove the joined front");
4101
4102        let listener = proxy
4103            .listeners
4104            .get(&token)
4105            .expect("listener must be present")
4106            .borrow();
4107        assert!(
4108            listener.get_tags(&key_joined).is_none(),
4109            "removing the joined front must clear its own tags entry"
4110        );
4111        assert_eq!(
4112            listener
4113                .get_tags(&key_split)
4114                .and_then(|t| t.tags.get("variant"))
4115                .map(String::as_str),
4116            Some("split"),
4117            "removing the joined front must not disturb its sibling's tags"
4118        );
4119    }
4120
4121    #[test]
4122    fn remove_sni_route_for_an_absent_exact_key_does_not_strip_a_sibling_wildcards_entry() {
4123        let mut listener = test_listener();
4124        listener
4125            .insert_sni_route(
4126                "*.example.com".to_owned(),
4127                vec![],
4128                "cluster-wildcard".to_owned(),
4129            )
4130            .expect("insert_sni_route must succeed for a valid test SNI");
4131
4132        // "a.example.com" was never inserted as its own route -- only the
4133        // wildcard catch-all exists. A remove targeting the exact host
4134        // (e.g. a stale `RemoveTcpFrontend` replayed from a hand-edited
4135        // `LoadState`) must be a no-op here, not reach into and strip the
4136        // WILDCARD's own catch-all entry.
4137        listener.remove_sni_route(
4138            "a.example.com".to_owned(),
4139            vec![],
4140            &"cluster-wildcard".to_owned(),
4141        );
4142
4143        let (_, wildcard_entries) = listener
4144            .sni_routes
4145            .domain_lookup(b"b.example.com", true)
4146            .expect(
4147                "the wildcard catch-all must survive a remove targeting an unrelated exact key",
4148            );
4149        assert_eq!(
4150            wildcard_entries,
4151            &vec![(AlpnMatcher::Any, "cluster-wildcard".to_owned())]
4152        );
4153    }
4154
4155    // ---- routing gate data invariant: never both cluster_id AND routes ----
4156
4157    #[test]
4158    fn a_no_sni_front_leaves_the_route_table_empty() {
4159        let mut listener = test_listener();
4160        listener.cluster_id = Some("legacy-catch-all".to_owned());
4161        assert!(
4162            listener.sni_routes.is_empty(),
4163            "a listener with only a no-SNI front must never populate sni_routes"
4164        );
4165    }
4166
4167    #[test]
4168    fn an_sni_scoped_front_leaves_cluster_id_unset() {
4169        let mut listener = test_listener();
4170        listener
4171            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".to_owned())
4172            .expect("insert_sni_route must succeed for a valid test SNI");
4173        assert!(
4174            listener.cluster_id.is_none(),
4175            "a listener with only SNI-scoped fronts must never populate the legacy cluster_id"
4176        );
4177    }
4178
4179    // ---- end-to-end through TcpProxy::add_tcp_front / remove_tcp_front ----
4180
4181    fn test_proxy() -> TcpProxy {
4182        let ServerParts {
4183            registry,
4184            sessions,
4185            pool,
4186            backends,
4187            ..
4188        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
4189        TcpProxy::new(registry, sessions, pool, backends)
4190    }
4191
4192    #[test]
4193    fn add_then_remove_sni_front_round_trips_through_tcp_proxy() {
4194        let mut proxy = test_proxy();
4195        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4196        let config = ListenerBuilder::new_tcp(address)
4197            .to_tcp(None)
4198            .expect("could not build listener config");
4199        let token = Token(0);
4200        proxy
4201            .add_listener(config, token)
4202            .expect("could not add listener");
4203
4204        let front = RequestTcpFrontend {
4205            cluster_id: "cluster-a".to_owned(),
4206            address,
4207            sni: Some("Example.COM".to_owned()),
4208            alpn: vec![],
4209            ..Default::default()
4210        };
4211        proxy
4212            .add_tcp_front(front.clone())
4213            .expect("add_tcp_front must succeed");
4214
4215        {
4216            let listener = proxy
4217                .listeners
4218                .get(&token)
4219                .expect("listener must be present")
4220                .borrow();
4221            assert!(listener.cluster_id.is_none());
4222            let (_, entries) = listener
4223                .sni_routes
4224                // Lowercased at insert time regardless of wire-form casing.
4225                .domain_lookup(b"example.com", true)
4226                .expect("example.com must be routable after add_tcp_front");
4227            assert_eq!(entries, &vec![(AlpnMatcher::Any, "cluster-a".to_owned())]);
4228        }
4229        assert_eq!(proxy.fronts.get("cluster-a"), Some(&token));
4230
4231        proxy
4232            .remove_tcp_front(front)
4233            .expect("remove_tcp_front must succeed");
4234
4235        {
4236            let listener = proxy
4237                .listeners
4238                .get(&token)
4239                .expect("listener must be present")
4240                .borrow();
4241            assert!(
4242                listener.sni_routes.is_empty(),
4243                "remove_tcp_front must leave no stranded route"
4244            );
4245        }
4246        assert_eq!(
4247            proxy.fronts.get("cluster-a"),
4248            None,
4249            "remove_tcp_front must undo add_tcp_front's self.fronts bookkeeping"
4250        );
4251    }
4252
4253    #[test]
4254    fn add_then_remove_legacy_no_sni_front_round_trips() {
4255        let mut proxy = test_proxy();
4256        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4257        let config = ListenerBuilder::new_tcp(address)
4258            .to_tcp(None)
4259            .expect("could not build listener config");
4260        let token = Token(0);
4261        proxy
4262            .add_listener(config, token)
4263            .expect("could not add listener");
4264
4265        let front = frontend("cluster-legacy", None, &[]);
4266        let front = RequestTcpFrontend { address, ..front };
4267        proxy
4268            .add_tcp_front(front.clone())
4269            .expect("add_tcp_front must succeed");
4270
4271        assert_eq!(
4272            proxy
4273                .listeners
4274                .get(&token)
4275                .expect("listener must be present")
4276                .borrow()
4277                .cluster_id,
4278            Some("cluster-legacy".to_owned())
4279        );
4280
4281        proxy
4282            .remove_tcp_front(front)
4283            .expect("remove_tcp_front must succeed");
4284
4285        assert_eq!(
4286            proxy
4287                .listeners
4288                .get(&token)
4289                .expect("listener must be present")
4290                .borrow()
4291                .cluster_id,
4292            None
4293        );
4294        assert_eq!(proxy.fronts.get("cluster-legacy"), None);
4295    }
4296
4297    // ---- add_tcp_front hard-rejects routing-corrupting requests --------
4298    //
4299    // Worker-side mirror of `command/src/config.rs`'s TOML config-load
4300    // invariants (sozu-proxy/sozu#1279 hardening): `AddTcpFrontend` can
4301    // reach the worker directly over the command socket, or via `LoadState`
4302    // replay, bypassing config.rs entirely, so `add_tcp_front` must defend
4303    // itself rather than rely on a debug-only assertion.
4304
4305    #[test]
4306    fn add_tcp_front_rejects_alpn_without_sni() {
4307        let mut proxy = test_proxy();
4308        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4309        let config = ListenerBuilder::new_tcp(address)
4310            .to_tcp(None)
4311            .expect("could not build listener config");
4312        proxy
4313            .add_listener(config, Token(0))
4314            .expect("could not add listener");
4315
4316        let front = frontend("cluster-a", None, &["h2"]);
4317        let front = RequestTcpFrontend { address, ..front };
4318        match proxy.add_tcp_front(front) {
4319            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4320            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4321        }
4322    }
4323
4324    #[test]
4325    fn add_tcp_front_rejects_no_sni_front_on_listener_with_sni_routes() {
4326        let mut proxy = test_proxy();
4327        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4328        let config = ListenerBuilder::new_tcp(address)
4329            .to_tcp(None)
4330            .expect("could not build listener config");
4331        proxy
4332            .add_listener(config, Token(0))
4333            .expect("could not add listener");
4334
4335        let sni_front = frontend("cluster-a", Some("example.com"), &[]);
4336        let sni_front = RequestTcpFrontend {
4337            address,
4338            ..sni_front
4339        };
4340        proxy
4341            .add_tcp_front(sni_front)
4342            .expect("the first, SNI-scoped frontend must be accepted");
4343
4344        let no_sni_front = frontend("cluster-b", None, &[]);
4345        let no_sni_front = RequestTcpFrontend {
4346            address,
4347            ..no_sni_front
4348        };
4349        match proxy.add_tcp_front(no_sni_front) {
4350            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4351            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4352        }
4353    }
4354
4355    #[test]
4356    fn add_tcp_front_rejects_sni_front_on_listener_with_no_sni_cluster() {
4357        let mut proxy = test_proxy();
4358        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4359        let config = ListenerBuilder::new_tcp(address)
4360            .to_tcp(None)
4361            .expect("could not build listener config");
4362        proxy
4363            .add_listener(config, Token(0))
4364            .expect("could not add listener");
4365
4366        let no_sni_front = frontend("cluster-a", None, &[]);
4367        let no_sni_front = RequestTcpFrontend {
4368            address,
4369            ..no_sni_front
4370        };
4371        proxy
4372            .add_tcp_front(no_sni_front)
4373            .expect("the first, no-SNI frontend must be accepted");
4374
4375        let sni_front = frontend("cluster-b", Some("example.com"), &[]);
4376        let sni_front = RequestTcpFrontend {
4377            address,
4378            ..sni_front
4379        };
4380        match proxy.add_tcp_front(sni_front) {
4381            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4382            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4383        }
4384    }
4385
4386    #[test]
4387    fn add_tcp_front_rejects_alpn_overlap_on_same_sni() {
4388        let mut proxy = test_proxy();
4389        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4390        let config = ListenerBuilder::new_tcp(address)
4391            .to_tcp(None)
4392            .expect("could not build listener config");
4393        proxy
4394            .add_listener(config, Token(0))
4395            .expect("could not add listener");
4396
4397        let first = frontend("cluster-a", Some("example.com"), &["h2"]);
4398        let first = RequestTcpFrontend { address, ..first };
4399        proxy
4400            .add_tcp_front(first)
4401            .expect("the first frontend must be accepted");
4402
4403        // Second frontend shares "h2" with the first on the same sni.
4404        let second = frontend("cluster-b", Some("example.com"), &["h2", "http/1.1"]);
4405        let second = RequestTcpFrontend { address, ..second };
4406        match proxy.add_tcp_front(second) {
4407            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4408            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4409        }
4410    }
4411
4412    #[test]
4413    fn add_tcp_front_rejects_duplicate_catch_all_on_same_sni() {
4414        let mut proxy = test_proxy();
4415        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4416        let config = ListenerBuilder::new_tcp(address)
4417            .to_tcp(None)
4418            .expect("could not build listener config");
4419        proxy
4420            .add_listener(config, Token(0))
4421            .expect("could not add listener");
4422
4423        let first = frontend("cluster-a", Some("example.com"), &[]);
4424        let first = RequestTcpFrontend { address, ..first };
4425        proxy
4426            .add_tcp_front(first)
4427            .expect("the first catch-all frontend must be accepted");
4428
4429        let second = frontend("cluster-b", Some("example.com"), &[]);
4430        let second = RequestTcpFrontend { address, ..second };
4431        match proxy.add_tcp_front(second) {
4432            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4433            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4434        }
4435    }
4436
4437    /// The valid, intended shape (sozu-proxy/sozu#1279's whole reason for
4438    /// existing) must still be accepted: disjoint, non-empty `alpn` lists
4439    /// on the same `(address, sni)`, and a catch-all alongside a
4440    /// specific-protocol entry.
4441    #[test]
4442    fn add_tcp_front_accepts_disjoint_alpn_and_catch_all_on_same_sni() {
4443        let mut proxy = test_proxy();
4444        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4445        let config = ListenerBuilder::new_tcp(address)
4446            .to_tcp(None)
4447            .expect("could not build listener config");
4448        proxy
4449            .add_listener(config, Token(0))
4450            .expect("could not add listener");
4451
4452        let h2 = frontend("cluster-h2", Some("example.com"), &["h2"]);
4453        let h2 = RequestTcpFrontend { address, ..h2 };
4454        proxy
4455            .add_tcp_front(h2)
4456            .expect("disjoint alpn frontend must be accepted");
4457
4458        let http11 = frontend("cluster-http11", Some("example.com"), &["http/1.1"]);
4459        let http11 = RequestTcpFrontend { address, ..http11 };
4460        proxy
4461            .add_tcp_front(http11)
4462            .expect("second disjoint alpn frontend must be accepted");
4463
4464        let catch_all = frontend("cluster-default", Some("example.com"), &[]);
4465        let catch_all = RequestTcpFrontend {
4466            address,
4467            ..catch_all
4468        };
4469        proxy
4470            .add_tcp_front(catch_all)
4471            .expect("a catch-all alongside specific-protocol entries must be accepted");
4472    }
4473
4474    // ---- tcp.sni_preread.active gauge accounting ----------------------
4475
4476    /// Read the current process-local `tcp.sni_preread.active` gauge,
4477    /// treating an absent key as 0. `dump_local_proxy_metrics` is a
4478    /// non-draining filter over the proxy `MetricsMap`, so repeated reads are
4479    /// side-effect free and the key is the raw metric name.
4480    fn sni_preread_active_gauge() -> i64 {
4481        use sozu_command::proto::command::filtered_metrics::Inner;
4482        crate::metrics::METRICS.with(|metrics| {
4483            metrics
4484                .borrow_mut()
4485                .dump_local_proxy_metrics()
4486                .get(names::tcp::sni_preread::ACTIVE)
4487                .and_then(|fm| fm.inner.as_ref())
4488                .and_then(|inner| match inner {
4489                    Inner::Gauge(v) => Some(*v as i64),
4490                    _ => None,
4491                })
4492                .unwrap_or(0)
4493        })
4494    }
4495
4496    #[test]
4497    fn entering_sni_preread_increments_the_active_gauge() {
4498        // Regression guard for the missing-`+1` gauge bug
4499        // (sozu-proxy/sozu#1279): `new_sni_preread` must bump
4500        // `tcp.sni_preread.active` by exactly one when a session ENTERS the
4501        // state, so each of the two `-1` decrements -- the "upgrade" exit in
4502        // `upgrade_sni_preread` and the "reject"/"teardown" exit in `close()`'s
4503        // `StateMarker::SniPreread` arm -- has a matching increment. Without
4504        // this `+1` the first `-1` underflows a fresh-zero gauge (clamped to 0,
4505        // ERROR-logged), pinning the gauge at 0 and rendering the e2e gauge
4506        // assertion vacuous.
4507        //
4508        // `METRICS` is a thread-local shared across unit tests on the same
4509        // worker thread, so this asserts the DELTA around one constructor call
4510        // (robust to any starting value), not an absolute reading. The
4511        // net-zero-per-session contract spans the full lifecycle (accept ->
4512        // live backend connect -> upgrade/teardown) and is the behavioural job
4513        // of the e2e gauge assertion
4514        // (`test_tcp_sni_reject_then_valid_connection_not_limited` in
4515        // `e2e/src/tests/tcp_sni_tests.rs`), not reproducible at this unit
4516        // level.
4517        let ServerParts {
4518            registry,
4519            sessions,
4520            pool,
4521            backends,
4522            ..
4523        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
4524
4525        let proxy = Rc::new(RefCell::new(TcpProxy::new(
4526            registry,
4527            sessions,
4528            pool.clone(),
4529            backends,
4530        )));
4531        let listener = Rc::new(RefCell::new(test_listener()));
4532
4533        let (front_buffer, back_buffer) = {
4534            let mut pool = pool.borrow_mut();
4535            (
4536                pool.checkout().expect("front buffer checkout must succeed"),
4537                pool.checkout().expect("back buffer checkout must succeed"),
4538            )
4539        };
4540
4541        // A non-blocking connect to a (likely unused) loopback port returns a
4542        // real `MioTcpStream` handle immediately, regardless of whether the
4543        // connection completes; `new_sni_preread` only reads `peer_addr()`.
4544        let socket = MioTcpStream::connect(
4545            format!("127.0.0.1:{}", provide_port())
4546                .parse()
4547                .expect("loopback address must parse"),
4548        )
4549        .expect("mio connect must return a socket handle");
4550
4551        let before = sni_preread_active_gauge();
4552        let session = TcpSession::new_sni_preread(
4553            back_buffer,
4554            Duration::from_secs(30),
4555            Duration::from_secs(30),
4556            front_buffer,
4557            Token(0),
4558            listener,
4559            proxy,
4560            socket,
4561            Duration::from_millis(0),
4562            Duration::from_secs(3),
4563            16384,
4564        );
4565        let after = sni_preread_active_gauge();
4566
4567        // The session is measured while still in `SniPreread`; hold it across
4568        // the read so no future `Drop` side effect could race the measurement.
4569        assert!(matches!(session.state, TcpStateMachine::SniPreread(_)));
4570
4571        assert_eq!(
4572            after - before,
4573            1,
4574            "entering the SniPreread state must increment tcp.sni_preread.active by exactly one"
4575        );
4576    }
4577
4578    // ---- absolute preread deadline + route-aware timeout (sozu-proxy/sozu#1290) ----
4579
4580    /// `frontend_timeout_resets_on_readable` must gate the reset on exactly
4581    /// one thing: whether the CURRENT state is an undecided `SniPreread`.
4582    /// Constructed directly (no socket I/O needed) since the predicate is a
4583    /// pure function of `&TcpStateMachine`.
4584    #[test]
4585    fn frontend_timeout_reset_is_gated_on_sni_preread_decision() {
4586        let mut pool = crate::pool::Pool::with_capacity(1, 1, 16 * 1024);
4587        let frontend_buffer = pool.checkout().expect("frontend buffer");
4588        let socket = MioTcpStream::connect(
4589            format!("127.0.0.1:{}", provide_port())
4590                .parse()
4591                .expect("loopback address must parse"),
4592        )
4593        .expect("mio connect must return a socket handle");
4594
4595        let undecided = TcpStateMachine::SniPreread(SniPreread::new(
4596            socket,
4597            Token(0),
4598            Ulid::generate(),
4599            frontend_buffer,
4600            16384,
4601        ));
4602        assert!(
4603            !frontend_timeout_resets_on_readable(&undecided),
4604            "an undecided SniPreread must not have its absolute deadline reset"
4605        );
4606
4607        // Every OTHER state resets normally -- represented here by
4608        // `ExpectProxyProtocol`, cheaply constructible without driving a
4609        // real route decision (unlike a ROUTED `SniPreread`, whose
4610        // `outcome` field has no test-only setter and can only become
4611        // `Some` by parsing a real ClientHello -- see
4612        // `frontend_timeout_restored_and_timeout_after_route_is_not_double_counted`
4613        // below for that scenario end-to-end).
4614        let socket2 = MioTcpStream::connect(
4615            format!("127.0.0.1:{}", provide_port())
4616                .parse()
4617                .expect("loopback address must parse"),
4618        )
4619        .expect("mio connect must return a socket handle");
4620        let container = crate::timer::TimeoutContainer::new_empty(Duration::from_secs(5));
4621        let other = TcpStateMachine::ExpectProxyProtocol(ExpectProxyProtocol::new(
4622            container,
4623            socket2,
4624            Token(1),
4625            Ulid::generate(),
4626        ));
4627        assert!(
4628            frontend_timeout_resets_on_readable(&other),
4629            "every non-preread-undecided state must keep resetting its frontend timeout"
4630        );
4631    }
4632
4633    /// Minimal single-record TLS ClientHello wire carrying only a
4634    /// `server_name` extension for `host` -- hand-built rather than reusing
4635    /// `tcp_preread::parser`'s test-only wire-building helpers (`mod
4636    /// parser` is private to `tcp_preread`, unreachable from this sibling
4637    /// module) to drive a REAL route decision through
4638    /// `TcpSession::readable()` for the regression test below.
4639    fn minimal_client_hello_wire(host: &str) -> Vec<u8> {
4640        let mut name_list = vec![0u8]; // name_type = host_name
4641        name_list.extend_from_slice(&(host.len() as u16).to_be_bytes());
4642        name_list.extend_from_slice(host.as_bytes());
4643        let mut sni_ext_data = Vec::new();
4644        sni_ext_data.extend_from_slice(&(name_list.len() as u16).to_be_bytes());
4645        sni_ext_data.extend_from_slice(&name_list);
4646        let mut sni_ext = Vec::new();
4647        sni_ext.extend_from_slice(&0x0000u16.to_be_bytes()); // server_name extension type
4648        sni_ext.extend_from_slice(&(sni_ext_data.len() as u16).to_be_bytes());
4649        sni_ext.extend_from_slice(&sni_ext_data);
4650
4651        let mut body = Vec::new();
4652        body.extend_from_slice(&[0x03, 0x03]); // legacy_version
4653        body.extend_from_slice(&[0u8; 32]); // random
4654        body.push(0); // session_id: empty
4655        body.extend_from_slice(&[0x00, 0x02, 0x13, 0x01]); // cipher_suites
4656        body.push(1); // compression_methods length
4657        body.push(0); // compression_method: null
4658        body.extend_from_slice(&(sni_ext.len() as u16).to_be_bytes()); // extensions block length
4659        body.extend_from_slice(&sni_ext);
4660
4661        let mut handshake = Vec::new();
4662        handshake.push(1u8); // msg_type = client_hello
4663        let hs_len = body.len() as u32;
4664        handshake.extend_from_slice(&hs_len.to_be_bytes()[1..4]);
4665        handshake.extend_from_slice(&body);
4666
4667        let mut record = Vec::new();
4668        record.push(22u8); // ContentType::handshake
4669        record.extend_from_slice(&[0x03, 0x03]); // legacy record version
4670        record.extend_from_slice(&(handshake.len() as u16).to_be_bytes());
4671        record.extend_from_slice(&handshake);
4672        record
4673    }
4674
4675    /// Read the current process-local `tcp.sni_preread.routed` counter.
4676    /// Same non-draining-read pattern as `sni_preread_active_gauge`, but
4677    /// for a `Count` metric instead of a `Gauge`.
4678    fn sni_preread_routed_count() -> i64 {
4679        use sozu_command::proto::command::filtered_metrics::Inner;
4680        crate::metrics::METRICS.with(|metrics| {
4681            metrics
4682                .borrow_mut()
4683                .dump_local_proxy_metrics()
4684                .get(names::tcp::sni_preread::ROUTED)
4685                .and_then(|fm| fm.inner.as_ref())
4686                .and_then(|inner| match inner {
4687                    Inner::Count(v) => Some(*v),
4688                    _ => None,
4689                })
4690                .unwrap_or(0)
4691        })
4692    }
4693
4694    /// End-to-end regression from the sozu-proxy/sozu#1290 review:
4695    ///
4696    /// (a) the moment a real ClientHello routes, `container_frontend_timeout`
4697    ///     must already carry the listener's configured `front_timeout` --
4698    ///     not just once `upgrade_sni_preread` eventually runs (which can be
4699    ///     one or more `ready()` cycles later, after the backend connects);
4700    /// (b) a front-timeout firing AFTER that route decision (backend connect
4701    ///     still pending) must not re-feed `Input::Timeout` into the
4702    ///     already-decided core: pre-fix, doing so replayed the SAME latched
4703    ///     `Output::Routed` through `SniPreread::handle_output`'s `Routed`
4704    ///     arm a second time, double-incrementing `tcp.sni_preread.routed`
4705    ///     (release) and tripping `debug_assert!(self.outcome.is_none(),
4706    ///     ...)` (debug -- this test runs in a debug build, so pre-fix it
4707    ///     panics here).
4708    #[test]
4709    fn frontend_timeout_restored_and_timeout_after_route_is_not_double_counted() {
4710        let ServerParts {
4711            registry,
4712            sessions,
4713            pool,
4714            backends,
4715            ..
4716        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
4717        let proxy = Rc::new(RefCell::new(TcpProxy::new(
4718            registry,
4719            sessions,
4720            pool.clone(),
4721            backends,
4722        )));
4723
4724        let mut bare_listener = test_listener();
4725        bare_listener
4726            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".to_owned())
4727            .expect("insert_sni_route must succeed for a valid test SNI");
4728        let configured_front_timeout = bare_listener.config.front_timeout;
4729        let listener = Rc::new(RefCell::new(bare_listener));
4730
4731        let std_listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind test listener");
4732        let addr = std_listener.local_addr().expect("listener local addr");
4733        let mut client = std::net::TcpStream::connect(addr).expect("connect test client");
4734        let (server, _) = std_listener.accept().expect("accept test server");
4735        server.set_nonblocking(true).expect("server nonblocking");
4736
4737        let (front_buffer, back_buffer) = {
4738            let mut pool = pool.borrow_mut();
4739            (
4740                pool.checkout().expect("front buffer checkout must succeed"),
4741                pool.checkout().expect("back buffer checkout must succeed"),
4742            )
4743        };
4744
4745        let mut session = TcpSession::new_sni_preread(
4746            back_buffer,
4747            Duration::from_secs(30),
4748            Duration::from_secs(30),
4749            front_buffer,
4750            Token(0),
4751            listener,
4752            proxy,
4753            MioTcpStream::from_std(server),
4754            Duration::from_millis(0),
4755            Duration::from_secs(3),
4756            16384,
4757        );
4758
4759        {
4760            use std::io::Write as _;
4761            client
4762                .write_all(&minimal_client_hello_wire("example.com"))
4763                .expect("write ClientHello");
4764            client.flush().ok();
4765        }
4766
4767        let routed_before = sni_preread_routed_count();
4768        for _ in 0..10 {
4769            if session.cluster_id.is_some() {
4770                break;
4771            }
4772            let _ = session.readable();
4773        }
4774        assert_eq!(
4775            session.cluster_id.as_deref(),
4776            Some("cluster-a"),
4777            "the session must have routed on a valid ClientHello for a configured SNI"
4778        );
4779        assert_eq!(
4780            sni_preread_routed_count() - routed_before,
4781            1,
4782            "routing must count tcp.sni_preread.routed exactly once"
4783        );
4784
4785        // (a) front_timeout is restored the moment routing succeeds.
4786        assert_eq!(
4787            session.container_frontend_timeout.duration(),
4788            Duration::from_secs(configured_front_timeout as u64),
4789            "the frontend timeout must already carry the configured front_timeout right after \
4790             routing, not only once the backend connects"
4791        );
4792
4793        // (b) a timeout firing after the route already latched must not
4794        // double-count tcp.sni_preread.routed, and (running in a debug
4795        // build) must not panic on SniPreread::handle_output's
4796        // `debug_assert!(self.outcome.is_none(), ...)`.
4797        let _ = session.timeout(Token(0));
4798        assert_eq!(
4799            sni_preread_routed_count() - routed_before,
4800            1,
4801            "a timeout firing after the route already latched must not double-count \
4802             tcp.sni_preread.routed"
4803        );
4804
4805        drop(client);
4806    }
4807}