tangled-cli 0.1.0

//! PDS request authentication.
//!
//! App-password sessions use plain `Authorization: Bearer` JWTs. OAuth
//! sessions are DPoP-bound: every request carries a fresh ES256-signed proof
//! (RFC 9449) and honors the server's `DPoP-Nonce` retry dance.

use std::sync::Mutex;

use anyhow::{anyhow, Result};
use base64::Engine;
use p256::ecdsa::signature::Signer;
use p256::ecdsa::{Signature, SigningKey};
use serde::de::DeserializeOwned;
use sha2::{Digest, Sha256};
use tangled_api::xrpc::{self, XrpcError};
use tangled_config::session::Session;

fn b64url(data: &[u8]) -> String {
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
}

/// Authentication mode for requests against the PDS.
pub enum PdsAuth {
    /// Unauthenticated (public endpoints).
    #[allow(dead_code)]
    None,
    /// App-password session JWT.
    Bearer(String),
    /// OAuth session with a DPoP-bound access token.
    Dpop(Box<DpopClient>),
}

impl PdsAuth {
    pub fn from_session(session: &Session) -> Result<Self> {
        match &session.oauth {
            Some(oauth) => Ok(PdsAuth::Dpop(Box::new(DpopClient::new(
                oauth.access_token.clone(),
                DpopKey::from_b64(&oauth.dpop_key)?,
            )))),
            None => Ok(PdsAuth::Bearer(session.access_jwt.clone())),
        }
    }

    pub async fn send<T: DeserializeOwned>(
        &self,
        rb: reqwest::RequestBuilder,
    ) -> Result<T> {
        crate::progress::with_spinner("Contacting PDS...", async {
            Ok(xrpc::parse(self.response(rb).await?).await?)
        })
        .await
    }

    pub async fn send_unit(&self, rb: reqwest::RequestBuilder) -> Result<()> {
        crate::progress::with_spinner("Contacting PDS...", async {
            Ok(xrpc::parse_unit(self.response(rb).await?).await?)
        })
        .await
    }

    pub async fn send_bytes(
        &self,
        rb: reqwest::RequestBuilder,
    ) -> Result<Vec<u8>> {
        crate::progress::with_spinner("Contacting PDS...", async {
            Ok(xrpc::parse_bytes(self.response(rb).await?).await?)
        })
        .await
    }

    async fn response(
        &self,
        rb: reqwest::RequestBuilder,
    ) -> Result<reqwest::Response> {
        match self {
            PdsAuth::None => Ok(rb.send().await.map_err(XrpcError::from)?),
            PdsAuth::Bearer(token) => Ok(rb
                .bearer_auth(token)
                .send()
                .await
                .map_err(XrpcError::from)?),
            PdsAuth::Dpop(client) => client.response(rb).await,
        }
    }
}

/// A P-256 key used to sign DPoP proofs.
pub struct DpopKey {
    signing: SigningKey,
}

impl DpopKey {
    pub fn generate() -> Self {
        Self {
            signing: SigningKey::random(&mut rand::thread_rng()),
        }
    }

    pub fn from_b64(encoded: &str) -> Result<Self> {
        let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .decode(encoded)
            .map_err(|e| anyhow!("invalid DPoP key encoding: {e}"))?;
        let signing = SigningKey::from_slice(&bytes)
            .map_err(|e| anyhow!("invalid DPoP key: {e}"))?;
        Ok(Self { signing })
    }

    pub fn to_b64(&self) -> String {
        b64url(&self.signing.to_bytes())
    }

    fn public_jwk(&self) -> serde_json::Value {
        let point = self.signing.verifying_key().to_encoded_point(false);
        serde_json::json!({
            "kty": "EC",
            "crv": "P-256",
            "x": b64url(point.x().expect("uncompressed point has x")),
            "y": b64url(point.y().expect("uncompressed point has y")),
        })
    }

    /// Builds a signed DPoP proof JWT for `method` on `url`.
    ///
    /// `nonce` is the latest server-provided DPoP nonce, and `access_token`
    /// is set (hashed into the `ath` claim) for resource-server requests.
    pub fn proof(
        &self,
        method: &str,
        url: &str,
        nonce: Option<&str>,
        access_token: Option<&str>,
    ) -> String {
        let header = serde_json::json!({
            "typ": "dpop+jwt",
            "alg": "ES256",
            "jwk": self.public_jwk(),
        });
        // htu must not include query or fragment.
        let htu = url.split(['?', '#']).next().unwrap_or(url);
        let mut jti = [0u8; 16];
        rand::RngCore::fill_bytes(&mut rand::thread_rng(), &mut jti);
        let mut payload = serde_json::json!({
            "jti": b64url(&jti),
            "htm": method,
            "htu": htu,
            "iat": chrono::Utc::now().timestamp(),
        });
        if let Some(nonce) = nonce {
            payload["nonce"] = serde_json::json!(nonce);
        }
        if let Some(token) = access_token {
            payload["ath"] =
                serde_json::json!(b64url(&Sha256::digest(token.as_bytes())));
        }

        let signing_input = format!(
            "{}.{}",
            b64url(header.to_string().as_bytes()),
            b64url(payload.to_string().as_bytes())
        );
        let signature: Signature = self.signing.sign(signing_input.as_bytes());
        format!("{}.{}", signing_input, b64url(&signature.to_bytes()))
    }
}

