1use std::{cell::RefCell, rc::Rc};
10
11use mio::{net::TcpStream, *};
12use nom::{Err, HexDisplay};
13use rusty_ulid::Ulid;
14use sozu_command::{
15 config::MAX_LOOP_ITERATIONS,
16 logging::{LogContext, ansi_palette},
17};
18
19use super::{header::ProxyAddr, parser::parse_v2_header};
20use crate::metrics::names;
21use crate::{
22 Protocol, Readiness, SessionMetrics, StateResult,
23 pool::Checkout,
24 protocol::{
25 SessionResult, SessionState,
26 pipe::{Pipe, WebSocketContext},
27 },
28 socket::{SocketHandler, SocketResult},
29 sozu_command::ready::Ready,
30 tcp::TcpListener,
31 timer::TimeoutContainer,
32};
33
34const V2_FIXED_HEADER_LEN: usize = 16;
35const MAX_HEADER_LEN: usize = 232;
36
37macro_rules! log_module_context {
42 () => {{
43 let (open, reset, _, _, _) = ansi_palette();
44 format!(
45 "{open}PROXY-EXPECT{reset}\t >>>",
46 open = open,
47 reset = reset
48 )
49 }};
50}
51
52macro_rules! log_context {
58 ($self:expr) => {{
59 let (open, reset, grey, gray, white) = ansi_palette();
60 format!(
61 "{gray}{ctx}{reset}\t{open}PROXY-EXPECT{reset}\t{grey}Session{reset}({gray}frontend{reset}={white}{frontend}{reset}, {gray}index{reset}={white}{index}{reset}, {gray}readiness{reset}={white}{readiness}{reset})\t >>>",
62 open = open,
63 reset = reset,
64 grey = grey,
65 gray = gray,
66 white = white,
67 ctx = $self.log_context(),
68 frontend = $self.frontend_token.0,
69 index = $self.index,
70 readiness = $self.frontend_readiness,
71 )
72 }};
73}
74
75#[derive(Clone, Copy)]
80pub enum HeaderLen {
81 V4,
82 V6,
83 Unix,
84}
85
86pub struct ExpectProxyProtocol<Front: SocketHandler> {
88 pub addresses: Option<ProxyAddr>,
89 pub container_frontend_timeout: TimeoutContainer,
90 frontend_buffer: [u8; MAX_HEADER_LEN],
91 pub frontend_readiness: Readiness,
92 pub frontend_token: Token,
93 pub frontend: Front,
94 read_target: usize,
95 index: usize,
96 pub request_id: Ulid,
97}
98
99impl<Front: SocketHandler> ExpectProxyProtocol<Front> {
100 pub fn new(
104 container_frontend_timeout: TimeoutContainer,
105 frontend: Front,
106 frontend_token: Token,
107 request_id: Ulid,
108 ) -> Self {
109 ExpectProxyProtocol {
110 addresses: None,
111 container_frontend_timeout,
112 frontend_buffer: [0; MAX_HEADER_LEN],
113 frontend_readiness: Readiness {
114 interest: Ready::READABLE | Ready::HUP | Ready::ERROR,
115 event: Ready::EMPTY,
116 },
117 frontend_token,
118 frontend,
119 read_target: V2_FIXED_HEADER_LEN,
120 index: 0,
121 request_id,
122 }
123 }
124
125 pub fn readable(&mut self, metrics: &mut SessionMetrics) -> SessionResult {
126 let total_len = self.read_target;
127
128 debug_assert!(
134 self.index <= total_len,
135 "read cursor must not exceed the current stage target"
136 );
137 debug_assert!(
138 total_len <= self.frontend_buffer.len(),
139 "stage target must fit the fixed proxy-protocol buffer"
140 );
141
142 let index_before = self.index;
143 let (sz, socket_result) = self
144 .frontend
145 .socket_read(&mut self.frontend_buffer[self.index..total_len]);
146 debug_assert!(
149 sz <= total_len - index_before,
150 "socket_read cannot return more bytes than the slice it was given"
151 );
152 trace!(
153 "{} read {} bytes and res={:?}, total_len = {}",
154 log_context!(self),
155 sz,
156 socket_result,
157 total_len
158 );
159
160 if sz > 0 {
161 self.index += sz;
162 debug_assert_eq!(
166 self.index,
167 index_before + sz,
168 "read cursor advances by exactly the bytes just read"
169 );
170 debug_assert!(
171 self.index <= self.frontend_buffer.len(),
172 "accumulated bytes must never exceed the fixed buffer bound"
173 );
174
175 count!(names::backend::BYTES_IN, sz as i64);
176 metrics.bin += sz;
177 } else {
178 debug_assert_eq!(
179 self.index, index_before,
180 "a non-positive read must leave the cursor unchanged"
181 );
182 self.frontend_readiness.event.remove(Ready::READABLE);
183 }
184
185 match socket_result {
186 SocketResult::Error => {
187 error!(
188 "{} front socket error, closing the connection (read {}, wrote {})",
189 log_context!(self),
190 metrics.bin,
191 metrics.bout
192 );
193 incr!(names::proxy_protocol::ERRORS);
194 self.frontend_readiness.reset();
195 return SessionResult::Close;
196 }
197 SocketResult::WouldBlock => {
198 self.frontend_readiness.event.remove(Ready::READABLE);
199 }
200 SocketResult::Closed => {
201 if self.index == 0 {
207 trace!(
208 "{} socket closed with 0 bytes, closing session",
209 log_context!(self)
210 );
211 return SessionResult::Close;
212 }
213 }
214 SocketResult::Continue => {}
215 }
216
217 match parse_v2_header(&self.frontend_buffer[..self.index]) {
218 Ok((rest, header)) => {
219 debug_assert!(
224 rest.len() <= self.index,
225 "parser remainder cannot exceed the accumulated input"
226 );
227 debug_assert!(
228 rest.is_empty(),
229 "exact-length reads must leave application payload in the socket"
230 );
231 trace!(
232 "{} got expect header: {:?}, rest.len() = {}",
233 log_context!(self),
234 header,
235 rest.len()
236 );
237 self.addresses = Some(header.addr);
238 SessionResult::Upgrade
239 }
240 Err(Err::Incomplete(_)) if self.index < self.read_target => SessionResult::Continue,
241 Err(Err::Incomplete(_)) if self.read_target == V2_FIXED_HEADER_LEN => {
242 let declared_len =
243 u16::from_be_bytes([self.frontend_buffer[14], self.frontend_buffer[15]])
244 as usize;
245 let declared_total = V2_FIXED_HEADER_LEN + declared_len;
246
247 if declared_total > self.frontend_buffer.len() {
248 error!(
249 "{} proxy protocol header declares {} bytes, exceeding maximum size ({} bytes), closing",
250 log_context!(self),
251 declared_total,
252 self.frontend_buffer.len()
253 );
254 incr!(names::proxy_protocol::ERRORS);
255 self.frontend_readiness.reset();
256 return SessionResult::Close;
257 }
258
259 if declared_total > V2_FIXED_HEADER_LEN {
260 self.read_target = declared_total;
261 return SessionResult::Continue;
262 }
263
264 error!(
265 "{} proxy protocol header is incomplete at its declared length ({} bytes), closing",
266 log_context!(self),
267 declared_total
268 );
269 incr!(names::proxy_protocol::ERRORS);
270 self.frontend_readiness.reset();
271 SessionResult::Close
272 }
273 Err(Err::Incomplete(_)) => {
274 error!(
275 "{} proxy protocol header is incomplete at its declared length ({} bytes), closing",
276 log_context!(self),
277 self.read_target
278 );
279 incr!(names::proxy_protocol::ERRORS);
280 self.frontend_readiness.reset();
281 SessionResult::Close
282 }
283 Err(Err::Error(e)) | Err(Err::Failure(e)) => {
284 error!(
285 "{} parse error, closing the connection:\n{}",
286 log_context!(self),
287 e.input.to_hex(16)
288 );
289 incr!(names::proxy_protocol::ERRORS);
290 self.frontend_readiness.reset();
291 SessionResult::Close
292 }
293 }
294 }
295
296 pub fn front_socket(&self) -> &TcpStream {
297 self.frontend.socket_ref()
298 }
299
300 pub fn into_pipe(
301 self,
302 front_buf: Checkout,
303 back_buf: Checkout,
304 backend_socket: Option<TcpStream>,
305 backend_token: Option<Token>,
306 listener: Rc<RefCell<TcpListener>>,
307 ) -> Pipe<Front, TcpListener> {
308 let addr = self
320 .addresses
321 .as_ref()
322 .and_then(|pa| pa.source())
323 .or_else(|| self.front_socket().peer_addr().ok());
324
325 let mut pipe = Pipe::new(
326 back_buf,
327 None,
328 backend_socket,
329 None,
330 None,
331 Some(self.container_frontend_timeout),
332 None,
333 front_buf,
334 self.frontend_token,
335 self.frontend,
336 listener,
337 Protocol::TCP,
338 self.request_id,
339 self.request_id,
340 addr,
341 WebSocketContext::Tcp,
342 );
343
344 pipe.frontend_readiness.event = self.frontend_readiness.event;
345
346 if let Some(backend_token) = backend_token {
347 pipe.set_back_token(backend_token);
348 }
349
350 pipe
351 }
352
353 pub fn log_context(&self) -> LogContext<'_> {
354 LogContext {
355 session_id: self.request_id,
356 request_id: None,
357 cluster_id: None,
358 backend_id: None,
359 }
360 }
361}
362
363impl<Front: SocketHandler> SessionState for ExpectProxyProtocol<Front> {
364 fn ready(
365 &mut self,
366 _session: Rc<RefCell<dyn crate::ProxySession>>,
367 _proxy: Rc<RefCell<dyn crate::L7Proxy>>,
368 metrics: &mut SessionMetrics,
369 ) -> SessionResult {
370 let mut counter = 0;
371
372 while counter < MAX_LOOP_ITERATIONS {
373 let frontend_interest = self.frontend_readiness.filter_interest();
374
375 trace!(
376 "{} {:?} -> None",
377 log_context!(self),
378 self.frontend_readiness
379 );
380
381 if self.frontend_readiness.event.is_hup() && !frontend_interest.is_readable() {
389 return SessionResult::Close;
390 }
391
392 if frontend_interest.is_empty() {
393 break;
394 }
395
396 if frontend_interest.is_readable() {
397 let session_result = self.readable(metrics);
398 if session_result != SessionResult::Continue {
399 return session_result;
400 }
401 }
402
403 if frontend_interest.is_error() {
404 error!("{} front error, disconnecting", log_context!(self));
405 self.frontend_readiness.interest = Ready::EMPTY;
406
407 return SessionResult::Close;
408 }
409
410 let counter_before = counter;
411 counter += 1;
412 debug_assert_eq!(counter, counter_before + 1, "loop counter advances by one");
416 debug_assert!(
417 counter <= MAX_LOOP_ITERATIONS,
418 "loop counter must stay within the iteration cap"
419 );
420 }
421
422 if counter >= MAX_LOOP_ITERATIONS {
423 error!(
424 "{} handling session went through {} iterations, there's a probable infinite loop bug, closing the connection",
425 log_context!(self),
426 MAX_LOOP_ITERATIONS
427 );
428 incr!(names::http::INFINITE_LOOP_ERROR);
429
430 self.print_state("");
431
432 return SessionResult::Close;
433 }
434
435 SessionResult::Continue
436 }
437
438 fn update_readiness(&mut self, token: Token, events: Ready) {
439 if self.frontend_token == token {
440 self.frontend_readiness.event |= events;
441 }
442 }
443
444 fn timeout(&mut self, token: Token, _metrics: &mut SessionMetrics) -> StateResult {
445 if self.frontend_token == token {
446 self.container_frontend_timeout.triggered();
447 return StateResult::CloseSession;
448 }
449
450 error!(
451 "{} got timeout for an invalid token: {:?}",
452 log_module_context!(),
453 token
454 );
455 StateResult::CloseSession
456 }
457
458 fn cancel_timeouts(&mut self) {
459 self.container_frontend_timeout.cancel();
460 }
461
462 fn print_state(&self, context: &str) {
463 error!(
464 "{} {} Session(Expect)\n\tFrontend:\n\t\ttoken: {:?}\treadiness: {:?}",
465 log_context!(self),
466 context,
467 self.frontend_token,
468 self.frontend_readiness
469 );
470 }
471}
472
473#[cfg(test)]
474mod expect_test {
475 use std::{
476 io::Write,
477 net::{IpAddr, Ipv4Addr, SocketAddr, TcpStream as StdTcpStream},
478 sync::{Arc, Barrier},
479 thread::{self, JoinHandle},
480 time::Duration,
481 };
482
483 use mio::net::TcpListener;
484 use rusty_ulid::Ulid;
485
486 use super::*;
487 use crate::protocol::proxy_protocol::header::*;
488
489 fn header_src() -> SocketAddr {
494 SocketAddr::new(IpAddr::V4(Ipv4Addr::new(125, 25, 10, 1)), 8080)
495 }
496
497 fn header_dst() -> SocketAddr {
498 SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 4, 5, 8)), 4200)
499 }
500
501 fn attributed_source_for(command: Command) -> Option<SocketAddr> {
512 setup_test_logger!();
513 let listener = TcpListener::bind("127.0.0.1:0".parse().expect("parse address error"))
514 .expect("could not bind the middleware listener");
515 let middleware_addr = listener
516 .local_addr()
517 .expect("the middleware listener must expose its address");
518 let barrier = Arc::new(Barrier::new(2));
519
520 let upfront = start_upfront_middleware(middleware_addr, barrier.clone(), command);
521 let addresses = start_middleware(listener, barrier);
522
523 upfront.join().expect("should join");
524
525 addresses.as_ref().and_then(ProxyAddr::source)
526 }
527
528 #[test]
529 fn middleware_should_receive_proxy_protocol_header_from_an_upfront_middleware() {
530 assert_eq!(
531 attributed_source_for(Command::Proxy),
532 Some(header_src()),
533 "a PROXY header still carries the real client through to the pipe phase"
534 );
535 }
536
537 #[test]
556 fn a_local_command_header_attributes_no_source_address() {
557 assert_eq!(
558 attributed_source_for(Command::Local),
559 None,
560 "a LOCAL header must attribute no source, so into_pipe falls back to peer_addr"
561 );
562 }
563
564 #[test]
565 fn oversized_declared_header_is_rejected_after_the_fixed_prelude() {
566 setup_test_logger!();
567 let listener = TcpListener::bind("127.0.0.1:0".parse().expect("parse address error"))
568 .expect("could not bind the middleware listener");
569 let middleware_addr = listener
570 .local_addr()
571 .expect("the middleware listener must expose its address");
572
573 let mut sender = StdTcpStream::connect(middleware_addr).expect("connect sender");
574 let mut wire = vec![
575 0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, 0x20, 0x00,
576 0x00, 0xD9, ];
578 wire.extend_from_slice(b"sentinel bytes must remain unread");
579 sender.write_all(&wire).expect("write oversized prelude");
580
581 let deadline = std::time::Instant::now() + Duration::from_secs(1);
582 let frontend = loop {
583 match listener.accept() {
584 Ok((stream, _)) => break stream,
585 Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
586 assert!(
587 std::time::Instant::now() < deadline,
588 "accept deadline elapsed"
589 );
590 thread::yield_now();
591 }
592 Err(error) => panic!("middleware accept failed: {error}"),
593 }
594 };
595
596 let mut metrics = SessionMetrics::new(None);
597 let mut expect = ExpectProxyProtocol::new(
598 TimeoutContainer::new(Duration::from_secs(10), Token(0)),
599 frontend,
600 Token(0),
601 Ulid::generate(),
602 );
603
604 assert_eq!(expect.readable(&mut metrics), SessionResult::Close);
605 assert_eq!(
606 metrics.bin, V2_FIXED_HEADER_LEN,
607 "the capacity check must happen before a second socket read"
608 );
609
610 drop(sender);
611 }
612
613 fn start_middleware(
615 upfront_middleware_conn_listener: TcpListener,
616 barrier: Arc<Barrier>,
617 ) -> Option<ProxyAddr> {
618 let session_stream;
619 barrier.wait();
620
621 loop {
623 if let Ok((stream, _addr)) = upfront_middleware_conn_listener.accept() {
624 session_stream = stream;
625 break;
626 }
627 }
628
629 let mut session_metrics = SessionMetrics::new(None);
630 let container_frontend_timeout = TimeoutContainer::new(Duration::from_secs(10), Token(0));
631 let mut expect_pp = ExpectProxyProtocol::new(
632 container_frontend_timeout,
633 session_stream,
634 Token(0),
635 Ulid::generate(),
636 );
637
638 let mut res = SessionResult::Continue;
639 while res == SessionResult::Continue {
640 res = expect_pp.readable(&mut session_metrics);
641 }
642
643 if res != SessionResult::Upgrade {
644 panic!("Should receive a complete proxy protocol header, res = {res:?}");
645 };
646
647 expect_pp.addresses
648 }
649
650 fn start_upfront_middleware(
652 next_middleware_addr: SocketAddr,
653 barrier: Arc<Barrier>,
654 command: Command,
655 ) -> JoinHandle<()> {
656 thread::spawn(move || {
657 let proxy_protocol = HeaderV2::new(command, header_src(), header_dst()).into_bytes();
658
659 barrier.wait();
660 match StdTcpStream::connect(next_middleware_addr) {
661 Ok(mut stream) => {
662 stream.write_all(&proxy_protocol).unwrap();
663 }
664 Err(e) => panic!("could not connect to the next middleware: {e}"),
665 };
666 })
667 }
668
669 struct UnusedSession;
673
674 impl crate::ProxySession for UnusedSession {
675 fn protocol(&self) -> Protocol {
676 unreachable!("the expect stage never asks the session")
677 }
678 fn ready(
679 &mut self,
680 _session: Rc<RefCell<dyn crate::ProxySession>>,
681 ) -> crate::SessionIsToBeClosed {
682 unreachable!("the expect stage never asks the session")
683 }
684 fn update_readiness(&mut self, _token: Token, _events: Ready) {
685 unreachable!("the expect stage never asks the session")
686 }
687 fn close(&mut self) {
688 unreachable!("the expect stage never asks the session")
689 }
690 fn timeout(&mut self, _t: Token) -> crate::SessionIsToBeClosed {
691 unreachable!("the expect stage never asks the session")
692 }
693 fn last_event(&self) -> std::time::Instant {
694 unreachable!("the expect stage never asks the session")
695 }
696 fn print_session(&self) {
697 unreachable!("the expect stage never asks the session")
698 }
699 fn frontend_token(&self) -> Token {
700 unreachable!("the expect stage never asks the session")
701 }
702 fn shutting_down(&mut self) -> crate::SessionIsToBeClosed {
703 unreachable!("the expect stage never asks the session")
704 }
705 }
706
707 struct UnusedProxy;
708
709 impl crate::L7Proxy for UnusedProxy {
710 fn kind(&self) -> sozu_command::proto::command::ListenerType {
711 unreachable!("the expect stage never asks the proxy")
712 }
713 fn register_socket(
714 &self,
715 _socket: &mut TcpStream,
716 _token: Token,
717 _interest: Interest,
718 ) -> Result<(), std::io::Error> {
719 unreachable!("the expect stage never asks the proxy")
720 }
721 fn add_session(&self, _session: Rc<RefCell<dyn crate::ProxySession>>) -> Token {
722 unreachable!("the expect stage never asks the proxy")
723 }
724 fn remove_session(&self, _token: Token) -> bool {
725 unreachable!("the expect stage never asks the proxy")
726 }
727 fn clusters(
728 &self,
729 ) -> &std::collections::HashMap<
730 sozu_command::state::ClusterId,
731 sozu_command::proto::command::Cluster,
732 > {
733 unreachable!("the expect stage never asks the proxy")
734 }
735 fn sessions(&self) -> Rc<RefCell<crate::server::SessionManager>> {
736 unreachable!("the expect stage never asks the proxy")
737 }
738 }
739
740 #[cfg(target_os = "linux")]
747 fn half_closed_expect(wire: &[u8]) -> (ExpectProxyProtocol<TcpStream>, StdTcpStream) {
748 use std::net::{Shutdown, TcpListener as StdTcpListener};
749
750 let listener = StdTcpListener::bind("127.0.0.1:0").expect("bind the expect listener");
751 let mut client = StdTcpStream::connect(
752 listener
753 .local_addr()
754 .expect("the expect listener must expose its address"),
755 )
756 .expect("connect the client");
757 let (accepted, _) = listener.accept().expect("accept the client");
758 accepted
759 .set_nonblocking(true)
760 .expect("accepted socket nonblocking");
761 client.write_all(wire).expect("write the request stream");
762 client
763 .shutdown(Shutdown::Write)
764 .expect("half-close the client write side");
765
766 let mut frontend = TcpStream::from_std(accepted);
767 let mut poll = Poll::new().expect("create poll");
768 poll.registry()
769 .register(&mut frontend, Token(0), Interest::READABLE)
770 .expect("register the accepted socket");
771 let mut events = Events::with_capacity(4);
772 poll.poll(&mut events, Some(Duration::from_secs(2)))
773 .expect("poll the accepted socket");
774 let ready = events
775 .iter()
776 .find(|event| event.token() == Token(0))
777 .map(Ready::from)
778 .expect("the queued request and FIN produce an epoll event");
779 assert!(
780 ready.is_readable() && ready.is_hup(),
781 "the queued bytes and FIN must coalesce into READABLE|HUP: {ready:?}"
782 );
783 poll.registry()
784 .deregister(&mut frontend)
785 .expect("deregister the accepted socket");
786
787 let mut expect = ExpectProxyProtocol::new(
788 TimeoutContainer::new(Duration::from_secs(10), Token(0)),
789 frontend,
790 Token(0),
791 Ulid::generate(),
792 );
793 expect.update_readiness(Token(0), ready);
794 (expect, client)
795 }
796
797 fn ready_once(expect: &mut ExpectProxyProtocol<TcpStream>) -> SessionResult {
798 let session: Rc<RefCell<dyn crate::ProxySession>> = Rc::new(RefCell::new(UnusedSession));
799 let proxy: Rc<RefCell<dyn crate::L7Proxy>> = Rc::new(RefCell::new(UnusedProxy));
800 let mut metrics = SessionMetrics::new(None);
801 expect.ready(session, proxy, &mut metrics)
802 }
803
804 #[cfg(target_os = "linux")]
814 #[test]
815 fn ready_parses_a_header_coalesced_with_the_client_half_close() {
816 let mut wire = HeaderV2::new(Command::Proxy, header_src(), header_dst()).into_bytes();
817 wire.extend_from_slice(b"GET / HTTP/1.1\r\n\r\n");
818 let (mut expect, _client) = half_closed_expect(&wire);
819
820 assert_eq!(ready_once(&mut expect), SessionResult::Upgrade);
821 assert_eq!(
822 expect.addresses.as_ref().and_then(ProxyAddr::source),
823 Some(header_src()),
824 "the coalesced header must be parsed before HUP is honoured"
825 );
826 assert!(
827 expect.frontend_readiness.event.is_hup(),
828 "HUP must survive the upgrade for the next stage to honour"
829 );
830 }
831
832 #[cfg(target_os = "linux")]
836 #[test]
837 fn ready_closes_a_half_close_without_a_complete_header() {
838 let (mut zero_bytes, _client) = half_closed_expect(&[]);
839 assert_eq!(ready_once(&mut zero_bytes), SessionResult::Close);
840
841 let header = HeaderV2::new(Command::Proxy, header_src(), header_dst()).into_bytes();
842 let (mut truncated, _client) = half_closed_expect(&header[..10]);
843 assert_eq!(ready_once(&mut truncated), SessionResult::Close);
844 assert_eq!(truncated.index, 10, "every readable byte must be consumed");
845 }
846}