ntp_usg-server 5.0.0

// Copyright 2026 U.S. Federal Government (in countries where recognized)
// SPDX-License-Identifier: Apache-2.0

//! NTS-KE (Key Establishment) server using the smol runtime (RFC 8915).
//!
//! Provides the same NTS-KE server functionality as [`crate::nts_ke_server`]
//! but using smol and futures-rustls instead of tokio and tokio-rustls.

use std::io;
use std::sync::{Arc, RwLock};
use std::time::Duration;

use futures_lite::io::{AsyncReadExt, AsyncWriteExt};
use tracing::{Instrument, debug};

use crate::nts_common::*;
pub use crate::nts_ke_server_common::NtsKeServerConfig;
use crate::nts_server_common::MasterKeyStore;

/// An NTS-KE server using the smol runtime.
pub struct NtsKeServer {
    tls_acceptor: futures_rustls::TlsAcceptor,
    key_store: Arc<RwLock<MasterKeyStore>>,
    ntp_server: Option<String>,
    ntp_port: Option<u16>,
    cookie_count: usize,
    listen_addr: String,
}

impl NtsKeServer {
    /// Create a new NTS-KE server from the given configuration and key store.
    pub fn new(
        config: NtsKeServerConfig,
        key_store: Arc<RwLock<MasterKeyStore>>,
    ) -> io::Result<Self> {
        let tls_config =
            crate::tls_config::nts_server_config(config.cert_chain, config.private_key)?;

        Ok(NtsKeServer {
            tls_acceptor: futures_rustls::TlsAcceptor::from(Arc::new(tls_config)),
            key_store,
            ntp_server: config.ntp_server,
            ntp_port: config.ntp_port,
            cookie_count: config.cookie_count,
            listen_addr: config.listen_addr,
        })
    }

    /// Run the NTS-KE server, accepting connections indefinitely.
    pub async fn run(&self) -> io::Result<()> {
        let listener = smol::net::TcpListener::bind(&self.listen_addr).await?;
        debug!("NTS-KE server (smol) listening on {}", self.listen_addr);

        loop {
            let (tcp_stream, peer_addr) = listener.accept().await?;
            debug!(peer = %peer_addr, "NTS-KE connection accepted");

            let acceptor = self.tls_acceptor.clone();
            let key_store = self.key_store.clone();
            let ntp_server = self.ntp_server.clone();
            let ntp_port = self.ntp_port;
            let cookie_count = self.cookie_count;

            smol::spawn(
                async move {
                    match acceptor.accept(tcp_stream).await {
                        Ok(tls_stream) => {
                            if let Err(e) = handle_nts_ke_connection(
                                tls_stream,
                                &key_store,
                                ntp_server.as_deref(),
                                ntp_port,
                                cookie_count,
                            )
                            .await
                            {
                                debug!(peer = %peer_addr, error = %e, "NTS-KE connection error");
                            }
                        }
                        Err(e) => {
                            debug!("TLS accept error from {}: {}", peer_addr, e);
                        }
                    }
                }
                .instrument(tracing::debug_span!(
                    "nts_ke_connection",
                    peer_addr = %peer_addr
                )),
            )
            .detach();
        }
    }

    /// Spawn a background task that rotates the cookie master key every
    /// `interval` (smol runtime).
    ///
    /// Rotation retires the current key (still usable for decryption during the
    /// store's grace period) and generates a fresh one, bounding how long any
    /// single key is in use. This is **opt-in**: without calling it, the master
    /// key is never rotated automatically. Pick an `interval` shorter than the
    /// grace period passed to [`MasterKeyStore::new`]. The returned task keeps
    /// rotation running; drop it to stop (or `.detach()` it to run forever).
    pub fn spawn_key_rotation(&self, interval: Duration) -> smol::Task<()> {
        let key_store = self.key_store.clone();
        smol::spawn(async move {
            loop {
                smol::Timer::after(interval).await;
                crate::nts_server_common::rotate_master_key(&key_store);
            }
        })
    }
}

/// Handle a single NTS-KE client connection (smol runtime).
async fn handle_nts_ke_connection(
    mut tls_stream: futures_rustls::server::TlsStream<smol::net::TcpStream>,
    key_store: &Arc<RwLock<MasterKeyStore>>,
    ntp_server: Option<&str>,
    ntp_port: Option<u16>,
    cookie_count: usize,
) -> io::Result<()> {
    // Read all client NTS-KE records until End of Message, capping the count
    // (see `MAX_KE_RECORDS`) to bound memory/CPU against a peer that never sends
    // End-of-Message.
    let mut records = Vec::new();
    loop {
        let record = read_ke_record_server(&mut tls_stream).await?;
        let is_eom = record.record_type == NTS_KE_END_OF_MESSAGE;
        records.push(record);
        if is_eom {
            break;
        }
        if records.len() >= crate::nts_ke_server_common::MAX_KE_RECORDS {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "NTS-KE request exceeded maximum record count",
            ));
        }
    }

    // Process records (shared logic: negotiate, export keys, generate cookies).
    let (_, tls_conn) = tls_stream.get_ref();
    let resp = crate::nts_ke_server_common::process_nts_ke_records(
        &records,
        tls_conn,
        key_store,
        ntp_server,
        ntp_port,
        cookie_count,
    )?;

    // Send response and close.
    tls_stream.write_all(&resp).await?;
    tls_stream.flush().await?;
    let _ = tls_stream.close().await;

    Ok(())
}

/// Read a single NTS-KE record from a server-side TLS stream (smol).
async fn read_ke_record_server(
    reader: &mut futures_rustls::server::TlsStream<smol::net::TcpStream>,
) -> io::Result<NtsKeRecord> {
    let mut hdr = [0u8; 4];
    reader.read_exact(&mut hdr).await?;
    let raw_type = u16::from_be_bytes([hdr[0], hdr[1]]);
    let body_length = u16::from_be_bytes([hdr[2], hdr[3]]);

    let critical = (raw_type & 0x8000) != 0;
    let record_type = raw_type & 0x7FFF;

    let mut body = vec![0u8; body_length as usize];
    reader.read_exact(&mut body).await?;

    Ok(NtsKeRecord {
        critical,
        record_type,
        body,
    })
}