Add simple HOSTID and port-number retention support

Cargo.toml

@@ -31,6 +31,7 @@ tokio-util = { version = "0.7.1", features = ["codec"] }
 tracing = "0.1.32"
 tracing-subscriber = "0.3.18"
 uuid = { version = "1.2.1", features = ["serde", "v4"] }
+socket2 = { version = "0.4.9" }
 
 [dev-dependencies]
 lazy_static = "1.4.0"

src/client.rs

@@ -1,5 +1,6 @@
 //! Client implementation for the `bore` service.
 
+use std::io::Read;
 use std::sync::Arc;
 
 use anyhow::{bail, Context, Result};
@@ -10,6 +11,7 @@ use uuid::Uuid;
 use crate::auth::Authenticator;
 use crate::shared::{
     proxy, ClientMessage, Delimited, ServerMessage, CONTROL_PORT, NETWORK_TIMEOUT,
+    BORE_KEEPINTERVAL, tcp_keepalive, parse_envvar_u64,
 };
 
 /// State structure for the client.
@@ -20,9 +22,12 @@ pub struct Client {
     /// Destination address of the server.
     to: String,
 
-    // Local host that is forwarded.
+    /// Local host that is forwarded.
     local_host: String,
 
+    /// Local host identity string.
+    host_id: String,
+
     /// Local port that is forwarded.
     local_port: u16,
 
@@ -33,24 +38,68 @@ pub struct Client {
     auth: Option<Authenticator>,
 }
 
+fn random_hostid(idlen: usize) -> String {
+    if idlen <= 4 {
+        if idlen == 0 { "".to_string() } else { std::iter::repeat_with(fastrand::alphanumeric).take(idlen).collect() }
+    } else {
+        let mut newid = String::with_capacity(idlen + 1);
+        let idlen = idlen - 4;
+        newid.push_str("TMP_");
+        let rid: String = std::iter::repeat_with(fastrand::alphanumeric).take(idlen).collect();
+        newid.push_str(&rid);
+        newid
+    }
+}
+
+fn read_hostid_fromfile(idlen: usize) -> String {
+    // Get the text file containing `BORE_HOSTID
+    let idfile: std::ffi::OsString = match std::env::var_os("BORE_IDFILE") {
+        Some(fpath) => fpath,
+        None => std::ffi::OsString::from("/tmp/bore_hostid.txt"),
+    };
+
+    let hfile = std::fs::OpenOptions::new().read(true)
+        .write(false).create(false).open(&idfile);
+    if hfile.is_err() {
+        return random_hostid(idlen);
+    }
+
+    let mut hfile = hfile.unwrap();
+    let mut idbuf = vec![0u8; idlen];
+    let rlen = hfile.read(&mut idbuf[..]).unwrap_or(0);
+    if rlen == 0 {
+        return random_hostid(idlen);
+    }
+
+    let idstr = String::from_utf8_lossy(&idbuf[..rlen]);
+    let hostid: &str = idstr.trim();
+    if hostid.is_empty() { random_hostid(idlen) } else { hostid.to_string() }
+}
+
 impl Client {
     /// Create a new client.
     pub async fn new(
         local_host: &str,
         local_port: u16,
+        id_str: &str,
         to: &str,
         port: u16,
         secret: Option<&str>,
     ) -> Result<Self> {
-        let mut stream = Delimited::new(connect_with_timeout(to, CONTROL_PORT).await?);
+        let kval = parse_envvar_u64(BORE_KEEPINTERVAL, 120);
+        let mut stream = Delimited::new(connect_with_timeout(to, CONTROL_PORT, kval).await?);
         let auth = secret.map(Authenticator::new);
         if let Some(auth) = &auth {
             auth.client_handshake(&mut stream).await?;
         }
 
-        stream.send(ClientMessage::Hello(port)).await?;
+        // Determine host ID for remote bore server
+        let hostid = if id_str.is_empty() { read_hostid_fromfile(16) } else { id_str.to_string() };
+        info!(hostid, "Using client IDString");
+
+        stream.send(ClientMessage::Hello(port, hostid.clone())).await?;
         let remote_port = match stream.recv_timeout().await? {
-            Some(ServerMessage::Hello(remote_port)) => remote_port,
+            Some(ServerMessage::Hello(remote_port, _)) => remote_port,
             Some(ServerMessage::Error(message)) => bail!("server error: {message}"),
             Some(ServerMessage::Challenge(_)) => {
                 bail!("server requires authentication, but no client secret was provided");
@@ -65,6 +114,7 @@ impl Client {
             conn: Some(stream),
             to: to.to_string(),
             local_host: local_host.to_string(),
+            host_id: hostid,
             local_port,
             remote_port,
             auth,
@@ -82,17 +132,18 @@ impl Client {
         let this = Arc::new(self);
         loop {
             match conn.recv().await? {
-                Some(ServerMessage::Hello(_)) => warn!("unexpected hello"),
+                Some(ServerMessage::Hello(_, _)) => warn!("unexpected hello"),
                 Some(ServerMessage::Challenge(_)) => warn!("unexpected challenge"),
                 Some(ServerMessage::Heartbeat) => (),
                 Some(ServerMessage::Connection(id)) => {
                     let this = Arc::clone(&this);
+                    let hostid: String = this.host_id.clone();
                     tokio::spawn(
                         async move {
-                            info!("new connection");
+                            info!(hostid, "new connection");
                             match this.handle_connection(id).await {
-                                Ok(_) => info!("connection exited"),
-                                Err(err) => warn!(%err, "connection exited with error"),
+                                Ok(_) => info!(hostid, "connection exited"),
+                                Err(err) => warn!(hostid, %err, "connection exited with error"),
                             }
                         }
                         .instrument(info_span!("proxy", %id)),
@@ -105,13 +156,14 @@ impl Client {
     }
 
     async fn handle_connection(&self, id: Uuid) -> Result<()> {
+        let kval = parse_envvar_u64("BORE_KEEPINTERVAL", 120);
         let mut remote_conn =
-            Delimited::new(connect_with_timeout(&self.to[..], CONTROL_PORT).await?);
+            Delimited::new(connect_with_timeout(&self.to[..], CONTROL_PORT, kval).await?);
         if let Some(auth) = &self.auth {
             auth.client_handshake(&mut remote_conn).await?;
         }
         remote_conn.send(ClientMessage::Accept(id)).await?;
-        let mut local_conn = connect_with_timeout(&self.local_host, self.local_port).await?;
+        let mut local_conn = connect_with_timeout(&self.local_host, self.local_port, kval).await?;
         let parts = remote_conn.into_parts();
         debug_assert!(parts.write_buf.is_empty(), "framed write buffer not empty");
         local_conn.write_all(&parts.read_buf).await?; // mostly of the cases, this will be empty
@@ -120,9 +172,9 @@ impl Client {
     }
 }
 
-async fn connect_with_timeout(to: &str, port: u16) -> Result<TcpStream> {
+async fn connect_with_timeout(to: &str, port: u16, keepival: u64) -> Result<TcpStream> {
     match timeout(NETWORK_TIMEOUT, TcpStream::connect((to, port))).await {
-        Ok(res) => res,
+        Ok(res) => if res.is_ok() { Ok(tcp_keepalive(res.unwrap(), 3, keepival)) } else { res },
         Err(err) => Err(err.into()),
     }
     .with_context(|| format!("could not connect to {to}:{port}"))

src/main.rs

@@ -23,6 +23,10 @@ enum Command {
         #[clap(short, long, value_name = "HOST", default_value = "localhost")]
         local_host: String,
 
+        /// The host ID-String for remote server
+        #[clap(short, long, env = "BORE_HOSTID", default_value = "")]
+        id_string: String,
+
         /// Address of the remote server to expose local ports to.
         #[clap(short, long, env = "BORE_SERVER")]
         to: String,
@@ -66,11 +70,12 @@ async fn run(command: Command) -> Result<()> {
         Command::Local {
             local_host,
             local_port,
+            id_string,
             to,
             port,
             secret,
         } => {
-            let client = Client::new(&local_host, local_port, &to, port, secret.as_deref()).await?;
+            let client = Client::new(&local_host, local_port, &id_string, &to, port, secret.as_deref()).await?;
             client.listen().await?;
         }
         Command::Server {