diff --git a/Cargo.lock b/Cargo.lock index 30c5250..a1e954d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1367,6 +1367,7 @@ dependencies = [ "rustls", "rusty_ulid 1.0.0", "serde", + "serde_json", "structopt", "tokio", "tracing", diff --git a/crates/ptth_quic/Cargo.toml b/crates/ptth_quic/Cargo.toml index 2175cc3..2e25ff4 100644 --- a/crates/ptth_quic/Cargo.toml +++ b/crates/ptth_quic/Cargo.toml @@ -21,6 +21,7 @@ rmp-serde = "1.1.1" rustls = "0.20.7" rusty_ulid = "1.0.0" serde = "1.0.151" +serde_json = "1.0.89" structopt = "0.3.26" tokio = { version = "1.23.0", features = ["full"] } tracing-subscriber = "0.3.16" diff --git a/crates/ptth_quic/src/executable_relay_server.rs b/crates/ptth_quic/src/executable_relay_server.rs index 0e2478f..4d1a898 100644 --- a/crates/ptth_quic/src/executable_relay_server.rs +++ b/crates/ptth_quic/src/executable_relay_server.rs @@ -35,7 +35,11 @@ pub async fn main (opt: Opt) -> anyhow::Result <()> tokio::fs::create_dir_all ("ptth_quic_output").await?; tokio::fs::write ("ptth_quic_output/quic_server.crt", &server_cert).await?; - let relay_state = Arc::new (RelayState::default ()); + let relay_state = RelayState::default (); + if let Err (e) = relay_state.reload_config ().await { + error! ("{:?}", e); + } + let relay_state = Arc::new (relay_state); let make_svc = { let relay_state = Arc::clone (&relay_state); @@ -55,6 +59,20 @@ pub async fn main (opt: Opt) -> anyhow::Result <()> let http_addr = SocketAddr::from (([0, 0, 0, 0], 4004)); let http_server = Server::bind (&http_addr); + let _task_reload_config = { + let relay_state = Arc::clone (&relay_state); + tokio::spawn (async move { + let mut interval = tokio::time::interval (std::time::Duration::from_secs (60)); + interval.set_missed_tick_behavior (tokio::time::MissedTickBehavior::Skip); + + loop { + interval.tick ().await; + + relay_state.reload_config ().await.ok (); + } + }) + }; + let task_quic_server = { let relay_state = Arc::clone (&relay_state); tokio::spawn (async move { @@ -172,12 +190,20 @@ struct RelayState { #[derive (Default)] struct Config { - ip_nicknames: HashMap , + ip_nicknames: BTreeMap <[u8; 4], String>, +} + +impl From for Config { + fn from (x: ConfigFile) -> Self { + Self { + ip_nicknames: x.ip_nicknames.into_iter ().collect (), + } + } } #[derive (Deserialize)] struct ConfigFile { - ip_nicknames: HashMap , + ip_nicknames: Vec <([u8; 4], String)>, } struct DirecState { @@ -223,7 +249,15 @@ struct P4State { } impl RelayState { - + async fn reload_config (&self) -> anyhow::Result <()> { + let s = tokio::fs::read_to_string ("config/ptth_quic_relay_server.json").await?; + let config: ConfigFile = serde_json::from_str (&s)?; + let config = Arc::new (Config::from (config)); + + self.config.store (config); + + Ok (()) + } } struct RequestP2ToP4 { @@ -300,11 +334,18 @@ async fn handle_quic_connection ( { let id = Ulid::generate (); + let config = relay_state.config.load (); + let remote_addr = conn.remote_address (); let ip_nickname = match remote_addr { SocketAddr::V4 (x) => { let ip = x.ip ().octets (); - "Unknown" + debug! ("{:?}", ip); + + match config.ip_nicknames.get (&ip) { + Some (nick) => nick.as_str (), + _ => "Unknown", + } }, _ => "Unknown, not IPv4", };