atproto-plc / src / builder.rs

//! Builder pattern for creating did:plc identifiers

use crate::crypto::SigningKey;
use crate::did::Did;
use crate::document::ServiceEndpoint;
use crate::encoding::{base32_encode, dag_cbor_encode, sha256};
use crate::error::{PlcError, Result};
use crate::operations::{Operation, UnsignedOperation};
use crate::validation::{
    validate_also_known_as, validate_rotation_keys, validate_services,
    validate_verification_methods,
};
use std::collections::HashMap;

/// Builder for creating new did:plc identifiers
///
/// # Examples
///
/// ```
/// use atproto_plc::{DidBuilder, SigningKey, ServiceEndpoint};
///
/// let rotation_key = SigningKey::generate_p256();
/// let signing_key = SigningKey::generate_k256();
///
/// let (did, operation, keys) = DidBuilder::new()
///     .add_rotation_key(rotation_key)
///     .add_verification_method("atproto".into(), signing_key)
///     .add_also_known_as("at://alice.example.com".into())
///     .add_service(
///         "atproto_pds".into(),
///         ServiceEndpoint::new(
///             "AtprotoPersonalDataServer".into(),
///             "https://pds.example.com".into(),
///         ),
///     )
///     .build()?;
///
/// println!("Created DID: {}", did);
/// # Ok::<(), atproto_plc::PlcError>(())
/// ```
pub struct DidBuilder {
    rotation_keys: Vec<SigningKey>,
    verification_methods: HashMap<String, SigningKey>,
    also_known_as: Vec<String>,
    services: HashMap<String, ServiceEndpoint>,
}

impl DidBuilder {
    /// Create a new DID builder
    pub fn new() -> Self {
        Self {
            rotation_keys: Vec::new(),
            verification_methods: HashMap::new(),
            also_known_as: Vec::new(),
            services: HashMap::new(),
        }
    }

    /// Add a rotation key (1-5 required, no duplicates)
    ///
    /// Rotation keys are used to sign operations and can be used to recover
    /// control of the DID within a 72-hour window.
    ///
    /// # Examples
    ///
    /// ```
    /// use atproto_plc::{DidBuilder, SigningKey};
    ///
    /// let key = SigningKey::generate_p256();
    /// let builder = DidBuilder::new().add_rotation_key(key);
    /// ```
    pub fn add_rotation_key(mut self, key: SigningKey) -> Self {
        self.rotation_keys.push(key);
        self
    }

    /// Add a verification method (max 10)
    ///
    /// Verification methods are cryptographic keys used for authentication
    /// and signing. In ATProto, these are typically used for signing posts
    /// and other records.
    ///
    /// # Examples
    ///
    /// ```
    /// use atproto_plc::{DidBuilder, SigningKey};
    ///
    /// let key = SigningKey::generate_k256();
    /// let builder = DidBuilder::new()
    ///     .add_verification_method("atproto".into(), key);
    /// ```
    pub fn add_verification_method(mut self, name: String, key: SigningKey) -> Self {
        self.verification_methods.insert(name, key);
        self
    }

    /// Add an also-known-as URI
    ///
    /// Also-known-as URIs are alternate identifiers for the same entity.
    /// In ATProto, this is typically the user's handle.
    ///
    /// # Examples
    ///
    /// ```
    /// use atproto_plc::DidBuilder;
    ///
    /// let builder = DidBuilder::new()
    ///     .add_also_known_as("at://alice.bsky.social".into());
    /// ```
    pub fn add_also_known_as(mut self, uri: String) -> Self {
        self.also_known_as.push(uri);
        self
    }

    /// Add a service endpoint
    ///
    /// Services are endpoints that provide functionality for the DID.
    /// In ATProto, this is typically the Personal Data Server (PDS).
    ///
    /// # Examples
    ///
    /// ```
    /// use atproto_plc::{DidBuilder, ServiceEndpoint};
    ///
    /// let builder = DidBuilder::new()
    ///     .add_service(
    ///         "atproto_pds".into(),
    ///         ServiceEndpoint::new(
    ///             "AtprotoPersonalDataServer".into(),
    ///             "https://pds.example.com".into(),
    ///         ),
    ///     );
    /// ```
    pub fn add_service(mut self, name: String, endpoint: ServiceEndpoint) -> Self {
        self.services.insert(name, endpoint);
        self
    }

