OHTTP keys should be rotated

This pr addresses #445. It implements OHTTP-key rotation to payjoin-mailroom
Mailroom operators can now decide the time interval for keys to be rotated.
Also if a key has expired, a 422 error is returned to clients. Clients can
handle they key-rotation via the cach-control header returned by the directory.

Author: zealsham

payjoin-mailroom/src/config.rs

@@ -12,6 +12,8 @@ pub struct Config {
     pub storage_dir: PathBuf,
     #[serde(deserialize_with = "deserialize_duration_secs")]
     pub timeout: Duration,
+    #[serde(deserialize_with = "deserialize_optional_duration_secs")]
+    pub ohttp_keys_max_age: Option<Duration>,
     pub v1: Option<V1Config>,
     #[cfg(feature = "telemetry")]
     pub telemetry: Option<TelemetryConfig>,
@@ -85,6 +87,7 @@ impl Default for Config {
             listener: "[::]:8080".parse().expect("valid default listener address"),
             storage_dir: PathBuf::from("./data"),
             timeout: Duration::from_secs(30),
+            ohttp_keys_max_age: Some(Duration::from_secs(7 * 24 * 60 * 60)),
             v1: None,
             #[cfg(feature = "telemetry")]
             telemetry: None,
@@ -104,17 +107,33 @@ where
     Ok(Duration::from_secs(secs))
 }
 
+fn deserialize_optional_duration_secs<'de, D>(deserializer: D) -> Result<Option<Duration>, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    let secs: Option<u64> = Option::deserialize(deserializer)?;
+    match secs {
+        None => Ok(None),
+        Some(0) => Err(<D::Error as serde::de::Error>::custom(
+            "ohttp_keys_max_age must be greater than 0 seconds when set",
+        )),
+        Some(s) => Ok(Some(Duration::from_secs(s))),
+    }
+}
+
 impl Config {
     pub fn new(
         listener: ListenerAddress,
         storage_dir: PathBuf,
         timeout: Duration,
+        ohttp_keys_max_age: Option<Duration>,
         v1: Option<V1Config>,
     ) -> Self {
         Self {
             listener,
             storage_dir,
             timeout,
+            ohttp_keys_max_age,
             v1,
             #[cfg(feature = "telemetry")]
             telemetry: None,

payjoin-mailroom/src/directory.rs

@@ -1,14 +1,18 @@
+use std::path::PathBuf;
 use std::pin::Pin;
 use std::str::FromStr;
+use std::sync::atomic::{AtomicU8, Ordering};
 use std::sync::Arc;
 use std::task::{Context, Poll};
+use std::time::{Duration, Instant};
 
 use anyhow::Result;
 use axum::body::{Body, Bytes};
-use axum::http::header::{HeaderValue, ACCESS_CONTROL_ALLOW_ORIGIN, CONTENT_TYPE};
+use axum::http::header::{HeaderValue, ACCESS_CONTROL_ALLOW_ORIGIN, CACHE_CONTROL, CONTENT_TYPE};
 use axum::http::{Method, Request, Response, StatusCode, Uri};
 use http_body_util::BodyExt;
 use payjoin::directory::{ShortId, ShortIdError, ENCAPSULATED_MESSAGE_BYTES};
+use tokio::sync::RwLock;
 use tracing::{debug, error, trace, warn};
 
 use crate::db::{Db, Error as DbError, SendableError};
@@ -28,6 +32,79 @@ const V1_VERSION_UNSUPPORTED_RES_JSON: &str =
 
 pub type BoxError = Box<dyn std::error::Error + Send + Sync>;
 
+// Two-slot OHTTP key set supporting rotation overlap.
+//
+// Key IDs alternate between 0 and 1. Both slots are always populated.
+// The current key is served to new clients; both slots are accepted
+// for decapsulation so that clients with a cached previous key still
+// work during the grace window after a switch.
+#[derive(Debug)]
+pub(crate) struct KeySlot {
+    pub(crate) server: ohttp::Server,
+    pub(crate) valid_from: Instant,
+}
+
+#[derive(Debug)]
+pub struct KeyRotatingServer {
+    keys: [RwLock<KeySlot>; 2],
+    current_key_id: AtomicU8,
+}
+
+impl KeyRotatingServer {
+    pub(crate) fn new(slot0: KeySlot, slot1: KeySlot, current_key_id: u8) -> Self {
+        assert!(current_key_id <= 1, "key_id must be 0 or 1");
+        Self {
+            keys: [RwLock::new(slot0), RwLock::new(slot1)],
+            current_key_id: AtomicU8::new(current_key_id),
+        }
+    }
+
+    pub fn current_key_id(&self) -> u8 { self.current_key_id.load(Ordering::Acquire) }
+
+    pub async fn current_valid_from(&self) -> Instant {
+        let id = self.current_key_id();
+        self.keys[id as usize].read().await.valid_from
+    }
+
+    pub fn next_key_id(&self) -> u8 { 1 - self.current_key_id() }
+
+    // Look up the server matching the key_id in an OHTTP message and
+    // decapsulate. The first byte of an OHTTP encapsulated request is the
+    // key identifier (RFC 9458 Section 4.3).
+    pub async fn decapsulate(
+        &self,
+        ohttp_body: &[u8],
+    ) -> std::result::Result<(Vec<u8>, ohttp::ServerResponse), ohttp::Error> {
+        let key_id = ohttp_body.first().copied().ok_or(ohttp::Error::Truncated)?;
+        match self.keys.get(key_id as usize) {
+            Some(slot) => slot.read().await.server.decapsulate(ohttp_body),
+            None => Err(ohttp::Error::KeyId),
+        }
+    }
+
+    // Encode the current key's config for serving to clients.
+    pub async fn encode_current(&self) -> std::result::Result<Vec<u8>, ohttp::Error> {
+        let id = self.current_key_id();
+        self.keys[id as usize].read().await.server.config().encode()
+    }
+
+    // Flip which key is advertised to new clients.
+    pub fn switch(&self) {
+        let old = self.current_key_id();
+        self.current_key_id.store(1 - old, Ordering::Release);
+    }
+
+    // Replace a slot with fresh key material.
+    pub async fn overwrite(&self, key_id: u8, server: ohttp::Server) {
+        assert!(key_id <= 1, "key_id must be 0 or 1");
+        *self.keys[key_id as usize].write().await = KeySlot { server, valid_from: Instant::now() };
+    }
+
+    pub async fn valid_from(&self, key_id: u8) -> Instant {
+        self.keys[key_id as usize].read().await.valid_from
+    }
+}
+
 /// Opaque blocklist of Bitcoin addresses stored as script pubkeys.
 ///
 /// Addresses are converted to `ScriptBuf` at parse time so that
@@ -91,7 +168,8 @@ fn parse_address_lines(text: &str) -> std::collections::HashSet<bitcoin::ScriptB
 #[derive(Clone)]
 pub struct Service<D: Db> {
     db: D,
-    ohttp: ohttp::Server,
+    ohttp: Arc<KeyRotatingServer>,
+    ohttp_keys_max_age: Option<Duration>,
     sentinel_tag: SentinelTag,
     v1: Option<V1>,
 }
@@ -117,10 +195,18 @@ where
 }
 
 impl<D: Db> Service<D> {
-    pub fn new(db: D, ohttp: ohttp::Server, sentinel_tag: SentinelTag, v1: Option<V1>) -> Self {
-        Self { db, ohttp, sentinel_tag, v1 }
+    pub fn new(
+        db: D,
+        ohttp: Arc<KeyRotatingServer>,
+        ohttp_keys_max_age: Option<Duration>,
+        sentinel_tag: SentinelTag,
+        v1: Option<V1>,
+    ) -> Self {
+        Self { db, ohttp, ohttp_keys_max_age, sentinel_tag, v1 }
     }
 
+    pub fn ohttp_key_set(&self) -> &Arc<KeyRotatingServer> { &self.ohttp }
+
     async fn serve_request<B>(&self, req: Request<B>) -> Result<Response<Body>>
     where
         B: axum::body::HttpBody<Data = Bytes> + Send + 'static,
@@ -200,10 +286,10 @@ impl<D: Db> Service<D> {
             .map_err(|e| HandlerError::BadRequest(anyhow::anyhow!(e.into())))?
             .to_bytes();
 
-        // Decapsulate OHTTP request
         let (bhttp_req, res_ctx) = self
             .ohttp
             .decapsulate(&ohttp_body)
+            .await
             .map_err(|e| HandlerError::OhttpKeyRejection(e.into()))?;
         let mut cursor = std::io::Cursor::new(bhttp_req);
         let req = bhttp::Message::read_bhttp(&mut cursor)
@@ -380,11 +466,22 @@ impl<D: Db> Service<D> {
     async fn get_ohttp_keys(&self) -> Result<Response<Body>, HandlerError> {
         let ohttp_keys = self
             .ohttp
-            .config()
-            .encode()
+            .encode_current()
+            .await
             .map_err(|e| HandlerError::InternalServerError(e.into()))?;
         let mut res = Response::new(full(ohttp_keys));
         res.headers_mut().insert(CONTENT_TYPE, HeaderValue::from_static("application/ohttp-keys"));
+        if let Some(max_age) = self.ohttp_keys_max_age {
+            let remaining = max_age.saturating_sub(self.ohttp.current_valid_from().await.elapsed());
+            res.headers_mut().insert(
+                CACHE_CONTROL,
+                HeaderValue::from_str(&format!(
+                    "public, s-maxage={}, immutable",
+                    remaining.as_secs()
+                ))
+                .expect("valid header value"),
+            );
+        }
         Ok(res)
     }
 
@@ -412,6 +509,56 @@ impl<D: Db> Service<D> {
     }
 }
 
+// Grace period after a switch during which the old key is still
+// accepted for decapsulation.
+const ROTATION_GRACE: Duration = Duration::from_secs(30);
+
+// Background task that rotates OHTTP keys on a fixed interval.
+//
+// Each cycle: sleep until the current key is about to expire, switch
+// to the other slot, wait out the grace period, then overwrite the
+// old slot with fresh key material so it is ready for the next cycle.
+pub fn spawn_key_rotation(keyset: Arc<KeyRotatingServer>, keys_dir: PathBuf, interval: Duration) {
+    tokio::spawn(async move {
+        loop {
+            let switch_delay = {
+                let valid_from = keyset.current_valid_from().await;
+                let switch_at = valid_from + interval - ROTATION_GRACE / 2;
+                switch_at.saturating_duration_since(Instant::now())
+            };
+            tokio::time::sleep(switch_delay).await;
+
+            let old_key_id = keyset.current_key_id();
+            keyset.switch();
+            let new_key_id = 1 - old_key_id;
+            tracing::info!("Switched OHTTP serving: key_id {old_key_id} -> {new_key_id}");
+
+            let overwrite_delay = {
+                let valid_from = keyset.valid_from(old_key_id).await;
+                let overwrite_at = valid_from + interval + ROTATION_GRACE;
+                overwrite_at.saturating_duration_since(Instant::now())
+            };
+            tokio::time::sleep(overwrite_delay).await;
+
+            let config = match crate::key_config::gen_ohttp_server_config_with_id(old_key_id) {
+                Ok(c) => c,
+                Err(e) => {
+                    tracing::error!("Failed to generate OHTTP key: {e}");
+                    continue;
+                }
+            };
+            let _ = tokio::fs::remove_file(keys_dir.join(format!("{old_key_id}.ikm"))).await;
+            if let Err(e) = crate::key_config::persist_key_config(&config, &keys_dir).await {
+                tracing::error!("Failed to persist OHTTP key: {e}");
+                continue;
+            }
+
+            keyset.overwrite(old_key_id, config.into_server()).await;
+            tracing::info!("Overwrote OHTTP key_id {old_key_id} with fresh material");
+        }
+    });
+}
+
 fn handle_peek<E: SendableError>(
     result: Result<Arc<Vec<u8>>, DbError<E>>,
     timeout_response: Response<Body>,
@@ -485,8 +632,8 @@ impl HandlerError {
             }
             HandlerError::OhttpKeyRejection(e) => {
                 const OHTTP_KEY_REJECTION_RES_JSON: &str = r#"{"type":"https://iana.org/assignments/http-problem-types#ohttp-key", "title": "key identifier unknown"}"#;
-                warn!("Bad request: Key configuration rejected: {}", e);
-                *res.status_mut() = StatusCode::BAD_REQUEST;
+                warn!("Key configuration rejected: {}", e);
+                *res.status_mut() = StatusCode::UNPROCESSABLE_ENTITY;
                 res.headers_mut()
                     .insert(CONTENT_TYPE, HeaderValue::from_static("application/problem+json"));
                 *res.body_mut() = full(OHTTP_KEY_REJECTION_RES_JSON);
@@ -592,9 +739,15 @@ mod tests {
     async fn test_service(v1: Option<V1>) -> Service<FilesDb> {
         let dir = tempfile::tempdir().expect("tempdir");
         let db = FilesDb::init(Duration::from_millis(100), dir.keep()).await.expect("db init");
-        let ohttp: ohttp::Server =
-            crate::key_config::gen_ohttp_server_config().expect("ohttp config").into();
-        Service::new(db, ohttp, SentinelTag::new([0u8; 32]), v1)
+        let now = Instant::now();
+        let c0 = crate::key_config::gen_ohttp_server_config_with_id(0).expect("ohttp config");
+        let c1 = crate::key_config::gen_ohttp_server_config_with_id(1).expect("ohttp config");
+        let keyset = Arc::new(KeyRotatingServer::new(
+            KeySlot { server: c0.into_server(), valid_from: now },
+            KeySlot { server: c1.into_server(), valid_from: now },
+            0,
+        ));
+        Service::new(db, keyset, None, SentinelTag::new([0u8; 32]), v1)
     }
 
     /// A valid ShortId encoded as bech32 for use in URL paths.
@@ -826,9 +979,15 @@ mod tests {
         let dir = tempfile::tempdir().expect("tempdir");
         let db = FilesDb::init(Duration::from_millis(100), dir.keep()).await.expect("db init");
         let db = MetricsDb::new(db, metrics);
-        let ohttp: ohttp::Server =
-            crate::key_config::gen_ohttp_server_config().expect("ohttp config").into();
-        let svc = Service::new(db, ohttp, SentinelTag::new([0u8; 32]), None);
+        let now = Instant::now();
+        let c0 = crate::key_config::gen_ohttp_server_config_with_id(0).expect("ohttp config");
+        let c1 = crate::key_config::gen_ohttp_server_config_with_id(1).expect("ohttp config");
+        let keyset = Arc::new(KeyRotatingServer::new(
+            KeySlot { server: c0.into_server(), valid_from: now },
+            KeySlot { server: c1.into_server(), valid_from: now },
+            0,
+        ));
+        let svc = Service::new(db, keyset, None, SentinelTag::new([0u8; 32]), None);
 
         let id = valid_short_id_path();
         let res = svc