fix(#368): Update the restrictions before each request

instead of only once per conections, to avoid using a stale
restriction config when multiple request arrive on the same tcp stream.

Author: erebe

Cargo.lock

@@ -21,6 +21,7 @@ futures-util = { version = "0.3.30" }
 hickory-resolver = { version = "0.24.1", features = ["tokio", "dns-over-https-rustls", "dns-over-rustls", "native-certs"] }
 ppp = { version = "2.2.0", features = [] }
 async-channel = { version = "2.3.1", features = [] }
+arc-swap = { version = "1.7.1", features = [] }
 
 # For config file parsing
 regex = { version = "1.11.0", default-features = false, features = ["std", "perf"] }

Cargo.toml

@@ -1,40 +1,40 @@
 use super::types::RestrictionsRules;
 use crate::restrictions::config_reloader::RestrictionsRulesReloaderState::{Config, Static};
 use anyhow::Context;
+use arc_swap::ArcSwap;
 use log::trace;
 use notify::{EventKind, RecommendedWatcher, Watcher};
 use parking_lot::Mutex;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::thread;
 use std::time::Duration;
-use tokio::sync::futures::Notified;
-use tokio::sync::Notify;
 use tracing::{error, info, warn};
 
 struct ConfigReloaderState {
     fs_watcher: Mutex<RecommendedWatcher>,
     config_path: PathBuf,
-    should_reload_config: Notify,
 }
 
+#[derive(Clone)]
 enum RestrictionsRulesReloaderState {
-    Static(Notify),
+    Static,
     Config(Arc<ConfigReloaderState>),
 }
 
 impl RestrictionsRulesReloaderState {
     fn fs_watcher(&self) -> &Mutex<RecommendedWatcher> {
         match self {
-            Static(_) => unreachable!(),
+            Static => unreachable!(),
             Config(this) => &this.fs_watcher,
         }
     }
 }
 
+#[derive(Clone)]
 pub struct RestrictionsRulesReloader {
     state: RestrictionsRulesReloaderState,
-    restrictions: Arc<RestrictionsRules>,
+    restrictions: Arc<ArcSwap<RestrictionsRules>>,
 }
 
 impl RestrictionsRulesReloader {
@@ -44,37 +44,40 @@ impl RestrictionsRulesReloader {
             config_path
         } else {
             return Ok(Self {
-                state: Static(Notify::new()),
-                restrictions: Arc::new(restrictions_rules),
+                state: Static,
+                restrictions: Arc::new(ArcSwap::from_pointee(restrictions_rules)),
             });
         };
-
-        let this = Arc::new(ConfigReloaderState {
-            fs_watcher: Mutex::new(notify::recommended_watcher(|_| {})?),
-            should_reload_config: Notify::new(),
-            config_path,
-        });
+        let reloader = Self {
+            state: Config(Arc::new(ConfigReloaderState {
+                fs_watcher: Mutex::new(notify::recommended_watcher(|_| {})?),
+                config_path,
+            })),
+            restrictions: Arc::new(ArcSwap::from_pointee(restrictions_rules)),
+        };
 
         info!("Starting to watch restriction config file for changes to reload them");
         let mut watcher = notify::recommended_watcher({
-            let this = Config(this.clone());
+            let reloader = reloader.clone();
 
-            move |event: notify::Result<notify::Event>| Self::handle_config_fs_event(&this, event)
+            move |event: notify::Result<notify::Event>| Self::handle_config_fs_event(&reloader, event)
         })
         .with_context(|| "Cannot create restriction config watcher")?;
 
-        watcher.watch(&this.config_path, notify::RecursiveMode::NonRecursive)?;
-        *this.fs_watcher.lock() = watcher;
+        match &reloader.state {
+            Static => {}
+            Config(cfg) => {
+                watcher.watch(&cfg.config_path, notify::RecursiveMode::NonRecursive)?;
+                *cfg.fs_watcher.lock() = watcher
+            }
+        }
 
-        Ok(Self {
-            state: Config(this),
-            restrictions: Arc::new(restrictions_rules),
-        })
+        Ok(reloader)
     }
 
-    pub fn reload_restrictions_config(&mut self) {
+    pub fn reload_restrictions_config(&self) {
         let restrictions = match &self.state {
-            Static(_) => return,
+            Static => return,
             Config(st) => match RestrictionsRules::from_config_file(&st.config_path) {
                 Ok(restrictions) => {
                     info!("Restrictions config file has been reloaded");
@@ -87,21 +90,14 @@ impl RestrictionsRulesReloader {
             },
         };
 
-        self.restrictions = Arc::new(restrictions);
+        self.restrictions.store(Arc::new(restrictions));
     }
 
-    pub const fn restrictions_rules(&self) -> &Arc<RestrictionsRules> {
+    pub const fn restrictions_rules(&self) -> &Arc<ArcSwap<RestrictionsRules>> {
         &self.restrictions
     }
 
-    pub fn reload_notifier(&self) -> Notified {
-        match &self.state {
-            Static(st) => st.notified(),
-            Config(st) => st.should_reload_config.notified(),
-        }
-    }
-
-    fn try_rewatch_config(this: RestrictionsRulesReloaderState, path: PathBuf) {
+    fn try_rewatch_config(this: RestrictionsRulesReloader, path: PathBuf) {
         thread::spawn(move || {
             while !path.exists() {
                 warn!(
@@ -110,7 +106,7 @@ impl RestrictionsRulesReloader {
                 );
                 thread::sleep(Duration::from_secs(10));
             }
-            let mut watcher = this.fs_watcher().lock();
+            let mut watcher = this.state.fs_watcher().lock();
             let _ = watcher.unwatch(&path);
             let Ok(_) = watcher
                 .watch(&path, notify::RecursiveMode::NonRecursive)
@@ -123,23 +119,20 @@ impl RestrictionsRulesReloader {
             };
             drop(watcher);
 
-            // Generate a fake event to force-reload the certificate
+            // Generate a fake event to force-reload the config
             let event = notify::Event {
                 kind: EventKind::Create(notify::event::CreateKind::Any),
                 paths: vec![path],
                 attrs: Default::default(),
             };
 
-            match &this {
-                Static(_) => Self::handle_config_fs_event(&this, Ok(event)),
-                Config(_) => Self::handle_config_fs_event(&this, Ok(event)),
-            }
+            Self::handle_config_fs_event(&this, Ok(event))
         });
     }
 
-    fn handle_config_fs_event(this: &RestrictionsRulesReloaderState, event: notify::Result<notify::Event>) {
-        let this = match this {
-            Static(_) => return,
+    fn handle_config_fs_event(reloader: &RestrictionsRulesReloader, event: notify::Result<notify::Event>) {
+        let this = match &reloader.state {
+            Static => return,
             Config(st) => st,
         };
 
@@ -159,11 +152,11 @@ impl RestrictionsRulesReloader {
         if let Some(path) = event.paths.iter().find(|p| p.ends_with(&this.config_path)) {
             match event.kind {
                 EventKind::Create(_) | EventKind::Modify(_) => {
-                    this.should_reload_config.notify_one();
+                    reloader.reload_restrictions_config();
                 }
                 EventKind::Remove(_) => {
                     warn!("Restriction config file has been removed, trying to re-set a watch for it");
-                    Self::try_rewatch_config(Config(this.clone()), path.to_path_buf());
+                    Self::try_rewatch_config(reloader.clone(), path.to_path_buf());
                 }
                 EventKind::Access(_) | EventKind::Other | EventKind::Any => {
                     trace!("Ignoring event {:?}", event);

src/restrictions/config_reloader.rs

@@ -4,23 +4,23 @@ use http_body_util::Either;
 use std::fmt;
 use std::fmt::{Debug, Formatter};
 
-use bytes::Bytes;
-use http_body_util::combinators::BoxBody;
-use std::net::SocketAddr;
-use std::path::PathBuf;
-use std::pin::Pin;
-use std::sync::{Arc, LazyLock};
-use std::time::Duration;
-
 use crate::protocols;
 use crate::tunnel::{try_to_sock_addr, LocalProtocol, RemoteAddr};
+use arc_swap::ArcSwap;
+use bytes::Bytes;
+use http_body_util::combinators::BoxBody;
 use hyper::body::Incoming;
 use hyper::server::conn::{http1, http2};
 use hyper::service::service_fn;
 use hyper::{http, Request, Response, StatusCode, Version};
 use hyper_util::rt::{TokioExecutor, TokioTimer};
 use parking_lot::Mutex;
 use socket2::SockRef;
+use std::net::SocketAddr;
+use std::path::PathBuf;
+use std::pin::Pin;
+use std::sync::{Arc, LazyLock};
+use std::time::Duration;
 
 use crate::protocols::dns::DnsResolver;
 use crate::protocols::tls;
@@ -37,7 +37,6 @@ use crate::tunnel::server::utils::{
 use crate::tunnel::tls_reloader::TlsReloader;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio::net::TcpListener;
-use tokio::select;
 use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer};
 use tokio_rustls::TlsAcceptor;
 use tracing::{error, info, span, warn, Instrument, Level, Span};
@@ -285,29 +284,41 @@ impl WsServer {
 
         // setup upgrade request handler
         let mk_websocket_upgrade_fn = |server: WsServer,
-                                       restrictions: Arc<RestrictionsRules>,
+                                       restrictions: Arc<ArcSwap<RestrictionsRules>>,
                                        restrict_path: Option<String>,
                                        client_addr: SocketAddr| {
             move |req: Request<Incoming>| {
-                ws_server_upgrade(server.clone(), restrictions.clone(), restrict_path.clone(), client_addr, req)
-                    .map::<anyhow::Result<_>, _>(Ok)
-                    .instrument(mk_span())
+                ws_server_upgrade(
+                    server.clone(),
+                    restrictions.load().clone(),
+                    restrict_path.clone(),
+                    client_addr,
+                    req,
+                )
+                .map::<anyhow::Result<_>, _>(Ok)
+                .instrument(mk_span())
             }
         };
 
         let mk_http_upgrade_fn = |server: WsServer,
-                                  restrictions: Arc<RestrictionsRules>,
+                                  restrictions: Arc<ArcSwap<RestrictionsRules>>,
                                   restrict_path: Option<String>,
                                   client_addr: SocketAddr| {
             move |req: Request<Incoming>| {
-                http_server_upgrade(server.clone(), restrictions.clone(), restrict_path.clone(), client_addr, req)
-                    .map::<anyhow::Result<_>, _>(Ok)
-                    .instrument(mk_span())
+                http_server_upgrade(
+                    server.clone(),
+                    restrictions.load().clone(),
+                    restrict_path.clone(),
+                    client_addr,
+                    req,
+                )
+                .map::<anyhow::Result<_>, _>(Ok)
+                .instrument(mk_span())
             }
         };
 
         let mk_auto_upgrade_fn = |server: WsServer,
-                                  restrictions: Arc<RestrictionsRules>,
+                                  restrictions: Arc<ArcSwap<RestrictionsRules>>,
                                   restrict_path: Option<String>,
                                   client_addr: SocketAddr| {
             move |req: Request<Incoming>| {
@@ -316,13 +327,13 @@ impl WsServer {
                 let restrict_path = restrict_path.clone();
                 async move {
                     if fastwebsockets::upgrade::is_upgrade_request(&req) {
-                        ws_server_upgrade(server.clone(), restrictions.clone(), restrict_path, client_addr, req)
+                        ws_server_upgrade(server.clone(), restrictions.load().clone(), restrict_path, client_addr, req)
                             .map::<anyhow::Result<_>, _>(Ok)
                             .await
                     } else if req.version() == Version::HTTP_2 {
                         http_server_upgrade(
                             server.clone(),
-                            restrictions.clone(),
+                            restrictions.load().clone(),
                             restrict_path.clone(),
                             client_addr,
                             req,
@@ -357,25 +368,11 @@ impl WsServer {
         };
 
         // Bind server and run forever to serve incoming connections.
-        let mut restrictions = RestrictionsRulesReloader::new(restrictions, self.config.restriction_config.clone())?;
-        let mut await_config_reload = Box::pin(restrictions.reload_notifier());
+        let restrictions = RestrictionsRulesReloader::new(restrictions, self.config.restriction_config.clone())?;
         let listener = TcpListener::bind(&self.config.bind).await?;
 
         loop {
-            let cnx = select! {
-                biased;
-
-                _ = &mut await_config_reload => {
-                    drop(await_config_reload);
-                    restrictions.reload_restrictions_config();
-                    await_config_reload = Box::pin(restrictions.reload_notifier());
-                    continue;
-                },
-
-                cnx = listener.accept() => { cnx }
-            };
-
-            let (stream, peer_addr) = match cnx {
+            let (stream, peer_addr) = match listener.accept().await {
                 Ok(ret) => ret,
                 Err(err) => {
                     warn!("Error while accepting connection {:?}", err);
@@ -423,7 +420,7 @@ impl WsServer {
                                 }
 
                                 let http_upgrade_fn =
-                                    mk_http_upgrade_fn(server, restrictions.clone(), restrict_path, peer_addr);
+                                    mk_http_upgrade_fn(server, restrictions, restrict_path, peer_addr);
                                 let con_fut = conn_builder.serve_connection(tls_stream, service_fn(http_upgrade_fn));
                                 if let Err(e) = con_fut.await {
                                     error!("Error while upgrading cnx to http: {:?}", e);
@@ -432,7 +429,7 @@ impl WsServer {
                             // websocket
                             _ => {
                                 let websocket_upgrade_fn =
-                                    mk_websocket_upgrade_fn(server, restrictions.clone(), restrict_path, peer_addr);
+                                    mk_websocket_upgrade_fn(server, restrictions, restrict_path, peer_addr);
                                 let conn_fut = http1::Builder::new()
                                     .timer(TokioTimer::new())
                                     // https://github.com/erebe/wstunnel/issues/358
@@ -460,7 +457,7 @@ impl WsServer {
                             conn_fut.http2().keep_alive_interval(ping);
                         }
 
-                        let websocket_upgrade_fn = mk_auto_upgrade_fn(server, restrictions.clone(), None, peer_addr);
+                        let websocket_upgrade_fn = mk_auto_upgrade_fn(server, restrictions, None, peer_addr);
                         let upgradable =
                             conn_fut.serve_connection_with_upgrades(stream, service_fn(websocket_upgrade_fn));