Harden release readiness gates

src/bin/dosh-server.rs

@@ -651,10 +651,10 @@ impl ServerState {
     /// Remove a client from whichever session holds it, keeping the index in
     /// sync. Returns true when a client was actually removed.
     fn remove_client_everywhere(&mut self, client_id: &[u8; 16]) -> bool {
-        if let Some(session_name) = self.client_index.remove(client_id) {
-            if let Some(session) = self.sessions.get_mut(&session_name) {
-                return session.clients.remove(client_id).is_some();
-            }
+        if let Some(session_name) = self.client_index.remove(client_id)
+            && let Some(session) = self.sessions.get_mut(&session_name)
+        {
+            return session.clients.remove(client_id).is_some();
         }
         false
     }
@@ -987,9 +987,9 @@ async fn handle_native_user_auth(
     let body = protocol::decrypt_body(packet, &pending.session_key, CLIENT_TO_SERVER)?;
     let req: NativeUserAuthBody = protocol::from_body(&body)?;
     if pending.client.requested_mode == "doctor" {
-        let check = {
+        let check_result = {
             let locked = state.lock().expect("server state poisoned");
-            if let Err(err) = verify_native_user_auth_from_config(
+            verify_native_user_auth_from_config(
                 &locked.config,
                 &pending.client,
                 &pending.server,
@@ -997,16 +997,7 @@ async fn handle_native_user_auth(
                 Some(peer.ip()),
             )
             .context("verify native user auth")
-            {
-                return send_reject_to_client(
-                    socket,
-                    peer,
-                    packet.header.conn_id,
-                    &format!("{err:#}"),
-                )
-                .await;
-            }
-            NativeAuthCheckOkBody {
+            .map(|_| NativeAuthCheckOkBody {
                 requested_user: pending.client.requested_user.clone(),
                 native_auth_enabled: locked.config.native_auth,
                 allow_tcp_forwarding: locked.config.allow_tcp_forwarding,
@@ -1014,6 +1005,18 @@ async fn handle_native_user_auth(
                 allow_agent_forwarding: locked.config.allow_agent_forwarding,
                 policy_flags: Vec::new(),
                 server_version: env!("CARGO_PKG_VERSION").to_string(),
+            })
+        };
+        let check = match check_result {
+            Ok(check) => check,
+            Err(err) => {
+                return send_reject_to_client(
+                    socket,
+                    peer,
+                    packet.header.conn_id,
+                    &format!("{err:#}"),
+                )
+                .await;
             }
         };
         let body = protocol::to_body(&check)?;
@@ -1065,17 +1068,19 @@ async fn handle_native_user_auth(
         });
     }
 
-    let attached: Result<(
-        [u8; 16],
-        [u8; 16],
-        Vec<u8>,
-        [u8; 32],
-        String,
-        String,
-        u64,
-        Vec<u8>,
-        Vec<ForwardingRequest>,
-    )> = {
+    struct NativeAttachResult {
+        client_id: [u8; 16],
+        key_id: [u8; 16],
+        attach_ticket: Vec<u8>,
+        attach_ticket_psk: [u8; 32],
+        session_name: String,
+        mode: String,
+        output_seq: u64,
+        snapshot: Vec<u8>,
+        requested_forwardings: Vec<ForwardingRequest>,
+    }
+
+    let attached: Result<NativeAttachResult> = {
         let mut locked = state.lock().expect("server state poisoned");
         verify_native_user_auth_from_config(
             &locked.config,
@@ -1150,7 +1155,7 @@ async fn handle_native_user_auth(
                     previous_session_key: None,
                 },
             );
-            Ok((
+            Ok(NativeAttachResult {
                 client_id,
                 key_id,
                 attach_ticket,
@@ -1159,22 +1164,12 @@ async fn handle_native_user_auth(
                 mode,
                 output_seq,
                 snapshot,
-                req.auth.requested_forwardings.clone(),
-            ))
+                requested_forwardings: req.auth.requested_forwardings.clone(),
+            })
         })
     };
 
-    let (
-        client_id,
-        key_id,
-        attach_ticket,
-        attach_ticket_psk,
-        session_name,
-        mode,
-        output_seq,
-        snapshot,
-        requested_forwardings,
-    ) = match attached {
+    let attached = match attached {
         Ok(attached) => attached,
         Err(err) => {
             // Auth/session setup failed: clean up the agent socket we bound early.
@@ -1185,11 +1180,12 @@ async fn handle_native_user_auth(
                 .await;
         }
     };
+    let client_id = attached.client_id;
     if let Err(err) = start_remote_forwards(
         state,
         socket,
         client_id,
-        requested_forwardings.clone(),
+        attached.requested_forwardings.clone(),
         pending.client.requested_session.clone(),
     )
     .await
@@ -1207,14 +1203,14 @@ async fn handle_native_user_auth(
 
     let ok = NativeAuthOk {
         client_id,
-        session: session_name,
-        mode,
+        session: attached.session_name,
+        mode: attached.mode,
         session_key: pending.session_key,
-        session_key_id: key_id,
-        attach_ticket,
-        attach_ticket_psk,
-        initial_seq: output_seq,
-        snapshot,
+        session_key_id: attached.key_id,
+        attach_ticket: attached.attach_ticket,
+        attach_ticket_psk: attached.attach_ticket_psk,
+        initial_seq: attached.output_seq,
+        snapshot: attached.snapshot,
         policy_flags: Vec::new(),
     };
     let body = protocol::to_body(&NativeAuthOkBody { ok })?;
@@ -1238,78 +1234,84 @@ async fn handle_bootstrap_attach(
     body: Vec<u8>,
 ) -> Result<()> {
     let req: BootstrapAttachRequest = protocol::from_body(&body)?;
-    let (client_id, key, key_id, session_name, mode, output_seq, snapshot) = {
+    type BootstrapAttachResult = ([u8; 16], [u8; 32], [u8; 16], String, String, u64, Vec<u8>);
+    let attached: Result<BootstrapAttachResult> = {
         let mut locked = state.lock().expect("server state poisoned");
         if !verify_bootstrap(&req.bootstrap, &locked.secret)? {
-            return send_reject(socket, peer, "invalid or expired bootstrap").await;
-        }
-        if !locked.sessions.contains_key(&req.bootstrap.session) && !locked.config.create_on_attach
+            Err(anyhow!("invalid or expired bootstrap"))
+        } else if !locked.sessions.contains_key(&req.bootstrap.session)
+            && !locked.config.create_on_attach
         {
-            return send_reject(socket, peer, "session does not exist").await;
-        }
-        locked.ensure_session(
-            &req.bootstrap.session,
-            req.cols,
-            req.rows,
-            &req.bootstrap.mode,
-            &req.requested_env,
-        )?;
-        let session = locked
-            .sessions
-            .get_mut(&req.bootstrap.session)
-            .expect("session exists");
-        if mode_allows_terminal_updates(&req.bootstrap.mode) {
-            apply_terminal_size(
-                session.pty.as_ref(),
-                &mut session.parser,
+            Err(anyhow!("session does not exist"))
+        } else {
+            locked.ensure_session(
+                &req.bootstrap.session,
                 req.cols,
                 req.rows,
+                &req.bootstrap.mode,
+                &req.requested_env,
             )?;
+            let session = locked
+                .sessions
+                .get_mut(&req.bootstrap.session)
+                .expect("session exists");
+            if mode_allows_terminal_updates(&req.bootstrap.mode) {
+                apply_terminal_size(
+                    session.pty.as_ref(),
+                    &mut session.parser,
+                    req.cols,
+                    req.rows,
+                )?;
+            }
+            let client_id = crypto::random_16();
+            let snapshot = screen_snapshot(session.parser.screen());
+            let output_seq = session.output_seq;
+            let bootstrap_session = req.bootstrap.session.clone();
+            locked.insert_client(
+                &bootstrap_session,
+                client_id,
+                ClientState {
+                    endpoint: peer,
+                    mode: req.bootstrap.mode.clone(),
+                    session_key: req.bootstrap.session_key,
+                    last_acked: output_seq,
+                    replay: ReplayWindow::default(),
+                    send_seq: 1,
+                    cols: req.cols,
+                    rows: req.rows,
+                    last_seen: Instant::now(),
+                    pending: VecDeque::new(),
+                    last_frame_sent: Instant::now() - Duration::from_secs(3600),
+                    paced_snapshot_due: false,
+                    allowed_forwardings: Vec::new(),
+                    stream_writers: HashMap::new(),
+                    opened_streams: HashSet::new(),
+                    stream_send_credit: HashMap::new(),
+                    stream_pending_data: HashMap::new(),
+                    stream_next_send_offset: HashMap::new(),
+                    stream_sent_data: HashMap::new(),
+                    stream_next_recv_offset: HashMap::new(),
+                    stream_recv_pending: HashMap::new(),
+                    epoch: 0,
+                    epoch_started: Instant::now(),
+                    epoch_packets: 0,
+                    previous_session_key: None,
+                },
+            );
+            Ok((
+                client_id,
+                req.bootstrap.session_key,
+                req.bootstrap.session_key_id,
+                req.bootstrap.session.clone(),
+                req.bootstrap.mode.clone(),
+                output_seq,
+                snapshot,
+            ))
         }
-        let client_id = crypto::random_16();
-        let snapshot = screen_snapshot(session.parser.screen());
-        let output_seq = session.output_seq;
-        let bootstrap_session = req.bootstrap.session.clone();
-        locked.insert_client(
-            &bootstrap_session,
-            client_id,
-            ClientState {
-                endpoint: peer,
-                mode: req.bootstrap.mode.clone(),
-                session_key: req.bootstrap.session_key,
-                last_acked: output_seq,
-                replay: ReplayWindow::default(),
-                send_seq: 1,
-                cols: req.cols,
-                rows: req.rows,
-                last_seen: Instant::now(),
-                pending: VecDeque::new(),
-                last_frame_sent: Instant::now() - Duration::from_secs(3600),
-                paced_snapshot_due: false,
-                allowed_forwardings: Vec::new(),
-                stream_writers: HashMap::new(),
-                opened_streams: HashSet::new(),
-                stream_send_credit: HashMap::new(),
-                stream_pending_data: HashMap::new(),
-                stream_next_send_offset: HashMap::new(),
-                stream_sent_data: HashMap::new(),
-                stream_next_recv_offset: HashMap::new(),
-                stream_recv_pending: HashMap::new(),
-                epoch: 0,
-                epoch_started: Instant::now(),
-                epoch_packets: 0,
-                previous_session_key: None,
-            },
-        );
-        (
-            client_id,
-            req.bootstrap.session_key,
-            req.bootstrap.session_key_id,
-            req.bootstrap.session.clone(),
-            req.bootstrap.mode.clone(),
-            output_seq,
-            snapshot,
-        )
+    };
+    let (client_id, key, key_id, session_name, mode, output_seq, snapshot) = match attached {
+        Ok(attached) => attached,
+        Err(err) => return send_reject(socket, peer, &err.to_string()).await,
     };
     let ok = AttachOk {
         client_id,
@@ -1341,24 +1343,29 @@ async fn handle_ticket_attach(
     body: Vec<u8>,
 ) -> Result<()> {
     let env: TicketAttachEnvelope = protocol::from_body(&body)?;
-    let (ticket, request_plain) = {
+    let opened = {
         let locked = state.lock().expect("server state poisoned");
         if !locked.config.allow_attach_tickets {
-            return send_reject(socket, peer, "attach tickets disabled").await;
+            Err(anyhow!("attach tickets disabled"))
+        } else {
+            let ticket = open_attach_ticket(&locked.secret, &env.ticket)?;
+            let request_key = crypto::hkdf32(
+                &ticket.psk,
+                &env.client_nonce,
+                b"dosh/ticket-attach-request/v1",
+            )?;
+            let request_plain = crypto::open(
+                &request_key,
+                &env.client_nonce,
+                b"dosh-ticket-attach-request-v1",
+                &env.ciphertext,
+            )?;
+            Ok((ticket, request_plain))
         }
-        let ticket = open_attach_ticket(&locked.secret, &env.ticket)?;
-        let request_key = crypto::hkdf32(
-            &ticket.psk,
-            &env.client_nonce,
-            b"dosh/ticket-attach-request/v1",
-        )?;
-        let request_plain = crypto::open(
-            &request_key,
-            &env.client_nonce,
-            b"dosh-ticket-attach-request-v1",
-            &env.ciphertext,
-        )?;
-        (ticket, request_plain)
+    };
+    let (ticket, request_plain) = match opened {
+        Ok(opened) => opened,
+        Err(err) => return send_reject(socket, peer, &err.to_string()).await,
     };
     let req: TicketAttachBody = protocol::from_body(&request_plain)?;
     if req.session != ticket.session || req.mode != ticket.mode {
@@ -1367,66 +1374,71 @@ async fn handle_ticket_attach(
 
     let session_key = crypto::random_32();
     let session_key_id = protocol::session_key_id(&session_key);
-    let (client_id, output_seq, snapshot) = {
+    let attached = {
         let mut locked = state.lock().expect("server state poisoned");
         if !locked.sessions.contains_key(&req.session) && !locked.config.create_on_attach {
-            return send_reject(socket, peer, "session does not exist").await;
-        }
-        locked.ensure_session(
-            &req.session,
-            req.cols,
-            req.rows,
-            &req.mode,
-            &req.requested_env,
-        )?;
-        let session = locked
-            .sessions
-            .get_mut(&req.session)
-            .expect("session exists");
-        if mode_allows_terminal_updates(&req.mode) {
-            apply_terminal_size(
-                session.pty.as_ref(),
-                &mut session.parser,
+            Err(anyhow!("session does not exist"))
+        } else {
+            locked.ensure_session(
+                &req.session,
                 req.cols,
                 req.rows,
+                &req.mode,
+                &req.requested_env,
             )?;
+            let session = locked
+                .sessions
+                .get_mut(&req.session)
+                .expect("session exists");
+            if mode_allows_terminal_updates(&req.mode) {
+                apply_terminal_size(
+                    session.pty.as_ref(),
+                    &mut session.parser,
+                    req.cols,
+                    req.rows,
+                )?;
+            }
+            let client_id = crypto::random_16();
+            let snapshot = screen_snapshot(session.parser.screen());
+            let output_seq = session.output_seq;
+            let ticket_session = req.session.clone();
+            locked.insert_client(
+                &ticket_session,
+                client_id,
+                ClientState {
+                    endpoint: peer,
+                    mode: req.mode.clone(),
+                    session_key,
+                    last_acked: output_seq,
+                    replay: ReplayWindow::default(),
+                    send_seq: 1,
+                    cols: req.cols,
+                    rows: req.rows,
+                    last_seen: Instant::now(),
+                    pending: VecDeque::new(),
+                    last_frame_sent: Instant::now() - Duration::from_secs(3600),
+                    paced_snapshot_due: false,
+                    allowed_forwardings: Vec::new(),
+                    stream_writers: HashMap::new(),
+                    opened_streams: HashSet::new(),
+                    stream_send_credit: HashMap::new(),
+                    stream_pending_data: HashMap::new(),
+                    stream_next_send_offset: HashMap::new(),
+                    stream_sent_data: HashMap::new(),
+                    stream_next_recv_offset: HashMap::new(),
+                    stream_recv_pending: HashMap::new(),
+                    epoch: 0,
+                    epoch_started: Instant::now(),
+                    epoch_packets: 0,
+                    previous_session_key: None,
+                },
+            );
+            Ok((client_id, output_seq, snapshot))
         }
-        let client_id = crypto::random_16();
-        let snapshot = screen_snapshot(session.parser.screen());
-        let output_seq = session.output_seq;
-        let ticket_session = req.session.clone();
-        locked.insert_client(
-            &ticket_session,
-            client_id,
-            ClientState {
-                endpoint: peer,
-                mode: req.mode.clone(),
-                session_key,
-                last_acked: output_seq,
-                replay: ReplayWindow::default(),
-                send_seq: 1,
-                cols: req.cols,
-                rows: req.rows,
-                last_seen: Instant::now(),
-                pending: VecDeque::new(),
-                last_frame_sent: Instant::now() - Duration::from_secs(3600),
-                paced_snapshot_due: false,
-                allowed_forwardings: Vec::new(),
-                stream_writers: HashMap::new(),
-                opened_streams: HashSet::new(),
-                stream_send_credit: HashMap::new(),
-                stream_pending_data: HashMap::new(),
-                stream_next_send_offset: HashMap::new(),
-                stream_sent_data: HashMap::new(),
-                stream_next_recv_offset: HashMap::new(),
-                stream_recv_pending: HashMap::new(),
-                epoch: 0,
-                epoch_started: Instant::now(),
-                epoch_packets: 0,
-                previous_session_key: None,
-            },
-        );
-        (client_id, output_seq, snapshot)
+    };
+    let (client_id, output_seq, snapshot) = match attached {
+        Ok(attached) => attached,
+        Err(err) => return send_reject(socket, peer, &err.to_string()).await,
     };
 
     let ok = AttachOk {