    /// Build and sign the genesis operation, returning the DID, operation, and keys
    ///
    /// This method:
    /// 1. Validates all inputs
    /// 2. Creates an unsigned genesis operation
    /// 3. Signs it with the first rotation key
    /// 4. Derives the DID from the signed operation's hash
    /// 5. Returns the DID, signed operation, and all keys for safekeeping
    ///
    /// # Errors
    ///
    /// Returns errors if:
    /// - No rotation keys provided
    /// - Too many rotation keys (>5)
    /// - Too many verification methods (>10)
    /// - Invalid URIs or service endpoints
    /// - Signing fails
    ///
    /// # Examples
    ///
    /// ```
    /// use atproto_plc::{DidBuilder, SigningKey};
    ///
    /// let rotation_key = SigningKey::generate_p256();
    ///
    /// let (did, operation, keys) = DidBuilder::new()
    ///     .add_rotation_key(rotation_key)
    ///     .build()?;
    ///
    /// assert!(did.as_str().starts_with("did:plc:"));
    /// # Ok::<(), atproto_plc::PlcError>(())
    /// ```
    pub fn build(self) -> Result<(Did, Operation, BuilderKeys)> {
        // Validate inputs
        if self.rotation_keys.is_empty() {
            return Err(PlcError::InvalidRotationKeys(
                "At least one rotation key is required".to_string(),
            ));
        }

        // Convert keys to did:key format for validation
        let rotation_key_strings: Vec<String> =
            self.rotation_keys.iter().map(|k| k.to_did_key()).collect();

        let verification_method_strings: HashMap<String, String> = self
            .verification_methods
            .iter()
            .map(|(name, key)| (name.clone(), key.to_did_key()))
            .collect();

        // Validate all fields
        validate_rotation_keys(&rotation_key_strings)?;
        validate_verification_methods(&verification_method_strings)?;
        validate_also_known_as(&self.also_known_as)?;
        validate_services(&self.services)?;

        // Create unsigned genesis operation
        let unsigned = UnsignedOperation::PlcOperation {
            rotation_keys: rotation_key_strings,
            verification_methods: verification_method_strings,
            also_known_as: self.also_known_as,
            services: self.services,
            prev: None, // Genesis has no previous operation
        };

        // Sign with the first rotation key
        let signed = unsigned.sign(&self.rotation_keys[0])?;

        // Derive DID from the signed operation
        let did = Self::derive_did(&signed)?;

        // Collect all keys to return
        let keys = BuilderKeys {
            rotation_keys: self.rotation_keys,
            verification_methods: self.verification_methods,
        };

        Ok((did, signed, keys))
    }

    /// Derive a DID from a signed genesis operation
    ///
    /// The DID is derived by:
    /// 1. Computing the CID of the signed operation
    /// 2. Taking the SHA-256 hash of the operation
    /// 3. Base32-encoding the hash
    /// 4. Taking the first 24 characters
    fn derive_did(operation: &Operation) -> Result<Did> {
        // Get the CID of the operation
        let _cid = operation.cid()?;

        // The DID is derived from the CID by taking the hash portion
        // For simplicity, we'll hash the entire serialized operation
        let serialized =
            dag_cbor_encode(operation).map_err(|e| PlcError::DagCborError(e.to_string()))?;

        let hash = sha256(&serialized);
        let encoded = base32_encode(&hash);

        // Take first 24 characters for the DID identifier
        let identifier = &encoded[..24.min(encoded.len())];

        Did::from_identifier(identifier)
    }
}

impl Default for DidBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Keys returned from the builder
///
/// These should be stored securely by the application.
/// The rotation keys can be used to update or recover the DID.
/// The verification methods are used for signing application data.
pub struct BuilderKeys {
    /// Rotation keys (private keys)
    pub rotation_keys: Vec<SigningKey>,

    /// Verification method keys (private keys)
    pub verification_methods: HashMap<String, SigningKey>,
}

impl BuilderKeys {
    /// Get a rotation key by index
    pub fn rotation_key(&self, index: usize) -> Option<&SigningKey> {
        self.rotation_keys.get(index)
    }

    /// Get a verification method key by name
    pub fn verification_method(&self, name: &str) -> Option<&SigningKey> {
        self.verification_methods.get(name)
    }

    /// Get the primary rotation key (first one)
    pub fn primary_rotation_key(&self) -> Option<&SigningKey> {
        self.rotation_key(0)
    }
}

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

    #[test]
    fn test_builder_basic() {
        let rotation_key = SigningKey::generate_p256();

        let (did, operation, keys) = DidBuilder::new()
            .add_rotation_key(rotation_key)
            .build()
            .unwrap();

        assert!(did.as_str().starts_with("did:plc:"));
        assert!(operation.is_genesis());
        assert_eq!(keys.rotation_keys.len(), 1);
    }

    #[test]
    fn test_builder_with_verification_methods() {
        let rotation_key = SigningKey::generate_p256();
        let signing_key = SigningKey::generate_k256();

        let (did, _, keys) = DidBuilder::new()
            .add_rotation_key(rotation_key)
            .add_verification_method("atproto".into(), signing_key)
            .build()
            .unwrap();

        assert!(did.as_str().starts_with("did:plc:"));
        assert_eq!(keys.verification_methods.len(), 1);
        assert!(keys.verification_method("atproto").is_some());
    }

    #[test]
    fn test_builder_with_services() {
        let rotation_key = SigningKey::generate_p256();

        let (did, _, _) = DidBuilder::new()
            .add_rotation_key(rotation_key)
            .add_service(
                "atproto_pds".into(),
                ServiceEndpoint::new(
                    "AtprotoPersonalDataServer".into(),
                    "https://pds.example.com".into(),
                ),
            )
            .build()
            .unwrap();

        assert!(did.as_str().starts_with("did:plc:"));
    }

    #[test]
    fn test_builder_no_rotation_keys() {
        let result = DidBuilder::new().build();
        assert!(result.is_err());
    }

    #[test]
    fn test_builder_too_many_rotation_keys() {
        let mut builder = DidBuilder::new();

        for _ in 0..6 {
            builder = builder.add_rotation_key(SigningKey::generate_p256());
        }

        assert!(builder.build().is_err());
    }

    #[test]
    fn test_builder_keys_access() {
        let rotation_key = SigningKey::generate_p256();
        let signing_key = SigningKey::generate_k256();

        let (_, _, keys) = DidBuilder::new()
            .add_rotation_key(rotation_key)
            .add_verification_method("atproto".into(), signing_key)
            .build()
            .unwrap();

        assert!(keys.primary_rotation_key().is_some());
        assert!(keys.verification_method("atproto").is_some());
        assert!(keys.verification_method("nonexistent").is_none());
    }
}