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    rc::Rc,
7    time::{Duration, Instant},
8};
9
10use mio::{
11    Interest, Registry, Token,
12    net::{TcpListener as MioTcpListener, TcpStream as MioTcpStream},
13};
14use rusty_ulid::Ulid;
15use socket2::SockRef;
16use sozu_command::{
17    ObjectKind,
18    config::{
19        DEFAULT_SNI_PREREAD_MAX_BYTES, DEFAULT_SNI_PREREAD_TIMEOUT, MAX_LOOP_ITERATIONS,
20        MIN_SNI_PREREAD_MAX_BYTES, validate_sni_pattern,
21    },
22    listener_key::ListenerKey,
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    load_balancing::affinity_key_from_ip,
35    pool::{Checkout, Pool},
36    protocol::mux::router::cluster_reads_affinity_key,
37    protocol::{
38        Pipe,
39        pipe::WebSocketContext,
40        proxy_protocol::{
41            expect::ExpectProxyProtocol, relay::RelayProxyProtocol, send::SendProxyProtocol,
42        },
43        tcp_preread::{AlpnMatcher, PrereadConfig, shell::SniPreread},
44    },
45    retry::RetryPolicy,
46    server::{ListenToken, SessionManager, push_event},
47    socket::{server_bind, stats::socket_rtt},
48    sozu_command::{
49        proto::command::{
50            Event, EventKind, ProxyProtocolConfig, RequestTcpFrontend, TcpListenerConfig,
51            UpdateTcpListenerConfig, WorkerRequest, WorkerResponse,
52        },
53        ready::Ready,
54        state::ClusterId,
55    },
56    timer::TimeoutContainer,
57};
58
59StateMachineBuilder! {
60    /// The various Stages of a TCP connection:
61    ///
62    /// 1. optional SniPreread (SNI-routed listeners only, sozu-proxy/sozu#1279)
63    /// 2. optional (ExpectProxyProtocol | SendProxyProtocol | RelayProxyProtocol)
64    /// 3. Pipe
65    enum TcpStateMachine {
66        Pipe(Pipe<MioTcpStream, TcpListener>),
67        SendProxyProtocol(SendProxyProtocol<MioTcpStream>),
68        RelayProxyProtocol(RelayProxyProtocol<MioTcpStream>),
69        ExpectProxyProtocol(ExpectProxyProtocol<MioTcpStream>),
70        SniPreread(SniPreread<MioTcpStream>),
71    }
72}
73
74/// This macro is defined uniquely in this module to help the tracking of kawa h1
75/// issues inside Sōzu. Colored output uses the unified log-context scheme:
76/// bold bright-white protocol label, light-grey `Session` keyword, gray keys
77/// and bright-white values.
78macro_rules! log_context {
79    ($self:expr) => {{
80        let (open, reset, grey, gray, white) = ansi_palette();
81        format!(
82            "{gray}{ctx}{reset}\t{open}TCP{reset}\t{grey}Session{reset}({gray}frontend{reset}={white}{frontend}{reset}, {gray}backend{reset}={white}{backend}{reset})\t >>>",
83            open = open,
84            reset = reset,
85            grey = grey,
86            gray = gray,
87            white = white,
88            ctx = $self.log_context(),
89            frontend = $self.frontend_token.0,
90            backend = $self
91                .backend_token
92                .map(|token| token.0.to_string())
93                .unwrap_or_else(|| "<none>".to_string()),
94        )
95    }};
96}
97
98/// Module-level prefix for log lines emitted from this file when no
99/// [`TcpSession`] is in scope. Produces a bold bright-white `TCP` label
100/// (uniform with the per-session `log_context!`) when the logger is in
101/// colored mode. Used by [`TcpProxy`] callbacks (notify, accept,
102/// create_session, soft_stop, hard_stop, status) and the `testing`
103/// helper module which own a listener/token map but have no
104/// `frontend_token` of their own.
105macro_rules! log_module_context {
106    () => {{
107        let (open, reset, _, _, _) = sozu_command::logging::ansi_palette();
108        format!("{open}TCP{reset}\t >>>", open = open, reset = reset)
109    }};
110}
111
112pub struct TcpSession {
113    backend_buffer: Option<Checkout>,
114    backend_connected: BackendConnectionStatus,
115    backend_id: Option<String>,
116    /// Address of the backend `connect_to_backend` last dialed, handed to the
117    /// `Pipe` for its access log (`Pipe::set_backend_address`), so the log
118    /// never asks the socket with `getpeername(2)`.
119    backend_address: Option<SocketAddr>,
120    backend_token: Option<Token>,
121    backend: Option<Rc<RefCell<Backend>>>,
122    cluster_id: Option<ClusterId>,
123    /// Configured lifetime of `cluster_id` captured when this session routes.
124    /// Kept on the session rather than the public `SessionMetrics` value so
125    /// delayed TCP and Pipe emissions cannot target a replacement cluster.
126    cluster_metrics_incarnation: crate::metrics::ClusterMetricsCapture,
127    configured_backend_timeout: Duration,
128    connection_attempt: u8,
129    container_backend_timeout: TimeoutContainer,
130    container_frontend_timeout: TimeoutContainer,
131    frontend_address: Option<SocketAddr>,
132    frontend_buffer: Option<Checkout>,
133    frontend_token: Token,
134    has_been_closed: SessionIsToBeClosed,
135    last_event: Instant,
136    listener: Rc<RefCell<TcpListener>>,
137    metrics: SessionMetrics,
138    proxy: Rc<RefCell<TcpProxy>>,
139    request_id: Ulid,
140    state: TcpStateMachine,
141    /// `true` once `connect_to_backend` has accounted this session
142    /// against the per-(cluster, source-IP) connection counter. Drives
143    /// the symmetric `untrack_all_cluster_ip` call in `close`. The flag
144    /// is per-session, not per-attempt: a TCP session has at most one
145    /// `(cluster, ip)` slot, so the SessionManager-side idempotency
146    /// already covers retries — this flag exists only to short-circuit
147    /// the close path's untrack when the feature is disabled or no
148    /// admit ever ran.
149    cluster_ip_tracked: bool,
150    /// SNI-preread routing result (sozu-proxy/sozu#1279), captured once by
151    /// `upgrade_sni_preread` for every `proxy_protocol` case and consumed
152    /// (`Option::take`) at the point the session actually reaches `Pipe`:
153    /// immediately in `build_pipe_from_preread` for
154    /// `Expect`/`Relay`/`None`, or one `ready()` cycle later in
155    /// `upgrade_send` for `SendHeader` (which transitions through
156    /// `SendProxyProtocol` first). `None` for every non-SNI-routed session.
157    routed_sni: Option<String>,
158    /// Paired with `routed_sni`: the client's first ALPN offer, mapped to a
159    /// known `&'static str` label (`"h2"` / `"http/1.1"`) for the access
160    /// log, or `None` if the client offered nothing recognized. Sōzu never
161    /// negotiates ALPN itself on the TCP passthrough path -- the backend
162    /// terminates TLS -- so this is informational (the client's
163    /// preference), not a negotiated value.
164    routed_alpn_label: Option<&'static str>,
165    /// Canonical access-log tags key for the MATCHED SNI/ALPN frontend
166    /// (`sni_tags_key`), rebuilt in `upgrade_sni_preread` from the route
167    /// decision's `matched_sni_pattern` + `matched_alpn`. `None` for every
168    /// non-SNI-routed session, whose tags stay keyed by the bare listener
169    /// address exactly as before SNI routing existed. Unlike `routed_sni`
170    /// it is never consumed: `log_request` reads it for the session-level
171    /// access log, and the `Pipe` receives a clone at upgrade time for the
172    /// post-upgrade log.
173    tags_key: Option<String>,
174}
175
176impl TcpSession {
177    #[allow(clippy::too_many_arguments)]
178    fn new(
179        backend_buffer: Checkout,
180        backend_id: Option<String>,
181        cluster_id: Option<ClusterId>,
182        configured_backend_timeout: Duration,
183        configured_connect_timeout: Duration,
184        configured_frontend_timeout: Duration,
185        frontend_buffer: Checkout,
186        frontend_token: Token,
187        listener: Rc<RefCell<TcpListener>>,
188        proxy_protocol: Option<ProxyProtocolConfig>,
189        proxy: Rc<RefCell<TcpProxy>>,
190        socket: MioTcpStream,
191        peer: SocketAddr,
192        wait_time: Duration,
193    ) -> TcpSession {
194        // The address `accept(2)` returned with `socket`.
195        let frontend_address = Some(peer);
196        let mut frontend_buffer_session = None;
197        let mut backend_buffer_session = None;
198
199        let request_id = Ulid::generate();
200
201        let container_frontend_timeout =
202            TimeoutContainer::new(configured_frontend_timeout, frontend_token);
203        let container_backend_timeout = TimeoutContainer::new_empty(configured_connect_timeout);
204        let cluster_metrics_incarnation = cluster_id
205            .as_deref()
206            .map_or_else(Default::default, crate::metrics::capture_cluster_metrics);
207
208        let state = match proxy_protocol {
209            Some(ProxyProtocolConfig::RelayHeader) => {
210                backend_buffer_session = Some(backend_buffer);
211                gauge_add!(names::protocol::PROXY_RELAY, 1);
212                TcpStateMachine::RelayProxyProtocol(RelayProxyProtocol::new(
213                    socket,
214                    frontend_token,
215                    request_id,
216                    None,
217                    frontend_buffer,
218                ))
219            }
220            Some(ProxyProtocolConfig::ExpectHeader) => {
221                frontend_buffer_session = Some(frontend_buffer);
222                backend_buffer_session = Some(backend_buffer);
223                gauge_add!(names::protocol::PROXY_EXPECT, 1);
224                TcpStateMachine::ExpectProxyProtocol(ExpectProxyProtocol::new(
225                    container_frontend_timeout.clone(),
226                    socket,
227                    frontend_token,
228                    request_id,
229                ))
230            }
231            Some(ProxyProtocolConfig::SendHeader) => {
232                frontend_buffer_session = Some(frontend_buffer);
233                backend_buffer_session = Some(backend_buffer);
234                gauge_add!(names::protocol::PROXY_SEND, 1);
235                TcpStateMachine::SendProxyProtocol(SendProxyProtocol::new(
236                    socket,
237                    frontend_token,
238                    request_id,
239                    None,
240                ))
241            }
242            None => {
243                gauge_add!(names::protocol::TCP, 1);
244                let mut pipe = Pipe::new(
245                    backend_buffer,
246                    backend_id.clone(),
247                    None,
248                    None,
249                    None,
250                    None,
251                    cluster_id.clone(),
252                    frontend_buffer,
253                    frontend_token,
254                    socket,
255                    listener.clone(),
256                    Protocol::TCP,
257                    request_id,
258                    request_id,
259                    frontend_address,
260                    WebSocketContext::Tcp,
261                );
262                pipe.set_cluster_id(cluster_id.clone());
263                pipe.set_cluster_metrics_incarnation(cluster_metrics_incarnation);
264                TcpStateMachine::Pipe(pipe)
265            }
266        };
267
268        let metrics = SessionMetrics::new(Some(wait_time));
269        //FIXME: timeout usage
270
271        TcpSession {
272            backend_buffer: backend_buffer_session,
273            backend_connected: BackendConnectionStatus::NotConnected,
274            backend_id,
275            backend_address: None,
276            backend_token: None,
277            backend: None,
278            cluster_id,
279            cluster_metrics_incarnation,
280            configured_backend_timeout,
281            connection_attempt: 0,
282            container_backend_timeout,
283            container_frontend_timeout,
284            frontend_address,
285            frontend_buffer: frontend_buffer_session,
286            frontend_token,
287            has_been_closed: false,
288            last_event: Instant::now(),
289            listener,
290            metrics,
291            proxy,
292            request_id,
293            state,
294            cluster_ip_tracked: false,
295            routed_sni: None,
296            routed_alpn_label: None,
297            tags_key: None,
298        }
299    }
300
301    /// Construct a session that starts in [`TcpStateMachine::SniPreread`]
302    /// instead of resolving a `proxy_protocol` up front -- the cluster (and
303    /// therefore the per-cluster `proxy_protocol`) is only known once
304    /// [`crate::protocol::tcp_preread::SniPrereadCore`] decides a route.
305    /// Mirrors [`Self::new`]'s tail; kept as a separate constructor rather
306    /// than folding a synthetic sentinel into `proxy_protocol:
307    /// Option<ProxyProtocolConfig>` (a proto-generated enum this crate does
308    /// not own).
309    #[allow(clippy::too_many_arguments)]
310    fn new_sni_preread(
311        backend_buffer: Checkout,
312        configured_backend_timeout: Duration,
313        configured_connect_timeout: Duration,
314        frontend_buffer: Checkout,
315        frontend_token: Token,
316        listener: Rc<RefCell<TcpListener>>,
317        proxy: Rc<RefCell<TcpProxy>>,
318        socket: MioTcpStream,
319        peer: SocketAddr,
320        wait_time: Duration,
321        preread_timeout: Duration,
322        effective_max_bytes: usize,
323    ) -> TcpSession {
324        // The address `accept(2)` returned with `socket`.
325        let frontend_address = Some(peer);
326        let request_id = Ulid::generate();
327
328        // Armed with the SHORT preread timeout directly (not the listener's
329        // configured front_timeout) -- `upgrade_sni_preread` restores the
330        // configured duration on the SAME container once routed, so there is
331        // exactly one `TimeoutContainer` for the frontend token throughout,
332        // never a diverging clone (a clone independently rearmed to a
333        // shorter duration would strand `TcpSession::readable`'s own
334        // unconditional `reset()` on a since-cancelled timer-wheel entry).
335        let container_frontend_timeout = TimeoutContainer::new(preread_timeout, frontend_token);
336        let container_backend_timeout = TimeoutContainer::new_empty(configured_connect_timeout);
337
338        let state = TcpStateMachine::SniPreread(SniPreread::new(
339            socket,
340            frontend_token,
341            request_id,
342            frontend_buffer,
343            effective_max_bytes,
344        ));
345
346        // Enter the `SniPreread` state: +1 the active gauge exactly once, and
347        // unconditionally, so every one of the two `-1` decrements has a
348        // matching increment. The gauge is decremented on precisely one of the
349        // two mutually-exclusive exits: the "upgrade" exit in
350        // `upgrade_sni_preread` (which first transitions `self.state` away from
351        // `SniPreread`, so `close()` cannot re-decrement), and the
352        // "reject"/"teardown" exit in `close()`'s `StateMarker::SniPreread`
353        // arm. A session therefore nets to 0 and never underflows.
354        gauge_add!(names::tcp::sni_preread::ACTIVE, 1);
355
356        let metrics = SessionMetrics::new(Some(wait_time));
357
358        TcpSession {
359            backend_buffer: Some(backend_buffer),
360            backend_connected: BackendConnectionStatus::NotConnected,
361            backend_id: None,
362            backend_address: None,
363            backend_token: None,
364            backend: None,
365            cluster_id: None,
366            cluster_metrics_incarnation: Default::default(),
367            configured_backend_timeout,
368            connection_attempt: 0,
369            container_backend_timeout,
370            container_frontend_timeout,
371            frontend_address,
372            frontend_buffer: None,
373            frontend_token,
374            has_been_closed: false,
375            last_event: Instant::now(),
376            listener,
377            metrics,
378            proxy,
379            request_id,
380            state,
381            cluster_ip_tracked: false,
382            routed_sni: None,
383            routed_alpn_label: None,
384            tags_key: None,
385        }
386    }
387
388    /// Source-IP for per-(cluster, source-IP) accounting.
389    ///
390    /// Prefer the parsed PROXY-v2 source from whichever upgrade phase is
391    /// in flight, then the post-upgrade `Pipe.session_address`, finally
392    /// the raw TCP `peer_addr` captured at session creation. The
393    /// `Pipe::session_address` itself is already PROXY-v2-aware after
394    /// `expect.rs::into_pipe` and `relay.rs::into_pipe`.
395    fn effective_session_address(&self) -> Option<SocketAddr> {
396        match &self.state {
397            TcpStateMachine::Pipe(pipe) => pipe.get_session_address(),
398            TcpStateMachine::ExpectProxyProtocol(epp) => {
399                epp.addresses.as_ref().and_then(|pa| pa.source())
400            }
401            TcpStateMachine::RelayProxyProtocol(rpp) => {
402                rpp.addresses.as_ref().and_then(|pa| pa.source())
403            }
404            TcpStateMachine::SniPreread(preread) => preread.outcome().and_then(|o| o.proxy_source),
405            TcpStateMachine::SendProxyProtocol(_) | TcpStateMachine::FailedUpgrade(_) => None,
406        }
407        .or(self.frontend_address)
408    }
409
410    fn log_request(&self) {
411        let listener = self.listener.borrow();
412        let context = self.log_context();
413        self.metrics
414            .register_end_of_session_for_incarnation(&context, self.cluster_metrics_incarnation);
415        // SNI-routed sessions carry the matched front's own tags key
416        // (`sni_tags_key`, stashed by `upgrade_sni_preread`); everything
417        // else keeps the historical bare-address key.
418        let address_key = TcpFrontendTagsKey::Address(*listener.get_addr()).to_string();
419        let tags_key = self.tags_key.as_deref().unwrap_or(&address_key);
420        info_access!(
421            on_failure: { incr!(names::access_logs::UNSENT) },
422            message: None,
423            context,
424            session_address: self.frontend_address,
425            backend_address: None,
426            protocol: "TCP",
427            endpoint: EndpointRecord::Tcp,
428            tags: listener.get_tags(tags_key),
429            client_rtt: socket_rtt(self.state.front_socket()),
430            server_rtt: None,
431            user_agent: None,
432            x_request_id: None,
433            // Sōzu never terminates TLS on the TCP path (the frontend is a
434            // raw `MioTcpStream`), so no negotiated version/cipher exists.
435            // A preread SNI/ALPN (SNI-routed listeners) is stamped on the
436            // `Pipe`'s own access log via `set_tls_metadata` at upgrade
437            // time; this pre-Pipe log site emits `None` for all four TLS
438            // fields and the parsed XFF chain.
439            tls_version: None,
440            tls_cipher: None,
441            tls_sni: None,
442            tls_alpn: None,
443            xff_chain: None,
444            service_time: self.metrics.service_time(),
445            response_time: self.metrics.backend_response_time(),
446            request_time: self.metrics.request_time(),
447            start_time_ns: self.metrics.start_wall_ns(),
448            bytes_in: self.metrics.bin,
449            bytes_out: self.metrics.bout,
450            otel: None,
451        );
452    }
453
454    fn front_hup(&mut self) -> SessionResult {
455        let listener = self.listener.borrow();
456        match &mut self.state {
457            TcpStateMachine::Pipe(pipe) => pipe.frontend_hup(&mut self.metrics),
458            // No access log here, mirroring `readable()`'s own error paths
459            // for the other pre-Pipe states (none of them call
460            // `log_request()` either): the shell itself decides silent vs.
461            // metered based on whether any bytes were ever received.
462            TcpStateMachine::SniPreread(preread) => {
463                let cfg = listener.preread_config(preread.effective_max_bytes());
464                preread.on_front_closed(&cfg);
465                SessionResult::Close
466            }
467            _ => {
468                self.log_request();
469                SessionResult::Close
470            }
471        }
472    }
473
474    fn back_hup(&mut self) -> SessionResult {
475        // `SniPreread` falls into the wildcard catch-all below (unconditional
476        // close + access log), same as Send/Relay/Expect: a backend HUP
477        // while still prereading is an ordinary connect-time failure with no
478        // preread-specific accounting to do (the core only ever reasons
479        // about frontend bytes).
480        match &mut self.state {
481            TcpStateMachine::Pipe(pipe) => pipe.backend_hup(&mut self.metrics),
482            _ => {
483                self.log_request();
484                SessionResult::Close
485            }
486        }
487    }
488
489    fn log_context(&self) -> LogContext<'_> {
490        LogContext {
491            session_id: self.request_id,
492            request_id: Some(self.request_id),
493            cluster_id: self.cluster_id.as_deref(),
494            backend_id: self.backend_id.as_deref(),
495        }
496    }
497
498    fn readable(&mut self) -> SessionResult {
499        // The absolute SNI-preread deadline (armed once, at session
500        // creation) must stand while undecided -- see
501        // `frontend_timeout_resets_on_readable`'s doc.
502        if frontend_timeout_resets_on_readable(&self.state)
503            && !self.container_frontend_timeout.reset()
504        {
505            error!(
506                "{} Could not reset frontend timeout on readable",
507                log_context!(self)
508            );
509        }
510        if self.backend_connected == BackendConnectionStatus::Connected
511            && !self.container_backend_timeout.reset()
512        {
513            error!(
514                "{} Could not reset backend timeout on readable",
515                log_context!(self)
516            );
517        }
518        let listener = self.listener.borrow();
519        let result = match &mut self.state {
520            TcpStateMachine::Pipe(pipe) => pipe.readable(&mut self.metrics),
521            TcpStateMachine::RelayProxyProtocol(pp) => pp.readable(&mut self.metrics),
522            TcpStateMachine::ExpectProxyProtocol(pp) => pp.readable(&mut self.metrics),
523            TcpStateMachine::SendProxyProtocol(_) => SessionResult::Continue,
524            TcpStateMachine::SniPreread(preread) => {
525                let cfg = listener.preread_config(preread.effective_max_bytes());
526                preread.readable(&mut self.metrics, &cfg)
527            }
528            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
529        };
530        drop(listener);
531
532        // Sync `cluster_id` the moment SNI preread lands a route, so
533        // `connect_to_backend`'s cluster source (`self.cluster_id.clone().or_else(...)`)
534        // sees it without waiting for a second dispatch.
535        if let TcpStateMachine::SniPreread(preread) = &self.state
536            && self.cluster_id.is_none()
537            && let Some(outcome) = preread.outcome()
538        {
539            self.cluster_id = Some(outcome.cluster.clone());
540            self.cluster_metrics_incarnation =
541                crate::metrics::capture_cluster_metrics(outcome.cluster.as_ref());
542            // Restore the listener's configured `front_timeout` THE MOMENT
543            // routing succeeds, not only once the backend connect completes
544            // (previously done only in `upgrade_sni_preread`, which can run
545            // one or more `ready()` cycles later): a slow-but-legitimate
546            // backend connect must be bounded by
547            // `front_timeout`/`connect_timeout`, never by the short
548            // `sni_preread_timeout` that only makes sense while a route
549            // decision is still pending (sozu-proxy/sozu#1290). This is
550            // also the point from which
551            // `frontend_timeout_resets_on_readable` starts resetting this
552            // container again on every future `readable()`.
553            self.container_frontend_timeout
554                .set_duration(Duration::from_secs(
555                    self.listener.borrow().config.front_timeout as u64,
556                ));
557        }
558
559        result
560    }
561
562    fn writable(&mut self) -> SessionResult {
563        match &mut self.state {
564            TcpStateMachine::Pipe(pipe) => pipe.writable(&mut self.metrics),
565            _ => SessionResult::Continue,
566        }
567    }
568
569    fn back_readable(&mut self) -> SessionResult {
570        if !self.container_frontend_timeout.reset() {
571            error!(
572                "{} Could not reset frontend timeout on back_readable",
573                log_context!(self)
574            );
575        }
576        if !self.container_backend_timeout.reset() {
577            error!(
578                "{} Could not reset backend timeout on back_readable",
579                log_context!(self)
580            );
581        }
582
583        match &mut self.state {
584            TcpStateMachine::Pipe(pipe) => pipe.backend_readable(&mut self.metrics),
585            _ => SessionResult::Continue,
586        }
587    }
588
589    fn back_writable(&mut self) -> SessionResult {
590        match &mut self.state {
591            TcpStateMachine::Pipe(pipe) => pipe.backend_writable(&mut self.metrics),
592            TcpStateMachine::RelayProxyProtocol(pp) => pp.back_writable(&mut self.metrics),
593            TcpStateMachine::SendProxyProtocol(pp) => pp.back_writable(&mut self.metrics),
594            // The FIRST backend-writable event while routed drives the
595            // upgrade out of `SniPreread` -- see
596            // `SniPreread::back_writable`'s doc and `upgrade_sni_preread`.
597            TcpStateMachine::SniPreread(preread) => preread.back_writable(),
598            TcpStateMachine::ExpectProxyProtocol(_) => SessionResult::Continue,
599            TcpStateMachine::FailedUpgrade(_) => {
600                unreachable!()
601            }
602        }
603    }
604
605    fn back_socket_mut(&mut self) -> Option<&mut MioTcpStream> {
606        match &mut self.state {
607            TcpStateMachine::Pipe(pipe) => pipe.back_socket_mut(),
608            TcpStateMachine::SendProxyProtocol(pp) => pp.back_socket_mut(),
609            TcpStateMachine::RelayProxyProtocol(pp) => pp.back_socket_mut(),
610            TcpStateMachine::SniPreread(preread) => preread.back_socket_mut(),
611            TcpStateMachine::ExpectProxyProtocol(_) => None,
612            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
613        }
614    }
615
616    pub fn upgrade(&mut self) -> SessionIsToBeClosed {
617        let new_state = match self.state.take() {
618            TcpStateMachine::SendProxyProtocol(spp) => self.upgrade_send(spp),
619            TcpStateMachine::RelayProxyProtocol(rpp) => self.upgrade_relay(rpp),
620            TcpStateMachine::ExpectProxyProtocol(epp) => self.upgrade_expect(epp),
621            TcpStateMachine::SniPreread(preread) => self.upgrade_sni_preread(preread),
622            TcpStateMachine::Pipe(_) => None,
623            TcpStateMachine::FailedUpgrade(_) => todo!(),
624        };
625
626        match new_state {
627            Some(state) => {
628                self.state = state;
629                false
630            } // The state stays FailedUpgrade, but the Session should be closed right after
631
632            None => true,
633        }
634    }
635
636    fn upgrade_send(
637        &mut self,
638        send_proxy_protocol: SendProxyProtocol<MioTcpStream>,
639    ) -> Option<TcpStateMachine> {
640        if self.backend_buffer.is_some() && self.frontend_buffer.is_some() {
641            let mut pipe = send_proxy_protocol.into_pipe(
642                self.frontend_buffer.take().unwrap(),
643                self.backend_buffer.take().unwrap(),
644                self.listener.clone(),
645            );
646
647            // `SendProxyProtocol::into_pipe` overwrites the whole readiness
648            // (clobbering the backend-writable interest `Pipe::new` armed for a
649            // non-empty inherited frontend buffer) and re-inserts only
650            // READABLE. That is harmless for a legacy `SendHeader` session
651            // (empty accumulator) but strands the coalesced payload tail
652            // carried in from `SniPreread`'s `SendHeader` branch -- the
653            // sozu-proxy/sozu#1279 close-before-flush truncation, where bytes
654            // already queued for the backend were dropped instead of flushed.
655            // Re-run the inherited-write arm by feeding the pipe's
656            // own current events back through `restore_readiness_events`; it is
657            // additive (never clears interest) and a no-op when the buffers are
658            // empty, so the legacy path is unaffected.
659            let frontend_event = pipe.frontend_readiness.event;
660            let backend_event = pipe.backend_readiness.event;
661            pipe.restore_readiness_events(frontend_event, backend_event);
662
663            pipe.set_cluster_id(self.cluster_id.clone());
664            pipe.set_cluster_metrics_incarnation(self.cluster_metrics_incarnation);
665            pipe.set_backend_address(self.backend_address);
666            // Only `Some` when this `SendProxyProtocol` was itself reached
667            // via `upgrade_sni_preread`'s `SendHeader` branch (sozu-proxy/sozu#1279)
668            // -- a legacy, non-SNI-routed `SendHeader` cluster never
669            // populates these fields, so this is a no-op for it.
670            if let Some(sni) = self.routed_sni.take() {
671                pipe.set_tls_metadata(None, None, Some(sni), self.routed_alpn_label.take());
672            }
673            // `None` for legacy sessions (keeps the bare-address tags
674            // lookup); the matched front's composed key for SNI-routed ones.
675            pipe.set_tags_key(self.tags_key.clone());
676            gauge_add!(names::protocol::PROXY_SEND, -1);
677            gauge_add!(names::protocol::TCP, 1);
678            return Some(TcpStateMachine::Pipe(pipe));
679        }
680
681        error!(
682            "{} Missing the frontend or backend buffer queue, we can't switch to a pipe",
683            log_context!(self)
684        );
685        None
686    }
687
688    fn upgrade_relay(&mut self, rpp: RelayProxyProtocol<MioTcpStream>) -> Option<TcpStateMachine> {
689        if self.backend_buffer.is_some() {
690            let mut pipe =
691                rpp.into_pipe(self.backend_buffer.take().unwrap(), self.listener.clone());
692            pipe.set_cluster_id(self.cluster_id.clone());
693            pipe.set_cluster_metrics_incarnation(self.cluster_metrics_incarnation);
694            pipe.set_backend_address(self.backend_address);
695            gauge_add!(names::protocol::PROXY_RELAY, -1);
696            gauge_add!(names::protocol::TCP, 1);
697            return Some(TcpStateMachine::Pipe(pipe));
698        }
699
700        error!(
701            "{} Missing the backend buffer queue, we can't switch to a pipe",
702            log_context!(self)
703        );
704        None
705    }
706
707    fn upgrade_expect(
708        &mut self,
709        epp: ExpectProxyProtocol<MioTcpStream>,
710    ) -> Option<TcpStateMachine> {
711        if self.frontend_buffer.is_some() && self.backend_buffer.is_some() {
712            let mut pipe = epp.into_pipe(
713                self.frontend_buffer.take().unwrap(),
714                self.backend_buffer.take().unwrap(),
715                None,
716                None,
717                self.listener.clone(),
718            );
719
720            pipe.set_cluster_id(self.cluster_id.clone());
721            pipe.set_cluster_metrics_incarnation(self.cluster_metrics_incarnation);
722            pipe.set_backend_address(self.backend_address);
723            gauge_add!(names::protocol::PROXY_EXPECT, -1);
724            gauge_add!(names::protocol::TCP, 1);
725            return Some(TcpStateMachine::Pipe(pipe));
726        }
727
728        error!(
729            "{} Missing the backend buffer queue, we can't switch to a pipe",
730            log_context!(self)
731        );
732        None
733    }
734
735    /// Dispatch out of [`TcpStateMachine::SniPreread`] once its backend has
736    /// connected, by the ROUTED cluster's `proxy_protocol` config:
737    ///
738    /// - `Some(SendHeader)` -> `SendProxyProtocol` synthesizes its OWN PPv2
739    ///   header for the backend, so any inbound PPv2 prefix this listener's
740    ///   `expect_proxy` preread already parsed (`content_offset` bytes) is
741    ///   dropped from the accumulator first: the wire order is `[synth
742    ///   PPv2][ClientHello...]`, never both headers back to back.
743    /// - `Some(ExpectHeader)` -> the inbound PPv2 prefix is consumed the
744    ///   same way (Sōzu terminates it locally; the backend never sees a
745    ///   PROXY header at all), then straight into `Pipe`.
746    /// - `Some(RelayHeader)` -> NO consume: the already-parsed inbound
747    ///   header bytes ARE the header this backend expects, replayed
748    ///   verbatim ahead of the ClientHello.
749    /// - `None` -> also consumed: a listener with `expect_proxy` but a
750    ///   `None`-proxy_protocol cluster still must not leak the stray
751    ///   inbound PPv2 prefix onto a backend that expects none -- the
752    ///   preread parsed those bytes for routing only, and a backend with
753    ///   no PROXY-protocol contract would read them as part of the TLS
754    ///   stream.
755    ///
756    /// `tcp.sni_preread.duration` is recorded on EVERY exit from this
757    /// function, including the four defensive early returns below: once
758    /// `preread.outcome()`/`preread`'s `SniPreread` value is consumed by the
759    /// `TcpStateMachine::FailedUpgrade`/`Pipe` transition its caller drives,
760    /// `close()`'s `StateMarker::SniPreread` arm can no longer reach a
761    /// `SniPreread` to read `started_at()` from, so recording later is not an
762    /// option. `tcp.sni_preread.active`, by contrast, is decremented exactly
763    /// once, on the "upgrade" exit named by the gauge's `-1 on every exit`
764    /// contract; the "reject"/"teardown" exits are each other's counterpart
765    /// in `SniPreread::handle_output` (metric only) and `TcpSession::close`'s
766    /// `StateMarker::SniPreread` arm (gauge).
767    /// Shared abort path for `upgrade_sni_preread`'s early-return guards:
768    /// logs `reason` through the same envelope as the rest of this module,
769    /// then records `tcp.sni_preread.duration` -- see the long comment on
770    /// `upgrade_sni_preread` for why that metric must fire on every exit --
771    /// before returning `None`.
772    fn abort_sni_preread_upgrade(
773        &self,
774        preread: &SniPreread<MioTcpStream>,
775        reason: &str,
776    ) -> Option<TcpStateMachine> {
777        error!("{} {}", log_context!(self), reason);
778        time!(
779            names::tcp::sni_preread::DURATION,
780            preread.started_at().elapsed().as_millis() as i64
781        );
782        None
783    }
784
785    fn upgrade_sni_preread(
786        &mut self,
787        mut preread: SniPreread<MioTcpStream>,
788    ) -> Option<TcpStateMachine> {
789        // Every early return below (a route decision missing, or the
790        // backend socket/token/buffer not yet wired) must happen BEFORE the
791        // `tcp.sni_preread.active` gauge is touched: `close()`'s
792        // `StateMarker::SniPreread` arm runs unconditionally whenever
793        // `self.state` is still (or, via `FailedUpgrade`, was last)
794        // `SniPreread` -- decrementing here AND there for the same session
795        // would underflow the gauge on this (defensive, should-never-happen)
796        // failure path. The gauge is deferred to `close()` on these paths,
797        // but the duration is NOT: it is recorded right before each `return
798        // None` below, symmetric with the success path's `time!` call.
799        let Some(outcome) = preread.outcome().cloned() else {
800            return self.abort_sni_preread_upgrade(
801                &preread,
802                "upgrade_sni_preread called before a route decision",
803            );
804        };
805        let Some(backend_socket) = preread.backend.take() else {
806            return self.abort_sni_preread_upgrade(
807                &preread,
808                "SNI preread upgrade with no backend socket set",
809            );
810        };
811        let Some(backend_token) = preread.backend_token else {
812            return self.abort_sni_preread_upgrade(
813                &preread,
814                "SNI preread upgrade with no backend token set",
815            );
816        };
817        let Some(back_buffer) = self.backend_buffer.take() else {
818            return self.abort_sni_preread_upgrade(
819                &preread,
820                "SNI preread upgrade with no backend buffer queued",
821            );
822        };
823
824        gauge_add!(names::tcp::sni_preread::ACTIVE, -1);
825        time!(
826            names::tcp::sni_preread::DURATION,
827            preread.started_at().elapsed().as_millis() as i64
828        );
829
830        // `TcpSession::readable` normally captured this route (and its
831        // metrics incarnation) already; only capture when it did not, so the
832        // incarnation stays the one observed at routing time.
833        if self.cluster_id.as_ref() != Some(&outcome.cluster) {
834            self.cluster_id = Some(outcome.cluster.clone());
835            self.cluster_metrics_incarnation =
836                crate::metrics::capture_cluster_metrics(outcome.cluster.as_ref());
837        }
838        // `container_frontend_timeout` is NOT restored here anymore: by the
839        // time this runs, `TcpSession::readable`'s route-capture block has
840        // already restored it to the listener's configured `front_timeout`
841        // the moment the route decision first became visible (potentially
842        // one or more `ready()` cycles before this upgrade, while the
843        // backend was still connecting) -- see that block's doc
844        // (sozu-proxy/sozu#1290). Restoring it again here
845        // would just re-arm the same duration a second time.
846        // Access-log tagging: stash the routed SNI/ALPN for
847        // whichever of the four `proxy_protocol` branches below eventually
848        // reaches `Pipe` -- immediately via `build_pipe_from_preread` for
849        // `Expect`/`Relay`/`None`, or one `ready()` cycle later via
850        // `upgrade_send` for `SendHeader` (see that method and the
851        // `routed_sni` field doc).
852        self.routed_sni = Some(outcome.sni.clone());
853        self.routed_alpn_label = known_alpn_label(&outcome.alpn);
854        // Rebuild the MATCHED front's tags key from the route decision's
855        // identity (`matched_sni_pattern` is the trie key — the configured
856        // pattern, not the client's concrete SNI — and `matched_alpn` the
857        // winning matcher), so the access log emits the tags of the front
858        // that actually routed this session, not whichever front was added
859        // last. Must compose the same key `add_tcp_front` stored — see
860        // `sni_tags_key`'s canonical-form doc.
861        self.tags_key = Some(sni_tags_key(
862            self.listener.borrow().get_addr(),
863            &outcome.matched_sni_pattern,
864            &alpn_matcher_protocols(&outcome.matched_alpn),
865        ));
866
867        let proxy_protocol = self
868            .proxy
869            .borrow()
870            .configs
871            .get(&outcome.cluster)
872            .and_then(|c| c.proxy_protocol);
873
874        let frontend_event = preread.frontend_readiness.event;
875        let backend_event = preread.backend_readiness.event;
876        let mut frontend_buffer = preread.frontend_buffer;
877        let frontend = preread.frontend;
878        let frontend_token = preread.frontend_token;
879        let request_id = preread.request_id;
880
881        match proxy_protocol {
882            Some(ProxyProtocolConfig::SendHeader) => {
883                frontend_buffer.consume(outcome.content_offset);
884                self.frontend_buffer = Some(frontend_buffer);
885                self.backend_buffer = Some(back_buffer);
886                gauge_add!(names::protocol::PROXY_SEND, 1);
887                let mut spp = SendProxyProtocol::new(
888                    frontend,
889                    frontend_token,
890                    request_id,
891                    Some(backend_socket),
892                );
893                spp.frontend_readiness.event = frontend_event;
894                spp.backend_readiness.event = backend_event;
895                spp.set_back_token(backend_token);
896                spp.set_back_connected(BackendConnectionStatus::Connected);
897                Some(TcpStateMachine::SendProxyProtocol(spp))
898            }
899            Some(ProxyProtocolConfig::ExpectHeader) | None => {
900                frontend_buffer.consume(outcome.content_offset);
901                Some(self.build_pipe_from_preread(
902                    back_buffer,
903                    frontend_buffer,
904                    frontend,
905                    frontend_token,
906                    frontend_event,
907                    backend_event,
908                    backend_socket,
909                    backend_token,
910                    request_id,
911                    outcome.proxy_source,
912                    outcome.cluster,
913                ))
914            }
915            Some(ProxyProtocolConfig::RelayHeader) => Some(self.build_pipe_from_preread(
916                back_buffer,
917                frontend_buffer,
918                frontend,
919                frontend_token,
920                frontend_event,
921                backend_event,
922                backend_socket,
923                backend_token,
924                request_id,
925                outcome.proxy_source,
926                outcome.cluster,
927            )),
928        }
929    }
930
931    #[allow(clippy::too_many_arguments)]
932    fn build_pipe_from_preread(
933        &mut self,
934        back_buffer: Checkout,
935        frontend_buffer: Checkout,
936        frontend: MioTcpStream,
937        frontend_token: Token,
938        frontend_event: Ready,
939        backend_event: Ready,
940        backend_socket: MioTcpStream,
941        backend_token: Token,
942        request_id: Ulid,
943        proxy_source: Option<SocketAddr>,
944        cluster_id: ClusterId,
945    ) -> TcpStateMachine {
946        let addr = proxy_source.or(self.frontend_address);
947        let mut pipe = Pipe::new(
948            back_buffer,
949            self.backend_id.clone(),
950            Some(backend_socket),
951            None,
952            None,
953            None,
954            Some(cluster_id),
955            frontend_buffer,
956            frontend_token,
957            frontend,
958            self.listener.clone(),
959            Protocol::TCP,
960            request_id,
961            request_id,
962            addr,
963            WebSocketContext::Tcp,
964        );
965        pipe.set_cluster_metrics_incarnation(self.cluster_metrics_incarnation);
966        // `Pipe::new` armed backend-writable for the inherited frontend
967        // accumulator (the ClientHello + any coalesced payload) via
968        // `arm_inherited_buffer_writes`. Restore the preread's readiness
969        // events through `restore_readiness_events` rather than a bare
970        // `pipe.frontend_readiness.event = …` / `pipe.backend_readiness.event
971        // = …` pair: it sets both `.event`s and THEN re-runs the inherited
972        // arm, so the synthetic backend-writable event survives even in the
973        // case where the restored `backend_event` does not itself carry
974        // WRITABLE (the byte-for-byte drain of the accumulator must not depend
975        // on that). The eventual flush-on-close of that accumulator is
976        // guaranteed by `Pipe::readable`'s half-close drain (sozu-proxy/sozu#1279).
977        pipe.restore_readiness_events(frontend_event, backend_event);
978        pipe.set_back_token(backend_token);
979        pipe.set_backend_address(self.backend_address);
980        // Access-log tagging: reaching `Pipe` straight from
981        // `SniPreread` (Expect/Relay/None) -- unlike `SendHeader`, which
982        // detours through `SendProxyProtocol` first (see `upgrade_send`).
983        if let Some(sni) = self.routed_sni.take() {
984            pipe.set_tls_metadata(None, None, Some(sni), self.routed_alpn_label.take());
985        }
986        // The matched front's composed tags key (always `Some` here — this
987        // is only reachable from `upgrade_sni_preread`, which just set it).
988        pipe.set_tags_key(self.tags_key.clone());
989        gauge_add!(names::protocol::TCP, 1);
990        TcpStateMachine::Pipe(pipe)
991    }
992
993    fn front_readiness(&mut self) -> &mut Readiness {
994        match &mut self.state {
995            TcpStateMachine::Pipe(pipe) => &mut pipe.frontend_readiness,
996            TcpStateMachine::SendProxyProtocol(pp) => &mut pp.frontend_readiness,
997            TcpStateMachine::RelayProxyProtocol(pp) => &mut pp.frontend_readiness,
998            TcpStateMachine::ExpectProxyProtocol(pp) => &mut pp.frontend_readiness,
999            TcpStateMachine::SniPreread(preread) => &mut preread.frontend_readiness,
1000            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
1001        }
1002    }
1003
1004    fn back_readiness(&mut self) -> Option<&mut Readiness> {
1005        match &mut self.state {
1006            TcpStateMachine::Pipe(pipe) => Some(&mut pipe.backend_readiness),
1007            TcpStateMachine::SendProxyProtocol(pp) => Some(&mut pp.backend_readiness),
1008            TcpStateMachine::RelayProxyProtocol(pp) => Some(&mut pp.backend_readiness),
1009            TcpStateMachine::SniPreread(preread) => Some(&mut preread.backend_readiness),
1010            TcpStateMachine::ExpectProxyProtocol(_) => None,
1011            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
1012        }
1013    }
1014
1015    fn set_back_socket(&mut self, socket: MioTcpStream) {
1016        match &mut self.state {
1017            TcpStateMachine::Pipe(pipe) => pipe.set_back_socket(socket),
1018            TcpStateMachine::SendProxyProtocol(pp) => pp.set_back_socket(socket),
1019            TcpStateMachine::RelayProxyProtocol(pp) => pp.set_back_socket(socket),
1020            TcpStateMachine::SniPreread(preread) => preread.set_back_socket(socket),
1021            TcpStateMachine::ExpectProxyProtocol(_) => {
1022                error!(
1023                    "{} We should not set the back socket for the expect proxy protocol",
1024                    log_context!(self)
1025                );
1026                panic!(
1027                    "{} We should not set the back socket for the expect proxy protocol",
1028                    log_context!(self)
1029                );
1030            }
1031            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
1032        }
1033    }
1034
1035    fn set_back_token(&mut self, token: Token) {
1036        // The frontend must own a token distinct from the backend's: the two
1037        // index different slab slots, so wiring the same token to both would
1038        // alias two sessions onto one slot.
1039        debug_assert_ne!(
1040            token, self.frontend_token,
1041            "backend token must differ from the frontend token"
1042        );
1043        self.backend_token = Some(token);
1044
1045        match &mut self.state {
1046            TcpStateMachine::Pipe(pipe) => pipe.set_back_token(token),
1047            TcpStateMachine::SendProxyProtocol(pp) => pp.set_back_token(token),
1048            TcpStateMachine::SniPreread(preread) => preread.set_back_token(token),
1049            TcpStateMachine::RelayProxyProtocol(pp) => pp.set_back_token(token),
1050            TcpStateMachine::ExpectProxyProtocol(_) => self.backend_token = Some(token),
1051            TcpStateMachine::FailedUpgrade(_) => unreachable!(),
1052        }
1053
1054        // Postcondition: the session now owns exactly the token it was asked
1055        // to register — every arm above (including the Expect arm, which only
1056        // stores the session-side token) leaves `backend_token == Some(token)`.
1057        debug_assert_eq!(
1058            self.backend_token,
1059            Some(token),
1060            "set_back_token must leave the session owning the registered token"
1061        );
1062    }
1063
1064    fn set_backend_id(&mut self, id: String) {
1065        self.backend_id = Some(id.clone());
1066        if let TcpStateMachine::Pipe(pipe) = &mut self.state {
1067            pipe.set_backend_id(Some(id));
1068        }
1069    }
1070
1071    /// Record the dialed backend's address; a session that is not a `Pipe`
1072    /// yet hands it over at upgrade time instead.
1073    fn set_backend_address(&mut self, address: SocketAddr) {
1074        self.backend_address = Some(address);
1075        if let TcpStateMachine::Pipe(pipe) = &mut self.state {
1076            pipe.set_backend_address(Some(address));
1077        }
1078    }
1079
1080    fn back_connected(&self) -> BackendConnectionStatus {
1081        self.backend_connected
1082    }
1083
1084    fn set_back_connected(&mut self, status: BackendConnectionStatus) {
1085        let last = self.backend_connected;
1086        // Transitioning INTO `Connected` bumps the backend-connection gauge by
1087        // exactly +1. Doing so from an already-`Connected` state would
1088        // double-count (gauge drift that only `close_backend`'s single -1
1089        // would later reconcile, leaving the gauge permanently +1). The
1090        // promotion always comes from a `Connecting` (the normal handshake
1091        // completion in `ready_inner`) — never from `Connected` itself.
1092        debug_assert!(
1093            status != BackendConnectionStatus::Connected
1094                || last != BackendConnectionStatus::Connected,
1095            "set_back_connected(Connected) must not run on an already-Connected backend (gauge would double-count)"
1096        );
1097        self.backend_connected = status;
1098
1099        // Postcondition: the requested status is now in effect.
1100        debug_assert_eq!(
1101            self.backend_connected, status,
1102            "set_back_connected must record the requested status"
1103        );
1104
1105        if status == BackendConnectionStatus::Connected {
1106            gauge_add!(names::backend::CONNECTIONS, 1);
1107            gauge_add!(
1108                names::backend::CONNECTIONS_PER_BACKEND,
1109                1,
1110                self.cluster_id.as_deref(),
1111                self.metrics.backend_id.as_deref(),
1112                self.cluster_metrics_incarnation
1113            );
1114
1115            // the back timeout was of connect_timeout duration before,
1116            // now that we're connected, move to backend_timeout duration
1117            self.container_backend_timeout
1118                .set_duration(self.configured_backend_timeout);
1119            self.container_frontend_timeout.reset();
1120
1121            if let TcpStateMachine::SendProxyProtocol(spp) = &mut self.state {
1122                spp.set_back_connected(BackendConnectionStatus::Connected);
1123            }
1124
1125            if let Some(backend) = self.backend.as_ref() {
1126                let mut backend = backend.borrow_mut();
1127
1128                if backend.retry_policy.is_down() {
1129                    incr!(
1130                        names::backend::UP,
1131                        self.cluster_id.as_deref(),
1132                        self.metrics.backend_id.as_deref(),
1133                        self.cluster_metrics_incarnation
1134                    );
1135                    gauge!(
1136                        names::backend::AVAILABLE,
1137                        1,
1138                        self.cluster_id.as_deref(),
1139                        self.metrics.backend_id.as_deref(),
1140                        self.cluster_metrics_incarnation
1141                    );
1142                    info!(
1143                        "{} backend server {} at {} is up",
1144                        log_context!(self),
1145                        backend.backend_id,
1146                        backend.address
1147                    );
1148                    push_event(Event {
1149                        kind: EventKind::BackendUp as i32,
1150                        backend_id: Some(backend.backend_id.to_owned()),
1151                        address: Some(backend.address.into()),
1152                        cluster_id: None,
1153                        metric_detail: None,
1154                    });
1155                }
1156
1157                if let BackendConnectionStatus::Connecting(start) = last {
1158                    let now = Instant::now();
1159                    backend.set_connection_time(now - start, now);
1160                }
1161
1162                //successful connection, rest failure counter
1163                backend.failures = 0;
1164                backend.retry_policy.succeed(Instant::now());
1165            }
1166        }
1167    }
1168
1169    fn remove_backend(&mut self) {
1170        if let Some(backend) = self.backend.take() {
1171            (*backend.borrow_mut()).dec_connections();
1172        }
1173
1174        self.backend_token = None;
1175
1176        // Postcondition: the backend handle and its token are torn down
1177        // together — neither may outlive the other (a dangling token would
1178        // leave a stale slab reference; a dangling handle would over-count
1179        // backend connections).
1180        debug_assert!(
1181            self.backend.is_none(),
1182            "remove_backend must release the backend handle"
1183        );
1184        debug_assert!(
1185            self.backend_token.is_none(),
1186            "remove_backend must clear the backend token"
1187        );
1188    }
1189
1190    fn fail_backend_connection(&mut self) {
1191        if let Some(backend) = self.backend.as_ref() {
1192            let backend = &mut *backend.borrow_mut();
1193            backend.failures += 1;
1194
1195            let already_unavailable = backend.retry_policy.is_down();
1196            // The TCP proxy is not simulated: it keeps the wall clock and the
1197            // thread-local generator at its own edge (#1684).
1198            backend.retry_policy.fail(Instant::now(), &mut rand::rng());
1199            incr!(
1200                names::backend::CONNECTIONS_ERROR,
1201                self.cluster_id.as_deref(),
1202                self.metrics.backend_id.as_deref(),
1203                self.cluster_metrics_incarnation
1204            );
1205            if !already_unavailable && backend.retry_policy.is_down() {
1206                error!(
1207                    "{} backend server {} at {} is down",
1208                    log_context!(self),
1209                    backend.backend_id,
1210                    backend.address
1211                );
1212                incr!(
1213                    names::backend::DOWN,
1214                    self.cluster_id.as_deref(),
1215                    self.metrics.backend_id.as_deref(),
1216                    self.cluster_metrics_incarnation
1217                );
1218                gauge!(
1219                    names::backend::AVAILABLE,
1220                    0,
1221                    self.cluster_id.as_deref(),
1222                    self.metrics.backend_id.as_deref(),
1223                    self.cluster_metrics_incarnation
1224                );
1225
1226                push_event(Event {
1227                    kind: EventKind::BackendDown as i32,
1228                    backend_id: Some(backend.backend_id.to_owned()),
1229                    address: Some(backend.address.into()),
1230                    cluster_id: None,
1231                    metric_detail: None,
1232                });
1233            }
1234        }
1235    }
1236
1237    pub fn test_back_socket(&mut self) -> SessionIsToBeClosed {
1238        match self.back_socket_mut() {
1239            Some(ref mut s) => {
1240                let mut tmp = [0u8; 1];
1241                let res = s.peek(&mut tmp[..]);
1242
1243                match res {
1244                    // if the socket is half open, it will report 0 bytes read (EOF)
1245                    Ok(0) => false,
1246                    Ok(_) => true,
1247                    Err(e) => matches!(e.kind(), std::io::ErrorKind::WouldBlock),
1248                }
1249            }
1250            None => false,
1251        }
1252    }
1253
1254    pub fn cancel_timeouts(&mut self) {
1255        self.container_frontend_timeout.cancel();
1256        self.container_backend_timeout.cancel();
1257    }
1258
1259    /// Full cross-field invariant sweep for the TCP session state machine.
1260    ///
1261    /// Run as a run-to-completion postcondition at the END of `ready()` (the
1262    /// only public entry point that drives the front/back token + readiness
1263    /// state machine). These are OUR-logic invariants — never reachable from
1264    /// hostile traffic — so a violation is a bug in Sōzu, not a malformed
1265    /// peer. Compiled out in release.
1266    #[cfg(debug_assertions)]
1267    fn check_invariants(&self) {
1268        // Connection-attempt budget: every retry path increments
1269        // `connection_attempt` and `connect_to_backend` refuses once the
1270        // counter reaches the cluster's `max_connection_attempts`, so the
1271        // value never exceeds the largest budget that can be configured (and
1272        // resets to 0 on success). The live budget itself is no bound here:
1273        // `SetMaxConnectionAttempts` or an `AddCluster` may lower it while a
1274        // session is retrying.
1275        debug_assert!(
1276            u32::from(self.connection_attempt)
1277                <= sozu_command::config::MAX_CONNECTION_ATTEMPTS_UPPER_BOUND,
1278            "connection_attempt ({}) must never exceed the largest attempt budget",
1279            self.connection_attempt,
1280        );
1281
1282        // Token ownership: a fully-connected backend always owns a backend
1283        // token (set by `set_back_token` during `connect_to_backend`, before
1284        // the status can ever flip to `Connected`). The `Connecting` phase is
1285        // deliberately excluded: there is a transient window inside
1286        // `connect_to_backend` where the status is `Connecting` but the token
1287        // has not been wired yet — that window never spans a `ready()`
1288        // boundary, so the postcondition still holds here.
1289        if self.backend_connected == BackendConnectionStatus::Connected {
1290            debug_assert!(
1291                self.backend_token.is_some(),
1292                "a Connected backend must own a backend token"
1293            );
1294        }
1295
1296        // A live backend handle implies the matching token is present: the
1297        // two are wired together in `connect_to_backend` and torn down
1298        // together in `remove_backend` (which clears the token) — they must
1299        // never drift apart. (For the pure-TCP proxy `backend` is currently
1300        // always `None`, so this is a guard against a future regression that
1301        // starts populating it without the token.)
1302        if self.backend.is_some() {
1303            debug_assert!(
1304                self.backend_token.is_some(),
1305                "a live backend handle must have a backend token"
1306            );
1307        }
1308
1309        // Once the session has been closed it is terminal: the backend has
1310        // been released and the per-(cluster, source-IP) slot untracked.
1311        if self.has_been_closed {
1312            debug_assert!(
1313                self.backend.is_none(),
1314                "a closed session must have released its backend handle"
1315            );
1316            debug_assert!(
1317                !self.cluster_ip_tracked,
1318                "a closed session must have untracked its (cluster, source-IP) slot"
1319            );
1320        }
1321    }
1322
1323    /// Attempt a fresh backend connect, exactly like `ready_inner`'s
1324    /// top-of-function gate -- but callable a second time from inside the
1325    /// dispatch loop. A `SniPreread` session's route decision can complete
1326    /// INSIDE `readable()`'s own dispatch (the SAME `ready_inner` call), and
1327    /// without a second attempt right after that dispatch the session would
1328    /// stall until an unrelated readiness event re-entered `ready_inner`.
1329    ///
1330    /// A no-op whenever `back_connected() != NotConnected` (already
1331    /// attempted, or backend already up), the state is a NOT-YET-ROUTED
1332    /// `SniPreread` (the cluster -- and therefore the backend to dial -- is
1333    /// unknown until `SniPrereadCore` decides), or the state is still
1334    /// `ExpectProxyProtocol` (see below), so this changes nothing for any
1335    /// other pre-existing state/path.
1336    ///
1337    /// `ExpectProxyProtocol` has NO backend side at all: `back_readiness`
1338    /// returns `None` for it and `set_back_socket` panics outright ("We
1339    /// should not set the back socket for the expect proxy protocol"). The
1340    /// backend is dialed from the `Pipe` that `upgrade_expect` installs once
1341    /// the inbound PROXY header has been parsed -- the same `ready()` pass
1342    /// re-enters `ready_inner` after the upgrade, so nothing is deferred
1343    /// beyond the header. Without this arm the top-of-`ready_inner` connect
1344    /// gate fired on the FIRST readiness event of a freshly accepted socket
1345    /// (the WRITABLE epoll reports before the client has sent a byte) and
1346    /// drove `connect_to_backend` straight into that panic, killing the
1347    /// worker and every other session on it.
1348    fn attempt_backend_connect_if_needed(
1349        &mut self,
1350        session: &Rc<RefCell<dyn ProxySession>>,
1351    ) -> Option<SessionResult> {
1352        if self.back_connected() != BackendConnectionStatus::NotConnected {
1353            return None;
1354        }
1355        if matches!(&self.state, TcpStateMachine::SniPreread(preread) if !preread.is_routed()) {
1356            return None;
1357        }
1358        if matches!(&self.state, TcpStateMachine::ExpectProxyProtocol(_)) {
1359            return None;
1360        }
1361
1362        let connection_result = self.connect_to_backend(session.clone());
1363        if let Err(err) = &connection_result {
1364            match err {
1365                // Already logged at warn! + metered at the retry-budget
1366                // gate in connect_to_backend; avoid double-emission.
1367                BackendConnectionError::MaxConnectionRetries(_) => trace!(
1368                    "{} Error connecting to backend: {}",
1369                    log_context!(self),
1370                    err
1371                ),
1372                _ => warn!(
1373                    "{} Error connecting to backend: {}",
1374                    log_context!(self),
1375                    err
1376                ),
1377            }
1378        }
1379        handle_connection_result(connection_result)
1380    }
1381
1382    fn ready_inner(&mut self, session: Rc<RefCell<dyn ProxySession>>) -> SessionResult {
1383        let mut counter = 0;
1384
1385        let back_connected = self.back_connected();
1386        if back_connected.is_connecting() {
1387            // A `Connecting` backend always carries a backend readiness:
1388            // `connect_to_backend` is the only writer of that status, and it
1389            // never runs for the two states `back_readiness` has none for
1390            // (`ExpectProxyProtocol`, guarded in
1391            // `attempt_backend_connect_if_needed`; `FailedUpgrade`, which
1392            // `close()` short-circuits at its `state.failed()` gate). Read it
1393            // once instead of unwrapping twice, and treat the impossible
1394            // `None` as the broken invariant it is: a backend handshake with
1395            // no readiness slot can never complete, so close rather than
1396            // panic the worker.
1397            let Some(back_event) = self.back_readiness().map(|readiness| readiness.event) else {
1398                error!(
1399                    "{} backend is connecting but the session state carries no backend readiness, closing",
1400                    log_context!(self)
1401                );
1402                return SessionResult::Close;
1403            };
1404            if back_event.is_hup() && !self.test_back_socket() {
1405                //retry connecting the backend
1406                debug!(
1407                    "{} error connecting to backend, trying again",
1408                    log_context!(self)
1409                );
1410                self.connection_attempt += 1;
1411                self.fail_backend_connection();
1412
1413                // trigger a backend reconnection
1414                self.close_backend();
1415                let connection_result = self.connect_to_backend(session.clone());
1416                if let Err(err) = &connection_result {
1417                    match err {
1418                        // Already logged at warn! + metered at the retry-budget
1419                        // gate in connect_to_backend; avoid double-emission.
1420                        BackendConnectionError::MaxConnectionRetries(_) => trace!(
1421                            "{} Error connecting to backend: {}",
1422                            log_context!(self),
1423                            err
1424                        ),
1425                        _ => warn!(
1426                            "{} Error connecting to backend: {}",
1427                            log_context!(self),
1428                            err
1429                        ),
1430                    }
1431                }
1432
1433                if let Some(state_result) = handle_connection_result(connection_result) {
1434                    return state_result;
1435                }
1436            } else if back_event != Ready::EMPTY {
1437                self.connection_attempt = 0;
1438                self.set_back_connected(BackendConnectionStatus::Connected);
1439            }
1440        } else if back_connected == BackendConnectionStatus::NotConnected
1441            && let Some(state_result) = self.attempt_backend_connect_if_needed(&session)
1442        {
1443            return state_result;
1444        }
1445
1446        // A client that half-closes right behind its PROXY-v2 header and
1447        // payload can deliver both as ONE `READABLE | HUP` event (Linux
1448        // coalesces `EPOLLIN | EPOLLRDHUP`). `front_hup` closes every
1449        // pre-`Pipe` state (`SniPreread` without an access log), so
1450        // honouring HUP first would drop a complete, already-queued header
1451        // (sozu-proxy/sozu#1823).
1452        // Drain the expect state's readable bytes first: a parsed header
1453        // returns `Upgrade` and the recursive `ready()` hands the still-set
1454        // HUP to `Pipe::frontend_hup`, which keeps the session alive while
1455        // request bytes remain in the kernel or the pipe's buffers. An empty
1456        // or truncated header
1457        // ends with READABLE cleared (`readable` removes it on every
1458        // zero-byte read) or a `Close`, so a dead client still closes,
1459        // through `readable`'s own zero-byte arm or the HUP check below.
1460        if let TcpStateMachine::ExpectProxyProtocol(_) = self.state {
1461            while (self.front_readiness().interest & self.front_readiness().event).is_readable()
1462                && self.front_readiness().event.is_hup()
1463                && counter < MAX_LOOP_ITERATIONS
1464            {
1465                let session_result = self.readable();
1466                if session_result != SessionResult::Continue {
1467                    return session_result;
1468                }
1469                counter += 1;
1470            }
1471        }
1472
1473        if self.front_readiness().event.is_hup() {
1474            let session_result = self.front_hup();
1475            if session_result != SessionResult::Continue {
1476                return session_result;
1477            }
1478            // `front_hup` drained in-flight request bytes and wants the
1479            // session kept alive (`Pipe::frontend_hup`'s in-flight branch):
1480            // the client already sent FIN, so under edge-triggered epoll no
1481            // further frontend event will ever arrive -- returning here
1482            // would stall the session forever waiting for a wake-up that
1483            // never comes. Clear the now-consumed HUP bit and fall through
1484            // into the loop below so `readable` (drains the kernel tail) and
1485            // `back_writable` (flushes `frontend_buffer`) can run
1486            // synchronously in this same pass, exactly how a backend HUP is
1487            // already handled inside the loop.
1488            self.front_readiness().event.remove(Ready::HUP);
1489        }
1490
1491        while counter < MAX_LOOP_ITERATIONS {
1492            let front_interest = self.front_readiness().interest & self.front_readiness().event;
1493            let back_interest = self
1494                .back_readiness()
1495                .map(|r| r.interest & r.event)
1496                .unwrap_or(Ready::EMPTY);
1497
1498            trace!(
1499                "{} Frontend interest({:?}) and backend interest({:?})",
1500                log_context!(self),
1501                front_interest,
1502                back_interest
1503            );
1504
1505            if front_interest == Ready::EMPTY && back_interest == Ready::EMPTY {
1506                break;
1507            }
1508
1509            if self
1510                .back_readiness()
1511                .map(|r| r.event.is_hup())
1512                .unwrap_or(false)
1513                && self.front_readiness().interest.is_writable()
1514                && !self.front_readiness().event.is_writable()
1515            {
1516                break;
1517            }
1518
1519            if front_interest.is_readable() {
1520                let session_result = self.readable();
1521                if session_result != SessionResult::Continue {
1522                    return session_result;
1523                }
1524                // A `SniPreread` route decision can complete INSIDE this
1525                // very `readable()` call; without a second attempt here the
1526                // session would stall until an unrelated readiness event
1527                // re-entered `ready_inner` to reach the top-of-function
1528                // connect gate. A no-op for every other state/backend
1529                // status (see `attempt_backend_connect_if_needed`'s guard).
1530                if let Some(state_result) = self.attempt_backend_connect_if_needed(&session) {
1531                    return state_result;
1532                }
1533            }
1534
1535            if back_interest.is_writable() {
1536                let session_result = self.back_writable();
1537                if session_result != SessionResult::Continue {
1538                    return session_result;
1539                }
1540            }
1541
1542            if back_interest.is_readable() {
1543                let session_result = self.back_readable();
1544                if session_result != SessionResult::Continue {
1545                    return session_result;
1546                }
1547            }
1548
1549            if front_interest.is_writable() {
1550                let session_result = self.writable();
1551                if session_result != SessionResult::Continue {
1552                    return session_result;
1553                }
1554            }
1555
1556            if back_interest.is_hup() {
1557                let session_result = self.back_hup();
1558                if session_result != SessionResult::Continue {
1559                    return session_result;
1560                }
1561            }
1562
1563            if front_interest.is_error() {
1564                error!(
1565                    "{} Frontend socket error, disconnecting",
1566                    log_context!(self)
1567                );
1568                self.front_readiness().interest = Ready::EMPTY;
1569                if let Some(r) = self.back_readiness() {
1570                    r.interest = Ready::EMPTY;
1571                }
1572
1573                return SessionResult::Close;
1574            }
1575
1576            if back_interest.is_error() && self.back_hup() == SessionResult::Close {
1577                self.front_readiness().interest = Ready::EMPTY;
1578                if let Some(r) = self.back_readiness() {
1579                    r.interest = Ready::EMPTY;
1580                }
1581
1582                error!("{} backend socket error, disconnecting", log_context!(self));
1583                return SessionResult::Close;
1584            }
1585
1586            counter += 1;
1587        }
1588
1589        if counter >= MAX_LOOP_ITERATIONS {
1590            error!(
1591                "{} Handling session went through {} iterations, there's a probable infinite loop bug, closing the connection",
1592                log_context!(self),
1593                MAX_LOOP_ITERATIONS
1594            );
1595
1596            incr!(names::tcp::INFINITE_LOOP_ERROR);
1597
1598            let front_interest = self.front_readiness().interest & self.front_readiness().event;
1599            let back_interest = self
1600                .back_readiness()
1601                .map(|r| r.interest & r.event)
1602                .unwrap_or(Ready::EMPTY);
1603
1604            let back = self.back_readiness().cloned();
1605
1606            error!(
1607                "{} readiness: front {:?} / back {:?} | front: {:?} | back: {:?} ",
1608                log_context!(self),
1609                self.front_readiness(),
1610                back,
1611                front_interest,
1612                back_interest
1613            );
1614
1615            self.print_session();
1616
1617            return SessionResult::Close;
1618        }
1619
1620        SessionResult::Continue
1621    }
1622
1623    /// TCP session closes its backend on its own, without defering this task to the state
1624    fn close_backend(&mut self) {
1625        // No `EPOLL_CTL_DEL` for the backend socket: its last `close(2)`
1626        // takes it out of the epoll set, and it comes before the next
1627        // `epoll_wait`, so no later result carries the freed token for it.
1628        // The result being dispatched may still hold one of its events, and
1629        // another session can take the freed key before that event's turn:
1630        // `Server::ready_if_still_owned` (`lib/src/server.rs`) drops it.
1631        // From `close()`, the server drops the session right after. From
1632        // `ready_inner`'s retry, `connect_to_backend` either replaces the
1633        // socket (`set_back_socket` drops the old one) or fails, and every
1634        // failure closes the session (`handle_connection_result`). No session
1635        // socket is ever duplicated (no `dup`, `try_clone` or worker fork).
1636        if let Some(token) = self.backend_token
1637            && self.back_socket_mut().is_some()
1638        {
1639            self.proxy
1640                .borrow()
1641                .sessions
1642                .borrow_mut()
1643                .slab
1644                .try_remove(token.0);
1645        }
1646        self.remove_backend();
1647
1648        let back_connected = self.back_connected();
1649        if back_connected != BackendConnectionStatus::NotConnected {
1650            if let Some(r) = self.back_readiness() {
1651                r.event = Ready::EMPTY;
1652            }
1653
1654            let log_context = log_context!(self);
1655            if let Some(sock) = self.back_socket_mut() {
1656                // TCP-only backend in the pure-TCP proxy: no outbound TLS
1657                // buffer to truncate, so `Shutdown::Both` is the right call.
1658                // If the TCP listener ever gains an inline TLS upgrade,
1659                // switch to `Shutdown::Write` here.
1660                if let Err(e) = sock.shutdown(Shutdown::Both)
1661                    && e.kind() != ErrorKind::NotConnected
1662                {
1663                    error!(
1664                        "{} Error closing back socket({:?}): {:?}",
1665                        log_context, sock, e
1666                    );
1667                }
1668            }
1669        }
1670
1671        // The -1 here pairs with the +1 in `set_back_connected(Connected)`:
1672        // we decrement the gauge exactly once, iff this session had actually
1673        // reached `Connected`. A `Connecting`/`NotConnected` backend never
1674        // bumped the gauge, so it must not decrement it either — that
1675        // asymmetry would underflow the gauge (a correctness bug, never a
1676        // rounding issue).
1677        if back_connected == BackendConnectionStatus::Connected {
1678            gauge_add!(names::backend::CONNECTIONS, -1);
1679            gauge_add!(
1680                names::backend::CONNECTIONS_PER_BACKEND,
1681                -1,
1682                self.cluster_id.as_deref(),
1683                self.metrics.backend_id.as_deref(),
1684                self.cluster_metrics_incarnation
1685            );
1686        }
1687
1688        self.set_back_connected(BackendConnectionStatus::NotConnected);
1689
1690        // Postcondition: the backend is fully torn down — `remove_backend`
1691        // cleared the token/handle above and the status is now `NotConnected`,
1692        // so a subsequent `connect_to_backend` starts from a clean slate.
1693        debug_assert_eq!(
1694            self.backend_connected,
1695            BackendConnectionStatus::NotConnected,
1696            "close_backend must leave the backend NotConnected"
1697        );
1698        debug_assert!(
1699            self.backend_token.is_none(),
1700            "close_backend must clear the backend token"
1701        );
1702    }
1703
1704    fn connect_to_backend(
1705        &mut self,
1706        session_rc: Rc<RefCell<dyn ProxySession>>,
1707    ) -> Result<BackendConnectAction, BackendConnectionError> {
1708        // Precondition: the retry budget can sit AT the ceiling (the gate
1709        // below converts that into `MaxConnectionRetries`) but the increment
1710        // in `ready_inner` must never have pushed it past the largest budget
1711        // that can be configured.
1712        debug_assert!(
1713            u32::from(self.connection_attempt)
1714                <= sozu_command::config::MAX_CONNECTION_ATTEMPTS_UPPER_BOUND,
1715            "connection_attempt ({}) overflowed the largest attempt budget before the retry gate",
1716            self.connection_attempt,
1717        );
1718
1719        // Prefer the SNI-routed cluster (set by `TcpSession::readable` once
1720        // `SniPrereadCore` decides) over the listener's legacy no-SNI
1721        // catch-all -- a listener never configures both (sozu-proxy/sozu#1279),
1722        // but this order is also simply correct for the routed case, where
1723        // `listener.cluster_id` is `None`.
1724        let cluster_id = self
1725            .cluster_id
1726            .clone()
1727            .or_else(|| self.listener.borrow().cluster_id.clone())
1728            .ok_or(BackendConnectionError::NotFound(ObjectKind::TcpCluster))?;
1729
1730        // Capture the metrics incarnation exactly once, when the session is
1731        // first routed. Every path that sets `cluster_id` (session creation,
1732        // SNI preread) captures alongside it, so a retry with `cluster_id`
1733        // already set keeps the original capture even when it is `None`.
1734        if self.cluster_id.is_none() {
1735            self.cluster_metrics_incarnation =
1736                crate::metrics::capture_cluster_metrics(cluster_id.as_ref());
1737            self.cluster_id = Some(cluster_id.clone());
1738        }
1739
1740        // The cluster's own budget when it sets one, the worker's otherwise
1741        // (sozu-proxy/sozu#1800).
1742        let max_attempts = {
1743            let proxy = self.proxy.borrow();
1744            let cluster_attempts = proxy
1745                .configs
1746                .get(&cluster_id)
1747                .and_then(|c| c.max_connection_attempts);
1748            proxy
1749                .sessions
1750                .borrow()
1751                .effective_max_connection_attempts(cluster_attempts)
1752        };
1753        if u32::from(self.connection_attempt) >= max_attempts {
1754            incr!(
1755                names::backend::CONNECT_RETRIES_EXHAUSTED,
1756                self.cluster_id.as_deref(),
1757                self.metrics.backend_id.as_deref(),
1758                self.cluster_metrics_incarnation
1759            );
1760            warn!(
1761                "{} Max connection attempt reached ({})",
1762                log_context!(self),
1763                self.connection_attempt
1764            );
1765            return Err(BackendConnectionError::MaxConnectionRetries(Some(
1766                cluster_id,
1767            )));
1768        }
1769
1770        if self.proxy.borrow().sessions.borrow().at_capacity() {
1771            return Err(BackendConnectionError::MaxSessionsMemory);
1772        }
1773
1774        // Per-(cluster, source-IP) connection limit gate (TCP). The
1775        // source IP comes from `effective_session_address`, which folds
1776        // a parsed PROXY-v2 source over the raw `peer_addr`. The mux's
1777        // Router does the same gate for HTTP/HTTPS sessions; here it
1778        // runs for raw TCP. Rejection produces a graceful TCP FIN via
1779        // `BackendConnectionError::TooManyConnectionsPerIp` →
1780        // `handle_connection_result` → `SessionResult::Close` — TCP has
1781        // no HTTP envelope to carry a 429 / `Retry-After`.
1782        let (cluster_max_connections_per_ip, cluster_max_connections_per_subnet, keyed) = self
1783            .proxy
1784            .borrow()
1785            .configs
1786            .get(&cluster_id)
1787            .map(|c| {
1788                (
1789                    c.max_connections_per_ip,
1790                    c.max_connections_per_subnet,
1791                    c.reads_affinity_key,
1792                )
1793            })
1794            .unwrap_or((None, None, false));
1795        if let Some(ip) = self.effective_session_address().map(|sa| sa.ip()) {
1796            let sessions_rc = self.proxy.borrow().sessions.clone();
1797            // Both caps, through the one combined gate the mux uses too,
1798            // so raw TCP and HTTP/HTTPS can never disagree about what
1799            // "admitted" means.
1800            let at_limit = sessions_rc.borrow().cluster_connection_at_limit(
1801                self.frontend_token,
1802                &cluster_id,
1803                &ip,
1804                cluster_max_connections_per_ip,
1805                cluster_max_connections_per_subnet,
1806            );
1807            if at_limit {
1808                debug!(
1809                    "{} per-(cluster, source-IP/subnet) limit hit for cluster {} from {}",
1810                    log_context!(self),
1811                    cluster_id,
1812                    ip
1813                );
1814                return Err(BackendConnectionError::TooManyConnectionsPerIp { cluster_id });
1815            }
1816            sessions_rc.borrow_mut().track_cluster_connection(
1817                self.frontend_token,
1818                cluster_id.clone(),
1819                ip,
1820                cluster_max_connections_per_subnet,
1821            );
1822            self.cluster_ip_tracked = true;
1823        }
1824
1825        // `HRW` and `MAGLEV` pin a TCP client on its source IP, and shuffle
1826        // sharding ranks its shard by it — the PROXY-v2 source when the
1827        // listener expects one, as for the gate above. A TCP session has no
1828        // request to read a header or cookie from, so the source IP is the
1829        // only key. Any other cluster ignores a key, so none is derived.
1830        let affinity_key = if keyed {
1831            self.effective_session_address()
1832                .map(|address| affinity_key_from_ip(address.ip()))
1833        } else {
1834            None
1835        };
1836        let (backend, mut stream) = self
1837            .proxy
1838            .borrow()
1839            .backends
1840            .borrow_mut()
1841            .backend_from_cluster_id(&cluster_id, affinity_key, Instant::now())
1842            .map_err(BackendConnectionError::Backend)?;
1843
1844        if let Err(e) = stream.set_nodelay(true) {
1845            error!(
1846                "{} Error setting nodelay on back socket({:?}): {:?}",
1847                log_context!(self),
1848                stream,
1849                e
1850            );
1851        }
1852        // Keep the selected backend's handle, as the mux keeps its own in
1853        // its `BackendRegistry`: `fail_backend_connection` and
1854        // `set_back_connected` charge it (failure counter, retry policy,
1855        // UP/DOWN transitions), and `remove_backend` releases the connection
1856        // `Backend::try_connect` counted on it.
1857        self.backend = Some(Rc::clone(&backend));
1858        self.backend_connected = BackendConnectionStatus::Connecting(Instant::now());
1859
1860        let back_token = {
1861            let proxy = self.proxy.borrow();
1862            let mut s = proxy.sessions.borrow_mut();
1863            let entry = s.slab.vacant_entry();
1864            let back_token = Token(entry.key());
1865            let _entry = entry.insert(session_rc.clone());
1866            back_token
1867        };
1868
1869        if let Err(e) = self.proxy.borrow().registry.register(
1870            &mut stream,
1871            back_token,
1872            Interest::READABLE | Interest::WRITABLE,
1873        ) {
1874            error!(
1875                "{} Error registering back socket({:?}): {:?}",
1876                log_context!(self),
1877                stream,
1878                e
1879            );
1880        }
1881
1882        self.container_backend_timeout.set(back_token);
1883
1884        self.set_back_token(back_token);
1885        self.set_back_socket(stream);
1886
1887        self.metrics.backend_id = Some(Rc::from(backend.borrow().backend_id.as_str()));
1888        self.metrics.backend_start();
1889        self.set_backend_id(backend.borrow().backend_id.clone());
1890        self.set_backend_address(backend.borrow().address);
1891
1892        // Postcondition of a successful New connect: the session is wired to
1893        // its freshly-registered backend token and the status reflects an
1894        // in-flight handshake (`Connecting`). The promotion to `Connected`
1895        // happens later in `ready_inner` once the socket signals writable.
1896        debug_assert!(
1897            self.backend_token.is_some(),
1898            "a New backend connection must own its backend token"
1899        );
1900        debug_assert!(
1901            self.backend_connected.is_connecting(),
1902            "a New backend connection must be in the Connecting state"
1903        );
1904
1905        Ok(BackendConnectAction::New)
1906    }
1907}
1908
1909impl ProxySession for TcpSession {
1910    fn close(&mut self) {
1911        if self.has_been_closed {
1912            return;
1913        }
1914
1915        // Past the idempotency guard the session is closing for the first
1916        // time: every gauge-restore / untrack below must run exactly once, so
1917        // re-entry on an already-closed session would double-decrement.
1918        debug_assert!(
1919            !self.has_been_closed,
1920            "close() body must only run on a not-yet-closed session"
1921        );
1922
1923        // TODO: the state should handle the timeouts
1924        trace!("{} Closing TCP session", log_context!(self));
1925        self.metrics.service_stop();
1926
1927        // Drain the per-(cluster, source-IP) accounting before any
1928        // early-return path below. The fail / non-fail close branches
1929        // both count, and the SessionManager-side untrack is idempotent
1930        // (no-op when the slot was never tracked) so this is safe even
1931        // when `cluster_ip_tracked` is false.
1932        if self.cluster_ip_tracked {
1933            self.proxy
1934                .borrow()
1935                .sessions
1936                .borrow_mut()
1937                .untrack_all_cluster_ip(self.frontend_token);
1938            self.cluster_ip_tracked = false;
1939        }
1940
1941        // Restore gauges. `SniPreread` is the "reject"/"teardown" half of
1942        // `tcp.sni_preread.active`'s "-1 on every exit" contract -- the
1943        // "upgrade" exit already decremented it in `upgrade_sni_preread`,
1944        // which also transitions `self.state` away from `SniPreread` before
1945        // `close()` can ever observe that marker again. A session that
1946        // reaches `close()` still marked `SniPreread` (directly, or via
1947        // `FailedUpgrade(SniPreread)` if `upgrade_sni_preread` itself failed)
1948        // therefore never had its gauge/duration accounted for yet.
1949        match self.state.marker() {
1950            StateMarker::Pipe => gauge_add!(names::protocol::TCP, -1),
1951            StateMarker::SendProxyProtocol => gauge_add!(names::protocol::PROXY_SEND, -1),
1952            StateMarker::RelayProxyProtocol => gauge_add!(names::protocol::PROXY_RELAY, -1),
1953            StateMarker::ExpectProxyProtocol => gauge_add!(names::protocol::PROXY_EXPECT, -1),
1954            StateMarker::SniPreread => {
1955                gauge_add!(names::tcp::sni_preread::ACTIVE, -1);
1956                if let TcpStateMachine::SniPreread(preread) = &self.state {
1957                    time!(
1958                        names::tcp::sni_preread::DURATION,
1959                        preread.started_at().elapsed().as_millis() as i64
1960                    );
1961                }
1962            }
1963        }
1964
1965        if self.state.failed() {
1966            match self.state.marker() {
1967                StateMarker::Pipe => incr!(names::tcp::UPGRADE_PIPE_FAILED),
1968                StateMarker::SendProxyProtocol => incr!(names::tcp::UPGRADE_SEND_FAILED),
1969                StateMarker::RelayProxyProtocol => incr!(names::tcp::UPGRADE_RELAY_FAILED),
1970                StateMarker::ExpectProxyProtocol => incr!(names::tcp::UPGRADE_EXPECT_FAILED),
1971                StateMarker::SniPreread => incr!(names::tcp::UPGRADE_SNI_PREREAD_FAILED),
1972            }
1973            return;
1974        }
1975
1976        self.cancel_timeouts();
1977
1978        let front_socket = self.state.front_socket();
1979        // TCP listener is plaintext at this layer — `Shutdown::Both` does not
1980        // truncate any TLS write buffer, so the canonical anti-pattern
1981        // (forces a TCP RST on the read direction, dropping in-flight bytes)
1982        // does not apply. Move to `Shutdown::Write` if a TLS upgrade ever
1983        // wraps this listener.
1984        if let Err(e) = front_socket.shutdown(Shutdown::Both) {
1985            // error 107 NotConnected can happen when was never fully connected, or was already disconnected due to error
1986            if e.kind() != ErrorKind::NotConnected {
1987                error!(
1988                    "{} Error shutting down front socket({:?}): {:?}",
1989                    log_context!(self),
1990                    front_socket,
1991                    e
1992                );
1993            }
1994        }
1995
1996        // Free the frontend slot. No `EPOLL_CTL_DEL`: the server drops the
1997        // session right after `close()`, and that last `close(2)` of the
1998        // front socket takes it out of the epoll set (see `close_backend`).
1999        self.proxy
2000            .borrow()
2001            .sessions
2002            .borrow_mut()
2003            .slab
2004            .try_remove(self.frontend_token.0);
2005
2006        self.close_backend();
2007        self.has_been_closed = true;
2008
2009        // Postcondition of the normal close path: the session is terminal and
2010        // every accounting slot has been released — `close_backend` cleared
2011        // the backend token, and the per-(cluster, source-IP) untrack above
2012        // reset the flag. The idempotency guard now short-circuits any repeat.
2013        debug_assert!(self.has_been_closed, "close() must mark the session closed");
2014        debug_assert!(
2015            self.backend_token.is_none(),
2016            "close() must leave no dangling backend token"
2017        );
2018        debug_assert!(
2019            !self.cluster_ip_tracked,
2020            "close() must untrack the (cluster, source-IP) slot"
2021        );
2022    }
2023
2024    fn timeout(&mut self, token: Token) -> SessionIsToBeClosed {
2025        // The frontend and backend slots are distinct tokens, so the two
2026        // dispatch arms below are mutually exclusive — a single token can
2027        // never match both. (Obsolete tokens matching neither are tolerated
2028        // and fall through to the `false` arm.)
2029        debug_assert!(
2030            self.backend_token != Some(self.frontend_token),
2031            "frontend and backend tokens must never collide"
2032        );
2033        if self.frontend_token == token {
2034            self.container_frontend_timeout.triggered();
2035            // The preread deadline firing always closes the session either
2036            // way (matches every other state's front-timeout behavior).
2037            // Route-aware: only feed `Input::Timeout` into the core while
2038            // still UNDECIDED, for its `tcp.sni_preread.rejected.fragmented`
2039            // metric + log side effect. Once a route has already latched
2040            // (backend connect still pending), this same front-timeout
2041            // firing is a plain "connect/upgrade took too long" close, NOT a
2042            // fresh preread verdict -- re-feeding `Input::Timeout` into an
2043            // already-decided core would just replay the SAME latched
2044            // `Output::Routed` through `SniPreread::handle_output`'s
2045            // `Routed` arm a SECOND time: double-incrementing
2046            // `tcp.sni_preread.routed` in release, and tripping its
2047            // `debug_assert!(self.outcome.is_none(), ...)` in debug
2048            // (sozu-proxy/sozu#1290).
2049            if let TcpStateMachine::SniPreread(preread) = &mut self.state {
2050                if preread.is_routed() {
2051                    debug!(
2052                        "{} frontend timeout while a routed SNI-preread session was still \
2053                         waiting on its backend connect",
2054                        log_context!(self)
2055                    );
2056                } else {
2057                    let listener = self.listener.borrow();
2058                    let cfg = listener.preread_config(preread.effective_max_bytes());
2059                    preread.on_timeout(&cfg);
2060                }
2061            }
2062            return true;
2063        }
2064        if self.backend_token == Some(token) {
2065            self.container_backend_timeout.triggered();
2066            return true;
2067        }
2068        // invalid token, obsolete timeout triggered
2069        false
2070    }
2071
2072    fn protocol(&self) -> Protocol {
2073        Protocol::TCP
2074    }
2075
2076    fn update_readiness(&mut self, token: Token, events: Ready) {
2077        trace!(
2078            "{} token {:?} got event {}",
2079            log_context!(self),
2080            token,
2081            super::ready_to_string(events)
2082        );
2083
2084        self.last_event = Instant::now();
2085        self.metrics.wait_start();
2086
2087        if self.frontend_token == token {
2088            self.front_readiness().event = self.front_readiness().event | events;
2089        } else if self.backend_token == Some(token)
2090            && let Some(r) = self.back_readiness()
2091        {
2092            r.event |= events;
2093        }
2094    }
2095
2096    fn ready(&mut self, session: Rc<RefCell<dyn ProxySession>>) -> SessionIsToBeClosed {
2097        self.metrics.service_start();
2098
2099        let session_result = self.ready_inner(session.clone());
2100
2101        let to_bo_closed = match session_result {
2102            SessionResult::Close => true,
2103            SessionResult::Continue => false,
2104            SessionResult::Upgrade => match self.upgrade() {
2105                false => self.ready(session),
2106                true => true,
2107            },
2108        };
2109
2110        self.metrics.service_stop();
2111
2112        // Run-to-completion postcondition: the front/back token + readiness
2113        // state machine must satisfy its cross-field invariants after every
2114        // `ready()` pass. Cfg-guarded so the call (and `check_invariants`
2115        // itself) is absent from release builds.
2116        #[cfg(debug_assertions)]
2117        self.check_invariants();
2118
2119        to_bo_closed
2120    }
2121
2122    fn shutting_down(&mut self) -> SessionIsToBeClosed {
2123        true
2124    }
2125
2126    fn last_event(&self) -> Instant {
2127        self.last_event
2128    }
2129
2130    fn print_session(&self) {
2131        let state: String = match &self.state {
2132            TcpStateMachine::ExpectProxyProtocol(_) => String::from("Expect"),
2133            TcpStateMachine::SendProxyProtocol(_) => String::from("Send"),
2134            TcpStateMachine::RelayProxyProtocol(_) => String::from("Relay"),
2135            TcpStateMachine::Pipe(_) => String::from("TCP"),
2136            TcpStateMachine::SniPreread(_) => String::from("SniPreread"),
2137            TcpStateMachine::FailedUpgrade(marker) => format!("FailedUpgrade({marker:?})"),
2138        };
2139
2140        let front_readiness = match &self.state {
2141            TcpStateMachine::ExpectProxyProtocol(expect) => Some(&expect.frontend_readiness),
2142            TcpStateMachine::SendProxyProtocol(send) => Some(&send.frontend_readiness),
2143            TcpStateMachine::RelayProxyProtocol(relay) => Some(&relay.frontend_readiness),
2144            TcpStateMachine::Pipe(pipe) => Some(&pipe.frontend_readiness),
2145            TcpStateMachine::SniPreread(preread) => Some(&preread.frontend_readiness),
2146            TcpStateMachine::FailedUpgrade(_) => None,
2147        };
2148
2149        let back_readiness = match &self.state {
2150            TcpStateMachine::SendProxyProtocol(send) => Some(&send.backend_readiness),
2151            TcpStateMachine::RelayProxyProtocol(relay) => Some(&relay.backend_readiness),
2152            TcpStateMachine::Pipe(pipe) => Some(&pipe.backend_readiness),
2153            TcpStateMachine::SniPreread(preread) => Some(&preread.backend_readiness),
2154            TcpStateMachine::ExpectProxyProtocol(_) => None,
2155            TcpStateMachine::FailedUpgrade(_) => None,
2156        };
2157
2158        error!(
2159            "\
2160{} Session ({:?})
2161\tFrontend:
2162\t\ttoken: {:?}\treadiness: {:?}
2163\tBackend:
2164\t\ttoken: {:?}\treadiness: {:?}\tstatus: {:?}\tcluster id: {:?}",
2165            log_context!(self),
2166            state,
2167            self.frontend_token,
2168            front_readiness,
2169            self.backend_token,
2170            back_readiness,
2171            self.backend_connected,
2172            self.cluster_id
2173        );
2174        error!("Metrics: {:?}", self.metrics);
2175    }
2176
2177    fn frontend_token(&self) -> Token {
2178        self.frontend_token
2179    }
2180}
2181
2182pub struct TcpListener {
2183    active: SessionIsToBeClosed,
2184    address: SocketAddr,
2185    cluster_id: Option<ClusterId>,
2186    config: TcpListenerConfig,
2187    listener: Option<MioTcpListener>,
2188    /// A socket this listener holds but has NOT registered.
2189    ///
2190    /// `activate()` parks an inherited socket here when `Registry::register`
2191    /// fails, instead of letting the `?` drop and close a descriptor
2192    /// `Listeners::get_*` has already removed from the SCM table (sozu#1342).
2193    /// The next activation reuses it rather than binding a second one.
2194    ///
2195    /// Deliberately NOT `listener`: every other consumer reads that field's `Some` as
2196    /// "registered and live" — `soft_stop` and `hard_stop` deregister it and
2197    /// report the failure, `give_back_listener` hands it back as an activated
2198    /// socket, `accept` accepts on it. A parked socket is none of those things,
2199    /// and putting it there made a failed registration answer `ENOENT` from
2200    /// `deregister` and fail the whole soft stop.
2201    parked_listener: Option<MioTcpListener>,
2202    /// Set by `activate()`, cleared by the first `accept()` that answers
2203    /// `WouldBlock`: while set, `accept()` sets `TCP_NODELAY` on each socket
2204    /// it returns, because a connection queued before `activate()` set the
2205    /// flag on the listener did not inherit it. See the field of the same
2206    /// name on `HttpListener` (`lib/src/http.rs`).
2207    nodelay_backlog: bool,
2208    /// SNI -> `(AlpnMatcher, ClusterId)` route table (sozu-proxy/sozu#1279).
2209    /// Populated by `add_tcp_front`/`remove_tcp_front` from
2210    /// `RequestTcpFrontend.sni`/`.alpn`; empty for a listener whose fronts
2211    /// are all no-SNI (the legacy `cluster_id` catch-all). A listener never
2212    /// mixes both (enforced at config load, `command/src/config.rs`), but
2213    /// `create_session`'s routing gate stays defensive and checks both.
2214    sni_routes: TrieNode<Vec<(AlpnMatcher, ClusterId)>>,
2215    tags: BTreeMap<String, CachedTags>,
2216    token: Token,
2217}
2218
2219impl ListenerHandler for TcpListener {
2220    fn get_addr(&self) -> &SocketAddr {
2221        &self.address
2222    }
2223
2224    fn get_tags(&self, key: &str) -> Option<&CachedTags> {
2225        self.tags.get(key)
2226    }
2227
2228    fn set_tags(&mut self, key: String, tags: Option<BTreeMap<String, String>>) {
2229        match tags {
2230            Some(tags) => self.tags.insert(key, CachedTags::new(tags)),
2231            None => self.tags.remove(&key),
2232        };
2233    }
2234
2235    fn protocol(&self) -> Protocol {
2236        Protocol::TCP
2237    }
2238
2239    fn public_address(&self) -> SocketAddr {
2240        self.config
2241            .public_address
2242            .map(|addr| addr.into())
2243            .unwrap_or(self.address)
2244    }
2245}
2246
2247impl TcpListener {
2248    /// This listener's identity: its address, plus the network interface its
2249    /// socket is bound to when it has one.
2250    pub fn key(&self) -> ListenerKey {
2251        ListenerKey::new(self.address, self.config.interface.as_deref())
2252    }
2253
2254    /// Whether this listener is the one `key` names.
2255    pub fn is(&self, key: &ListenerKey) -> bool {
2256        key.matches(&self.address, self.config.interface.as_deref())
2257    }
2258
2259    fn new(config: TcpListenerConfig, token: Token) -> Result<TcpListener, ListenerError> {
2260        Ok(TcpListener {
2261            cluster_id: None,
2262            listener: None,
2263            parked_listener: None,
2264            nodelay_backlog: false,
2265            token,
2266            address: config.address.into(),
2267            config,
2268            active: false,
2269            sni_routes: TrieNode::root(),
2270            tags: BTreeMap::new(),
2271        })
2272    }
2273
2274    /// Validate that a worker can build this TCP listener configuration WITHOUT
2275    /// constructing the full listener or binding a socket. TCP listener
2276    /// construction has no fallible config today (no rustls context, no answer
2277    /// templates; a bad bind surfaces later as an `ActivateListener` failure),
2278    /// so this currently always succeeds. It exists for surface parity with the
2279    /// HTTP/HTTPS validators the main process calls before committing an
2280    /// `Add*Listener` to `ConfigState` and fanning it out (sozu#1301), and is
2281    /// the hook for any future TCP-config validation.
2282    pub fn validate_config(_config: &TcpListenerConfig) -> Result<(), ListenerError> {
2283        Ok(())
2284    }
2285
2286    /// Build the [`PrereadConfig`] this listener's `SniPreread` sessions
2287    /// feed to [`crate::protocol::tcp_preread::SniPrereadCore::handle_input`].
2288    /// `routes`/`inbound_proxy`/`timeout` come straight from the listener
2289    /// config; `effective_max_bytes` is session-specific (already clamped to
2290    /// that session's buffer capacity at construction, see `create_session`),
2291    /// so it is passed in rather than re-derived here.
2292    fn preread_config(&self, effective_max_bytes: usize) -> PrereadConfig<'_> {
2293        PrereadConfig {
2294            routes: &self.sni_routes,
2295            inbound_proxy: self.config.expect_proxy,
2296            max_bytes: effective_max_bytes,
2297            timeout: Duration::from_secs(u64::from(
2298                self.config
2299                    .sni_preread_timeout
2300                    .unwrap_or(DEFAULT_SNI_PREREAD_TIMEOUT),
2301            )),
2302            accept_wildcard: true,
2303        }
2304    }
2305
2306    /// Validate an incoming `AddTcpFrontend` against this listener's
2307    /// CURRENT routing state before any mutation, mirroring
2308    /// `command/src/config.rs`'s TOML config-load TCP SNI/ALPN invariants
2309    /// (sozu-proxy/sozu#1279):
2310    ///
2311    /// - `alpn` set with no `sni`: the worker's no-SNI catch-all path never
2312    ///   consults `alpn`, so the protocol list would silently never be
2313    ///   enforced (mirrors `ConfigError::AlpnWithoutSni`).
2314    /// - a no-SNI frontend added to a listener that already has SNI-scoped
2315    ///   routes, or an SNI-scoped frontend added to a listener that already
2316    ///   has a no-SNI catch-all cluster (mirrors
2317    ///   `ConfigError::TcpListenerMixesSniAndNoSni`).
2318    /// - an ALPN protocol, or a catch-all (empty `alpn`), that overlaps an
2319    ///   existing route already registered for the same `(address, sni)`
2320    ///   (mirrors `ConfigError::TcpFrontendAlpnOverlap` /
2321    ///   `TcpFrontendMultipleAlpnCatchAll`).
2322    ///
2323    /// Config-load already rejects all of these shapes for requests built
2324    /// from a TOML file, but `AddTcpFrontend` can also arrive directly over
2325    /// the command socket, or via `LoadState` replay of a hand-edited or
2326    /// stale state file, bypassing config.rs entirely -- the worker must
2327    /// not silently corrupt its own routing table when that happens.
2328    fn validate_new_tcp_front(&self, front: &RequestTcpFrontend) -> Result<(), ProxyError> {
2329        let reject = |reason: String| {
2330            Err(ProxyError::InvalidTcpFrontend {
2331                address: self.address,
2332                reason,
2333            })
2334        };
2335
2336        match &front.sni {
2337            None => {
2338                if !front.alpn.is_empty() {
2339                    return reject(format!(
2340                        "alpn = {:?} set without sni: alpn only matches within an SNI-scoped \
2341                         preread, so a frontend without sni would silently ignore its alpn list",
2342                        front.alpn
2343                    ));
2344                }
2345                if !self.sni_routes.is_empty() {
2346                    return reject(
2347                        "a no-SNI frontend cannot be added to a listener that already has \
2348                         SNI-scoped routes"
2349                            .to_string(),
2350                    );
2351                }
2352            }
2353            Some(sni) => {
2354                if self.cluster_id.is_some() {
2355                    return reject(
2356                        "an SNI-scoped frontend cannot be added to a listener that already has \
2357                         a no-SNI catch-all cluster"
2358                            .to_string(),
2359                    );
2360                }
2361
2362                // SNI SHAPE check: delegate to the SAME validator config-load
2363                // uses (`sozu_command::config::validate_sni_pattern`) rather
2364                // than a hand-rolled partial check. A direct `AddTcpFrontend`
2365                // over the command socket, or a `LoadState` replay, bypasses
2366                // config.rs entirely, so the worker boundary must enforce
2367                // the identical rule -- including the checks a bare '/'/'*'
2368                // scan used to miss: empty string, non-ASCII, and any empty
2369                // label (leading/trailing/consecutive dots). An unvalidated
2370                // leading-empty-label pattern like `.example.com` would
2371                // otherwise reach `insert_sni_route` ->
2372                // `pattern_trie::insert_recursive`'s RELEASE-mode
2373                // `assert_ne!(partial_key, &b""[..])` and crash the worker.
2374                let normalized_sni = match validate_sni_pattern(sni) {
2375                    Ok(normalized) => normalized,
2376                    Err(config_error) => {
2377                        return reject(format!(
2378                            "sni {sni:?} failed SNI shape validation: {config_error}"
2379                        ));
2380                    }
2381                };
2382
2383                // Same key as `insert_sni_route`'s own lookup, so this checks
2384                // against exactly the entries the new route would be appended
2385                // alongside. This is bookkeeping over the key's OWN node, not
2386                // the routing lookup (`preread_config` keeps `true` for
2387                // that), so `accept_wildcard` must make the lookup EXACT for
2388                // both key shapes:
2389                //
2390                // - literal key -> `false`. With `true`, a literal key with
2391                //   no child yet (e.g. `a.example.com` when only
2392                //   `*.example.com` exists) falls back to the sibling
2393                //   wildcard's entries (`pattern_trie.rs`'s `lookup`
2394                //   wildcard-fallback branch), misattributing the wildcard's
2395                //   catch-all to the exact key (falsely rejecting a
2396                //   legitimate exact catch-all) — and symmetrically,
2397                //   `insert_sni_route`/`remove_sni_route` would corrupt the
2398                //   WILDCARD's `Vec` instead of touching a distinct
2399                //   exact-key node.
2400                // - wildcard key -> `true`. Unlike `lookup_mut` (which
2401                //   short-circuits `partial_key == b"*"` before consulting
2402                //   `accept_wildcard`, so insert/remove stay on `false`),
2403                //   the immutable `lookup` reaches a wildcard entry ONLY
2404                //   through the fallback branch; with `false` an existing
2405                //   `*.example.com` entry is invisible here and a duplicate
2406                //   catch-all / overlapping-ALPN wildcard front bypasses
2407                //   validation. For a wildcard key, `true` IS the exact
2408                //   self-lookup: the traversal descends the key's own
2409                //   literal ancestry, and `insert` never creates a literal
2410                //   `*` child that could shadow the node's `wildcard` slot.
2411                //
2412                // `starts_with(b"*.")` matches every wildcard shape
2413                // config-load can emit: `command/src/config.rs`'s
2414                // `validate_sni_pattern` only admits a single leading `*.`
2415                // label (any other `*` placement is rejected), and a plain
2416                // hostname key traverses literal children in both `lookup`
2417                // and `insert_recursive`. A bare `*` key (config-load
2418                // rejects it; only a bypassing IPC request can carry one)
2419                // is the known gap: `lookup_mut`'s short-circuit maps it to
2420                // the wildcard slot while this immutable self-lookup (flag
2421                // `false`) cannot see an existing entry there, so duplicate
2422                // bare-`*` fronts are not detected -- sibling keys stay
2423                // untouched either way.
2424                //
2425                // `normalized_sni` (not a fresh `sni.to_ascii_lowercase()`)
2426                // is used here: `validate_sni_pattern` already returned the
2427                // canonical lowercased form above, and reusing it keeps this
2428                // function's notion of "the key" identical to what
2429                // `insert_sni_route`/`remove_sni_route` compute from the
2430                // same input.
2431                let key = normalized_sni.into_bytes();
2432                let accept_wildcard_for_self_lookup = key.starts_with(b"*.");
2433                if let Some((_, existing)) = self
2434                    .sni_routes
2435                    .domain_lookup(&key, accept_wildcard_for_self_lookup)
2436                {
2437                    let new_is_catch_all = front.alpn.is_empty();
2438                    for (matcher, _cluster_id) in existing {
2439                        match matcher {
2440                            AlpnMatcher::Any if new_is_catch_all => {
2441                                return reject(format!(
2442                                    "sni {sni:?} already has a catch-all (empty alpn) \
2443                                     frontend: at most one frontend per (address, sni) may \
2444                                     omit alpn"
2445                                ));
2446                            }
2447                            AlpnMatcher::OneOf(protocols) => {
2448                                if let Some(overlap) = front
2449                                    .alpn
2450                                    .iter()
2451                                    .find(|protocol| protocols.contains(protocol.as_bytes()))
2452                                {
2453                                    return reject(format!(
2454                                        "sni {sni:?} already has a frontend matching ALPN \
2455                                         protocol {overlap:?}: ALPN matchers for the same \
2456                                         (address, sni) must not overlap"
2457                                    ));
2458                                }
2459                            }
2460                            AlpnMatcher::Any => {}
2461                        }
2462                    }
2463                }
2464            }
2465        }
2466
2467        Ok(())
2468    }
2469
2470    /// Add one `(AlpnMatcher, ClusterId)` entry to this listener's SNI route
2471    /// table, appending to the SNI key's existing `Vec` if one is already
2472    /// present rather than clobbering it (multiple ALPN-scoped fronts can
2473    /// share the same SNI). `sni` is defensively lowercased: the
2474    /// SNI-preread core normalizes (lowercase, no trailing dot) before ever
2475    /// looking a route up, so the table key must already be in that form.
2476    ///
2477    /// Returns `Err` when the underlying trie insert reports
2478    /// [`InsertResult::Failed`] (a malformed key, e.g. an empty label) —
2479    /// the caller (`TcpProxy::add_tcp_front`) must propagate this rather
2480    /// than silently keep going with an add that reports success while the
2481    /// route table never gained a usable entry. With `validate_new_tcp_front`
2482    /// enforcing the shared SNI shape validator before this ever runs, a
2483    /// `Failed` result should be unreachable here — the `debug_assert_ne!`
2484    /// below asserts OUR OWN validated invariant, not raw wire/IPC input,
2485    /// and is a loud Sōzu-internal-bug canary in debug builds; the `Err`
2486    /// path is the release-mode graceful fallback if that invariant is ever
2487    /// violated by a future bug.
2488    fn insert_sni_route(
2489        &mut self,
2490        sni: String,
2491        alpn: Vec<String>,
2492        cluster_id: ClusterId,
2493    ) -> Result<(), ProxyError> {
2494        let (key, matcher) = route_key_and_matcher(&sni, alpn);
2495        // `accept_wildcard: false` — see `validate_new_tcp_front`'s comment:
2496        // an exact key must never fall back to a sibling wildcard's entry,
2497        // or this push would corrupt the WILDCARD's route `Vec` instead of
2498        // creating this key's own node.
2499        match self.sni_routes.domain_lookup_mut(&key, false) {
2500            Some((_, entries)) => {
2501                entries.push((matcher, cluster_id));
2502                Ok(())
2503            }
2504            None => {
2505                let insert_result = self
2506                    .sni_routes
2507                    .domain_insert(key, vec![(matcher, cluster_id)]);
2508                debug_assert_ne!(
2509                    insert_result,
2510                    InsertResult::Failed,
2511                    "insert_sni_route's key must already be validated by validate_new_tcp_front"
2512                );
2513                if insert_result == InsertResult::Failed {
2514                    error!(
2515                        "{} SNI route insert failed for {:?} despite passing shape \
2516                         validation -- rejecting to avoid a silently dead route",
2517                        log_module_context!(),
2518                        sni
2519                    );
2520                    return Err(ProxyError::InvalidTcpFrontend {
2521                        address: self.address,
2522                        reason: format!(
2523                            "internal error: the route table rejected sni {sni:?} despite \
2524                             passing shape validation"
2525                        ),
2526                    });
2527                }
2528                Ok(())
2529            }
2530        }
2531    }
2532
2533    /// Symmetric counterpart to [`Self::insert_sni_route`]: removes the
2534    /// matching `(AlpnMatcher, ClusterId)` entry, then `domain_remove`s the
2535    /// SNI key itself once its `Vec` is empty (no stranded empty entries in
2536    /// the trie).
2537    fn remove_sni_route(&mut self, sni: String, alpn: Vec<String>, cluster_id: &str) {
2538        let (key, matcher) = route_key_and_matcher(&sni, alpn);
2539        // `accept_wildcard: false` — same reasoning as `insert_sni_route`:
2540        // removing the exact key must never reach into and strip a sibling
2541        // wildcard's entries.
2542        if let Some((_, entries)) = self.sni_routes.domain_lookup_mut(&key, false) {
2543            entries.retain(|(m, c)| !(*m == matcher && **c == *cluster_id));
2544            if entries.is_empty() {
2545                self.sni_routes.domain_remove(&key);
2546            }
2547        }
2548    }
2549
2550    pub fn activate(
2551        &mut self,
2552        registry: &Registry,
2553        tcp_listener: Option<MioTcpListener>,
2554    ) -> Result<Token, ProxyError> {
2555        if self.active {
2556            return Ok(self.token);
2557        }
2558        let address: SocketAddr = self.config.address.into();
2559
2560        // Reuse a socket a previous attempt parked here after a failed
2561        // registration, rather than binding a second one. Taking it
2562        // unconditionally is what keeps the two from coexisting: if the caller
2563        // ALSO handed one over, the parked socket is disposed of deliberately
2564        // below instead of being silently overwritten.
2565        let parked = self.parked_listener.take();
2566        let mut listener = match (tcp_listener, parked) {
2567            (Some(inherited), Some(parked)) => {
2568                // Not reachable while `Listeners::get_*` removes the entry it
2569                // hands over and `scm_listeners` is written once: a parked
2570                // socket means this address's entry was already consumed, so
2571                // the caller cannot have another. Neither of those is enforced
2572                // here, so dispose of the loser explicitly and loudly rather
2573                // than letting an assignment close it — that silent close is
2574                // sozu#1342 itself.
2575                warn!(
2576                    "{} closing the socket parked on the TCP listener for {}: an \
2577                     inherited socket arrived for the same address and is the one being \
2578                     activated",
2579                    log_module_context!(),
2580                    address
2581                );
2582                drop(parked);
2583                inherited
2584            }
2585            (Some(inherited), None) => inherited,
2586            (None, Some(parked)) => parked,
2587            (None, None) => server_bind(address, self.config.interface.as_deref())
2588                .map_err(|e| ProxyError::BindToSocket(address, e))?,
2589        };
2590
2591        // Once per listener, on whichever socket won above — freshly bound,
2592        // inherited over SCM_RIGHTS, or parked. Every socket accepted from now
2593        // on inherits it; see `nodelay_backlog` for the ones already queued.
2594        if let Err(e) = SockRef::from(&listener).set_tcp_nodelay(true) {
2595            error!(
2596                "{} error setting nodelay on listen socket({:?}): {:?}",
2597                log_module_context!(),
2598                listener,
2599                e
2600            );
2601        }
2602
2603        let registration = registry
2604            .register(&mut listener, self.token, Interest::READABLE)
2605            .map_err(ProxyError::RegisterListener);
2606        if let Err(error) = registration {
2607            // Park the socket instead of letting the `?` drop it. `register`
2608            // is fallible (EEXIST for a descriptor still registered because
2609            // `notify_deactivate_listener` only logs a failed deregister,
2610            // EBADF, ENOMEM), and a socket still held in a local is dropped —
2611            // and closed — on the way out. For an SCM-inherited socket that is
2612            // the sozu#1342 close in a second place: `Listeners::get_*` has
2613            // already removed the entry, so nothing else owns the descriptor.
2614            // It goes to `parked_listener`, never to `listener`, because every other
2615            // consumer of that field reads `Some` as "registered and live".
2616            self.parked_listener = Some(listener);
2617            return Err(error);
2618        }
2619
2620        self.listener = Some(listener);
2621        self.nodelay_backlog = true;
2622        self.active = true;
2623        Ok(self.token)
2624    }
2625
2626    /// Apply a partial-update patch to this TCP listener's live configuration.
2627    ///
2628    /// Fields absent in the patch (i.e. `None`) are preserved unchanged.
2629    pub fn update_config(&mut self, patch: &UpdateTcpListenerConfig) -> Result<(), ListenerError> {
2630        if let Some(v) = patch.public_address {
2631            self.config.public_address = Some(v);
2632        }
2633        if let Some(v) = patch.expect_proxy {
2634            self.config.expect_proxy = v;
2635        }
2636        if let Some(v) = patch.front_timeout {
2637            self.config.front_timeout = v;
2638        }
2639        if let Some(v) = patch.back_timeout {
2640            self.config.back_timeout = v;
2641        }
2642        if let Some(v) = patch.connect_timeout {
2643            self.config.connect_timeout = v;
2644        }
2645        Ok(())
2646    }
2647}
2648
2649fn handle_connection_result(
2650    connection_result: Result<BackendConnectAction, BackendConnectionError>,
2651) -> Option<SessionResult> {
2652    match connection_result {
2653        // reuse connection or send a default answer, we can continue
2654        Ok(BackendConnectAction::Reuse) => None,
2655        Ok(BackendConnectAction::New) | Ok(BackendConnectAction::Replace) => {
2656            // we must wait for an event
2657            Some(SessionResult::Continue)
2658        }
2659        Err(_) => {
2660            // in case of BackendConnectionError::Backend(BackendError::ConnectionFailures(..))
2661            // we may want to retry instead of closing
2662            Some(SessionResult::Close)
2663        }
2664    }
2665}
2666
2667/// `min(listener.config.sni_preread_max_bytes, frontend_buffer.capacity())`,
2668/// floored at [`MIN_SNI_PREREAD_MAX_BYTES`] -- the SNI-preread core has no
2669/// independent backstop of its own (`SniPrereadCore::handle_input` trusts
2670/// `PrereadConfig::max_bytes` entirely), so the shell must never hand it a
2671/// cap the checked-out buffer cannot actually hold (the `min`), NOR a cap so
2672/// small the preread read is zero-length and spins until the loop guard (the
2673/// `max`). Config-load rejects a sub-floor `sni_preread_max_bytes` loudly
2674/// (`ConfigError::SniPrereadMaxBytesTooSmall`), but a `0` knob from a direct
2675/// `sozu listener tcp add` CLI/IPC request, or a stale `LoadState`
2676/// replay, bypasses that check and reaches the worker -- this floor degrades
2677/// it to the 5-byte TLS-record-header minimum instead, killing the spin for
2678/// EVERY config source at the single point of use. `buffer_capacity` is
2679/// always `>= MIN_SNI_PREREAD_MAX_BYTES` in practice (buffers are KB-sized),
2680/// so the `min` never fights the `max`.
2681fn effective_sni_preread_max_bytes(configured: Option<u32>, buffer_capacity: usize) -> usize {
2682    (configured.unwrap_or(DEFAULT_SNI_PREREAD_MAX_BYTES) as usize)
2683        .min(buffer_capacity)
2684        .max(MIN_SNI_PREREAD_MAX_BYTES as usize)
2685}
2686
2687/// Whether `TcpSession::readable`'s frontend-timeout reset should fire for
2688/// the CURRENT state. `false` only while the session is an UNDECIDED
2689/// `SniPreread`: the preread deadline armed at session creation
2690/// (`new_sni_preread`) is an ABSOLUTE budget, not a per-fragment idle timer
2691/// -- resetting it on every `readable()` event would let a client
2692/// trickling one byte just before each expiry hold the session (and both
2693/// its checked-out buffers) open far past the configured
2694/// `sni_preread_timeout` (sozu-proxy/sozu#1290). Every
2695/// other state -- a `SniPreread` that has already routed, or any state
2696/// reached after it -- resets normally: once routed, the frontend timeout
2697/// reverts to being a genuine per-read idle timer (see the `front_timeout`
2698/// restore in `readable()`'s route-capture block).
2699fn frontend_timeout_resets_on_readable(state: &TcpStateMachine) -> bool {
2700    !matches!(
2701        state,
2702        TcpStateMachine::SniPreread(preread) if preread.outcome().is_none()
2703    )
2704}
2705
2706/// Access-log ALPN tag for an SNI-routed TCP session: the client's FIRST
2707/// offered protocol (client preference order, matching
2708/// `SniPrereadCore::route`'s own routing precedence), mapped to a known
2709/// `&'static str` label -- the same two labels `https.rs`'s own ALPN
2710/// negotiation records (`"h2"` / `"http/1.1"`) -- so `tcp.sni_preread`
2711/// sessions and terminated-TLS sessions chart under the same values.
2712/// `None` for an empty offer or an unrecognized protocol: Sōzu never
2713/// terminates TLS on this path, so this is the client's stated preference,
2714/// not a negotiated outcome.
2715fn known_alpn_label(offered: &[Vec<u8>]) -> Option<&'static str> {
2716    match offered.first().map(Vec::as_slice) {
2717        Some(b"h2") => Some("h2"),
2718        Some(b"http/1.1") => Some("http/1.1"),
2719        _ => None,
2720    }
2721}
2722
2723/// Worker-internal, collision-free identity for a TCP frontend's access-log
2724/// tags entry (sozu-proxy/sozu#1290). ALPN protocol
2725/// identifiers are opaque RFC 7301 byte strings -- nothing forbids a `,` or
2726/// `|` inside one -- so a naive `sorted_alpn.join(",")` string key let two
2727/// legal, DISJOINT fronts collide: a single protocol `"a,b"` and the pair
2728/// `["a", "b"]` both joined to the identical string `"a,b"`, so the second
2729/// `add_tcp_front` silently clobbered the first front's tags entry.
2730///
2731/// `TcpListener::tags: BTreeMap<String, CachedTags>` and
2732/// `Pipe::set_tags_key(Option<String>)` (`lib/src/protocol/pipe.rs`) both
2733/// key on a plain `String` -- crossing either boundary still needs one --
2734/// so this type does not replace that storage; it is the SINGLE place that
2735/// composes the string, via [`Self::fmt`]'s length-prefixed per-protocol
2736/// encoding, so no choice of in-band separator can be confused with a
2737/// protocol-name boundary the way a bare `join(",")` could.
2738#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
2739enum TcpFrontendTagsKey {
2740    /// Legacy no-SNI catch-all front: keyed by the bare listener address,
2741    /// exactly as before SNI routing existed.
2742    Address(SocketAddr),
2743    /// An SNI-scoped front. `alpn` is sorted (and deduped, since a plain
2744    /// `Vec<String>` cannot enforce uniqueness the way
2745    /// [`AlpnMatcher::OneOf`]'s `BTreeSet` does) so two constructions from
2746    /// the same logical protocol set always compare equal regardless of
2747    /// the caller's original order.
2748    Sni {
2749        address: SocketAddr,
2750        sni: String,
2751        alpn: Vec<String>,
2752    },
2753}
2754
2755impl TcpFrontendTagsKey {
2756    /// The identity triple `command/src/state.rs`'s `add_tcp_frontend`
2757    /// deduplicates on: `sni` is lowercased to match `insert_sni_route`'s
2758    /// trie key, `alpn` sorted + deduped so `add_tcp_front`'s operator
2759    /// order and the route-time rebuild's [`AlpnMatcher::OneOf`] `BTreeSet`
2760    /// order (`alpn_matcher_protocols`) always agree.
2761    fn sni(address: SocketAddr, sni: &str, alpn: &[String]) -> Self {
2762        let mut alpn: Vec<String> = alpn.to_vec();
2763        alpn.sort_unstable();
2764        alpn.dedup();
2765        TcpFrontendTagsKey::Sni {
2766            address,
2767            sni: sni.to_ascii_lowercase(),
2768            alpn,
2769        }
2770    }
2771}
2772
2773impl std::fmt::Display for TcpFrontendTagsKey {
2774    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2775        match self {
2776            TcpFrontendTagsKey::Address(address) => write!(f, "{address}"),
2777            TcpFrontendTagsKey::Sni { address, sni, alpn } => {
2778                write!(f, "{address}|{sni}|")?;
2779                for protocol in alpn {
2780                    // `<len>:<bytes>` per protocol: the length prefix is
2781                    // CHECKED against the following bytes, never scanned
2782                    // for, so a `,`/`:` embedded in one protocol name can
2783                    // never be misread as a boundary between two --
2784                    // unlike the old bare `join(",")`.
2785                    write!(f, "{}:{protocol}", protocol.len())?;
2786                }
2787                Ok(())
2788            }
2789        }
2790    }
2791}
2792
2793/// Canonical access-log tags key for an SNI-scoped TCP frontend: one
2794/// listener can carry many SNI/ALPN fronts, each with its own `tags`, so
2795/// keying tags by the bare listener address (the pre-SNI behavior, kept
2796/// verbatim for no-SNI fronts) would let the LAST added front clobber every
2797/// sibling and let removing ANY front clear tags for the whole address. See
2798/// [`TcpFrontendTagsKey`] for the collision-free composition.
2799fn sni_tags_key(address: &SocketAddr, sni: &str, alpn: &[String]) -> String {
2800    TcpFrontendTagsKey::sni(*address, sni, alpn).to_string()
2801}
2802
2803/// The matched [`AlpnMatcher`]'s protocols as strings, for rebuilding the
2804/// [`sni_tags_key`] a route decision maps to: `Any` is the empty-`alpn`
2805/// catch-all front, `OneOf` yields its (BTreeSet-sorted) protocol list.
2806/// `from_utf8_lossy` is defensive — matcher protocols originate from
2807/// config/IPC `String`s, never raw network bytes.
2808fn alpn_matcher_protocols(matcher: &AlpnMatcher) -> Vec<String> {
2809    match matcher {
2810        AlpnMatcher::Any => Vec::new(),
2811        AlpnMatcher::OneOf(set) => set
2812            .iter()
2813            .map(|protocol| String::from_utf8_lossy(protocol).into_owned())
2814            .collect(),
2815    }
2816}
2817
2818/// Shared `(trie key, ALPN matcher)` construction for
2819/// [`TcpListener::insert_sni_route`] and [`TcpListener::remove_sni_route`]:
2820/// `sni` lowercased into the trie key, `alpn` collapsed to
2821/// [`AlpnMatcher::Any`] when empty (catch-all) or [`AlpnMatcher::OneOf`]
2822/// otherwise. `alpn` is taken by value since neither caller needs it again
2823/// afterward.
2824fn route_key_and_matcher(sni: &str, alpn: Vec<String>) -> (Vec<u8>, AlpnMatcher) {
2825    let key = sni.to_ascii_lowercase().into_bytes();
2826    let matcher = if alpn.is_empty() {
2827        AlpnMatcher::Any
2828    } else {
2829        AlpnMatcher::OneOf(alpn.into_iter().map(String::into_bytes).collect())
2830    };
2831    (key, matcher)
2832}
2833
2834#[derive(Debug)]
2835pub struct ClusterConfiguration {
2836    proxy_protocol: Option<ProxyProtocolConfig>,
2837    /// Whether selection reads the client's affinity key: the policy is
2838    /// `HRW` or `MAGLEV`, or the cluster shards. `TcpSession::connect_to_backend`
2839    /// derives the key only then.
2840    reads_affinity_key: bool,
2841    /// Per-cluster override of the global per-(cluster, source-IP)
2842    /// connection limit. `None` inherits the global default,
2843    /// `Some(0)` is explicit "unlimited", `Some(n > 0)` overrides.
2844    /// Resolved against `SessionManager::effective_max_connections_per_ip`
2845    /// at admit time in `connect_to_backend`.
2846    pub max_connections_per_ip: Option<u64>,
2847    /// Per-cluster override of the global per-(cluster, source-SUBNET)
2848    /// connection limit. Same three-state semantics as
2849    /// `max_connections_per_ip` above, resolved against
2850    /// `SessionManager::effective_max_connections_per_subnet` at admit
2851    /// time. An independent second cap: both must admit.
2852    pub max_connections_per_subnet: Option<u64>,
2853    /// Per-cluster override of the global backend connection attempt
2854    /// budget, resolved against
2855    /// `SessionManager::effective_max_connection_attempts` in
2856    /// `connect_to_backend`. `None` inherits the global value.
2857    pub max_connection_attempts: Option<u32>,
2858}
2859
2860pub struct TcpProxy {
2861    fronts: HashMap<ClusterId, Token>,
2862    backends: Rc<RefCell<BackendMap>>,
2863    listeners: HashMap<Token, Rc<RefCell<TcpListener>>>,
2864    configs: HashMap<ClusterId, ClusterConfiguration>,
2865    registry: Registry,
2866    sessions: Rc<RefCell<SessionManager>>,
2867    pool: Rc<RefCell<Pool>>,
2868}
2869
2870impl TcpProxy {
2871    pub fn new(
2872        registry: Registry,
2873        sessions: Rc<RefCell<SessionManager>>,
2874        pool: Rc<RefCell<Pool>>,
2875        backends: Rc<RefCell<BackendMap>>,
2876    ) -> TcpProxy {
2877        TcpProxy {
2878            backends,
2879            listeners: HashMap::new(),
2880            configs: HashMap::new(),
2881            fronts: HashMap::new(),
2882            registry,
2883            sessions,
2884            pool,
2885        }
2886    }
2887
2888    pub fn add_listener(
2889        &mut self,
2890        config: TcpListenerConfig,
2891        token: Token,
2892    ) -> Result<Token, ProxyError> {
2893        if self.listeners.contains_key(&token) {
2894            return Err(ProxyError::ListenerAlreadyPresent);
2895        }
2896        let mut tcp_listener = TcpListener::new(config, token).map_err(ProxyError::AddListener)?;
2897        // TCP frontends name their listener by address alone, so a listener
2898        // added next to one already on its address (on another interface)
2899        // takes over the routes that one serves: every listener on an address
2900        // holds the same ones, whenever it was added (sozu-proxy/sozu#719).
2901        if let Some(sibling) = self
2902            .listeners_at(tcp_listener.address)
2903            .ok()
2904            .and_then(|listeners| listeners.into_iter().next())
2905        {
2906            let sibling = sibling.borrow();
2907            tcp_listener.cluster_id.clone_from(&sibling.cluster_id);
2908            tcp_listener.sni_routes = sibling.sni_routes.clone();
2909            tcp_listener.tags = sibling.tags.clone();
2910        }
2911        match self.listeners.entry(token) {
2912            Entry::Vacant(entry) => {
2913                entry.insert(Rc::new(RefCell::new(tcp_listener)));
2914                Ok(token)
2915            }
2916            _ => Err(ProxyError::ListenerAlreadyPresent),
2917        }
2918    }
2919
2920    pub fn remove_listener(&mut self, address: &ListenerKey) -> SessionIsToBeClosed {
2921        let len = self.listeners.len();
2922
2923        self.listeners.retain(|_, l| !l.borrow().is(address));
2924        self.listeners.len() < len
2925    }
2926
2927    /// What this proxy would do with a socket handed to
2928    /// [`activate_listener`](Self::activate_listener) for `addr`. See
2929    /// [`crate::InheritedSocketFate`].
2930    ///
2931    /// The event loop is single-threaded, so nothing can change the listener's
2932    /// `active` flag between this answer and the `activate_listener` call that
2933    /// acts on it.
2934    pub fn inherited_socket_fate(&self, addr: &ListenerKey) -> crate::InheritedSocketFate {
2935        use crate::InheritedSocketFate;
2936        match self
2937            .listeners
2938            .values()
2939            .find(|listener| listener.borrow().is(addr))
2940        {
2941            None => InheritedSocketFate::Unclaimed,
2942            Some(listener) if listener.borrow().active => InheritedSocketFate::Refused,
2943            Some(_) => InheritedSocketFate::Adopted,
2944        }
2945    }
2946
2947    pub fn activate_listener(
2948        &self,
2949        addr: &ListenerKey,
2950        tcp_listener: Option<MioTcpListener>,
2951    ) -> Result<Token, ProxyError> {
2952        let listener = self
2953            .listeners
2954            .values()
2955            .find(|listener| listener.borrow().is(addr))
2956            .ok_or(ProxyError::NoListenerFound(addr.clone()))?;
2957
2958        listener.borrow_mut().activate(&self.registry, tcp_listener)
2959    }
2960
2961    /// The slab token reserved for the listener at `address`, if this proxy
2962    /// holds one.
2963    ///
2964    /// A listener owns exactly one slab slot for its whole `AddListener` ->
2965    /// `RemoveListener` lifetime, and no `remove_listener` implementation
2966    /// touches the session slab. `Server` therefore reads the token here
2967    /// BEFORE dropping the listener, so that reserved slot is released exactly
2968    /// once, at the end of the lifetime.
2969    pub fn listener_token(&self, address: &ListenerKey) -> Option<Token> {
2970        self.listeners
2971            .iter()
2972            .find(|(_, listener)| listener.borrow().is(address))
2973            .map(|(token, _)| *token)
2974    }
2975
2976    /// The no-SNI cluster the listener at `address` routes to, for tests
2977    /// that drive the proxy through `Server::notify_proxys` (`lib/src/server.rs`).
2978    #[cfg(test)]
2979    pub(crate) fn listener_cluster(&self, address: &ListenerKey) -> Option<ClusterId> {
2980        let token = self.listener_token(address)?;
2981        self.listeners[&token].borrow().cluster_id.clone()
2982    }
2983
2984    pub fn give_back_listeners(&mut self) -> Vec<(ListenerKey, MioTcpListener)> {
2985        self.listeners
2986            .values()
2987            .filter_map(|listener| {
2988                let mut owned = listener.borrow_mut();
2989                if let Some(listener) = owned.listener.take() {
2990                    // Reset `active` so a subsequent `activate()` re-binds
2991                    // instead of short-circuiting on the stale flag.
2992                    owned.active = false;
2993                    return Some((owned.key(), listener));
2994                }
2995
2996                None
2997            })
2998            .collect()
2999    }
3000
3001    pub fn give_back_listener(
3002        &mut self,
3003        address: &ListenerKey,
3004    ) -> Result<(Token, MioTcpListener), ProxyError> {
3005        let listener = self
3006            .listeners
3007            .values()
3008            .find(|listener| listener.borrow().is(address))
3009            .ok_or(ProxyError::NoListenerFound(address.clone()))?;
3010
3011        let mut owned = listener.borrow_mut();
3012
3013        let taken_listener = owned
3014            .listener
3015            .take()
3016            .ok_or(ProxyError::UnactivatedListener)?;
3017
3018        // Reset `active` so a subsequent `activate()` re-binds instead of
3019        // short-circuiting on the stale flag.
3020        owned.active = false;
3021
3022        Ok((owned.token, taken_listener))
3023    }
3024
3025    /// Apply a partial-update patch to the identified TCP listener.
3026    pub fn update_listener(&mut self, patch: UpdateTcpListenerConfig) -> Result<(), ProxyError> {
3027        let address: SocketAddr = patch.address.into();
3028        let listener = self
3029            .listeners
3030            .values()
3031            .find(|l| l.borrow().address == address)
3032            .ok_or(ProxyError::NoListenerFound(address.into()))?;
3033        listener
3034            .borrow_mut()
3035            .update_config(&patch)
3036            .map_err(|listener_error| ProxyError::ListenerActivation {
3037                address,
3038                listener_error,
3039            })
3040    }
3041
3042    /// Every listener on `address`, whatever its interface. A TCP frontend
3043    /// names its listener by address alone, so it applies to all the
3044    /// listeners sharing that address on different interfaces.
3045    fn listeners_at(
3046        &self,
3047        address: SocketAddr,
3048    ) -> Result<Vec<Rc<RefCell<TcpListener>>>, ProxyError> {
3049        let mut listeners: Vec<_> = self
3050            .listeners
3051            .iter()
3052            .filter(|(_, listener)| listener.borrow().address == address)
3053            .map(|(token, listener)| (*token, listener.clone()))
3054            .collect();
3055        // Token order, so a failure is reported the same way on every run.
3056        listeners.sort_by_key(|(token, _)| *token);
3057        let listeners: Vec<_> = listeners
3058            .into_iter()
3059            .map(|(_, listener)| listener)
3060            .collect();
3061        if listeners.is_empty() {
3062            return Err(ProxyError::NoListenerFound(address.into()));
3063        }
3064        Ok(listeners)
3065    }
3066
3067    pub fn add_tcp_front(&mut self, front: RequestTcpFrontend) -> Result<(), ProxyError> {
3068        let address = front.address.into();
3069        let listeners = self.listeners_at(address)?;
3070
3071        // Hard-reject a request that would corrupt this listener's SNI/ALPN
3072        // routing invariants BEFORE any mutation below. Config-load
3073        // (`command/src/config.rs`, sozu-proxy/sozu#1279) already rejects
3074        // the same shapes for TOML-sourced requests, but `AddTcpFrontend`
3075        // can also arrive directly over the command socket, or via
3076        // `LoadState` replay of a hand-edited/stale state file, bypassing
3077        // config.rs entirely. Every listener on the address is checked before
3078        // any of them is touched.
3079        for listener in &listeners {
3080            listener.borrow().validate_new_tcp_front(&front)?;
3081        }
3082
3083        let cluster_id = ClusterId::from(front.cluster_id);
3084        for listener in listeners {
3085            let mut listener = listener.borrow_mut();
3086            self.fronts.insert(cluster_id.clone(), listener.token);
3087
3088            match &front.sni {
3089                Some(sni) => {
3090                    // Per-frontend tags key: many SNI/ALPN fronts share one
3091                    // listener, so the bare-address key (kept for no-SNI fronts
3092                    // below) would clobber siblings — see `sni_tags_key`.
3093                    listener.set_tags(
3094                        sni_tags_key(&address, sni, &front.alpn),
3095                        Some(front.tags.clone()),
3096                    );
3097                    listener.insert_sni_route(
3098                        sni.to_owned(),
3099                        front.alpn.clone(),
3100                        cluster_id.clone(),
3101                    )?;
3102                }
3103                None => {
3104                    listener.set_tags(
3105                        TcpFrontendTagsKey::Address(address).to_string(),
3106                        Some(front.tags.clone()),
3107                    );
3108                    listener.cluster_id = Some(cluster_id.clone());
3109                }
3110            }
3111
3112            // POST: the mixing invariant must hold after every successful add —
3113            // `validate_new_tcp_front` is the enforcement point above, this is
3114            // the cheap live re-check that it actually held.
3115            debug_assert!(
3116                listener.cluster_id.is_none() || listener.sni_routes.is_empty(),
3117                "a TCP listener must never mix a no-SNI catch-all cluster with SNI-scoped routes"
3118            );
3119        }
3120
3121        Ok(())
3122    }
3123
3124    pub fn remove_tcp_front(&mut self, front: RequestTcpFrontend) -> Result<(), ProxyError> {
3125        let address = front.address.into();
3126
3127        for listener in self.listeners_at(address)? {
3128            let mut listener = listener.borrow_mut();
3129            match &front.sni {
3130                Some(sni) => {
3131                    // Clear ONLY this front's own tags entry (`sni_tags_key`) —
3132                    // the pre-SNI bare-address removal here used to strip tags
3133                    // for every sibling front on the listener.
3134                    listener.set_tags(sni_tags_key(&address, sni, &front.alpn), None);
3135                    listener.remove_sni_route(
3136                        sni.to_owned(),
3137                        front.alpn.clone(),
3138                        &front.cluster_id,
3139                    );
3140                    self.fronts.remove(front.cluster_id.as_str());
3141                }
3142                None => {
3143                    listener.set_tags(TcpFrontendTagsKey::Address(address).to_string(), None);
3144                    if let Some(cluster_id) = listener.cluster_id.take() {
3145                        self.fronts.remove(&cluster_id);
3146                    }
3147                }
3148            }
3149        }
3150
3151        Ok(())
3152    }
3153}
3154
3155impl ProxyConfiguration for TcpProxy {
3156    fn notify(&mut self, message: WorkerRequest) -> WorkerResponse {
3157        let request_type = match message.content.request_type {
3158            Some(t) => t,
3159            None => return WorkerResponse::error(message.id, "Empty request"),
3160        };
3161        match request_type {
3162            RequestType::AddTcpFrontend(front) => {
3163                if let Err(err) = self.add_tcp_front(front) {
3164                    return WorkerResponse::error(message.id, err);
3165                }
3166
3167                WorkerResponse::ok(message.id)
3168            }
3169            RequestType::RemoveTcpFrontend(front) => {
3170                if let Err(err) = self.remove_tcp_front(front) {
3171                    return WorkerResponse::error(message.id, err);
3172                }
3173
3174                WorkerResponse::ok(message.id)
3175            }
3176            RequestType::SoftStop(_) => {
3177                info!(
3178                    "{} {} processing soft shutdown",
3179                    log_module_context!(),
3180                    message.id
3181                );
3182                let listeners = std::mem::take(&mut self.listeners);
3183                for l in listeners.values() {
3184                    l.borrow_mut()
3185                        .listener
3186                        .take()
3187                        .map(|mut sock| self.registry.deregister(&mut sock));
3188                }
3189                WorkerResponse::processing(message.id)
3190            }
3191            RequestType::HardStop(_) => {
3192                info!("{} {} hard shutdown", log_module_context!(), message.id);
3193                for (_, l) in std::mem::take(&mut self.listeners) {
3194                    l.borrow_mut()
3195                        .listener
3196                        .take()
3197                        .map(|mut sock| self.registry.deregister(&mut sock));
3198                }
3199                WorkerResponse::ok(message.id)
3200            }
3201            RequestType::Status(_) => {
3202                info!("{} {} status", log_module_context!(), message.id);
3203                WorkerResponse::ok(message.id)
3204            }
3205            RequestType::AddCluster(cluster) => {
3206                let config = ClusterConfiguration {
3207                    proxy_protocol: cluster
3208                        .proxy_protocol
3209                        .and_then(|n| ProxyProtocolConfig::try_from(n).ok()),
3210                    reads_affinity_key: cluster_reads_affinity_key(&cluster),
3211                    max_connections_per_ip: cluster.max_connections_per_ip,
3212                    max_connections_per_subnet: cluster.max_connections_per_subnet,
3213                    max_connection_attempts: cluster.max_connection_attempts,
3214                };
3215                self.configs.insert(cluster.cluster_id.into(), config);
3216                WorkerResponse::ok(message.id)
3217            }
3218            RequestType::RemoveCluster(cluster_id) => {
3219                self.configs.remove(cluster_id.as_str());
3220                WorkerResponse::ok(message.id)
3221            }
3222            RequestType::RemoveListener(remove) => {
3223                if !self.remove_listener(&remove.listener_key()) {
3224                    WorkerResponse::error(
3225                        message.id,
3226                        format!("no TCP listener to remove at address {:?}", remove.address),
3227                    )
3228                } else {
3229                    WorkerResponse::ok(message.id)
3230                }
3231            }
3232            command => {
3233                debug!(
3234                    "{} {} unsupported message for TCP proxy, ignoring {:?}",
3235                    log_module_context!(),
3236                    message.id,
3237                    command
3238                );
3239                WorkerResponse::error(message.id, "unsupported message")
3240            }
3241        }
3242    }
3243
3244    fn accept(&mut self, token: ListenToken) -> Result<(MioTcpStream, SocketAddr), AcceptError> {
3245        let internal_token = Token(token.0);
3246        if let Some(listener) = self.listeners.get(&internal_token) {
3247            let mut listener = listener.borrow_mut();
3248            let accepted = match &listener.listener {
3249                Some(tcp_listener) => tcp_listener.accept(),
3250                None => return Err(AcceptError::IoError),
3251            };
3252            match accepted {
3253                Ok((frontend_sock, peer)) => {
3254                    if listener.nodelay_backlog
3255                        && let Err(e) = frontend_sock.set_nodelay(true)
3256                    {
3257                        error!(
3258                            "{} error setting nodelay on front socket({:?}): {:?}",
3259                            log_module_context!(),
3260                            frontend_sock,
3261                            e
3262                        );
3263                    }
3264                    Ok((frontend_sock, peer))
3265                }
3266                Err(e) if e.kind() == ErrorKind::WouldBlock => {
3267                    // The backlog queued before `activate()` is drained.
3268                    listener.nodelay_backlog = false;
3269                    Err(AcceptError::WouldBlock)
3270                }
3271                Err(e) => {
3272                    error!("{} accept() IO error: {:?}", log_module_context!(), e);
3273                    Err(AcceptError::IoError)
3274                }
3275            }
3276        } else {
3277            Err(AcceptError::IoError)
3278        }
3279    }
3280
3281    fn create_session(
3282        &mut self,
3283        mut frontend_sock: MioTcpStream,
3284        peer: SocketAddr,
3285        token: ListenToken,
3286        wait_time: Duration,
3287        proxy: Rc<RefCell<Self>>,
3288    ) -> Result<(), AcceptError> {
3289        let listener_token = Token(token.0);
3290
3291        let listener = self
3292            .listeners
3293            .get(&listener_token)
3294            .ok_or(AcceptError::IoError)?;
3295
3296        let owned = listener.borrow();
3297        let mut pool = self.pool.borrow_mut();
3298
3299        let (front_buffer, back_buffer) = match (pool.checkout(), pool.checkout()) {
3300            (Some(fb), Some(bb)) => (fb, bb),
3301            _ => {
3302                error!("{} could not get buffers from pool", log_module_context!());
3303                error!(
3304                    "{} Buffer capacity has been reached, stopping to accept new connections for now",
3305                    log_module_context!()
3306                );
3307                gauge!(names::accept_queue::BACKPRESSURE, 1);
3308                self.sessions.borrow_mut().can_accept = false;
3309
3310                return Err(AcceptError::BufferCapacityReached);
3311            }
3312        };
3313
3314        // A listener may route either by a legacy no-SNI catch-all cluster
3315        // OR by SNI-scoped routes (never both -- enforced at config load,
3316        // sozu-proxy/sozu#1279); reject only when NEITHER is configured.
3317        if owned.cluster_id.is_none() && owned.sni_routes.is_empty() {
3318            error!(
3319                "{} listener at address {:?} has no linked cluster",
3320                log_module_context!(),
3321                owned.address
3322            );
3323            return Err(AcceptError::IoError);
3324        }
3325
3326        // No `set_nodelay` here: the socket has `TCP_NODELAY` from its
3327        // listener, or from `TcpProxy::accept` if it was queued before the
3328        // listener set it.
3329
3330        let mut session_manager = self.sessions.borrow_mut();
3331        let entry = session_manager.slab.vacant_entry();
3332        let frontend_token = Token(entry.key());
3333
3334        if let Err(register_error) = self.registry.register(
3335            &mut frontend_sock,
3336            frontend_token,
3337            Interest::READABLE | Interest::WRITABLE,
3338        ) {
3339            error!(
3340                "{} error registering front socket({:?}): {:?}",
3341                log_module_context!(),
3342                frontend_sock,
3343                register_error
3344            );
3345            return Err(AcceptError::RegisterError);
3346        }
3347
3348        let session = if !owned.sni_routes.is_empty() {
3349            // Routing decides the cluster post-accept; the effective
3350            // preread cap can never exceed what the checked-out buffer can
3351            // actually hold, regardless of the configured knob.
3352            let effective_max_bytes = effective_sni_preread_max_bytes(
3353                owned.config.sni_preread_max_bytes,
3354                front_buffer.capacity(),
3355            );
3356            let preread_timeout = Duration::from_secs(u64::from(
3357                owned
3358                    .config
3359                    .sni_preread_timeout
3360                    .unwrap_or(DEFAULT_SNI_PREREAD_TIMEOUT),
3361            ));
3362            TcpSession::new_sni_preread(
3363                back_buffer,
3364                Duration::from_secs(owned.config.back_timeout as u64),
3365                Duration::from_secs(owned.config.connect_timeout as u64),
3366                front_buffer,
3367                frontend_token,
3368                listener.clone(),
3369                proxy,
3370                frontend_sock,
3371                peer,
3372                wait_time,
3373                preread_timeout,
3374                effective_max_bytes,
3375            )
3376        } else {
3377            let proxy_protocol = self
3378                .configs
3379                .get(owned.cluster_id.as_ref().unwrap())
3380                .and_then(|c| c.proxy_protocol);
3381            TcpSession::new(
3382                back_buffer,
3383                None,
3384                owned.cluster_id.clone(),
3385                Duration::from_secs(owned.config.back_timeout as u64),
3386                Duration::from_secs(owned.config.connect_timeout as u64),
3387                Duration::from_secs(owned.config.front_timeout as u64),
3388                front_buffer,
3389                frontend_token,
3390                listener.clone(),
3391                proxy_protocol,
3392                proxy,
3393                frontend_sock,
3394                peer,
3395                wait_time,
3396            )
3397        };
3398        incr!(names::tcp::REQUESTS);
3399
3400        let session = Rc::new(RefCell::new(session));
3401        entry.insert(session);
3402
3403        Ok(())
3404    }
3405}
3406
3407pub mod testing {
3408    use crate::testing::*;
3409
3410    /// This is not directly used by Sōzu but is available for example and testing purposes
3411    pub fn start_tcp_worker(
3412        config: TcpListenerConfig,
3413        max_buffers: usize,
3414        buffer_size: usize,
3415        channel: ProxyChannel,
3416    ) -> anyhow::Result<()> {
3417        let address = config.address.into();
3418
3419        let ServerParts {
3420            event_loop,
3421            registry,
3422            sessions,
3423            pool,
3424            backends,
3425            client_scm_socket: _,
3426            server_scm_socket,
3427            server_config,
3428        } = prebuild_server(max_buffers, buffer_size, true)?;
3429
3430        let token = {
3431            let mut sessions = sessions.borrow_mut();
3432            let entry = sessions.slab.vacant_entry();
3433            let key = entry.key();
3434            let _ = entry.insert(Rc::new(RefCell::new(ListenSession {
3435                protocol: Protocol::TCPListen,
3436            })));
3437            Token(key)
3438        };
3439
3440        let mut proxy = TcpProxy::new(registry, sessions.clone(), pool.clone(), backends.clone());
3441        proxy
3442            .add_listener(config, token)
3443            .with_context(|| "Failed at creating adding the listener")?;
3444        proxy
3445            .activate_listener(&address, None)
3446            .with_context(|| "Failed at creating activating the listener")?;
3447
3448        let mut server = Server::new(
3449            event_loop,
3450            channel,
3451            server_scm_socket,
3452            sessions,
3453            pool,
3454            backends,
3455            None,
3456            None,
3457            Some(proxy),
3458            server_config,
3459            None,
3460            false,
3461        )
3462        .with_context(|| "Failed at creating server")?;
3463
3464        debug!("{} starting event loop", log_module_context!());
3465        server.run();
3466        debug!("{} ending event loop", log_module_context!());
3467        Ok(())
3468    }
3469}
3470
3471#[cfg(test)]
3472mod tests {
3473    use std::{
3474        io::{Read, Write},
3475        net::{Shutdown, TcpListener, TcpStream},
3476        str,
3477        sync::{
3478            Arc, Barrier,
3479            atomic::{AtomicBool, Ordering},
3480        },
3481        thread,
3482        time::Duration,
3483    };
3484
3485    use sozu_command::{
3486        channel::Channel,
3487        config::ListenerBuilder,
3488        proto::command::{
3489            LoadBalancingParams, RequestTcpFrontend, SocketAddress, SoftStop, WorkerRequest,
3490            WorkerResponse, request::RequestType,
3491        },
3492    };
3493
3494    use super::testing::start_tcp_worker;
3495    use crate::testing::*;
3496
3497    /*
3498    #[test]
3499    #[cfg(target_pointer_width = "64")]
3500    fn size_test() {
3501      assert_size!(Pipe<mio::net::TcpStream>, 224);
3502      assert_size!(SendProxyProtocol<mio::net::TcpStream>, 144);
3503      assert_size!(RelayProxyProtocol<mio::net::TcpStream>, 152);
3504      assert_size!(ExpectProxyProtocol<mio::net::TcpStream>, 520);
3505      assert_size!(State, 528);
3506      // fails depending on the platform?
3507      //assert_size!(Session, 808);
3508    }*/
3509
3510    #[test]
3511    fn round_trip() {
3512        setup_test_logger!();
3513        let barrier = Arc::new(Barrier::new(2));
3514        let test_finished = Arc::new(AtomicBool::new(false));
3515
3516        let front_port1 = provide_port();
3517        let front_port2 = provide_port();
3518
3519        let backend_port = start_server(barrier.clone(), test_finished.clone());
3520        let mut command =
3521            start_proxy(backend_port, front_port1, front_port2).expect("Could not start proxy");
3522        barrier.wait();
3523
3524        thread::scope(|_s| {
3525            let front_addr = format!("127.0.0.1:{front_port1}");
3526
3527            let mut s1 = TcpStream::connect(&front_addr).expect("could not connect");
3528            s1.set_read_timeout(Some(Duration::from_secs(5)))
3529                .expect("could not set read timeout on s1");
3530
3531            let s3 = TcpStream::connect(&front_addr).expect("could not connect");
3532
3533            let mut s2 = TcpStream::connect(&front_addr).expect("could not connect");
3534            s2.set_read_timeout(Some(Duration::from_secs(5)))
3535                .expect("could not set read timeout on s2");
3536
3537            s1.write_all(b"hello ").expect("could not write to s1");
3538            println!("s1 sent");
3539
3540            s2.write_all(b"pouet pouet").expect("could not write to s2");
3541            println!("s2 sent");
3542
3543            let mut res = [0; 128];
3544            s1.write_all(b"coucou").expect("could not write to s1");
3545
3546            s3.shutdown(Shutdown::Both).expect("could not shutdown s3");
3547
3548            let sz2 = s2
3549                .read(&mut res[..])
3550                .expect("could not read from socket s2");
3551            println!("s2 received {:?}", str::from_utf8(&res[..sz2]));
3552            assert_eq!(&res[..sz2], &b"pouet pouet"[..]);
3553
3554            // Read in a loop: a single read() on a TCP stream is not
3555            // guaranteed to return all echoed data if the second write's
3556            // round trip (client → proxy → backend → proxy → client) is
3557            // still in flight when we poll.
3558            let expected = b"hello coucou";
3559            let mut total = 0;
3560            while total < expected.len() {
3561                let sz = s1
3562                    .read(&mut res[total..])
3563                    .expect("could not read from socket s1");
3564                assert!(sz > 0, "connection closed before receiving all data");
3565                total += sz;
3566            }
3567            println!(
3568                "s1 received again({}): {:?}",
3569                total,
3570                str::from_utf8(&res[..total])
3571            );
3572            assert_eq!(&res[..total], &expected[..]);
3573
3574            // Signal the echo server to stop
3575            test_finished.store(true, Ordering::Relaxed);
3576
3577            // Send SoftStop to the sozu worker so server.run() exits cleanly
3578            command
3579                .write_message(&WorkerRequest {
3580                    id: "ID_SOFTSTOP".to_owned(),
3581                    content: RequestType::SoftStop(SoftStop {}).into(),
3582                })
3583                .expect("could not send SoftStop to sozu worker");
3584        });
3585    }
3586
3587    /// Start an echo server on an ephemeral port.
3588    /// Returns the port the server is listening on.
3589    fn start_server(barrier: Arc<Barrier>, test_finished: Arc<AtomicBool>) -> u16 {
3590        let listener =
3591            TcpListener::bind("127.0.0.1:0").expect("could not bind echo server listener");
3592        let port = listener
3593            .local_addr()
3594            .expect("could not get echo server local address")
3595            .port();
3596
3597        listener
3598            .set_nonblocking(true)
3599            .expect("could not set echo server listener to non-blocking");
3600
3601        thread::spawn(move || {
3602            barrier.wait();
3603            let mut count: u8 = 0;
3604            loop {
3605                match listener.accept() {
3606                    Ok((mut stream, _)) => {
3607                        let finished = test_finished.clone();
3608                        thread::spawn(move || {
3609                            println!("got a new client: {count}");
3610                            stream
3611                                .set_read_timeout(Some(Duration::from_secs(2)))
3612                                .expect("could not set read timeout on echo client");
3613                            let mut buf = [0; 128];
3614                            loop {
3615                                match stream.read(&mut buf[..]) {
3616                                    Ok(0) => break,
3617                                    Ok(sz) => {
3618                                        println!(
3619                                            "ECHO[{count}] got \"{:?}\"",
3620                                            str::from_utf8(&buf[..sz])
3621                                        );
3622                                        stream
3623                                            .write_all(&buf[..sz])
3624                                            .expect("could not echo data back");
3625                                    }
3626                                    Err(ref e)
3627                                        if e.kind() == std::io::ErrorKind::WouldBlock
3628                                            || e.kind() == std::io::ErrorKind::TimedOut =>
3629                                    {
3630                                        if finished.load(Ordering::Relaxed) {
3631                                            println!("backend server stopping (client handler)");
3632                                            break;
3633                                        }
3634                                    }
3635                                    Err(_) => break,
3636                                }
3637                            }
3638                        });
3639                        count = count.wrapping_add(1);
3640                    }
3641                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
3642                        if test_finished.load(Ordering::Relaxed) {
3643                            println!("backend server stopping (accept loop)");
3644                            break;
3645                        }
3646                        thread::sleep(Duration::from_millis(50));
3647                    }
3648                    Err(e) => {
3649                        println!("connection failed: {e:?}");
3650                    }
3651                }
3652            }
3653        });
3654
3655        port
3656    }
3657
3658    /// Start a sozu TCP proxy worker with the given backend and frontend ports.
3659    fn start_proxy(
3660        backend_port: u16,
3661        front_port1: u16,
3662        front_port2: u16,
3663    ) -> anyhow::Result<Channel<WorkerRequest, WorkerResponse>> {
3664        let config = ListenerBuilder::new_tcp(SocketAddress::new_v4(127, 0, 0, 1, front_port1))
3665            .to_tcp(None)
3666            .expect("could not create listener config");
3667
3668        let (mut command, channel) =
3669            Channel::generate(1000, 10000).with_context(|| "should create a channel")?;
3670        let _jg = thread::spawn(move || {
3671            setup_test_logger!();
3672            start_tcp_worker(config, 100, 16384, channel).expect("could not start the tcp server");
3673        });
3674
3675        command
3676            .blocking()
3677            .expect("could not set command channel to blocking");
3678        {
3679            let front = RequestTcpFrontend {
3680                cluster_id: "yolo".to_owned(),
3681                address: SocketAddress::new_v4(127, 0, 0, 1, front_port1),
3682                ..Default::default()
3683            };
3684            let backend = sozu_command_lib::response::Backend {
3685                cluster_id: "yolo".into(),
3686                backend_id: "yolo-0".to_owned(),
3687                address: SocketAddress::new_v4(127, 0, 0, 1, backend_port).into(),
3688                load_balancing_parameters: Some(LoadBalancingParams::default()),
3689                sticky_id: None,
3690                backup: None,
3691            };
3692
3693            command
3694                .write_message(&WorkerRequest {
3695                    id: "ID_YOLO1".to_owned(),
3696                    content: RequestType::AddTcpFrontend(front).into(),
3697                })
3698                .expect("could not send AddTcpFrontend for front1");
3699            command
3700                .write_message(&WorkerRequest {
3701                    id: "ID_YOLO2".to_owned(),
3702                    content: RequestType::AddBackend(backend.to_add_backend()).into(),
3703                })
3704                .expect("could not send AddBackend for front1");
3705        }
3706        {
3707            let front = RequestTcpFrontend {
3708                cluster_id: "yolo".to_owned(),
3709                address: SocketAddress::new_v4(127, 0, 0, 1, front_port2),
3710                ..Default::default()
3711            };
3712            let backend = sozu_command::response::Backend {
3713                cluster_id: "yolo".into(),
3714                backend_id: "yolo-0".to_owned(),
3715                address: SocketAddress::new_v4(127, 0, 0, 1, backend_port).into(),
3716                load_balancing_parameters: Some(LoadBalancingParams::default()),
3717                sticky_id: None,
3718                backup: None,
3719            };
3720            command
3721                .write_message(&WorkerRequest {
3722                    id: "ID_YOLO3".to_owned(),
3723                    content: RequestType::AddTcpFrontend(front).into(),
3724                })
3725                .expect("could not send AddTcpFrontend for front2");
3726            command
3727                .write_message(&WorkerRequest {
3728                    id: "ID_YOLO4".to_owned(),
3729                    content: RequestType::AddBackend(backend.to_add_backend()).into(),
3730                })
3731                .expect("could not send AddBackend for front2");
3732        }
3733
3734        for _ in 0..4 {
3735            println!(
3736                "read_message: {:?}",
3737                command
3738                    .read_message()
3739                    .with_context(|| "could not read message")?
3740            );
3741        }
3742
3743        Ok(command)
3744    }
3745}
3746
3747/// Unit coverage for the SNI-preread routing shell added for
3748/// sozu-proxy/sozu#1279: the route-table mutations (`add_tcp_front` /
3749/// `remove_tcp_front`), the `AlpnMatcher` mapping, the effective preread
3750/// cap, and the routing gate's data invariants. None of this needs a live
3751/// socket or event loop -- it is a separate module (rather than nested in
3752/// the `tests` module above) purely to avoid that module's `use
3753/// std::net::TcpListener` import shadowing `super::TcpListener` (this
3754/// crate's listener struct).
3755#[cfg(test)]
3756mod sni_routing_tests {
3757    use sozu_command::{config::ListenerBuilder, proto::command::SocketAddress};
3758
3759    use super::*;
3760    use crate::testing::{ServerParts, prebuild_server, provide_port};
3761
3762    fn test_listener() -> TcpListener {
3763        let config = ListenerBuilder::new_tcp(SocketAddress::new_v4(127, 0, 0, 1, provide_port()))
3764            .to_tcp(None)
3765            .expect("could not build a TcpListenerConfig for the test");
3766        TcpListener::new(config, Token(0)).expect("could not build a bare TcpListener for the test")
3767    }
3768
3769    fn frontend(cluster_id: &str, sni: Option<&str>, alpn: &[&str]) -> RequestTcpFrontend {
3770        RequestTcpFrontend {
3771            cluster_id: cluster_id.to_owned(),
3772            address: SocketAddress::new_v4(127, 0, 0, 1, provide_port()),
3773            sni: sni.map(str::to_owned),
3774            alpn: alpn.iter().map(|p| p.to_string()).collect(),
3775            ..Default::default()
3776        }
3777    }
3778
3779    // ---- effective_sni_preread_max_bytes ------------------------------
3780
3781    #[test]
3782    fn effective_max_bytes_falls_back_to_default_when_unconfigured() {
3783        assert_eq!(
3784            effective_sni_preread_max_bytes(None, 65536),
3785            DEFAULT_SNI_PREREAD_MAX_BYTES as usize
3786        );
3787    }
3788
3789    #[test]
3790    fn effective_max_bytes_is_the_min_of_knob_and_capacity() {
3791        assert_eq!(effective_sni_preread_max_bytes(Some(8192), 16384), 8192);
3792        assert_eq!(effective_sni_preread_max_bytes(Some(32768), 16384), 16384);
3793        assert_eq!(effective_sni_preread_max_bytes(Some(16384), 16384), 16384);
3794    }
3795
3796    #[test]
3797    fn effective_max_bytes_never_below_the_floor() {
3798        // A `sni_preread_max_bytes = 0` knob reaching the worker from a
3799        // direct `sozu listener tcp add`/`update` CLI/IPC request (or a stale
3800        // LoadState replay) bypasses config.rs's loud MIN_SNI_PREREAD_MAX_BYTES
3801        // load-time reject. Without the floor the shell would issue
3802        // zero-length preread reads and spin until the loop guard closes each
3803        // session; the floor degrades a sub-minimum knob to the 5-byte
3804        // TLS-record-header minimum instead.
3805        assert_eq!(
3806            effective_sni_preread_max_bytes(Some(0), 16384),
3807            MIN_SNI_PREREAD_MAX_BYTES as usize,
3808            "a 0 knob must degrade to the floor, never 0 (would spin the preread)"
3809        );
3810        assert_eq!(
3811            effective_sni_preread_max_bytes(Some(3), 16384),
3812            MIN_SNI_PREREAD_MAX_BYTES as usize,
3813            "any sub-floor knob must be raised to the floor"
3814        );
3815        // The floor itself, and anything above it, are respected unchanged.
3816        assert_eq!(
3817            effective_sni_preread_max_bytes(Some(MIN_SNI_PREREAD_MAX_BYTES), 16384),
3818            MIN_SNI_PREREAD_MAX_BYTES as usize
3819        );
3820        assert_eq!(
3821            effective_sni_preread_max_bytes(Some(MIN_SNI_PREREAD_MAX_BYTES + 1), 16384),
3822            (MIN_SNI_PREREAD_MAX_BYTES + 1) as usize
3823        );
3824    }
3825
3826    // ---- known_alpn_label (access-log tagging) -------------------------
3827
3828    #[test]
3829    fn known_alpn_label_picks_the_clients_first_offer() {
3830        assert_eq!(
3831            known_alpn_label(&[b"h2".to_vec(), b"http/1.1".to_vec()]),
3832            Some("h2")
3833        );
3834        assert_eq!(
3835            known_alpn_label(&[b"http/1.1".to_vec(), b"h2".to_vec()]),
3836            Some("http/1.1"),
3837            "client preference order must win, not a fixed h2-first priority"
3838        );
3839    }
3840
3841    #[test]
3842    fn known_alpn_label_is_none_for_empty_or_unrecognized_offers() {
3843        assert_eq!(known_alpn_label(&[]), None);
3844        assert_eq!(known_alpn_label(&[b"spdy/1".to_vec()]), None);
3845    }
3846
3847    // ---- AlpnMatcher mapping + route-table add/remove symmetry --------
3848
3849    #[test]
3850    fn empty_alpn_maps_to_any_non_empty_maps_to_one_of() {
3851        let mut listener = test_listener();
3852        listener
3853            .insert_sni_route("example.com".to_owned(), vec![], "cluster-any".into())
3854            .expect("insert_sni_route must succeed for a valid test SNI");
3855        listener
3856            .insert_sni_route(
3857                "h2.example.com".to_owned(),
3858                vec!["h2".to_owned(), "http/1.1".to_owned()],
3859                "cluster-h2".into(),
3860            )
3861            .expect("insert_sni_route must succeed for a valid test SNI");
3862
3863        let (_, any_entries) = listener
3864            .sni_routes
3865            .domain_lookup(b"example.com", true)
3866            .expect("example.com must be routable");
3867        assert_eq!(any_entries, &vec![(AlpnMatcher::Any, "cluster-any".into())]);
3868
3869        let (_, h2_entries) = listener
3870            .sni_routes
3871            .domain_lookup(b"h2.example.com", true)
3872            .expect("h2.example.com must be routable");
3873        assert_eq!(
3874            h2_entries,
3875            &vec![(
3876                AlpnMatcher::OneOf([b"h2".to_vec(), b"http/1.1".to_vec()].into_iter().collect()),
3877                "cluster-h2".into()
3878            )]
3879        );
3880    }
3881
3882    #[test]
3883    fn insert_sni_route_appends_under_the_same_sni() {
3884        let mut listener = test_listener();
3885        listener
3886            .insert_sni_route(
3887                "example.com".to_owned(),
3888                vec!["h2".to_owned()],
3889                "cluster-h2".into(),
3890            )
3891            .expect("insert_sni_route must succeed for a valid test SNI");
3892        listener
3893            .insert_sni_route("example.com".to_owned(), vec![], "cluster-default".into())
3894            .expect("insert_sni_route must succeed for a valid test SNI");
3895
3896        let (_, entries) = listener
3897            .sni_routes
3898            .domain_lookup(b"example.com", true)
3899            .expect("example.com must be routable");
3900        assert_eq!(entries.len(), 2, "both fronts must share the SNI's Vec");
3901    }
3902
3903    #[test]
3904    fn remove_sni_route_drops_only_the_matching_entry() {
3905        let mut listener = test_listener();
3906        listener
3907            .insert_sni_route(
3908                "example.com".to_owned(),
3909                vec!["h2".to_owned()],
3910                "cluster-h2".into(),
3911            )
3912            .expect("insert_sni_route must succeed for a valid test SNI");
3913        listener
3914            .insert_sni_route("example.com".to_owned(), vec![], "cluster-default".into())
3915            .expect("insert_sni_route must succeed for a valid test SNI");
3916
3917        listener.remove_sni_route(
3918            "example.com".to_owned(),
3919            vec!["h2".to_owned()],
3920            "cluster-h2",
3921        );
3922
3923        let (_, entries) = listener
3924            .sni_routes
3925            .domain_lookup(b"example.com", true)
3926            .expect("example.com must still be routable via the remaining entry");
3927        assert_eq!(
3928            entries,
3929            &vec![(AlpnMatcher::Any, "cluster-default".into())],
3930            "removing one entry must not disturb the other"
3931        );
3932    }
3933
3934    #[test]
3935    fn remove_sni_route_empties_the_trie_key_when_the_last_entry_goes() {
3936        let mut listener = test_listener();
3937        listener
3938            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".into())
3939            .expect("insert_sni_route must succeed for a valid test SNI");
3940        assert!(!listener.sni_routes.is_empty());
3941
3942        listener.remove_sni_route("example.com".to_owned(), vec![], "cluster-a");
3943
3944        assert!(
3945            listener.sni_routes.is_empty(),
3946            "domain_remove must run once the SNI's Vec empties, leaving no stranded key"
3947        );
3948        assert!(
3949            listener
3950                .sni_routes
3951                .domain_lookup(b"example.com", true)
3952                .is_none()
3953        );
3954    }
3955
3956    #[test]
3957    fn remove_sni_route_on_an_absent_sni_is_a_harmless_no_op() {
3958        let mut listener = test_listener();
3959        listener
3960            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".into())
3961            .expect("insert_sni_route must succeed for a valid test SNI");
3962
3963        // Removing a route for a SNI that was never inserted must not panic
3964        // and must not disturb the existing route.
3965        listener.remove_sni_route("other.example.net".to_owned(), vec![], "cluster-a");
3966
3967        assert!(
3968            listener
3969                .sni_routes
3970                .domain_lookup(b"example.com", true)
3971                .is_some()
3972        );
3973    }
3974
3975    // ---- exact-key bookkeeping must not fall back to a sibling wildcard
3976    // (route-table corruption caught in sozu-proxy/sozu#1290 review) ----
3977
3978    #[test]
3979    fn insert_sni_route_creates_a_distinct_node_for_an_exact_key_over_a_sibling_wildcard() {
3980        let mut listener = test_listener();
3981        // Wildcard catch-all first.
3982        listener
3983            .insert_sni_route(
3984                "*.example.com".to_owned(),
3985                vec![],
3986                "cluster-wildcard".into(),
3987            )
3988            .expect("insert_sni_route must succeed for a valid test SNI");
3989        // Exact ALPN-scoped route for one specific subdomain.
3990        listener
3991            .insert_sni_route(
3992                "a.example.com".to_owned(),
3993                vec!["h2".to_owned()],
3994                "cluster-a-h2".into(),
3995            )
3996            .expect("insert_sni_route must succeed for a valid test SNI");
3997
3998        // The exact key must have gotten its OWN trie node -- with
3999        // `accept_wildcard: true` this lookup would instead fall back to
4000        // (and the insert above would have corrupted) the wildcard's node,
4001        // since no literal `a` child existed yet at insert time.
4002        let (_, a_entries) = listener
4003            .sni_routes
4004            .domain_lookup(b"a.example.com", false)
4005            .expect("a.example.com must have a distinct exact-key node");
4006        assert_eq!(
4007            a_entries,
4008            &vec![(
4009                AlpnMatcher::OneOf([b"h2".to_vec()].into_iter().collect()),
4010                "cluster-a-h2".into()
4011            )],
4012            "the exact key's own Vec must hold only its own entry, not the wildcard's"
4013        );
4014
4015        // Any OTHER subdomain must still resolve to ONLY the wildcard, via
4016        // the same `accept_wildcard: true` lookup the routing path uses
4017        // (`preread_config`) -- it must never see `a.example.com`'s h2 route.
4018        let (_, b_entries) = listener
4019            .sni_routes
4020            .domain_lookup(b"b.example.com", true)
4021            .expect("b.example.com must fall back to the wildcard catch-all");
4022        assert_eq!(
4023            b_entries,
4024            &vec![(AlpnMatcher::Any, "cluster-wildcard".into())],
4025            "an unrelated subdomain must see ONLY the wildcard's entry"
4026        );
4027    }
4028
4029    #[test]
4030    fn validate_new_tcp_front_accepts_an_exact_catch_all_sibling_of_a_wildcard_catch_all() {
4031        let mut listener = test_listener();
4032        listener
4033            .insert_sni_route(
4034                "*.example.com".to_owned(),
4035                vec![],
4036                "cluster-wildcard".into(),
4037            )
4038            .expect("insert_sni_route must succeed for a valid test SNI");
4039
4040        // An exact catch-all for one subdomain must be accepted: it is a
4041        // SIBLING of the wildcard's catch-all, not a duplicate of it. With
4042        // `accept_wildcard: true` this lookup would wrongly find the
4043        // wildcard's own `AlpnMatcher::Any` entry and reject it as "already
4044        // has a catch-all".
4045        let front = frontend("cluster-a", Some("a.example.com"), &[]);
4046        assert!(
4047            listener.validate_new_tcp_front(&front).is_ok(),
4048            "an exact catch-all must be accepted when only a SIBLING wildcard has a catch-all"
4049        );
4050    }
4051
4052    #[test]
4053    fn validate_new_tcp_front_rejects_a_duplicate_wildcard_catch_all() {
4054        let mut listener = test_listener();
4055        listener
4056            .insert_sni_route(
4057                "*.example.com".to_owned(),
4058                vec![],
4059                "cluster-wildcard".into(),
4060            )
4061            .expect("insert_sni_route must succeed for a valid test SNI");
4062
4063        // A SECOND catch-all for the SAME wildcard key is an ambiguous
4064        // duplicate and must be rejected. The immutable trie `lookup` has no
4065        // literal-`*` short-circuit (only `lookup_mut` does), so a plain
4066        // `accept_wildcard: false` self-lookup never sees the existing
4067        // wildcard entry and waves the duplicate through -- which
4068        // `insert_sni_route` (lookup_mut, short-circuit present) would then
4069        // happily append.
4070        let front = frontend("cluster-dup", Some("*.example.com"), &[]);
4071        assert!(
4072            listener.validate_new_tcp_front(&front).is_err(),
4073            "a second catch-all on the same wildcard SNI must be rejected as a duplicate"
4074        );
4075    }
4076
4077    #[test]
4078    fn validate_new_tcp_front_rejects_overlapping_alpn_on_the_same_wildcard() {
4079        let mut listener = test_listener();
4080        listener
4081            .insert_sni_route(
4082                "*.example.com".to_owned(),
4083                vec!["h2".to_owned()],
4084                "cluster-wildcard-h2".into(),
4085            )
4086            .expect("insert_sni_route must succeed for a valid test SNI");
4087
4088        // Same wildcard key, overlapping ALPN protocol: ambiguous, must be
4089        // rejected (same bypass as the duplicate catch-all above).
4090        let front = frontend("cluster-dup", Some("*.example.com"), &["h2"]);
4091        assert!(
4092            listener.validate_new_tcp_front(&front).is_err(),
4093            "an overlapping ALPN matcher on the same wildcard SNI must be rejected"
4094        );
4095    }
4096
4097    #[test]
4098    fn validate_new_tcp_front_accepts_a_disjoint_alpn_addition_on_the_same_wildcard() {
4099        let mut listener = test_listener();
4100        listener
4101            .insert_sni_route(
4102                "*.example.com".to_owned(),
4103                vec![],
4104                "cluster-wildcard".into(),
4105            )
4106            .expect("insert_sni_route must succeed for a valid test SNI");
4107
4108        // A non-overlapping ALPN-scoped addition alongside the wildcard's
4109        // catch-all stays legal -- the wildcard-aware self-lookup must not
4110        // over-reject.
4111        let front = frontend("cluster-h2", Some("*.example.com"), &["h2"]);
4112        assert!(
4113            listener.validate_new_tcp_front(&front).is_ok(),
4114            "a disjoint ALPN addition on the same wildcard SNI must be accepted"
4115        );
4116    }
4117
4118    /// The worker boundary must reject malformed SNI SHAPES that a direct
4119    /// `AddTcpFrontend` (command socket) or `LoadState` replay could carry
4120    /// past config.rs's `validate_sni_pattern`: a `/.../` label would be
4121    /// inserted into the `pattern_trie` as a REGEX route and a misplaced
4122    /// `*` as an unintended wildcard — silently widening routing.
4123    #[test]
4124    fn validate_new_tcp_front_rejects_malformed_sni_shapes() {
4125        let listener = test_listener();
4126
4127        for bad in [
4128            "/[a-z]+/.example.com",
4129            "foo/bar.example.com",
4130            "a.*.example.com",
4131            "*.*.example.com",
4132            "*",
4133        ] {
4134            let front = frontend("cluster-a", Some(bad), &[]);
4135            assert!(
4136                listener.validate_new_tcp_front(&front).is_err(),
4137                "malformed SNI shape {bad:?} must be rejected at the worker boundary"
4138            );
4139        }
4140
4141        // The two documented-legal shapes must still pass.
4142        for good in ["a.example.com", "*.example.com"] {
4143            let front = frontend("cluster-a", Some(good), &[]);
4144            assert!(
4145                listener.validate_new_tcp_front(&front).is_ok(),
4146                "legal SNI shape {good:?} must still be accepted"
4147            );
4148        }
4149    }
4150
4151    /// `validate_new_tcp_front`'s OLD hand-rolled shape check only rejected
4152    /// `/` and a misplaced `*`, letting an empty label (leading dot,
4153    /// trailing dot, consecutive dots), an empty string, or a non-ASCII
4154    /// pattern reach `insert_sni_route` -> `pattern_trie::insert_recursive`,
4155    /// whose RELEASE-mode `assert_ne!(partial_key, &b""[..])` panics the
4156    /// worker on a leading-empty-label key like `.example.com`. This is the
4157    /// worker-boundary counterpart to `config.rs`'s `validate_sni_pattern`
4158    /// tests: every shape the shared validator rejects at config-load must
4159    /// also be rejected here, for `AddTcpFrontend`/`LoadState` requests that
4160    /// bypass config.rs entirely.
4161    #[test]
4162    fn add_tcp_front_enforces_the_shared_sni_validator_at_the_worker_boundary() {
4163        let mut proxy = test_proxy();
4164        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4165        let config = ListenerBuilder::new_tcp(address)
4166            .to_tcp(None)
4167            .expect("could not build listener config");
4168        let token = Token(0);
4169        proxy
4170            .add_listener(config, token)
4171            .expect("could not add listener");
4172
4173        for bad in [
4174            ".example.com",     // leading empty label -- the pattern_trie-crashing shape
4175            "example.com.",     // trailing empty label
4176            "a..b.example.com", // consecutive dots -- empty middle label
4177            "",                 // empty pattern
4178            "exämple.com",      // non-ASCII
4179        ] {
4180            let front = frontend("cluster-a", Some(bad), &[]);
4181            let front = RequestTcpFrontend { address, ..front };
4182            match proxy.add_tcp_front(front) {
4183                Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4184                other => panic!("malformed sni {bad:?} must be rejected, got {other:?}"),
4185            }
4186
4187            let listener = proxy
4188                .listeners
4189                .get(&token)
4190                .expect("listener must be present")
4191                .borrow();
4192            assert!(
4193                listener.sni_routes.is_empty(),
4194                "a rejected sni {bad:?} must leave sni_routes empty"
4195            );
4196            assert!(
4197                listener.cluster_id.is_none(),
4198                "a rejected sni {bad:?} must leave cluster_id unset"
4199            );
4200        }
4201
4202        // A space IS valid ASCII, is not '*'/'/', and splits into non-empty
4203        // labels ("exa", "mple.com") -- none of the shared validator's rules
4204        // (empty pattern, non-ASCII, misplaced '*', empty label) cover "not a
4205        // valid hostname character" in general, so this shape is ACCEPTED,
4206        // not rejected. Documenting the validator's actual behavior rather
4207        // than assuming a hostname-shaped string with a space would be
4208        // caught too.
4209        let space_front = frontend("cluster-space", Some("exa mple.com"), &[]);
4210        let space_front = RequestTcpFrontend {
4211            address,
4212            ..space_front
4213        };
4214        assert!(
4215            proxy.add_tcp_front(space_front).is_ok(),
4216            "the shared SNI validator does not reject an embedded space -- \
4217             it is valid ASCII with no empty label"
4218        );
4219    }
4220
4221    // ---- per-frontend access-log tags keying (sozu-proxy/sozu#1290) ----
4222
4223    #[test]
4224    fn sni_fronts_keep_distinct_tags_under_distinct_keys() {
4225        let mut proxy = test_proxy();
4226        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4227        let config = ListenerBuilder::new_tcp(address)
4228            .to_tcp(None)
4229            .expect("could not build listener config");
4230        let token = Token(0);
4231        proxy
4232            .add_listener(config, token)
4233            .expect("could not add listener");
4234
4235        let front_a = RequestTcpFrontend {
4236            cluster_id: "cluster-a".to_owned(),
4237            address,
4238            sni: Some("a.example.com".to_owned()),
4239            alpn: vec![],
4240            tags: std::collections::BTreeMap::from([("team".to_owned(), "alpha".to_owned())]),
4241        };
4242        let front_b = RequestTcpFrontend {
4243            cluster_id: "cluster-b".to_owned(),
4244            address,
4245            sni: Some("b.example.com".to_owned()),
4246            alpn: vec!["h2".to_owned()],
4247            tags: std::collections::BTreeMap::from([("team".to_owned(), "beta".to_owned())]),
4248        };
4249        proxy
4250            .add_tcp_front(front_a)
4251            .expect("add_tcp_front A must succeed");
4252        proxy
4253            .add_tcp_front(front_b)
4254            .expect("add_tcp_front B must succeed");
4255
4256        let std_address: SocketAddr = address.into();
4257        let key_a = sni_tags_key(&std_address, "a.example.com", &[]);
4258        let key_b = sni_tags_key(&std_address, "b.example.com", &["h2".to_owned()]);
4259
4260        let listener = proxy
4261            .listeners
4262            .get(&token)
4263            .expect("listener must be present")
4264            .borrow();
4265        let tags_a = listener
4266            .get_tags(&key_a)
4267            .expect("front A's tags must live under its own composed key");
4268        assert_eq!(
4269            tags_a.tags.get("team").map(String::as_str),
4270            Some("alpha"),
4271            "front A's tags must survive front B's add, not be clobbered by it"
4272        );
4273        let tags_b = listener
4274            .get_tags(&key_b)
4275            .expect("front B's tags must live under its own composed key");
4276        assert_eq!(tags_b.tags.get("team").map(String::as_str), Some("beta"));
4277    }
4278
4279    #[test]
4280    fn removing_one_sni_front_clears_only_its_own_tags() {
4281        let mut proxy = test_proxy();
4282        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4283        let config = ListenerBuilder::new_tcp(address)
4284            .to_tcp(None)
4285            .expect("could not build listener config");
4286        let token = Token(0);
4287        proxy
4288            .add_listener(config, token)
4289            .expect("could not add listener");
4290
4291        let front_a = RequestTcpFrontend {
4292            cluster_id: "cluster-a".to_owned(),
4293            address,
4294            sni: Some("a.example.com".to_owned()),
4295            alpn: vec![],
4296            tags: std::collections::BTreeMap::from([("team".to_owned(), "alpha".to_owned())]),
4297        };
4298        let front_b = RequestTcpFrontend {
4299            cluster_id: "cluster-b".to_owned(),
4300            address,
4301            sni: Some("b.example.com".to_owned()),
4302            alpn: vec!["h2".to_owned()],
4303            tags: std::collections::BTreeMap::from([("team".to_owned(), "beta".to_owned())]),
4304        };
4305        proxy
4306            .add_tcp_front(front_a.clone())
4307            .expect("add_tcp_front A must succeed");
4308        proxy
4309            .add_tcp_front(front_b)
4310            .expect("add_tcp_front B must succeed");
4311
4312        proxy
4313            .remove_tcp_front(front_a)
4314            .expect("remove_tcp_front A must succeed");
4315
4316        let std_address: SocketAddr = address.into();
4317        let key_a = sni_tags_key(&std_address, "a.example.com", &[]);
4318        let key_b = sni_tags_key(&std_address, "b.example.com", &["h2".to_owned()]);
4319
4320        let listener = proxy
4321            .listeners
4322            .get(&token)
4323            .expect("listener must be present")
4324            .borrow();
4325        assert!(
4326            listener.get_tags(&key_a).is_none(),
4327            "removing front A must clear its own tags entry"
4328        );
4329        assert!(
4330            listener.get_tags(&key_b).is_some(),
4331            "removing front A must NOT clear sibling front B's tags"
4332        );
4333    }
4334
4335    /// A session routed to front A must look tags up under front A's key:
4336    /// the key `add_tcp_front` stores and the key `upgrade_sni_preread`
4337    /// rebuilds from the route decision (`matched_sni_pattern` +
4338    /// `matched_alpn`) must be identical, regardless of the operator's
4339    /// original ALPN order or SNI casing, and must never collide with the
4340    /// bare-address key used by no-SNI fronts.
4341    #[test]
4342    fn route_time_tags_key_rebuild_matches_the_add_time_key() {
4343        let address: SocketAddr = "127.0.0.1:9000".parse().expect("test address");
4344
4345        // Wildcard front, operator wrote mixed case + reverse ALPN order.
4346        let add_time = sni_tags_key(
4347            &address,
4348            "*.Example.COM",
4349            &["http/1.1".to_owned(), "h2".to_owned()],
4350        );
4351        let matcher =
4352            AlpnMatcher::OneOf([b"h2".to_vec(), b"http/1.1".to_vec()].into_iter().collect());
4353        let route_time = sni_tags_key(
4354            &address,
4355            "*.example.com", // matched_sni_pattern: the lowercased trie key
4356            &alpn_matcher_protocols(&matcher),
4357        );
4358        assert_eq!(
4359            add_time, route_time,
4360            "add-time and route-time keys must agree for the same front"
4361        );
4362
4363        // Catch-all front: empty alpn at add time <-> AlpnMatcher::Any.
4364        assert_eq!(
4365            sni_tags_key(&address, "a.example.com", &[]),
4366            sni_tags_key(
4367                &address,
4368                "a.example.com",
4369                &alpn_matcher_protocols(&AlpnMatcher::Any)
4370            )
4371        );
4372
4373        // A composed key never collides with the bare-address key.
4374        assert_ne!(add_time, address.to_string());
4375    }
4376
4377    /// Regression from the sozu-proxy/sozu#1290 review: ALPN protocol
4378    /// identifiers are opaque byte strings -- nothing forbids a `,` inside
4379    /// one -- so a SINGLE protocol `"a,b"` and the DISJOINT pair `["a",
4380    /// "b"]` are two legal, distinct `AlpnMatcher`s on the same
4381    /// `(address, sni)`, yet the naive `sorted_alpn.join(",")` key collapses
4382    /// both to the literal string `"a,b"`.
4383    #[test]
4384    fn alpn_sets_differing_only_by_an_embedded_separator_get_distinct_tags_keys() {
4385        let address: SocketAddr = "127.0.0.1:9001".parse().expect("test address");
4386        let key_joined = sni_tags_key(&address, "example.com", &["a,b".to_owned()]);
4387        let key_split = sni_tags_key(&address, "example.com", &["a".to_owned(), "b".to_owned()]);
4388        assert_ne!(
4389            key_joined, key_split,
4390            "a single \"a,b\" protocol must not collide with the disjoint [\"a\", \"b\"] pair"
4391        );
4392    }
4393
4394    /// End-to-end through `TcpProxy`: both fronts are accepted (they are
4395    /// genuinely disjoint `AlpnMatcher`s, so `validate_new_tcp_front`'s
4396    /// overlap check does not reject either), each keeps its OWN tags under
4397    /// its own composed key, and removing one leaves the other's tags
4398    /// intact.
4399    #[test]
4400    fn alpn_sets_differing_only_by_an_embedded_separator_keep_distinct_tags() {
4401        let mut proxy = test_proxy();
4402        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4403        let config = ListenerBuilder::new_tcp(address)
4404            .to_tcp(None)
4405            .expect("could not build listener config");
4406        let token = Token(0);
4407        proxy
4408            .add_listener(config, token)
4409            .expect("could not add listener");
4410
4411        let front_joined = RequestTcpFrontend {
4412            cluster_id: "cluster-joined".to_owned(),
4413            address,
4414            sni: Some("example.com".to_owned()),
4415            alpn: vec!["a,b".to_owned()],
4416            tags: std::collections::BTreeMap::from([("variant".to_owned(), "joined".to_owned())]),
4417        };
4418        let front_split = RequestTcpFrontend {
4419            cluster_id: "cluster-split".to_owned(),
4420            address,
4421            sni: Some("example.com".to_owned()),
4422            alpn: vec!["a".to_owned(), "b".to_owned()],
4423            tags: std::collections::BTreeMap::from([("variant".to_owned(), "split".to_owned())]),
4424        };
4425
4426        proxy
4427            .add_tcp_front(front_joined.clone())
4428            .expect("the single \"a,b\" protocol front must be accepted");
4429        proxy.add_tcp_front(front_split.clone()).expect(
4430            "the disjoint [\"a\", \"b\"] front must be accepted -- it is NOT the same ALPN \
4431                 set as [\"a,b\"]",
4432        );
4433
4434        let std_address: SocketAddr = address.into();
4435        let key_joined = sni_tags_key(&std_address, "example.com", &["a,b".to_owned()]);
4436        let key_split = sni_tags_key(
4437            &std_address,
4438            "example.com",
4439            &["a".to_owned(), "b".to_owned()],
4440        );
4441        assert_ne!(key_joined, key_split);
4442
4443        {
4444            let listener = proxy
4445                .listeners
4446                .get(&token)
4447                .expect("listener must be present")
4448                .borrow();
4449            assert_eq!(
4450                listener
4451                    .get_tags(&key_joined)
4452                    .and_then(|t| t.tags.get("variant"))
4453                    .map(String::as_str),
4454                Some("joined"),
4455                "the joined front's tags must live under its own composed key"
4456            );
4457            assert_eq!(
4458                listener
4459                    .get_tags(&key_split)
4460                    .and_then(|t| t.tags.get("variant"))
4461                    .map(String::as_str),
4462                Some("split"),
4463                "the split front's tags must live under its own DISTINCT composed key"
4464            );
4465        }
4466
4467        proxy
4468            .remove_tcp_front(front_joined)
4469            .expect("remove the joined front");
4470
4471        let listener = proxy
4472            .listeners
4473            .get(&token)
4474            .expect("listener must be present")
4475            .borrow();
4476        assert!(
4477            listener.get_tags(&key_joined).is_none(),
4478            "removing the joined front must clear its own tags entry"
4479        );
4480        assert_eq!(
4481            listener
4482                .get_tags(&key_split)
4483                .and_then(|t| t.tags.get("variant"))
4484                .map(String::as_str),
4485            Some("split"),
4486            "removing the joined front must not disturb its sibling's tags"
4487        );
4488    }
4489
4490    #[test]
4491    fn remove_sni_route_for_an_absent_exact_key_does_not_strip_a_sibling_wildcards_entry() {
4492        let mut listener = test_listener();
4493        listener
4494            .insert_sni_route(
4495                "*.example.com".to_owned(),
4496                vec![],
4497                "cluster-wildcard".into(),
4498            )
4499            .expect("insert_sni_route must succeed for a valid test SNI");
4500
4501        // "a.example.com" was never inserted as its own route -- only the
4502        // wildcard catch-all exists. A remove targeting the exact host
4503        // (e.g. a stale `RemoveTcpFrontend` replayed from a hand-edited
4504        // `LoadState`) must be a no-op here, not reach into and strip the
4505        // WILDCARD's own catch-all entry.
4506        listener.remove_sni_route("a.example.com".to_owned(), vec![], "cluster-wildcard");
4507
4508        let (_, wildcard_entries) = listener
4509            .sni_routes
4510            .domain_lookup(b"b.example.com", true)
4511            .expect(
4512                "the wildcard catch-all must survive a remove targeting an unrelated exact key",
4513            );
4514        assert_eq!(
4515            wildcard_entries,
4516            &vec![(AlpnMatcher::Any, "cluster-wildcard".into())]
4517        );
4518    }
4519
4520    // ---- routing gate data invariant: never both cluster_id AND routes ----
4521
4522    #[test]
4523    fn a_no_sni_front_leaves_the_route_table_empty() {
4524        let mut listener = test_listener();
4525        listener.cluster_id = Some("legacy-catch-all".into());
4526        assert!(
4527            listener.sni_routes.is_empty(),
4528            "a listener with only a no-SNI front must never populate sni_routes"
4529        );
4530    }
4531
4532    #[test]
4533    fn an_sni_scoped_front_leaves_cluster_id_unset() {
4534        let mut listener = test_listener();
4535        listener
4536            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".into())
4537            .expect("insert_sni_route must succeed for a valid test SNI");
4538        assert!(
4539            listener.cluster_id.is_none(),
4540            "a listener with only SNI-scoped fronts must never populate the legacy cluster_id"
4541        );
4542    }
4543
4544    // ---- end-to-end through TcpProxy::add_tcp_front / remove_tcp_front ----
4545
4546    fn test_proxy() -> TcpProxy {
4547        let ServerParts {
4548            registry,
4549            sessions,
4550            pool,
4551            backends,
4552            ..
4553        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
4554        TcpProxy::new(registry, sessions, pool, backends)
4555    }
4556
4557    /// Two listeners on one address, one bound to `lo` and one to no
4558    /// interface, coexist in a worker: each is activated, handed back and
4559    /// removed under its own (address, interface) key, each socket carries its
4560    /// own `SO_BINDTODEVICE`, and a frontend — which names its listener by
4561    /// address alone — reaches both.
4562    #[cfg(target_os = "linux")]
4563    #[test]
4564    fn listeners_sharing_an_address_on_different_interfaces_coexist() {
4565        use socket2::SockRef;
4566
4567        let mut proxy = test_proxy();
4568        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4569        let bare = ListenerKey::from(address);
4570        let on_lo = ListenerKey::new(address, Some("lo"));
4571        for (token, interface) in [(Token(0), None), (Token(1), Some("lo"))] {
4572            let config = ListenerBuilder::new_tcp(address)
4573                .with_interface(interface)
4574                .to_tcp(None)
4575                .expect("could not build listener config");
4576            proxy
4577                .add_listener(config, token)
4578                .expect("a listener on another interface must not collide");
4579        }
4580        assert_eq!(proxy.listener_token(&bare), Some(Token(0)));
4581        assert_eq!(proxy.listener_token(&on_lo), Some(Token(1)));
4582        assert_eq!(
4583            proxy.listener_token(&ListenerKey::new(address, Some("wg0"))),
4584            None,
4585            "an interface nobody added names no listener"
4586        );
4587
4588        proxy
4589            .add_tcp_front(RequestTcpFrontend {
4590                cluster_id: "cluster-a".to_owned(),
4591                address,
4592                ..Default::default()
4593            })
4594            .expect("add_tcp_front must succeed");
4595        for token in [Token(0), Token(1)] {
4596            assert_eq!(
4597                proxy.listeners[&token].borrow().cluster_id.as_deref(),
4598                Some("cluster-a"),
4599                "the frontend must reach the listener at {token:?}"
4600            );
4601        }
4602
4603        for key in [&bare, &on_lo] {
4604            if let Err(error) = proxy.activate_listener(key, None) {
4605                let message = error.to_string();
4606                if message.contains("Operation not permitted") {
4607                    eprintln!(
4608                        "SKIPPED listeners_sharing_an_address_on_different_interfaces_coexist: \
4609                         this environment may not bind a socket to an interface: {message}"
4610                    );
4611                    return;
4612                }
4613                panic!("activating {key}: {message}");
4614            }
4615        }
4616        let mut handed_back = proxy.give_back_listeners();
4617        handed_back.sort_by(|(a, _), (b, _)| a.cmp(b));
4618        let keys: Vec<&ListenerKey> = handed_back.iter().map(|(key, _)| key).collect();
4619        assert_eq!(keys, vec![&bare, &on_lo]);
4620        for (key, socket) in &handed_back {
4621            assert_eq!(
4622                SockRef::from(socket)
4623                    .device()
4624                    .expect("getsockopt(SO_BINDTODEVICE)"),
4625                key.interface
4626                    .as_ref()
4627                    .map(|interface| interface.as_bytes().to_vec()),
4628                "the socket handed back for {key} must be bound to its interface"
4629            );
4630        }
4631
4632        proxy.remove_listener(&on_lo);
4633        assert_eq!(proxy.listener_token(&on_lo), None);
4634        assert_eq!(
4635            proxy.listener_token(&bare),
4636            Some(Token(0)),
4637            "removing the lo listener must keep the one without interface"
4638        );
4639    }
4640
4641    /// A TCP listener added on an address that already routes to a cluster
4642    /// takes that route over from the listener already there — the no-SNI
4643    /// cluster and the SNI routes alike — so it forwards from its first
4644    /// connection (sozu-proxy/sozu#719).
4645    #[test]
4646    fn a_listener_added_next_to_a_sibling_takes_over_its_routes() {
4647        let mut proxy = test_proxy();
4648        let catch_all = SocketAddress::new_v4(0, 0, 0, 0, provide_port());
4649        let by_sni = SocketAddress::new_v4(0, 0, 0, 0, provide_port());
4650        for (token, address, sni) in [
4651            (Token(0), catch_all, None),
4652            (Token(2), by_sni, Some("example.com".to_owned())),
4653        ] {
4654            let config = ListenerBuilder::new_tcp(address)
4655                .to_tcp(None)
4656                .expect("could not build listener config");
4657            proxy.add_listener(config, token).expect("add the listener");
4658            proxy
4659                .add_tcp_front(RequestTcpFrontend {
4660                    cluster_id: "cluster-a".to_owned(),
4661                    address,
4662                    sni,
4663                    ..Default::default()
4664                })
4665                .expect("add_tcp_front must succeed");
4666        }
4667        for (token, address) in [(Token(1), catch_all), (Token(3), by_sni)] {
4668            let config = ListenerBuilder::new_tcp(address)
4669                .with_interface(Some("lo"))
4670                .to_tcp(None)
4671                .expect("could not build listener config");
4672            proxy.add_listener(config, token).expect("add the sibling");
4673        }
4674
4675        assert_eq!(
4676            proxy.listeners[&Token(1)].borrow().cluster_id.as_deref(),
4677            Some("cluster-a"),
4678            "the listener added after the frontend must route to its cluster"
4679        );
4680        let sibling = proxy.listeners[&Token(3)].borrow();
4681        let (_, entries) = sibling
4682            .sni_routes
4683            .domain_lookup(b"example.com", true)
4684            .expect("the listener added after the SNI frontend must route it");
4685        assert_eq!(entries, &vec![(AlpnMatcher::Any, "cluster-a".into())]);
4686    }
4687
4688    /// Traffic that arrives through an interface goes to the listener bound to
4689    /// it, not to the listener on the same address bound to none: Linux
4690    /// prefers the device-bound socket, and the two never share a
4691    /// `SO_REUSEPORT` group. Here a bare `0.0.0.0:P` and a `lo`-bound
4692    /// `0.0.0.0:P` both listen, a client connects through `127.0.0.1`, and
4693    /// only the `lo` listener has a connection to accept.
4694    #[cfg(target_os = "linux")]
4695    #[test]
4696    fn traffic_through_an_interface_reaches_the_listener_bound_to_it() {
4697        let mut proxy = test_proxy();
4698        let port = provide_port();
4699        let address = SocketAddress::new_v4(0, 0, 0, 0, port);
4700        for (token, interface) in [(Token(0), None), (Token(1), Some("lo"))] {
4701            let config = ListenerBuilder::new_tcp(address)
4702                .with_interface(interface)
4703                .to_tcp(None)
4704                .expect("could not build listener config");
4705            proxy.add_listener(config, token).expect("add the listener");
4706            let key = ListenerKey::new(address, interface);
4707            if let Err(error) = proxy.activate_listener(&key, None) {
4708                let message = error.to_string();
4709                if message.contains("Operation not permitted") {
4710                    eprintln!(
4711                        "SKIPPED traffic_through_an_interface_reaches_the_listener_bound_to_it: \
4712                         this environment may not bind a socket to an interface: {message}"
4713                    );
4714                    return;
4715                }
4716                panic!("activating {key}: {message}");
4717            }
4718        }
4719
4720        // Several connections: with the binding ignored, both sockets would
4721        // share one `SO_REUSEPORT` group and each connection would land on
4722        // either of them.
4723        const CONNECTIONS: usize = 8;
4724        let _clients: Vec<_> = (0..CONNECTIONS)
4725            .map(|_| {
4726                std::net::TcpStream::connect(("127.0.0.1", port))
4727                    .expect("connect through the loopback interface")
4728            })
4729            .collect();
4730        // The handshakes are complete once `connect` returns; poll briefly in
4731        // case an accept queue entry is not visible yet.
4732        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
4733        let mut accepted = 0;
4734        while accepted < CONNECTIONS {
4735            match proxy.accept(ListenToken(1)) {
4736                Ok(_) => accepted += 1,
4737                Err(AcceptError::WouldBlock) if std::time::Instant::now() < deadline => {
4738                    std::thread::sleep(std::time::Duration::from_millis(10));
4739                }
4740                Err(error) => panic!(
4741                    "the lo listener must accept every connection, got {accepted} of \
4742                     {CONNECTIONS}: {error:?}"
4743                ),
4744            }
4745        }
4746        assert!(
4747            matches!(proxy.accept(ListenToken(0)), Err(AcceptError::WouldBlock)),
4748            "the listener bound to no interface must not receive loopback traffic"
4749        );
4750    }
4751
4752    #[test]
4753    fn add_then_remove_sni_front_round_trips_through_tcp_proxy() {
4754        let mut proxy = test_proxy();
4755        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4756        let config = ListenerBuilder::new_tcp(address)
4757            .to_tcp(None)
4758            .expect("could not build listener config");
4759        let token = Token(0);
4760        proxy
4761            .add_listener(config, token)
4762            .expect("could not add listener");
4763
4764        let front = RequestTcpFrontend {
4765            cluster_id: "cluster-a".to_owned(),
4766            address,
4767            sni: Some("Example.COM".to_owned()),
4768            alpn: vec![],
4769            ..Default::default()
4770        };
4771        proxy
4772            .add_tcp_front(front.clone())
4773            .expect("add_tcp_front must succeed");
4774
4775        {
4776            let listener = proxy
4777                .listeners
4778                .get(&token)
4779                .expect("listener must be present")
4780                .borrow();
4781            assert!(listener.cluster_id.is_none());
4782            let (_, entries) = listener
4783                .sni_routes
4784                // Lowercased at insert time regardless of wire-form casing.
4785                .domain_lookup(b"example.com", true)
4786                .expect("example.com must be routable after add_tcp_front");
4787            assert_eq!(entries, &vec![(AlpnMatcher::Any, "cluster-a".into())]);
4788        }
4789        assert_eq!(proxy.fronts.get("cluster-a"), Some(&token));
4790
4791        proxy
4792            .remove_tcp_front(front)
4793            .expect("remove_tcp_front must succeed");
4794
4795        {
4796            let listener = proxy
4797                .listeners
4798                .get(&token)
4799                .expect("listener must be present")
4800                .borrow();
4801            assert!(
4802                listener.sni_routes.is_empty(),
4803                "remove_tcp_front must leave no stranded route"
4804            );
4805        }
4806        assert_eq!(
4807            proxy.fronts.get("cluster-a"),
4808            None,
4809            "remove_tcp_front must undo add_tcp_front's self.fronts bookkeeping"
4810        );
4811    }
4812
4813    #[test]
4814    fn add_then_remove_legacy_no_sni_front_round_trips() {
4815        let mut proxy = test_proxy();
4816        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4817        let config = ListenerBuilder::new_tcp(address)
4818            .to_tcp(None)
4819            .expect("could not build listener config");
4820        let token = Token(0);
4821        proxy
4822            .add_listener(config, token)
4823            .expect("could not add listener");
4824
4825        let front = frontend("cluster-legacy", None, &[]);
4826        let front = RequestTcpFrontend { address, ..front };
4827        proxy
4828            .add_tcp_front(front.clone())
4829            .expect("add_tcp_front must succeed");
4830
4831        assert_eq!(
4832            proxy
4833                .listeners
4834                .get(&token)
4835                .expect("listener must be present")
4836                .borrow()
4837                .cluster_id,
4838            Some("cluster-legacy".into())
4839        );
4840
4841        proxy
4842            .remove_tcp_front(front)
4843            .expect("remove_tcp_front must succeed");
4844
4845        assert_eq!(
4846            proxy
4847                .listeners
4848                .get(&token)
4849                .expect("listener must be present")
4850                .borrow()
4851                .cluster_id,
4852            None
4853        );
4854        assert_eq!(proxy.fronts.get("cluster-legacy"), None);
4855    }
4856
4857    // ---- add_tcp_front hard-rejects routing-corrupting requests --------
4858    //
4859    // Worker-side mirror of `command/src/config.rs`'s TOML config-load
4860    // invariants (sozu-proxy/sozu#1279 hardening): `AddTcpFrontend` can
4861    // reach the worker directly over the command socket, or via `LoadState`
4862    // replay, bypassing config.rs entirely, so `add_tcp_front` must defend
4863    // itself rather than rely on a debug-only assertion.
4864
4865    #[test]
4866    fn add_tcp_front_rejects_alpn_without_sni() {
4867        let mut proxy = test_proxy();
4868        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4869        let config = ListenerBuilder::new_tcp(address)
4870            .to_tcp(None)
4871            .expect("could not build listener config");
4872        proxy
4873            .add_listener(config, Token(0))
4874            .expect("could not add listener");
4875
4876        let front = frontend("cluster-a", None, &["h2"]);
4877        let front = RequestTcpFrontend { address, ..front };
4878        match proxy.add_tcp_front(front) {
4879            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4880            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4881        }
4882    }
4883
4884    #[test]
4885    fn add_tcp_front_rejects_no_sni_front_on_listener_with_sni_routes() {
4886        let mut proxy = test_proxy();
4887        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4888        let config = ListenerBuilder::new_tcp(address)
4889            .to_tcp(None)
4890            .expect("could not build listener config");
4891        proxy
4892            .add_listener(config, Token(0))
4893            .expect("could not add listener");
4894
4895        let sni_front = frontend("cluster-a", Some("example.com"), &[]);
4896        let sni_front = RequestTcpFrontend {
4897            address,
4898            ..sni_front
4899        };
4900        proxy
4901            .add_tcp_front(sni_front)
4902            .expect("the first, SNI-scoped frontend must be accepted");
4903
4904        let no_sni_front = frontend("cluster-b", None, &[]);
4905        let no_sni_front = RequestTcpFrontend {
4906            address,
4907            ..no_sni_front
4908        };
4909        match proxy.add_tcp_front(no_sni_front) {
4910            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4911            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4912        }
4913    }
4914
4915    #[test]
4916    fn add_tcp_front_rejects_sni_front_on_listener_with_no_sni_cluster() {
4917        let mut proxy = test_proxy();
4918        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4919        let config = ListenerBuilder::new_tcp(address)
4920            .to_tcp(None)
4921            .expect("could not build listener config");
4922        proxy
4923            .add_listener(config, Token(0))
4924            .expect("could not add listener");
4925
4926        let no_sni_front = frontend("cluster-a", None, &[]);
4927        let no_sni_front = RequestTcpFrontend {
4928            address,
4929            ..no_sni_front
4930        };
4931        proxy
4932            .add_tcp_front(no_sni_front)
4933            .expect("the first, no-SNI frontend must be accepted");
4934
4935        let sni_front = frontend("cluster-b", Some("example.com"), &[]);
4936        let sni_front = RequestTcpFrontend {
4937            address,
4938            ..sni_front
4939        };
4940        match proxy.add_tcp_front(sni_front) {
4941            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4942            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4943        }
4944    }
4945
4946    #[test]
4947    fn add_tcp_front_rejects_alpn_overlap_on_same_sni() {
4948        let mut proxy = test_proxy();
4949        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4950        let config = ListenerBuilder::new_tcp(address)
4951            .to_tcp(None)
4952            .expect("could not build listener config");
4953        proxy
4954            .add_listener(config, Token(0))
4955            .expect("could not add listener");
4956
4957        let first = frontend("cluster-a", Some("example.com"), &["h2"]);
4958        let first = RequestTcpFrontend { address, ..first };
4959        proxy
4960            .add_tcp_front(first)
4961            .expect("the first frontend must be accepted");
4962
4963        // Second frontend shares "h2" with the first on the same sni.
4964        let second = frontend("cluster-b", Some("example.com"), &["h2", "http/1.1"]);
4965        let second = RequestTcpFrontend { address, ..second };
4966        match proxy.add_tcp_front(second) {
4967            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4968            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4969        }
4970    }
4971
4972    #[test]
4973    fn add_tcp_front_rejects_duplicate_catch_all_on_same_sni() {
4974        let mut proxy = test_proxy();
4975        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
4976        let config = ListenerBuilder::new_tcp(address)
4977            .to_tcp(None)
4978            .expect("could not build listener config");
4979        proxy
4980            .add_listener(config, Token(0))
4981            .expect("could not add listener");
4982
4983        let first = frontend("cluster-a", Some("example.com"), &[]);
4984        let first = RequestTcpFrontend { address, ..first };
4985        proxy
4986            .add_tcp_front(first)
4987            .expect("the first catch-all frontend must be accepted");
4988
4989        let second = frontend("cluster-b", Some("example.com"), &[]);
4990        let second = RequestTcpFrontend { address, ..second };
4991        match proxy.add_tcp_front(second) {
4992            Err(ProxyError::InvalidTcpFrontend { .. }) => {}
4993            other => panic!("expected InvalidTcpFrontend, got {other:?}"),
4994        }
4995    }
4996
4997    /// The valid, intended shape (sozu-proxy/sozu#1279's whole reason for
4998    /// existing) must still be accepted: disjoint, non-empty `alpn` lists
4999    /// on the same `(address, sni)`, and a catch-all alongside a
5000    /// specific-protocol entry.
5001    #[test]
5002    fn add_tcp_front_accepts_disjoint_alpn_and_catch_all_on_same_sni() {
5003        let mut proxy = test_proxy();
5004        let address = SocketAddress::new_v4(127, 0, 0, 1, provide_port());
5005        let config = ListenerBuilder::new_tcp(address)
5006            .to_tcp(None)
5007            .expect("could not build listener config");
5008        proxy
5009            .add_listener(config, Token(0))
5010            .expect("could not add listener");
5011
5012        let h2 = frontend("cluster-h2", Some("example.com"), &["h2"]);
5013        let h2 = RequestTcpFrontend { address, ..h2 };
5014        proxy
5015            .add_tcp_front(h2)
5016            .expect("disjoint alpn frontend must be accepted");
5017
5018        let http11 = frontend("cluster-http11", Some("example.com"), &["http/1.1"]);
5019        let http11 = RequestTcpFrontend { address, ..http11 };
5020        proxy
5021            .add_tcp_front(http11)
5022            .expect("second disjoint alpn frontend must be accepted");
5023
5024        let catch_all = frontend("cluster-default", Some("example.com"), &[]);
5025        let catch_all = RequestTcpFrontend {
5026            address,
5027            ..catch_all
5028        };
5029        proxy
5030            .add_tcp_front(catch_all)
5031            .expect("a catch-all alongside specific-protocol entries must be accepted");
5032    }
5033
5034    // ---- tcp.sni_preread.active gauge accounting ----------------------
5035
5036    /// Read the current process-local `tcp.sni_preread.active` gauge,
5037    /// treating an absent key as 0. `dump_local_proxy_metrics` is a
5038    /// non-draining filter over the proxy `MetricsMap`, so repeated reads are
5039    /// side-effect free and the key is the raw metric name.
5040    fn sni_preread_active_gauge() -> i64 {
5041        use sozu_command::proto::command::filtered_metrics::Inner;
5042        crate::metrics::METRICS.with(|metrics| {
5043            metrics
5044                .borrow_mut()
5045                .dump_local_proxy_metrics()
5046                .get(names::tcp::sni_preread::ACTIVE)
5047                .and_then(|fm| fm.inner.as_ref())
5048                .and_then(|inner| match inner {
5049                    Inner::Gauge(v) => Some(*v as i64),
5050                    _ => None,
5051                })
5052                .unwrap_or(0)
5053        })
5054    }
5055
5056    #[test]
5057    fn entering_sni_preread_increments_the_active_gauge() {
5058        // Regression guard for the missing-`+1` gauge bug
5059        // (sozu-proxy/sozu#1279): `new_sni_preread` must bump
5060        // `tcp.sni_preread.active` by exactly one when a session ENTERS the
5061        // state, so each of the two `-1` decrements -- the "upgrade" exit in
5062        // `upgrade_sni_preread` and the "reject"/"teardown" exit in `close()`'s
5063        // `StateMarker::SniPreread` arm -- has a matching increment. Without
5064        // this `+1` the first `-1` underflows a fresh-zero gauge (clamped to 0,
5065        // ERROR-logged), pinning the gauge at 0 and rendering the e2e gauge
5066        // assertion vacuous.
5067        //
5068        // `METRICS` is a thread-local shared across unit tests on the same
5069        // worker thread, so this asserts the DELTA around one constructor call
5070        // (robust to any starting value), not an absolute reading. The
5071        // net-zero-per-session contract spans the full lifecycle (accept ->
5072        // live backend connect -> upgrade/teardown) and is the behavioural job
5073        // of the e2e gauge assertion
5074        // (`test_tcp_sni_per_ip_limiter_rejects_second_then_admits_after_release`
5075        // in `e2e/src/tests/tcp_sni_tests.rs`), not reproducible at this unit
5076        // level.
5077        let ServerParts {
5078            registry,
5079            sessions,
5080            pool,
5081            backends,
5082            ..
5083        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
5084
5085        let proxy = Rc::new(RefCell::new(TcpProxy::new(
5086            registry,
5087            sessions,
5088            pool.clone(),
5089            backends,
5090        )));
5091        let listener = Rc::new(RefCell::new(test_listener()));
5092
5093        let (front_buffer, back_buffer) = {
5094            let mut pool = pool.borrow_mut();
5095            (
5096                pool.checkout().expect("front buffer checkout must succeed"),
5097                pool.checkout().expect("back buffer checkout must succeed"),
5098            )
5099        };
5100
5101        // A non-blocking connect to a (likely unused) loopback port returns a
5102        // real `MioTcpStream` handle immediately, regardless of whether the
5103        // connection completes; `new_sni_preread` reads nothing from it.
5104        let peer: SocketAddr = format!("127.0.0.1:{}", provide_port())
5105            .parse()
5106            .expect("loopback address must parse");
5107        let socket = MioTcpStream::connect(peer).expect("mio connect must return a socket handle");
5108
5109        let before = sni_preread_active_gauge();
5110        let session = TcpSession::new_sni_preread(
5111            back_buffer,
5112            Duration::from_secs(30),
5113            Duration::from_secs(30),
5114            front_buffer,
5115            Token(0),
5116            listener,
5117            proxy,
5118            socket,
5119            peer,
5120            Duration::from_millis(0),
5121            Duration::from_secs(3),
5122            16384,
5123        );
5124        let after = sni_preread_active_gauge();
5125
5126        // The session is measured while still in `SniPreread`; hold it across
5127        // the read so no future `Drop` side effect could race the measurement.
5128        assert!(matches!(session.state, TcpStateMachine::SniPreread(_)));
5129
5130        assert_eq!(
5131            after - before,
5132            1,
5133            "entering the SniPreread state must increment tcp.sni_preread.active by exactly one"
5134        );
5135    }
5136
5137    // ---- absolute preread deadline + route-aware timeout (sozu-proxy/sozu#1290) ----
5138
5139    /// `frontend_timeout_resets_on_readable` must gate the reset on exactly
5140    /// one thing: whether the CURRENT state is an undecided `SniPreread`.
5141    /// Constructed directly (no socket I/O needed) since the predicate is a
5142    /// pure function of `&TcpStateMachine`.
5143    #[test]
5144    fn frontend_timeout_reset_is_gated_on_sni_preread_decision() {
5145        let mut pool = crate::pool::Pool::with_capacity(1, 1, 16 * 1024);
5146        let frontend_buffer = pool.checkout().expect("frontend buffer");
5147        let socket = MioTcpStream::connect(
5148            format!("127.0.0.1:{}", provide_port())
5149                .parse()
5150                .expect("loopback address must parse"),
5151        )
5152        .expect("mio connect must return a socket handle");
5153
5154        let undecided = TcpStateMachine::SniPreread(SniPreread::new(
5155            socket,
5156            Token(0),
5157            Ulid::generate(),
5158            frontend_buffer,
5159            16384,
5160        ));
5161        assert!(
5162            !frontend_timeout_resets_on_readable(&undecided),
5163            "an undecided SniPreread must not have its absolute deadline reset"
5164        );
5165
5166        // Every OTHER state resets normally -- represented here by
5167        // `ExpectProxyProtocol`, cheaply constructible without driving a
5168        // real route decision (unlike a ROUTED `SniPreread`, whose
5169        // `outcome` field has no test-only setter and can only become
5170        // `Some` by parsing a real ClientHello -- see
5171        // `frontend_timeout_restored_and_timeout_after_route_is_not_double_counted`
5172        // below for that scenario end-to-end).
5173        let socket2 = MioTcpStream::connect(
5174            format!("127.0.0.1:{}", provide_port())
5175                .parse()
5176                .expect("loopback address must parse"),
5177        )
5178        .expect("mio connect must return a socket handle");
5179        let container = crate::timer::TimeoutContainer::new_empty(Duration::from_secs(5));
5180        let other = TcpStateMachine::ExpectProxyProtocol(ExpectProxyProtocol::new(
5181            container,
5182            socket2,
5183            Token(1),
5184            Ulid::generate(),
5185        ));
5186        assert!(
5187            frontend_timeout_resets_on_readable(&other),
5188            "every non-preread-undecided state must keep resetting its frontend timeout"
5189        );
5190    }
5191
5192    /// Minimal single-record TLS ClientHello wire carrying only a
5193    /// `server_name` extension for `host` -- hand-built rather than reusing
5194    /// `tcp_preread::parser`'s test-only wire-building helpers (`mod
5195    /// parser` is private to `tcp_preread`, unreachable from this sibling
5196    /// module) to drive a REAL route decision through
5197    /// `TcpSession::readable()` for the regression test below.
5198    fn minimal_client_hello_wire(host: &str) -> Vec<u8> {
5199        let mut name_list = vec![0u8]; // name_type = host_name
5200        name_list.extend_from_slice(&(host.len() as u16).to_be_bytes());
5201        name_list.extend_from_slice(host.as_bytes());
5202        let mut sni_ext_data = Vec::new();
5203        sni_ext_data.extend_from_slice(&(name_list.len() as u16).to_be_bytes());
5204        sni_ext_data.extend_from_slice(&name_list);
5205        let mut sni_ext = Vec::new();
5206        sni_ext.extend_from_slice(&0x0000u16.to_be_bytes()); // server_name extension type
5207        sni_ext.extend_from_slice(&(sni_ext_data.len() as u16).to_be_bytes());
5208        sni_ext.extend_from_slice(&sni_ext_data);
5209
5210        let mut body = Vec::new();
5211        body.extend_from_slice(&[0x03, 0x03]); // legacy_version
5212        body.extend_from_slice(&[0u8; 32]); // random
5213        body.push(0); // session_id: empty
5214        body.extend_from_slice(&[0x00, 0x02, 0x13, 0x01]); // cipher_suites
5215        body.push(1); // compression_methods length
5216        body.push(0); // compression_method: null
5217        body.extend_from_slice(&(sni_ext.len() as u16).to_be_bytes()); // extensions block length
5218        body.extend_from_slice(&sni_ext);
5219
5220        let mut handshake = Vec::new();
5221        handshake.push(1u8); // msg_type = client_hello
5222        let hs_len = body.len() as u32;
5223        handshake.extend_from_slice(&hs_len.to_be_bytes()[1..4]);
5224        handshake.extend_from_slice(&body);
5225
5226        let mut record = Vec::new();
5227        record.push(22u8); // ContentType::handshake
5228        record.extend_from_slice(&[0x03, 0x03]); // legacy record version
5229        record.extend_from_slice(&(handshake.len() as u16).to_be_bytes());
5230        record.extend_from_slice(&handshake);
5231        record
5232    }
5233
5234    /// Read the current process-local `tcp.sni_preread.routed` counter.
5235    /// Same non-draining-read pattern as `sni_preread_active_gauge`, but
5236    /// for a `Count` metric instead of a `Gauge`.
5237    fn sni_preread_routed_count() -> i64 {
5238        use sozu_command::proto::command::filtered_metrics::Inner;
5239        crate::metrics::METRICS.with(|metrics| {
5240            metrics
5241                .borrow_mut()
5242                .dump_local_proxy_metrics()
5243                .get(names::tcp::sni_preread::ROUTED)
5244                .and_then(|fm| fm.inner.as_ref())
5245                .and_then(|inner| match inner {
5246                    Inner::Count(v) => Some(*v),
5247                    _ => None,
5248                })
5249                .unwrap_or(0)
5250        })
5251    }
5252
5253    /// End-to-end regression from the sozu-proxy/sozu#1290 review:
5254    ///
5255    /// (a) the moment a real ClientHello routes, `container_frontend_timeout`
5256    ///     must already carry the listener's configured `front_timeout` --
5257    ///     not just once `upgrade_sni_preread` eventually runs (which can be
5258    ///     one or more `ready()` cycles later, after the backend connects);
5259    /// (b) a front-timeout firing AFTER that route decision (backend connect
5260    ///     still pending) must not re-feed `Input::Timeout` into the
5261    ///     already-decided core: pre-fix, doing so replayed the SAME latched
5262    ///     `Output::Routed` through `SniPreread::handle_output`'s `Routed`
5263    ///     arm a second time, double-incrementing `tcp.sni_preread.routed`
5264    ///     (release) and tripping `debug_assert!(self.outcome.is_none(),
5265    ///     ...)` (debug -- this test runs in a debug build, so pre-fix it
5266    ///     panics here).
5267    #[test]
5268    fn frontend_timeout_restored_and_timeout_after_route_is_not_double_counted() {
5269        let ServerParts {
5270            registry,
5271            sessions,
5272            pool,
5273            backends,
5274            ..
5275        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
5276        let proxy = Rc::new(RefCell::new(TcpProxy::new(
5277            registry,
5278            sessions,
5279            pool.clone(),
5280            backends,
5281        )));
5282
5283        let mut bare_listener = test_listener();
5284        bare_listener
5285            .insert_sni_route("example.com".to_owned(), vec![], "cluster-a".into())
5286            .expect("insert_sni_route must succeed for a valid test SNI");
5287        let configured_front_timeout = bare_listener.config.front_timeout;
5288        let listener = Rc::new(RefCell::new(bare_listener));
5289
5290        let std_listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind test listener");
5291        let addr = std_listener.local_addr().expect("listener local addr");
5292        let mut client = std::net::TcpStream::connect(addr).expect("connect test client");
5293        let (server, server_peer) = std_listener.accept().expect("accept test server");
5294        server.set_nonblocking(true).expect("server nonblocking");
5295
5296        let (front_buffer, back_buffer) = {
5297            let mut pool = pool.borrow_mut();
5298            (
5299                pool.checkout().expect("front buffer checkout must succeed"),
5300                pool.checkout().expect("back buffer checkout must succeed"),
5301            )
5302        };
5303
5304        let mut session = TcpSession::new_sni_preread(
5305            back_buffer,
5306            Duration::from_secs(30),
5307            Duration::from_secs(30),
5308            front_buffer,
5309            Token(0),
5310            listener,
5311            proxy,
5312            MioTcpStream::from_std(server),
5313            server_peer,
5314            Duration::from_millis(0),
5315            Duration::from_secs(3),
5316            16384,
5317        );
5318
5319        {
5320            use std::io::Write as _;
5321            client
5322                .write_all(&minimal_client_hello_wire("example.com"))
5323                .expect("write ClientHello");
5324            client.flush().ok();
5325        }
5326
5327        let routed_before = sni_preread_routed_count();
5328        for _ in 0..10 {
5329            if session.cluster_id.is_some() {
5330                break;
5331            }
5332            let _ = session.readable();
5333        }
5334        assert_eq!(
5335            session.cluster_id.as_deref(),
5336            Some("cluster-a"),
5337            "the session must have routed on a valid ClientHello for a configured SNI"
5338        );
5339        assert_eq!(
5340            sni_preread_routed_count() - routed_before,
5341            1,
5342            "routing must count tcp.sni_preread.routed exactly once"
5343        );
5344
5345        // (a) front_timeout is restored the moment routing succeeds.
5346        assert_eq!(
5347            session.container_frontend_timeout.duration(),
5348            Duration::from_secs(configured_front_timeout as u64),
5349            "the frontend timeout must already carry the configured front_timeout right after \
5350             routing, not only once the backend connects"
5351        );
5352
5353        // (b) a timeout firing after the route already latched must not
5354        // double-count tcp.sni_preread.routed, and (running in a debug
5355        // build) must not panic on SniPreread::handle_output's
5356        // `debug_assert!(self.outcome.is_none(), ...)`.
5357        let _ = session.timeout(Token(0));
5358        assert_eq!(
5359            sni_preread_routed_count() - routed_before,
5360            1,
5361            "a timeout firing after the route already latched must not double-count \
5362             tcp.sni_preread.routed"
5363        );
5364
5365        drop(client);
5366    }
5367
5368    /// A complete PROXY-v2 `PROXY` command header for IPv4, the exact shape
5369    /// `ExpectProxyProtocol::readable` accepts (28 bytes: 12-byte magic,
5370    /// version+command, family, address-block length, then the two addresses
5371    /// and ports).
5372    fn proxy_protocol_v2_ipv4_header() -> [u8; 28] {
5373        [
5374            0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54,
5375            0x0A, // PROXY-v2 magic
5376            0x21, // version 2, command PROXY
5377            0x11, // AF_INET over STREAM
5378            0x00, 0x0C, // address block length = 12
5379            127, 0, 0, 1, // source address
5380            127, 0, 0, 1, // destination address
5381            0x1F, 0x90, // source port 8080
5382            0x10, 0x68, // destination port 4200
5383        ]
5384    }
5385
5386    /// A live `ExpectHeader` session plus everything that must outlive it:
5387    /// dropping either listener would break the connect the tests drive.
5388    struct ExpectProxyFixture {
5389        session: Rc<RefCell<TcpSession>>,
5390        /// The same session, type-erased -- `ready()` takes the handle the
5391        /// server dispatches with, and `connect_to_backend` inserts a clone
5392        /// of it under the backend token.
5393        proxy_session: Rc<RefCell<dyn ProxySession>>,
5394        frontend_token: Token,
5395        client: std::net::TcpStream,
5396        _frontend_listener: std::net::TcpListener,
5397        _backend_listener: std::net::TcpListener,
5398    }
5399
5400    /// Build a session for a TCP cluster configured with `expect_proxy = true`
5401    /// (`ProxyProtocolConfig::ExpectHeader`), wired to a real bound backend and a real accepted
5402    /// frontend socket, and registered in the slab exactly like
5403    /// `create_session` does (the frontend slot is reserved BEFORE the
5404    /// session exists, so `connect_to_backend`'s backend token lands on a
5405    /// different key -- `set_back_token` asserts the two differ).
5406    fn expect_proxy_fixture() -> ExpectProxyFixture {
5407        let ServerParts {
5408            registry,
5409            sessions,
5410            pool,
5411            backends,
5412            ..
5413        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
5414
5415        let backend_listener =
5416            std::net::TcpListener::bind("127.0.0.1:0").expect("bind test backend");
5417        let backend_address = backend_listener
5418            .local_addr()
5419            .expect("test backend local addr");
5420        backends.borrow_mut().add_backend(
5421            "cluster-expect",
5422            Backend::new("cluster-expect-1", backend_address, None, None, None),
5423        );
5424
5425        let proxy = Rc::new(RefCell::new(TcpProxy::new(
5426            registry,
5427            sessions.clone(),
5428            pool.clone(),
5429            backends,
5430        )));
5431        let listener = Rc::new(RefCell::new(test_listener()));
5432
5433        let frontend_listener =
5434            std::net::TcpListener::bind("127.0.0.1:0").expect("bind test frontend listener");
5435        let frontend_address = frontend_listener
5436            .local_addr()
5437            .expect("test frontend local addr");
5438        let client = std::net::TcpStream::connect(frontend_address).expect("connect test client");
5439        let (frontend, frontend_peer) = frontend_listener
5440            .accept()
5441            .expect("accept the client connection");
5442        frontend
5443            .set_nonblocking(true)
5444            .expect("frontend socket nonblocking");
5445
5446        let (frontend_buffer, backend_buffer) = {
5447            let mut pool = pool.borrow_mut();
5448            (
5449                pool.checkout().expect("front buffer checkout must succeed"),
5450                pool.checkout().expect("back buffer checkout must succeed"),
5451            )
5452        };
5453
5454        let frontend_token = {
5455            let mut session_manager = sessions.borrow_mut();
5456            let entry = session_manager.slab.vacant_entry();
5457            let token = Token(entry.key());
5458            entry.insert(Rc::new(RefCell::new(crate::server::ListenSession {
5459                protocol: Protocol::TCPListen,
5460            })));
5461            token
5462        };
5463
5464        let session = Rc::new(RefCell::new(TcpSession::new(
5465            backend_buffer,
5466            None,
5467            Some("cluster-expect".into()),
5468            Duration::from_secs(30),
5469            Duration::from_secs(30),
5470            Duration::from_secs(30),
5471            frontend_buffer,
5472            frontend_token,
5473            listener,
5474            Some(ProxyProtocolConfig::ExpectHeader),
5475            proxy,
5476            MioTcpStream::from_std(frontend),
5477            frontend_peer,
5478            Duration::from_millis(0),
5479        )));
5480        let proxy_session: Rc<RefCell<dyn ProxySession>> = session.clone();
5481        sessions.borrow_mut().slab[frontend_token.0] = proxy_session.clone();
5482
5483        assert!(
5484            matches!(
5485                session.borrow().state,
5486                TcpStateMachine::ExpectProxyProtocol(_)
5487            ),
5488            "an ExpectHeader cluster must start its session in the expect state"
5489        );
5490
5491        ExpectProxyFixture {
5492            session,
5493            proxy_session,
5494            frontend_token,
5495            client,
5496            _frontend_listener: frontend_listener,
5497            _backend_listener: backend_listener,
5498        }
5499    }
5500
5501    /// A TCP cluster configured with `expect_proxy = true`
5502    /// (`ProxyProtocolConfig::ExpectHeader`) starts its sessions
5503    /// in [`TcpStateMachine::ExpectProxyProtocol`], a state with NO backend
5504    /// side at all: `back_readiness` returns `None` and `set_back_socket`
5505    /// panics outright ("We should not set the back socket for the expect
5506    /// proxy protocol"). The backend may only be dialed once the inbound
5507    /// header has been parsed and `upgrade_expect` has swapped in a `Pipe`.
5508    ///
5509    /// `ready_inner`'s top-of-function connect gate, however, fires on the
5510    /// FIRST readiness event of a freshly accepted socket -- including the
5511    /// WRITABLE epoll reports for a socket that has sent no byte yet. Without
5512    /// the `ExpectProxyProtocol` arm of `attempt_backend_connect_if_needed`'s
5513    /// guard, completing the TCP handshake was enough to drive
5514    /// `connect_to_backend` into that panic, killing the worker and every
5515    /// other session on it.
5516    ///
5517    /// To SEE THIS RED: delete the `TcpStateMachine::ExpectProxyProtocol(_)`
5518    /// arm from `attempt_backend_connect_if_needed`'s early-return guard.
5519    /// The first `ready()` below then panics inside `connect_to_backend` ->
5520    /// `set_back_socket` before any assertion is reached.
5521    #[test]
5522    fn expect_proxy_dials_the_backend_only_after_the_header_is_parsed() {
5523        use std::io::Write as _;
5524
5525        let mut fixture = expect_proxy_fixture();
5526        let (session, proxy_session, frontend_token) = (
5527            fixture.session.clone(),
5528            fixture.proxy_session.clone(),
5529            fixture.frontend_token,
5530        );
5531
5532        // The handshake alone: epoll reports the accepted socket WRITABLE
5533        // before the client has sent a single byte.
5534        session
5535            .borrow_mut()
5536            .update_readiness(frontend_token, Ready::WRITABLE);
5537        let closed = session.borrow_mut().ready(proxy_session.clone());
5538        assert!(
5539            !closed,
5540            "a session still waiting for its PROXY header must stay open"
5541        );
5542        {
5543            let session = session.borrow();
5544            assert!(
5545                matches!(session.state, TcpStateMachine::ExpectProxyProtocol(_)),
5546                "the session must still be expecting its PROXY header"
5547            );
5548            assert_eq!(
5549                session.back_connected(),
5550                BackendConnectionStatus::NotConnected,
5551                "no backend may be dialed while the PROXY header is still unparsed"
5552            );
5553            assert!(
5554                session.backend_token.is_none(),
5555                "no backend token may be wired while the PROXY header is still unparsed"
5556            );
5557        }
5558
5559        // The header lands: NOW the session upgrades out of the expect state
5560        // and that very same connect gate is allowed to dial.
5561        fixture
5562            .client
5563            .write_all(&proxy_protocol_v2_ipv4_header())
5564            .expect("write the PROXY-v2 header");
5565        fixture.client.flush().ok();
5566        session
5567            .borrow_mut()
5568            .update_readiness(frontend_token, Ready::READABLE);
5569        let _ = session.borrow_mut().ready(proxy_session.clone());
5570        {
5571            let session = session.borrow();
5572            assert!(
5573                !matches!(session.state, TcpStateMachine::ExpectProxyProtocol(_)),
5574                "a parsed PROXY header must move the session out of the expect state"
5575            );
5576            assert_ne!(
5577                session.back_connected(),
5578                BackendConnectionStatus::NotConnected,
5579                "the backend must be dialed once the header has been parsed"
5580            );
5581            assert!(
5582                session.backend_token.is_some(),
5583                "a dialed backend must own its backend token"
5584            );
5585        }
5586    }
5587
5588    /// Queue `wire` on the fixture's client, half-close its write side, then
5589    /// observe the accepted frontend socket through a real epoll instance
5590    /// registered only AFTER both the bytes and the FIN are queued.
5591    ///
5592    /// On Linux a `shutdown(SHUT_WR)` behind pending bytes is reported as one
5593    /// event carrying `EPOLLIN | EPOLLRDHUP`, which `Ready::from(&Event)`
5594    /// maps to `READABLE | HUP` without `WRITE_CLOSED`: the peer ended its
5595    /// request stream but still owns a usable response stream. The asserts
5596    /// pin that coalesced shape, so the tests that call this exercise the
5597    /// real kernel event rather than a synthetic readiness word.
5598    #[cfg(target_os = "linux")]
5599    fn half_close_and_observe_frontend(fixture: &mut ExpectProxyFixture, wire: &[u8]) -> Ready {
5600        use std::{io::Write as _, net::Shutdown, os::fd::AsRawFd, time::Duration};
5601
5602        use mio::{Events, Interest, Poll, Token, unix::SourceFd};
5603
5604        fixture
5605            .client
5606            .write_all(wire)
5607            .expect("write the request stream");
5608        fixture
5609            .client
5610            .shutdown(Shutdown::Write)
5611            .expect("half-close the client write side");
5612
5613        let frontend_fd = fixture.session.borrow().state.front_socket().as_raw_fd();
5614        let mut source = SourceFd(&frontend_fd);
5615        let mut poll = Poll::new().expect("create frontend poll");
5616        let observed_token = Token(7);
5617        poll.registry()
5618            .register(
5619                &mut source,
5620                observed_token,
5621                Interest::READABLE | Interest::WRITABLE,
5622            )
5623            .expect("register accepted frontend socket");
5624        let mut events = Events::with_capacity(4);
5625        poll.poll(&mut events, Some(Duration::from_secs(2)))
5626            .expect("poll accepted frontend socket");
5627        let ready = events
5628            .iter()
5629            .find(|event| event.token() == observed_token)
5630            .map(Ready::from)
5631            .expect("the queued request and FIN produce an epoll event");
5632        assert!(
5633            ready.is_readable(),
5634            "the queued FIN must be reported readable: {ready:?}"
5635        );
5636        assert!(
5637            ready.is_hup(),
5638            "the peer FIN must be reported as HUP: {ready:?}"
5639        );
5640        assert!(
5641            !ready.is_write_closed(),
5642            "EPOLLRDHUP leaves the response direction open: {ready:?}"
5643        );
5644        ready
5645    }
5646
5647    /// A real Linux `shutdown(SHUT_WR)` can report the final request bytes
5648    /// and `EPOLLRDHUP` in the same epoll event (`READABLE | HUP`, see
5649    /// `half_close_and_observe_frontend`). An `ExpectProxyProtocol` session
5650    /// must consume the complete PROXY-v2 header before acting on HUP,
5651    /// upgrade to `Pipe`, drain the payload already queued behind the
5652    /// header, and forward it to the backend (sozu-proxy/sozu#1823).
5653    ///
5654    /// To SEE THIS RED: delete the `ExpectProxyProtocol` drain that
5655    /// `TcpSession::ready_inner` runs before its frontend-HUP check. HUP is
5656    /// then handled first, `front_hup`'s wildcard arm closes the expect state
5657    /// before the already-buffered header is parsed, and the `!closed`
5658    /// assertion fails.
5659    #[cfg(target_os = "linux")]
5660    #[test]
5661    fn expect_proxy_half_close_drains_payload_to_the_backend() {
5662        use std::{io::Read as _, time::Duration};
5663
5664        let mut fixture = expect_proxy_fixture();
5665
5666        let payload = b"request-after-proxy-header";
5667        let mut wire = proxy_protocol_v2_ipv4_header().to_vec();
5668        wire.extend_from_slice(payload);
5669        let ready = half_close_and_observe_frontend(&mut fixture, &wire);
5670
5671        fixture
5672            .session
5673            .borrow_mut()
5674            .update_readiness(fixture.frontend_token, ready);
5675        let closed = fixture
5676            .session
5677            .borrow_mut()
5678            .ready(fixture.proxy_session.clone());
5679        assert!(
5680            !closed,
5681            "READABLE|HUP must parse the header and retain the response direction"
5682        );
5683        assert!(
5684            !matches!(
5685                fixture.session.borrow().state,
5686                TcpStateMachine::ExpectProxyProtocol(_)
5687            ),
5688            "the complete header must upgrade the expect state"
5689        );
5690
5691        let (mut backend, _) = fixture
5692            ._backend_listener
5693            .accept()
5694            .expect("the parsed header triggers the backend dial");
5695        backend
5696            .set_read_timeout(Some(Duration::from_secs(2)))
5697            .expect("set backend read timeout");
5698        backend
5699            .set_write_timeout(Some(Duration::from_secs(2)))
5700            .expect("set backend write timeout");
5701
5702        let backend_token = fixture
5703            .session
5704            .borrow()
5705            .backend_token
5706            .expect("the upgraded session owns a backend token");
5707        // The backend connects and the payload drained behind the header is
5708        // flushed to it. Whether the pipe then stays open for a response is
5709        // `Pipe::check_connections`' half-close policy (it keeps a
5710        // frontend-closed pipe only while request bytes are in flight), the
5711        // same for every TCP session with or without a PROXY header, so it is
5712        // deliberately not asserted here.
5713        fixture
5714            .session
5715            .borrow_mut()
5716            .update_readiness(backend_token, Ready::WRITABLE);
5717        let _ = fixture
5718            .session
5719            .borrow_mut()
5720            .ready(fixture.proxy_session.clone());
5721
5722        let mut backend_received = vec![0; payload.len()];
5723        backend
5724            .read_exact(&mut backend_received)
5725            .expect("backend receives the payload queued before FIN");
5726        assert_eq!(backend_received, payload);
5727    }
5728
5729    /// The negative space of the drain above: a bare TCP health check
5730    /// (connect, zero bytes, FIN) observed as `READABLE | HUP` must still
5731    /// close at once and never dial the backend. The drain reads EOF with
5732    /// nothing accumulated, and `ExpectProxyProtocol::readable`'s zero-byte
5733    /// `SocketResult::Closed` arm closes the session -- the same outcome as
5734    /// when the FIN arrives in its own event.
5735    #[cfg(target_os = "linux")]
5736    #[test]
5737    fn expect_proxy_zero_byte_half_close_closes_without_dialing() {
5738        let mut fixture = expect_proxy_fixture();
5739        let ready = half_close_and_observe_frontend(&mut fixture, &[]);
5740
5741        fixture
5742            .session
5743            .borrow_mut()
5744            .update_readiness(fixture.frontend_token, ready);
5745        let closed = fixture
5746            .session
5747            .borrow_mut()
5748            .ready(fixture.proxy_session.clone());
5749        assert!(
5750            closed,
5751            "a zero-byte client half-close must close the session"
5752        );
5753        let session = fixture.session.borrow();
5754        assert_eq!(
5755            session.back_connected(),
5756            BackendConnectionStatus::NotConnected,
5757            "a zero-byte health check must never dial the backend"
5758        );
5759        assert!(
5760            session.backend_token.is_none(),
5761            "a zero-byte health check must never wire a backend token"
5762        );
5763    }
5764
5765    /// A client that sends only part of a PROXY-v2 header and then FIN can
5766    /// never complete it: once the drain has consumed every readable byte,
5767    /// the frontend HUP must still close the session instead of leaving it
5768    /// to the frontend timeout, and no backend may be dialed.
5769    #[cfg(target_os = "linux")]
5770    #[test]
5771    fn expect_proxy_partial_header_half_close_closes_without_dialing() {
5772        let mut fixture = expect_proxy_fixture();
5773        let header = proxy_protocol_v2_ipv4_header();
5774        let ready = half_close_and_observe_frontend(&mut fixture, &header[..10]);
5775
5776        fixture
5777            .session
5778            .borrow_mut()
5779            .update_readiness(fixture.frontend_token, ready);
5780        let closed = fixture
5781            .session
5782            .borrow_mut()
5783            .ready(fixture.proxy_session.clone());
5784        assert!(
5785            closed,
5786            "an incomplete PROXY header followed by FIN must close the session"
5787        );
5788        let session = fixture.session.borrow();
5789        assert_eq!(
5790            session.back_connected(),
5791            BackendConnectionStatus::NotConnected,
5792            "an incomplete PROXY header must never dial the backend"
5793        );
5794    }
5795
5796    /// Structural safety for `ready_inner`'s `Connecting` branch: it reads
5797    /// the backend readiness, which `back_readiness` has none of for
5798    /// `ExpectProxyProtocol` (and `unreachable!`s on for `FailedUpgrade`).
5799    ///
5800    /// No input sequence reaches this pairing -- `connect_to_backend` is the
5801    /// only writer of `Connecting` and can no longer run in the expect state
5802    /// (see the test above), and no transition ever enters
5803    /// `ExpectProxyProtocol` from another state -- so the state is
5804    /// synthesized here directly. The branch must degrade to a logged close
5805    /// rather than take the whole worker down with it: a backend handshake
5806    /// with no readiness slot can never complete.
5807    ///
5808    /// To SEE THIS RED: restore either `self.back_readiness().unwrap()` in
5809    /// `ready_inner`'s `Connecting` branch -- the `None` this test installs
5810    /// then panics with "called `Option::unwrap()` on a `None` value".
5811    #[test]
5812    fn connecting_backend_without_readiness_closes_instead_of_panicking() {
5813        let fixture = expect_proxy_fixture();
5814        let (session, proxy_session) = (fixture.session.clone(), fixture.proxy_session.clone());
5815
5816        // The invariant break itself: `Connecting` while the state still has
5817        // no backend side.
5818        session.borrow_mut().backend_connected =
5819            BackendConnectionStatus::Connecting(Instant::now());
5820        assert!(
5821            session.borrow_mut().back_readiness().is_none(),
5822            "the expect state must have no backend readiness -- the premise of this test"
5823        );
5824
5825        let closed = session.borrow_mut().ready(proxy_session.clone());
5826        assert!(
5827            closed,
5828            "a connecting backend with no readiness slot must close the session"
5829        );
5830    }
5831
5832    /// Drive an `ExpectHeader` fixture through its PROXY header, which
5833    /// upgrades it to `Pipe` and dials the backend in that same `ready()`.
5834    fn expect_proxy_fixture_dialed() -> ExpectProxyFixture {
5835        use std::io::Write as _;
5836
5837        let mut fixture = expect_proxy_fixture();
5838        fixture
5839            .client
5840            .write_all(&proxy_protocol_v2_ipv4_header())
5841            .expect("write the PROXY-v2 header");
5842        fixture
5843            .session
5844            .borrow_mut()
5845            .update_readiness(fixture.frontend_token, Ready::READABLE);
5846        let proxy_session = fixture.proxy_session.clone();
5847        let _ = fixture.session.borrow_mut().ready(proxy_session);
5848        {
5849            let session = fixture.session.borrow();
5850            assert!(
5851                matches!(session.state, TcpStateMachine::Pipe(_)),
5852                "precondition: the parsed header upgrades the session to Pipe"
5853            );
5854            assert!(
5855                session.backend_token.is_some(),
5856                "precondition: the Pipe dialed its backend"
5857            );
5858        }
5859        fixture
5860    }
5861
5862    /// The `Pipe`'s access log reports the address `connect_to_backend`
5863    /// dialed, handed over by `TcpSession::set_backend_address`, rather than
5864    /// a `getpeername(2)` on the backend socket.
5865    ///
5866    /// To SEE THIS RED: remove the `self.set_backend_address(...)` call from
5867    /// `connect_to_backend`. The pipe then reports `None`.
5868    #[test]
5869    fn a_dialed_pipe_reports_the_backend_address_it_was_given() {
5870        let fixture = expect_proxy_fixture_dialed();
5871        let backend_address = fixture
5872            ._backend_listener
5873            .local_addr()
5874            .expect("test backend local addr");
5875        let session = fixture.session.borrow();
5876        let TcpStateMachine::Pipe(pipe) = &session.state else {
5877            unreachable!("expect_proxy_fixture_dialed checked the Pipe state");
5878        };
5879        assert_eq!(
5880            pipe.get_backend_address(),
5881            Some(backend_address),
5882            "the pipe must report the backend address the session dialed"
5883        );
5884    }
5885
5886    /// Closing a TCP session issues no `EPOLL_CTL_DEL` for its front or
5887    /// back socket: the last `close(2)` of each, when the server drops the
5888    /// session right after `close()`, takes it out of the epoll set.
5889    ///
5890    /// The test reads the kernel's own epoll table (`/proc/self/fdinfo`), so
5891    /// it tells "never deregistered" from "removed by the close".
5892    ///
5893    /// To SEE THIS RED: put back
5894    /// `let _ = self.proxy.borrow().registry.deregister(&mut mio::unix::SourceFd(&front_socket.as_raw_fd()));`
5895    /// in `TcpSession::close` (or the matching backend call in
5896    /// `close_backend`). The first assertion after `close()` then fails.
5897    #[cfg(target_os = "linux")]
5898    #[test]
5899    fn closed_tcp_sessions_leave_their_sockets_to_close() {
5900        use std::os::fd::AsRawFd as _;
5901
5902        use crate::protocol::mux::test_support::epoll_watches;
5903
5904        let fixture = expect_proxy_fixture_dialed();
5905        let session = fixture.session.clone();
5906        let frontend_token = fixture.frontend_token;
5907        let registry = session
5908            .borrow()
5909            .proxy
5910            .borrow()
5911            .registry
5912            .try_clone()
5913            .expect("the test registry must clone");
5914
5915        // `create_session` registers the front socket; the fixture builds the
5916        // session directly, so register it the same way here.
5917        let (front_fd, back_fd) = {
5918            let mut session = session.borrow_mut();
5919            let front_fd = session.state.front_socket().as_raw_fd();
5920            registry
5921                .register(
5922                    &mut mio::unix::SourceFd(&front_fd),
5923                    frontend_token,
5924                    Interest::READABLE | Interest::WRITABLE,
5925                )
5926                .expect("register the front socket");
5927            let back_fd = session
5928                .back_socket_mut()
5929                .expect("the dialed Pipe owns a backend socket")
5930                .as_raw_fd();
5931            (front_fd, back_fd)
5932        };
5933        assert!(
5934            epoll_watches(&registry, front_fd) && epoll_watches(&registry, back_fd),
5935            "precondition: both sockets are in the epoll set"
5936        );
5937
5938        // What `Server::shut_down_sessions_by_frontend_tokens` does.
5939        let removed = session
5940            .borrow()
5941            .proxy
5942            .borrow()
5943            .sessions
5944            .borrow_mut()
5945            .slab
5946            .remove(frontend_token.0);
5947        session.borrow_mut().close();
5948        assert!(
5949            epoll_watches(&registry, front_fd),
5950            "TcpSession::close must not deregister a front socket it is about \
5951             to close: the close removes it, the EPOLL_CTL_DEL is a wasted syscall"
5952        );
5953        assert!(
5954            epoll_watches(&registry, back_fd),
5955            "TcpSession::close_backend must not deregister a backend socket it \
5956             is about to close: the close removes it"
5957        );
5958
5959        drop(removed);
5960        drop(session);
5961        drop(fixture);
5962        assert!(
5963            !epoll_watches(&registry, front_fd),
5964            "dropping the session closes the front socket's only descriptor, \
5965             which must take it out of the epoll set"
5966        );
5967        assert!(
5968            !epoll_watches(&registry, back_fd),
5969            "dropping the session closes the backend socket's only descriptor, \
5970             which must take it out of the epoll set"
5971        );
5972    }
5973
5974    /// A session dialed from a state that is not `Pipe` yet, plus everything
5975    /// that must outlive it.
5976    struct DialFixture {
5977        session: Rc<RefCell<TcpSession>>,
5978        proxy_session: Rc<RefCell<dyn ProxySession>>,
5979        frontend_token: Token,
5980        backend_address: SocketAddr,
5981        client: std::net::TcpStream,
5982        _frontend_listener: std::net::TcpListener,
5983        _backend_listener: std::net::TcpListener,
5984    }
5985
5986    /// Build a session for cluster `cluster-dial`, wired to a real bound
5987    /// backend and a real accepted frontend socket, registered in the slab the
5988    /// way `expect_proxy_fixture` does. `build` picks the constructor, hence
5989    /// the starting state; `sni_route` routes `example.com` to the cluster.
5990    fn dial_fixture(
5991        sni_route: bool,
5992        build: impl FnOnce(
5993            Checkout,
5994            Checkout,
5995            Token,
5996            Rc<RefCell<TcpListener>>,
5997            Rc<RefCell<TcpProxy>>,
5998            MioTcpStream,
5999            SocketAddr,
6000        ) -> TcpSession,
6001    ) -> DialFixture {
6002        let ServerParts {
6003            registry,
6004            sessions,
6005            pool,
6006            backends,
6007            ..
6008        } = prebuild_server(16, 16384, false).expect("could not prebuild a test server");
6009
6010        let backend_listener =
6011            std::net::TcpListener::bind("127.0.0.1:0").expect("bind test backend");
6012        let backend_address = backend_listener
6013            .local_addr()
6014            .expect("test backend local addr");
6015        backends.borrow_mut().add_backend(
6016            "cluster-dial",
6017            Backend::new("cluster-dial-1", backend_address, None, None, None),
6018        );
6019
6020        let proxy = Rc::new(RefCell::new(TcpProxy::new(
6021            registry,
6022            sessions.clone(),
6023            pool.clone(),
6024            backends,
6025        )));
6026        let mut bare_listener = test_listener();
6027        if sni_route {
6028            bare_listener
6029                .insert_sni_route("example.com".to_owned(), vec![], "cluster-dial".into())
6030                .expect("insert_sni_route must succeed for a valid test SNI");
6031        }
6032        let listener = Rc::new(RefCell::new(bare_listener));
6033
6034        let frontend_listener =
6035            std::net::TcpListener::bind("127.0.0.1:0").expect("bind test frontend listener");
6036        let client = std::net::TcpStream::connect(
6037            frontend_listener
6038                .local_addr()
6039                .expect("test frontend local addr"),
6040        )
6041        .expect("connect test client");
6042        let (frontend, frontend_peer) = frontend_listener
6043            .accept()
6044            .expect("accept the client connection");
6045        frontend
6046            .set_nonblocking(true)
6047            .expect("frontend socket nonblocking");
6048
6049        let (frontend_buffer, backend_buffer) = {
6050            let mut pool = pool.borrow_mut();
6051            (
6052                pool.checkout().expect("front buffer checkout must succeed"),
6053                pool.checkout().expect("back buffer checkout must succeed"),
6054            )
6055        };
6056
6057        let frontend_token = {
6058            let mut session_manager = sessions.borrow_mut();
6059            let entry = session_manager.slab.vacant_entry();
6060            let token = Token(entry.key());
6061            entry.insert(Rc::new(RefCell::new(crate::server::ListenSession {
6062                protocol: Protocol::TCPListen,
6063            })));
6064            token
6065        };
6066
6067        let session = Rc::new(RefCell::new(build(
6068            backend_buffer,
6069            frontend_buffer,
6070            frontend_token,
6071            listener,
6072            proxy,
6073            MioTcpStream::from_std(frontend),
6074            frontend_peer,
6075        )));
6076        let proxy_session: Rc<RefCell<dyn ProxySession>> = session.clone();
6077        sessions.borrow_mut().slab[frontend_token.0] = proxy_session.clone();
6078
6079        DialFixture {
6080            session,
6081            proxy_session,
6082            frontend_token,
6083            backend_address,
6084            client,
6085            _frontend_listener: frontend_listener,
6086            _backend_listener: backend_listener,
6087        }
6088    }
6089
6090    /// A legacy (non-SNI) session whose cluster uses `proxy_protocol`; its
6091    /// first `ready()` dials the backend while it is still in its
6092    /// PROXY-protocol state.
6093    fn proxy_protocol_session_dialed(proxy_protocol: ProxyProtocolConfig) -> DialFixture {
6094        let fixture = dial_fixture(
6095            false,
6096            |back_buffer, front_buffer, token, listener, proxy, socket, peer| {
6097                TcpSession::new(
6098                    back_buffer,
6099                    None,
6100                    Some("cluster-dial".into()),
6101                    Duration::from_secs(30),
6102                    Duration::from_secs(30),
6103                    Duration::from_secs(30),
6104                    front_buffer,
6105                    token,
6106                    listener,
6107                    Some(proxy_protocol),
6108                    proxy,
6109                    socket,
6110                    peer,
6111                    Duration::from_millis(0),
6112                )
6113            },
6114        );
6115        fixture
6116            .session
6117            .borrow_mut()
6118            .update_readiness(fixture.frontend_token, Ready::WRITABLE);
6119        let closed = fixture
6120            .session
6121            .borrow_mut()
6122            .ready(fixture.proxy_session.clone());
6123        assert!(!closed, "precondition: the dial keeps the session open");
6124        fixture
6125    }
6126
6127    #[test]
6128    fn static_tcp_route_keeps_its_capture_after_cluster_replacement() {
6129        let fixture = dial_fixture(
6130            false,
6131            |back_buffer, front_buffer, token, listener, proxy, socket, peer| {
6132                TcpSession::new(
6133                    back_buffer,
6134                    None,
6135                    Some("cluster-dial".into()),
6136                    Duration::from_secs(30),
6137                    Duration::from_secs(30),
6138                    Duration::from_secs(30),
6139                    front_buffer,
6140                    token,
6141                    listener,
6142                    None,
6143                    proxy,
6144                    socket,
6145                    peer,
6146                    Duration::from_millis(0),
6147                )
6148            },
6149        );
6150        let initial_capture = fixture.session.borrow().cluster_metrics_incarnation;
6151        assert!(
6152            matches!(
6153                initial_capture,
6154                crate::metrics::ClusterMetricsCapture::Active(_)
6155            ),
6156            "precondition: the static route captured an implicit active identity",
6157        );
6158
6159        crate::metrics::METRICS.with(|metrics| {
6160            let mut metrics = metrics.borrow_mut();
6161            metrics.remove_cluster("cluster-dial");
6162            metrics.add_cluster("cluster-dial");
6163            let replacement = metrics
6164                .cluster_incarnation("cluster-dial")
6165                .expect("replacement cluster identity");
6166            assert_ne!(
6167                initial_capture,
6168                crate::metrics::ClusterMetricsCapture::Active(replacement),
6169            );
6170        });
6171        fixture
6172            .session
6173            .borrow_mut()
6174            .update_readiness(fixture.frontend_token, Ready::WRITABLE);
6175        let closed = fixture
6176            .session
6177            .borrow_mut()
6178            .ready(fixture.proxy_session.clone());
6179        assert!(!closed, "the backend dial must keep the session open");
6180        assert_eq!(
6181            fixture.session.borrow().cluster_metrics_incarnation,
6182            initial_capture,
6183            "a retry after replacement must keep the identity captured at static routing time",
6184        );
6185        crate::metrics::METRICS.with(|metrics| metrics.borrow_mut().remove_cluster("cluster-dial"));
6186    }
6187
6188    #[test]
6189    fn sni_tcp_route_keeps_its_capture_after_cluster_replacement() {
6190        use std::io::Write as _;
6191
6192        let mut fixture = dial_fixture(
6193            true,
6194            |back_buffer, front_buffer, token, listener, proxy, socket, peer| {
6195                TcpSession::new_sni_preread(
6196                    back_buffer,
6197                    Duration::from_secs(30),
6198                    Duration::from_secs(30),
6199                    front_buffer,
6200                    token,
6201                    listener,
6202                    proxy,
6203                    socket,
6204                    peer,
6205                    Duration::from_millis(0),
6206                    Duration::from_secs(3),
6207                    16384,
6208                )
6209            },
6210        );
6211        fixture
6212            .client
6213            .write_all(&minimal_client_hello_wire("example.com"))
6214            .expect("write ClientHello");
6215        fixture.client.flush().ok();
6216        for _ in 0..10 {
6217            if fixture.session.borrow().cluster_id.is_some() {
6218                break;
6219            }
6220            let _ = fixture.session.borrow_mut().readable();
6221        }
6222        assert_eq!(
6223            fixture.session.borrow().cluster_id.as_deref(),
6224            Some("cluster-dial")
6225        );
6226        let initial_capture = fixture.session.borrow().cluster_metrics_incarnation;
6227        assert!(
6228            matches!(
6229                initial_capture,
6230                crate::metrics::ClusterMetricsCapture::Active(_)
6231            ),
6232            "precondition: SNI routing captured an implicit active identity",
6233        );
6234
6235        crate::metrics::METRICS.with(|metrics| {
6236            let mut metrics = metrics.borrow_mut();
6237            metrics.remove_cluster("cluster-dial");
6238            metrics.add_cluster("cluster-dial");
6239            let replacement = metrics
6240                .cluster_incarnation("cluster-dial")
6241                .expect("replacement cluster identity");
6242            assert_ne!(
6243                initial_capture,
6244                crate::metrics::ClusterMetricsCapture::Active(replacement),
6245            );
6246        });
6247        let closed = fixture
6248            .session
6249            .borrow_mut()
6250            .ready(fixture.proxy_session.clone());
6251        assert!(
6252            !closed,
6253            "the SNI-routed backend dial must keep the session open"
6254        );
6255        assert_eq!(
6256            fixture.session.borrow().cluster_metrics_incarnation,
6257            initial_capture,
6258            "SNI upgrade and retry after replacement must keep the route-time identity",
6259        );
6260        crate::metrics::METRICS.with(|metrics| metrics.borrow_mut().remove_cluster("cluster-dial"));
6261    }
6262
6263    /// Upgrade a dialed, not-yet-`Pipe` session and assert the `Pipe` reports
6264    /// the backend address the session dialed: the address its access log
6265    /// carries.
6266    fn assert_upgraded_pipe_reports_the_dialed_address(fixture: &DialFixture) {
6267        {
6268            let session = fixture.session.borrow();
6269            assert!(
6270                !matches!(session.state, TcpStateMachine::Pipe(_)),
6271                "precondition: the backend was dialed before the Pipe existed"
6272            );
6273            assert!(
6274                session.backend_token.is_some(),
6275                "precondition: the session dialed its backend"
6276            );
6277        }
6278        let closed = fixture.session.borrow_mut().upgrade();
6279        assert!(!closed, "the upgrade to Pipe must succeed");
6280        let session = fixture.session.borrow();
6281        let TcpStateMachine::Pipe(pipe) = &session.state else {
6282            panic!("the upgrade must install a Pipe");
6283        };
6284        assert_eq!(
6285            pipe.get_backend_address(),
6286            Some(fixture.backend_address),
6287            "the Pipe must report the backend address dialed before the upgrade"
6288        );
6289    }
6290
6291    /// `SendProxyProtocol` dials before it becomes a `Pipe`: `upgrade_send`
6292    /// must hand the dialed address over.
6293    ///
6294    /// To SEE THIS RED: make `TcpSession::set_backend_address` record nothing,
6295    /// or pass `None` to `pipe.set_backend_address` in `upgrade_send`.
6296    #[test]
6297    fn a_send_proxy_protocol_pipe_reports_the_dialed_backend_address() {
6298        let fixture = proxy_protocol_session_dialed(ProxyProtocolConfig::SendHeader);
6299        assert!(
6300            matches!(
6301                fixture.session.borrow().state,
6302                TcpStateMachine::SendProxyProtocol(_)
6303            ),
6304            "precondition: a SendHeader cluster starts in the send state"
6305        );
6306        assert_upgraded_pipe_reports_the_dialed_address(&fixture);
6307    }
6308
6309    /// `RelayProxyProtocol` dials before it becomes a `Pipe`: `upgrade_relay`
6310    /// must hand the dialed address over.
6311    ///
6312    /// To SEE THIS RED: make `TcpSession::set_backend_address` record nothing,
6313    /// or pass `None` to `pipe.set_backend_address` in `upgrade_relay`.
6314    #[test]
6315    fn a_relay_proxy_protocol_pipe_reports_the_dialed_backend_address() {
6316        let fixture = proxy_protocol_session_dialed(ProxyProtocolConfig::RelayHeader);
6317        assert!(
6318            matches!(
6319                fixture.session.borrow().state,
6320                TcpStateMachine::RelayProxyProtocol(_)
6321            ),
6322            "precondition: a RelayHeader cluster starts in the relay state"
6323        );
6324        assert_upgraded_pipe_reports_the_dialed_address(&fixture);
6325    }
6326
6327    /// An SNI-routed session dials from `SniPreread`: `upgrade_sni_preread`
6328    /// builds the `Pipe` through `build_pipe_from_preread`, which must hand
6329    /// the dialed address over.
6330    ///
6331    /// To SEE THIS RED: make `TcpSession::set_backend_address` record nothing,
6332    /// or pass `None` to `pipe.set_backend_address` in
6333    /// `build_pipe_from_preread`.
6334    #[test]
6335    fn an_sni_preread_pipe_reports_the_dialed_backend_address() {
6336        use std::io::Write as _;
6337
6338        let mut fixture = dial_fixture(
6339            true,
6340            |back_buffer, front_buffer, token, listener, proxy, socket, peer| {
6341                TcpSession::new_sni_preread(
6342                    back_buffer,
6343                    Duration::from_secs(30),
6344                    Duration::from_secs(30),
6345                    front_buffer,
6346                    token,
6347                    listener,
6348                    proxy,
6349                    socket,
6350                    peer,
6351                    Duration::from_millis(0),
6352                    Duration::from_secs(3),
6353                    16384,
6354                )
6355            },
6356        );
6357        fixture
6358            .client
6359            .write_all(&minimal_client_hello_wire("example.com"))
6360            .expect("write ClientHello");
6361        fixture.client.flush().ok();
6362        for _ in 0..10 {
6363            if fixture.session.borrow().cluster_id.is_some() {
6364                break;
6365            }
6366            let _ = fixture.session.borrow_mut().readable();
6367        }
6368        assert_eq!(
6369            fixture.session.borrow().cluster_id.as_deref(),
6370            Some("cluster-dial"),
6371            "precondition: the ClientHello routes the session"
6372        );
6373        let closed = fixture
6374            .session
6375            .borrow_mut()
6376            .ready(fixture.proxy_session.clone());
6377        assert!(!closed, "precondition: the dial keeps the session open");
6378        assert!(
6379            matches!(
6380                fixture.session.borrow().state,
6381                TcpStateMachine::SniPreread(_)
6382            ),
6383            "precondition: the routed session dials while still in SniPreread"
6384        );
6385        assert_upgraded_pipe_reports_the_dialed_address(&fixture);
6386    }
6387}
6388
6389/// Coverage for the `RelayProxyProtocol` -> `Pipe` handoff. It lives here
6390/// rather than in `relay.rs`'s own `relay_test` because `into_pipe` demands an
6391/// `Rc<RefCell<TcpListener>>` and every field of that struct — plus its
6392/// constructor — is private to this module, so the relay's test module cannot
6393/// build one. A separate module (rather than the `tests` module above) avoids
6394/// that module's `use std::net::TcpListener` shadowing `super::TcpListener`.
6395#[cfg(test)]
6396mod relay_upgrade_tests {
6397    use std::net::TcpListener as StdTcpListener;
6398
6399    use sozu_command::{config::ListenerBuilder, proto::command::SocketAddress};
6400
6401    use super::*;
6402    use crate::{pool::Pool, testing::provide_port};
6403
6404    fn test_listener() -> TcpListener {
6405        let config = ListenerBuilder::new_tcp(SocketAddress::new_v4(127, 0, 0, 1, provide_port()))
6406            .to_tcp(None)
6407            .expect("could not build a TcpListenerConfig for the test");
6408        TcpListener::new(config, Token(0)).expect("could not build a bare TcpListener for the test")
6409    }
6410
6411    fn connected_pair() -> (std::net::TcpStream, std::net::TcpStream) {
6412        let listener = StdTcpListener::bind("127.0.0.1:0").expect("bind the test listener");
6413        let address = listener
6414            .local_addr()
6415            .expect("the test listener must expose its address");
6416        let local = std::net::TcpStream::connect(address).expect("connect the test pair");
6417        let (peer, _addr) = listener.accept().expect("accept the test pair");
6418        local
6419            .set_nonblocking(true)
6420            .expect("the local end must be non-blocking");
6421        peer.set_nonblocking(true)
6422            .expect("the peer end must be non-blocking");
6423        (local, peer)
6424    }
6425
6426    /// `Pipe::new` ends with `arm_inherited_buffer_writes`, which arms the
6427    /// backend WRITABLE readiness when the inherited `frontend_buffer` is
6428    /// non-empty. `build_pipe_from_preread` carries a comment stating outright
6429    /// that a bare `pipe.backend_readiness.event = ..` pair must NOT be used
6430    /// in its place, because the drain of that buffer must not depend on the
6431    /// restored event carrying WRITABLE itself — and `upgrade_send` re-runs the
6432    /// arm for the same reason (sozu-proxy/sozu#1279's close-before-flush
6433    /// truncation). `RelayProxyProtocol::into_pipe` was the one pipe-entry path
6434    /// still using the bare pair. That was harmless only while the relay's
6435    /// buffer was always empty at upgrade time; now that the relay leaves
6436    /// pipelined client payload in it, this path joins the class those comments
6437    /// describe, and the surviving wake-up must be by construction rather than
6438    /// by the accident that `back_writable` happens to be entered with WRITABLE
6439    /// already set.
6440    ///
6441    /// To SEE THIS RED: in `relay.rs::into_pipe`, replace the
6442    /// `pipe.restore_readiness_events(..)` call with the bare pair
6443    /// `pipe.frontend_readiness.event = self.frontend_readiness.event;` /
6444    /// `pipe.backend_readiness.event = self.backend_readiness.event;`.
6445    #[test]
6446    fn into_pipe_rearms_the_backend_write_for_pipelined_payload() {
6447        let listener = Rc::new(RefCell::new(test_listener()));
6448        let (frontend, _front_peer) = connected_pair();
6449        let (backend, _backend_peer) = connected_pair();
6450
6451        let mut pool = Pool::with_capacity(2, 2, 4096);
6452        let mut frontend_buffer = pool.checkout().expect("frontend buffer");
6453        let backend_buffer = pool.checkout().expect("backend buffer");
6454        // What the relay leaves behind once the header is forwarded: the
6455        // payload the client pipelined into the same read.
6456        let payload = b"pipelined-payload";
6457        frontend_buffer.space()[..payload.len()].copy_from_slice(payload);
6458        frontend_buffer.fill(payload.len());
6459
6460        let mut relay = RelayProxyProtocol::new(
6461            MioTcpStream::from_std(frontend),
6462            Token(0),
6463            Ulid::generate(),
6464            Some(MioTcpStream::from_std(backend)),
6465            frontend_buffer,
6466        );
6467        // `into_pipe` never reads this, but a relay only ever upgrades with a
6468        // fully forwarded header, so leave the state coherent.
6469        relay.header_size = Some(28);
6470        relay.set_back_token(Token(1));
6471        // The case the `build_pipe_from_preread` comment names explicitly: the
6472        // readiness word restored from the previous state need not itself carry
6473        // WRITABLE, and the inherited drain must not depend on it.
6474        relay.backend_readiness.event = Ready::EMPTY;
6475        relay.frontend_readiness.event = Ready::EMPTY;
6476
6477        let pipe = relay.into_pipe(backend_buffer, listener);
6478
6479        assert!(
6480            pipe.backend_readiness.event.is_writable(),
6481            "a non-empty inherited frontend buffer must leave a backend WRITABLE event queued, or the pipelined payload is never flushed"
6482        );
6483        assert!(
6484            pipe.backend_readiness.interest.is_writable(),
6485            "the inherited-write arm must also hold the backend WRITABLE interest"
6486        );
6487    }
6488}