/// Sends requests authenticated with a DPoP-bound OAuth access token.
pub struct DpopClient {
    access_token: String,
    key: DpopKey,
    nonce: Mutex<Option<String>>,
}

impl DpopClient {
    pub fn new(access_token: String, key: DpopKey) -> Self {
        Self {
            access_token,
            key,
            nonce: Mutex::new(None),
        }
    }

    async fn response(
        &self,
        rb: reqwest::RequestBuilder,
    ) -> Result<reqwest::Response> {
        let mut retried = false;
        loop {
            let attempt = rb.try_clone().ok_or_else(|| {
                anyhow!("request body is not cloneable for DPoP retry")
            })?;
            let (client, request) = attempt.build_split();
            let mut request = request.map_err(XrpcError::from)?;

            let nonce = self.nonce.lock().unwrap().clone();
            let proof = self.key.proof(
                request.method().as_str(),
                request.url().as_str(),
                nonce.as_deref(),
                Some(&self.access_token),
            );
            let headers = request.headers_mut();
            headers.insert(
                reqwest::header::AUTHORIZATION,
                format!("DPoP {}", self.access_token).parse()?,
            );
            headers.insert("DPoP", proof.parse()?);

            let response =
                client.execute(request).await.map_err(XrpcError::from)?;

            let new_nonce = response
                .headers()
                .get("dpop-nonce")
                .and_then(|v| v.to_str().ok())
                .map(str::to_string);
            if let Some(n) = new_nonce.clone() {
                *self.nonce.lock().unwrap() = Some(n);
            }

            if response.status() == reqwest::StatusCode::UNAUTHORIZED
                && !retried
                && new_nonce.is_some()
                && wants_new_nonce(&response)
            {
                retried = true;
                continue;
            }
            return Ok(response);
        }
    }
}

fn wants_new_nonce(res: &reqwest::Response) -> bool {
    res.headers()
        .get(reqwest::header::WWW_AUTHENTICATE)
        .and_then(|v| v.to_str().ok())
        .is_some_and(|v| v.contains("use_dpop_nonce"))
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A DPoP request must retry exactly once when the server demands a
    /// nonce, carrying `Authorization: DPoP` plus a proof header, and store
    /// the server-provided nonce for the retry.
    #[tokio::test]
    async fn dpop_send_retries_with_server_nonce() {
        let mut server = mockito::Server::new_async().await;
        let mock = server
            .mock("GET", "/xrpc/test.method")
            .match_header("authorization", "DPoP token123")
            .match_header("dpop", mockito::Matcher::Regex("^eyJ".into()))
            .with_status(401)
            .with_header("DPoP-Nonce", "nonce-1")
            .with_header("WWW-Authenticate", "DPoP error=\"use_dpop_nonce\"")
            .with_body(r#"{"error":"use_dpop_nonce"}"#)
            .expect(2)
            .create_async()
            .await;

        let client =
            DpopClient::new("token123".to_string(), DpopKey::generate());
        let rb = reqwest::Client::new()
            .get(format!("{}/xrpc/test.method", server.url()));

        // The server keeps answering 401, so after one nonce retry the call
        // fails; the mock's expect(2) proves exactly one retry happened.
        let response = client.response(rb).await.unwrap();
        assert_eq!(response.status(), 401);
        mock.assert_async().await;
        assert_eq!(
            client.nonce.lock().unwrap().as_deref(),
            Some("nonce-1"),
            "server nonce should be stored for future requests"
        );
    }

    #[test]
    fn dpop_key_round_trips_through_base64() {
        let key = DpopKey::generate();
        let restored = DpopKey::from_b64(&key.to_b64()).unwrap();
        assert_eq!(key.to_b64(), restored.to_b64());
    }

    #[test]
    fn dpop_proof_is_a_signed_jwt_with_expected_claims() {
        use base64::Engine;
        let key = DpopKey::generate();
        let proof = key.proof(
            "POST",
            "https://pds.example/xrpc/com.atproto.repo.createRecord?x=1",
            Some("nonce-abc"),
            Some("access-token"),
        );
        let parts: Vec<&str> = proof.split('.').collect();
        assert_eq!(parts.len(), 3);
        let decode = |s: &str| {
            serde_json::from_slice::<serde_json::Value>(
                &base64::engine::general_purpose::URL_SAFE_NO_PAD
                    .decode(s)
                    .unwrap(),
            )
            .unwrap()
        };
        let header = decode(parts[0]);
        assert_eq!(header["typ"], "dpop+jwt");
        assert_eq!(header["alg"], "ES256");
        assert_eq!(header["jwk"]["crv"], "P-256");
        let payload = decode(parts[1]);
        assert_eq!(payload["htm"], "POST");
        // htu must strip the query string.
        assert_eq!(
            payload["htu"],
            "https://pds.example/xrpc/com.atproto.repo.createRecord"
        );
        assert_eq!(payload["nonce"], "nonce-abc");
        assert!(payload["ath"].is_string());
    }
}