1use crate::metrics::ReverseMetrics;
2use crate::{
3 redact_auth, relay_bidirectional_with_timeout, server_auth_handshake, ControlState,
4 ProtocolError,
5};
6use std::net::{IpAddr, SocketAddr};
7use std::sync::atomic::{AtomicU32, Ordering};
8use std::sync::Arc;
9use std::time::Instant;
10use tokio::net::{TcpListener, TcpStream};
11use tokio::sync::mpsc;
12use tokio::task::JoinSet;
13use tokio_util::sync::CancellationToken;
14use tracing::{debug, error, info, warn};
15
16#[derive(Debug, Clone)]
21pub struct ReverseServerConfig {
22 pub control_bind: SocketAddr,
24 pub external_bind: Option<SocketAddr>,
26 pub auth_username: Option<String>,
28 pub auth_password: Option<String>,
30 pub max_control_connections: u32,
32 pub read_timeout_ms: u64,
34 pub allow_bind: Option<Vec<SocketAddr>>,
38 pub max_listeners_per_client: u32,
42 pub max_streams_per_listener: u32,
44 pub max_pending_external: u32,
47}
48
49impl Default for ReverseServerConfig {
50 fn default() -> Self {
51 Self {
52 control_bind: "127.0.0.1:0".parse().unwrap(),
53 external_bind: None,
54 auth_username: None,
55 auth_password: None,
56 max_control_connections: 256,
57 read_timeout_ms: 300_000,
58 allow_bind: None,
59 max_listeners_per_client: 1,
60 max_streams_per_listener: 1024,
61 max_pending_external: 1024,
62 }
63 }
64}
65
66impl ReverseServerConfig {
67 pub fn is_bind_allowed(&self, addr: SocketAddr) -> bool {
71 match &self.allow_bind {
72 None => true,
73 Some(list) if list.is_empty() => true,
74 Some(list) => list.iter().any(|allowed| same_bind(allowed, &addr)),
75 }
76 }
77
78 pub fn is_loopback(addr: SocketAddr) -> bool {
80 match addr.ip() {
81 IpAddr::V4(v4) => v4.is_loopback(),
82 IpAddr::V6(v6) => v6.is_loopback(),
83 }
84 }
85
86 pub fn validate(&self) -> Result<(), ProtocolError> {
94 if let Some(external) = self.external_bind {
95 if !Self::is_loopback(external) {
100 let has_auth = self.auth_username.as_deref().is_some_and(|s| !s.is_empty())
101 && self.auth_password.as_deref().is_some_and(|s| !s.is_empty());
102 let has_allowlist = matches!(&self.allow_bind, Some(list) if !list.is_empty());
103 if !has_auth {
104 return Err(ProtocolError::ConfigInvalid(format!(
105 "reverse server external_bind={external} is non-loopback but no \
106 authentication is configured; set auth_username/auth_password or \
107 bind to loopback"
108 )));
109 }
110 if !has_allowlist {
111 return Err(ProtocolError::ConfigInvalid(format!(
112 "reverse server external_bind={external} is non-loopback but \
113 allow_bind is empty; configure an explicit allowlist"
114 )));
115 }
116 }
117 }
118 Ok(())
119 }
120}
121
122fn same_bind(a: &SocketAddr, b: &SocketAddr) -> bool {
123 a.port() == b.port()
124 && match (a.ip(), b.ip()) {
125 (IpAddr::V4(a4), IpAddr::V4(b4)) => a4 == b4,
126 (IpAddr::V6(a6), IpAddr::V6(b6)) => a6 == b6,
127 _ => false,
128 }
129}
130
131#[derive(Debug, Default)]
133pub struct ReverseServerState {
134 pub active_control: AtomicU32,
136 pub active_streams: AtomicU32,
138 pub pending_external: AtomicU32,
140 pub denied_bind: AtomicU32,
142 pub dropped_stream_limit: AtomicU32,
144 pub dropped_pending_limit: AtomicU32,
146}
147
148impl ReverseServerState {
149 pub fn snapshot(&self) -> ReverseServerStateSnapshot {
151 ReverseServerStateSnapshot {
152 active_control: self.active_control.load(Ordering::Relaxed),
153 active_streams: self.active_streams.load(Ordering::Relaxed),
154 pending_external: self.pending_external.load(Ordering::Relaxed),
155 denied_bind: self.denied_bind.load(Ordering::Relaxed),
156 dropped_stream_limit: self.dropped_stream_limit.load(Ordering::Relaxed),
157 dropped_pending_limit: self.dropped_pending_limit.load(Ordering::Relaxed),
158 }
159 }
160}
161
162#[derive(Debug, Clone, serde::Serialize)]
164pub struct ReverseServerStateSnapshot {
165 pub active_control: u32,
166 pub active_streams: u32,
167 pub pending_external: u32,
168 pub denied_bind: u32,
169 pub dropped_stream_limit: u32,
170 pub dropped_pending_limit: u32,
171}
172
173pub struct ReverseServer {
179 config: ReverseServerConfig,
180 cancel: CancellationToken,
181 metrics: Option<Arc<ReverseMetrics>>,
182 state: Arc<ReverseServerState>,
183}
184
185impl ReverseServer {
186 pub fn new(config: ReverseServerConfig) -> Self {
187 Self {
188 config,
189 cancel: CancellationToken::new(),
190 metrics: None,
191 state: Arc::new(ReverseServerState::default()),
192 }
193 }
194
195 pub fn set_metrics(&mut self, metrics: Arc<ReverseMetrics>) {
197 self.metrics = Some(metrics);
198 }
199
200 pub fn state_handle(&self) -> Arc<ReverseServerState> {
202 self.state.clone()
203 }
204
205 pub fn cancel_token(&self) -> CancellationToken {
207 self.cancel.clone()
208 }
209
210 async fn bind_external_listener(
213 config: &ReverseServerConfig,
214 state: &ReverseServerState,
215 ) -> Result<Option<TcpListener>, ProtocolError> {
216 let external_bind = match config.external_bind {
217 Some(addr) => addr,
218 None => return Ok(None),
219 };
220 if !config.is_bind_allowed(external_bind) {
221 state.denied_bind.fetch_add(1, Ordering::Relaxed);
222 return Err(ProtocolError::BindDenied(external_bind));
223 }
224 let listener = TcpListener::bind(external_bind).await?;
225 let addr = listener.local_addr()?;
226 info!(addr = %addr, "reverse server listening for external clients");
227 Ok(Some(listener))
228 }
229
230 pub async fn run(self) -> Result<(), ProtocolError> {
232 self.config.validate()?;
236
237 if self.config.auth_username.is_none() || self.config.auth_password.is_none() {
238 warn!(
239 control_bind = %self.config.control_bind,
240 "reverse server control channel has no authentication configured"
241 );
242 }
243
244 if let Some(external_bind) = self.config.external_bind {
246 if !self.config.is_bind_allowed(external_bind) {
247 self.state.denied_bind.fetch_add(1, Ordering::Relaxed);
248 return Err(ProtocolError::BindDenied(external_bind));
249 }
250 }
251
252 let control_listener = TcpListener::bind(&self.config.control_bind).await?;
253 let control_addr = control_listener.local_addr()?;
254 info!(addr = %control_addr, "reverse server listening for control connections");
255
256 let external_listener = Self::bind_external_listener(&self.config, &self.state).await?;
257
258 let config = Arc::new(self.config);
259 let cancel = self.cancel.clone();
260 let state = self.state.clone();
261 let metrics = self.metrics.clone();
262
263 let (control_tx, control_rx) = mpsc::unbounded_channel::<ControlStream>();
265
266 let config_clone = config.clone();
268 let cancel_clone = cancel.clone();
269 let control_tx_clone = control_tx.clone();
270 let metrics_clone = metrics.clone();
271 let state_clone = state.clone();
272 let control_task = tokio::spawn(async move {
273 Self::accept_control_connections(
274 control_listener,
275 config_clone,
276 cancel_clone,
277 control_tx_clone,
278 metrics_clone,
279 state_clone,
280 )
281 .await;
282 });
283
284 let external_task = if let Some(external_listener) = external_listener {
286 let config_clone = config.clone();
287 let cancel_clone = cancel.clone();
288 let metrics_clone = metrics.clone();
289 let state_clone = state.clone();
290 Some(tokio::spawn(async move {
291 Self::accept_external_clients(
292 external_listener,
293 config_clone,
294 cancel_clone,
295 control_rx,
296 metrics_clone,
297 state_clone,
298 )
299 .await;
300 }))
301 } else {
302 let state_clone = state.clone();
307 let metrics_clone = metrics.clone();
308 let cancel_clone = cancel.clone();
309 Some(tokio::spawn(async move {
310 let mut control_rx = control_rx;
311 loop {
312 tokio::select! {
313 Some(ctrl) = control_rx.recv() => {
314 debug!(
315 control_peer = %ctrl.peer_addr,
316 "dropping control connection: no external listener"
317 );
318 drop(ctrl.stream);
319 state_clone.active_control.fetch_sub(1, Ordering::Relaxed);
320 if let Some(m) = metrics_clone.as_deref() {
321 m.record_control_closed();
322 }
323 }
324 _ = cancel_clone.cancelled() => break,
325 }
326 }
327 }))
328 };
329
330 cancel.cancelled().await;
332 let drain_start = Instant::now();
333 info!("reverse server shutting down, draining active streams");
334
335 let _ = control_task.await;
338 if let Some(task) = external_task {
339 let _ = task.await;
340 }
341 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
342 let drain_ms = drain_start.elapsed().as_millis() as u64;
343 if let Some(ref m) = metrics {
344 m.record_drain(drain_ms);
345 }
346 info!(drain_ms, "reverse server drain complete");
347 Ok(())
348 }
349
350 async fn accept_control_connections(
352 listener: TcpListener,
353 config: Arc<ReverseServerConfig>,
354 cancel: CancellationToken,
355 control_tx: mpsc::UnboundedSender<ControlStream>,
356 metrics: Option<Arc<ReverseMetrics>>,
357 state: Arc<ReverseServerState>,
358 ) {
359 loop {
360 tokio::select! {
361 result = listener.accept() => {
362 match result {
363 Ok((stream, peer_addr)) => {
364 let prev = state.active_control.fetch_add(1, Ordering::AcqRel);
367 if prev >= config.max_control_connections {
368 state.active_control.fetch_sub(1, Ordering::Relaxed);
369 warn!(
370 peer = %peer_addr,
371 max = config.max_control_connections,
372 "rejecting control connection: max reached"
373 );
374 if let Some(ref m) = metrics {
375 m.record_control_rejected(peer_addr, "max_control_connections");
376 }
377 drop(stream);
378 continue;
379 }
380
381 let config = config.clone();
382 let control_tx = control_tx.clone();
383 let metrics = metrics.clone();
384 let state = state.clone();
385 tokio::spawn(async move {
386 if let Err(e) = Self::handle_control_connection(
387 stream,
388 peer_addr,
389 config,
390 control_tx,
391 metrics.as_deref(),
392 state.clone(),
393 ).await {
394 state.active_control.fetch_sub(1, Ordering::Relaxed);
395 debug!(peer = %peer_addr, error = %e, "control connection handler error");
396 }
397 });
398 }
399 Err(e) => {
400 error!(error = %e, "failed to accept control connection");
401 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
402 }
403 }
404 }
405 _ = cancel.cancelled() => {
406 break;
407 }
408 }
409 }
410 }
411
412 async fn handle_control_connection(
414 mut stream: TcpStream,
415 peer_addr: SocketAddr,
416 config: Arc<ReverseServerConfig>,
417 control_tx: mpsc::UnboundedSender<ControlStream>,
418 metrics: Option<&ReverseMetrics>,
419 state: Arc<ReverseServerState>,
420 ) -> Result<(), ProtocolError> {
421 info!(peer = %peer_addr, state = ?ControlState::Connecting, "new control connection");
422
423 let redacted = if config.auth_username.is_some() && config.auth_password.is_some() {
425 let authenticating_start = Instant::now();
426 let result = server_auth_handshake(
427 &mut stream,
428 config.auth_username.as_deref(),
429 config.auth_password.as_deref(),
430 )
431 .await;
432 let elapsed = authenticating_start.elapsed().as_millis() as u64;
433
434 match result {
435 Ok(redacted) => {
436 info!(
437 peer = %peer_addr,
438 auth = %redacted,
439 duration_ms = elapsed,
440 state = ?ControlState::Authenticating,
441 "control connection authenticated"
442 );
443 if let Some(m) = metrics {
444 m.record_control_accepted(peer_addr);
445 m.record_state_duration(ControlState::Authenticating, elapsed);
446 }
447 Some(redacted)
448 }
449 Err(e) => {
450 warn!(
451 peer = %peer_addr,
452 error = %e,
453 duration_ms = elapsed,
454 state = ?ControlState::Authenticating,
455 "control connection auth failed"
456 );
457 if let Some(m) = metrics {
458 m.record_auth_failure(peer_addr, &e.to_string());
459 }
460 return Err(e);
461 }
462 }
463 } else {
464 crate::write_handshake_accept(&mut stream).await?;
466 info!(
467 peer = %peer_addr,
468 state = ?ControlState::Authenticating,
469 "control connection accepted (no auth)"
470 );
471 if let Some(m) = metrics {
472 m.record_control_accepted(peer_addr);
473 }
474 None
475 };
476
477 let ctrl = ControlStream {
478 stream,
479 peer_addr,
480 redacted_auth: redacted,
481 };
482 if control_tx.send(ctrl).is_err() {
483 state.active_control.fetch_sub(1, Ordering::Relaxed);
484 if let Some(m) = metrics {
485 m.record_control_closed();
486 }
487 warn!(peer = %peer_addr, "control channel closed, cannot add to pool");
488 }
489
490 Ok(())
491 }
492
493 async fn accept_external_clients(
495 listener: TcpListener,
496 config: Arc<ReverseServerConfig>,
497 cancel: CancellationToken,
498 mut control_rx: mpsc::UnboundedReceiver<ControlStream>,
499 metrics: Option<Arc<ReverseMetrics>>,
500 state: Arc<ReverseServerState>,
501 ) {
502 let mut relay_tasks = JoinSet::new();
503
504 loop {
505 tokio::select! {
506 result = listener.accept() => {
507 match result {
508 Ok((external_stream, peer_addr)) => {
509 match state.active_streams.fetch_update(
510 Ordering::AcqRel,
511 Ordering::Acquire,
512 |current| {
513 (current < config.max_streams_per_listener)
514 .then_some(current + 1)
515 },
516 ) {
517 Ok(_) => {}
518 Err(current) => {
519 warn!(
520 peer = %peer_addr,
521 active = current,
522 max = config.max_streams_per_listener,
523 "dropping external client: max_streams_per_listener reached"
524 );
525 state.dropped_stream_limit.fetch_add(1, Ordering::Relaxed);
526 drop(external_stream);
527 continue;
528 }
529 }
530
531 match state.pending_external.fetch_update(
532 Ordering::AcqRel,
533 Ordering::Acquire,
534 |current| {
535 (current < config.max_pending_external)
536 .then_some(current + 1)
537 },
538 ) {
539 Ok(_) => {}
540 Err(current) => {
541 state.active_streams.fetch_sub(1, Ordering::Release);
542 warn!(
543 peer = %peer_addr,
544 pending = current,
545 max = config.max_pending_external,
546 "dropping external client: max_pending_external reached"
547 );
548 state.dropped_pending_limit.fetch_add(1, Ordering::Relaxed);
549 drop(external_stream);
550 continue;
551 }
552 }
553 let control = tokio::select! {
555 control = control_rx.recv() => control,
556 _ = cancel.cancelled() => {
557 state.pending_external.fetch_sub(1, Ordering::Release);
558 state.active_streams.fetch_sub(1, Ordering::Release);
559 drop(external_stream);
560 break;
561 }
562 };
563 match control {
564 Some(control) => {
565 state.pending_external.fetch_sub(1, Ordering::Release);
566 let metrics = metrics.clone();
567 let state = state.clone();
568 let idle_timeout = (config.read_timeout_ms > 0).then(|| {
569 std::time::Duration::from_millis(config.read_timeout_ms)
570 });
571 state.active_control.fetch_sub(1, Ordering::Relaxed);
572 relay_tasks.spawn(async move {
573 info!(
574 peer = %peer_addr,
575 control_peer = %control.peer_addr,
576 "relaying external client through control connection"
577 );
578 if let Some(m) = metrics.as_deref() {
579 m.record_stream_opened();
580 m.record_state_duration(ControlState::Ready, 0);
581 }
582 let relay_result = relay_bidirectional_with_timeout(
583 external_stream,
584 control.stream,
585 idle_timeout,
586 )
587 .await;
588 match relay_result {
589 Ok(()) => {
590 debug!(peer = %peer_addr, "relay finished cleanly");
591 }
592 Err(e) => {
593 debug!(peer = %peer_addr, error = %e, "relay ended");
594 }
595 }
596 if let Some(m) = metrics.as_deref() {
597 m.record_stream_closed(0);
598 m.record_control_closed();
599 }
600 state.active_streams.fetch_sub(1, Ordering::Release);
601 debug!(peer = %peer_addr, "relay finished");
602 });
603 }
604 None => {
605 state.pending_external.fetch_sub(1, Ordering::Release);
606 state.active_streams.fetch_sub(1, Ordering::Release);
607 warn!(peer = %peer_addr, "no control connections available, rejecting external client");
608 drop(external_stream);
609 }
610 }
611 }
612 Err(e) => {
613 error!(error = %e, "failed to accept external client");
614 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
615 }
616 }
617 }
618 _ = cancel.cancelled() => {
619 break;
620 }
621 }
622 }
623
624 relay_tasks.abort_all();
625 while relay_tasks.join_next().await.is_some() {}
626 }
627
628 pub fn shutdown(&self) {
630 self.cancel.cancel();
631 }
632}
633
634pub struct ControlStream {
637 pub stream: TcpStream,
638 pub peer_addr: SocketAddr,
639 pub redacted_auth: Option<String>,
640}
641
642pub fn format_auth_redacted(auth: &str) -> String {
644 redact_auth(auth)
645}
646
647#[cfg(test)]
648mod tests {
649 use super::*;
650
651 #[test]
652 fn is_bind_allowed_with_none() {
653 let cfg = ReverseServerConfig {
654 allow_bind: None,
655 ..Default::default()
656 };
657 assert!(cfg.is_bind_allowed("127.0.0.1:8080".parse().unwrap()));
658 }
659
660 #[test]
661 fn is_bind_allowed_with_empty() {
662 let cfg = ReverseServerConfig {
663 allow_bind: Some(vec![]),
664 ..Default::default()
665 };
666 assert!(cfg.is_bind_allowed("127.0.0.1:8080".parse().unwrap()));
667 }
668
669 #[test]
670 fn is_bind_allowed_match() {
671 let cfg = ReverseServerConfig {
672 allow_bind: Some(vec!["127.0.0.1:8080".parse().unwrap()]),
673 ..Default::default()
674 };
675 assert!(cfg.is_bind_allowed("127.0.0.1:8080".parse().unwrap()));
676 }
677
678 #[test]
679 fn is_bind_allowed_mismatch() {
680 let cfg = ReverseServerConfig {
681 allow_bind: Some(vec!["127.0.0.1:8080".parse().unwrap()]),
682 ..Default::default()
683 };
684 assert!(!cfg.is_bind_allowed("0.0.0.0:8080".parse().unwrap()));
685 assert!(!cfg.is_bind_allowed("127.0.0.1:9090".parse().unwrap()));
686 }
687
688 #[test]
689 fn state_snapshot_round_trip() {
690 let s = ReverseServerState::default();
691 s.active_control.fetch_add(3, Ordering::Relaxed);
692 s.active_streams.fetch_add(2, Ordering::Relaxed);
693 s.pending_external.fetch_add(1, Ordering::Relaxed);
694 s.denied_bind.fetch_add(1, Ordering::Relaxed);
695 s.dropped_stream_limit.fetch_add(4, Ordering::Relaxed);
696 s.dropped_pending_limit.fetch_add(5, Ordering::Relaxed);
697 let snap = s.snapshot();
698 assert_eq!(snap.active_control, 3);
699 assert_eq!(snap.active_streams, 2);
700 assert_eq!(snap.pending_external, 1);
701 assert_eq!(snap.denied_bind, 1);
702 assert_eq!(snap.dropped_stream_limit, 4);
703 assert_eq!(snap.dropped_pending_limit, 5);
704 }
705
706 #[test]
707 fn format_auth_redacted_basic() {
708 assert_eq!(format_auth_redacted("user:pass"), "user:****");
709 }
710
711 #[test]
712 fn same_bind_v4() {
713 let a: SocketAddr = "127.0.0.1:8080".parse().unwrap();
714 let b: SocketAddr = "127.0.0.1:8080".parse().unwrap();
715 assert!(same_bind(&a, &b));
716 }
717
718 #[test]
719 fn same_bind_different_port() {
720 let a: SocketAddr = "127.0.0.1:8080".parse().unwrap();
721 let b: SocketAddr = "127.0.0.1:9090".parse().unwrap();
722 assert!(!same_bind(&a, &b));
723 }
724
725 #[test]
726 fn validate_loopback_ok() {
727 let cfg = ReverseServerConfig {
728 control_bind: "127.0.0.1:0".parse().unwrap(),
729 external_bind: Some("127.0.0.1:0".parse().unwrap()),
730 ..Default::default()
731 };
732 assert!(cfg.validate().is_ok());
733 }
734
735 #[test]
736 fn validate_no_external_bind_ok() {
737 let cfg = ReverseServerConfig {
738 control_bind: "127.0.0.1:0".parse().unwrap(),
739 external_bind: None,
740 ..Default::default()
741 };
742 assert!(cfg.validate().is_ok());
743 }
744
745 #[test]
746 fn validate_non_loopback_without_auth_rejected() {
747 let cfg = ReverseServerConfig {
748 control_bind: "127.0.0.1:0".parse().unwrap(),
749 external_bind: Some("0.0.0.0:9000".parse().unwrap()),
750 auth_username: None,
751 auth_password: None,
752 ..Default::default()
753 };
754 let err = cfg.validate().unwrap_err();
755 assert!(
756 matches!(err, ProtocolError::ConfigInvalid(_)),
757 "got: {err:?}"
758 );
759 }
760
761 #[test]
762 fn validate_non_loopback_with_auth_but_no_allowlist_rejected() {
763 let cfg = ReverseServerConfig {
764 control_bind: "127.0.0.1:0".parse().unwrap(),
765 external_bind: Some("0.0.0.0:9000".parse().unwrap()),
766 auth_username: Some("user".to_string()),
767 auth_password: Some("pass".to_string()),
768 allow_bind: None,
769 ..Default::default()
770 };
771 let err = cfg.validate().unwrap_err();
772 assert!(
773 matches!(err, ProtocolError::ConfigInvalid(_)),
774 "got: {err:?}"
775 );
776 }
777
778 #[test]
779 fn validate_non_loopback_with_auth_and_allowlist_ok() {
780 let cfg = ReverseServerConfig {
781 control_bind: "127.0.0.1:0".parse().unwrap(),
782 external_bind: Some("0.0.0.0:9000".parse().unwrap()),
783 auth_username: Some("user".to_string()),
784 auth_password: Some("pass".to_string()),
785 allow_bind: Some(vec!["0.0.0.0:9000".parse().unwrap()]),
786 ..Default::default()
787 };
788 assert!(cfg.validate().is_ok());
789 }
790
791 #[test]
792 fn validate_ipv6_loopback_ok() {
793 let cfg = ReverseServerConfig {
794 control_bind: "127.0.0.1:0".parse().unwrap(),
795 external_bind: Some("[::1]:9000".parse().unwrap()),
796 ..Default::default()
797 };
798 assert!(cfg.validate().is_ok());
799 }
800
801 #[test]
802 fn validate_ipv6_non_loopback_without_auth_rejected() {
803 let cfg = ReverseServerConfig {
804 control_bind: "127.0.0.1:0".parse().unwrap(),
805 external_bind: Some("[2001:db8::1]:9000".parse().unwrap()),
806 ..Default::default()
807 };
808 let err = cfg.validate().unwrap_err();
809 assert!(
810 matches!(err, ProtocolError::ConfigInvalid(_)),
811 "got: {err:?}"
812 );
813 }
814}