regy 0.1.0

use std::{
    collections::HashSet,
    fmt, fs,
    io::{self, Read, Write},
    os::unix::fs::{MetadataExt, OpenOptionsExt, PermissionsExt},
    path::{Path, PathBuf},
    str::FromStr,
};

use chrono::{DateTime, Duration, Utc};
use data_encoding::BASE64URL_NOPAD;
use ring::hmac;
use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _};
use subtle::ConstantTimeEq;
use thiserror::Error;
use uuid::Uuid;

use super::{
    ClientBinding,
    identity::{SecretBytes, ServerIdentity, ensure_private_parent, sync_directory},
};

const CLIENT_HMAC_DOMAIN: &[u8] = b"regy-pc-agent-client-v1";
const SCHEMA_VERSION: u32 = 3;
const MAX_CLIENTS: usize = 100;
const MAX_CLIENTS_FILE_BYTES: u64 = 1024 * 1024;
const PRIVATE_FILE_MODE: u32 = 0o600;
const LAST_USE_WRITE_INTERVAL: Duration = Duration::hours(1);
const LOCK_OPEN_ATTEMPTS: usize = 8;

#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)]
pub(crate) enum PairingStoreError {
    #[error("secure randomness unavailable")]
    Randomness,
    #[error("server identity storage is invalid or unavailable")]
    IdentityStorage,
    #[error("approved client storage is unavailable")]
    Storage,
    #[error("approved client storage is malformed or unsupported")]
    Corrupt,
    #[error("approved client label is invalid")]
    InvalidLabel,
    #[error("approved client limit reached")]
    ClientLimit,
    #[error("approved client record is duplicated")]
    Duplicate,
    #[error("approved client binding does not match")]
    BindingMismatch,
    #[error("approved client binding is invalid")]
    InvalidBinding,
    #[error("approved client update was published but durability is uncertain")]
    CommitUncertain,
}

#[derive(Clone, PartialEq, Eq, Hash)]
pub(crate) struct ClientId(String);

impl ClientId {
    fn random() -> Self {
        Self(Uuid::new_v4().hyphenated().to_string())
    }

    pub(crate) fn as_str(&self) -> &str {
        &self.0
    }
}

impl FromStr for ClientId {
    type Err = &'static str;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        let parsed = Uuid::parse_str(value).map_err(|_| "client ID must be a canonical UUID")?;
        if parsed.hyphenated().to_string() != value {
            return Err("client ID must be a canonical lowercase hyphenated UUID");
        }
        Ok(Self(value.to_owned()))
    }
}

impl fmt::Display for ClientId {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl fmt::Debug for ClientId {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.debug_tuple("ClientId").field(&self.0).finish()
    }
}

impl Serialize for ClientId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&self.0)
    }
}

impl<'de> Deserialize<'de> for ClientId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        String::deserialize(deserializer)?
            .parse()
            .map_err(D::Error::custom)
    }
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub(crate) struct ApprovedClient {
    pub id: ClientId,
    pub label: String,
    pub binding: ClientBinding,
    pub created_at: DateTime<Utc>,
    pub last_used_at: DateTime<Utc>,
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct StoredClient {
    id: ClientId,
    label: String,
    credential_hmac: String,
    binding: ClientBinding,
    created_at: DateTime<Utc>,
    last_used_at: DateTime<Utc>,
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct ClientRecords {
    schema_version: u32,
    clients: Vec<StoredClient>,
}

#[derive(Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct SchemaTwoClientRecords {
    schema_version: u32,
    clients: Vec<SchemaTwoStoredClient>,
}

#[derive(Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct SchemaTwoStoredClient {
    id: ClientId,
    label: String,
    credential_hmac: String,
    binding: SchemaTwoClientBinding,
    created_at: DateTime<Utc>,
    last_used_at: DateTime<Utc>,
}

#[derive(Clone, Deserialize, PartialEq, Eq)]
#[serde(tag = "transport", rename_all = "snake_case", deny_unknown_fields)]
enum SchemaTwoClientBinding {
    LocalWebSocket,
    IrohEndpoint { endpoint_id: String },
}

#[derive(Deserialize)]
struct SchemaHeader {
    schema_version: u32,
}

#[derive(Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct LegacyStoredClient {
    id: ClientId,
    label: String,
    credential_hmac: String,
    created_at: DateTime<Utc>,
    last_used_at: DateTime<Utc>,
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct LegacyClientRecords {
    schema_version: u32,
    clients: Vec<LegacyStoredClient>,
}

impl Default for ClientRecords {
    fn default() -> Self {
        Self {
            schema_version: SCHEMA_VERSION,
            clients: Vec::new(),
        }
    }
}

pub(crate) trait ClientStore {
    fn approve(
        &self,
        label: &str,
        credential: &SecretBytes<32>,
        binding: &ClientBinding,
        now: DateTime<Utc>,
    ) -> Result<ApprovedClient, PairingStoreError>;

    fn authenticate(
        &self,
        credential: &SecretBytes<32>,
        binding: &ClientBinding,
        now: DateTime<Utc>,
    ) -> Result<Option<ApprovedClient>, PairingStoreError>;

    fn list(&self) -> Result<Vec<ApprovedClient>, PairingStoreError>;

    fn revoke(&self, id: &ClientId) -> Result<bool, PairingStoreError>;
}

#[derive(Clone)]
pub(crate) struct FileClientStore {
    path: PathBuf,
    identity: ServerIdentity,
    sync_parent: fn(&Path) -> io::Result<()>,
}

impl FileClientStore {
    pub(crate) fn new(path: impl Into<PathBuf>, identity: ServerIdentity) -> Self {
        Self {
            path: path.into(),
            identity,
            sync_parent: sync_directory,
        }
    }

    #[cfg(test)]
    pub(crate) fn with_parent_sync(
        path: impl Into<PathBuf>,
        identity: ServerIdentity,
        sync_parent: fn(&Path) -> io::Result<()>,
    ) -> Self {
        Self {
            path: path.into(),
            identity,
            sync_parent,
        }
    }

    fn credential_hmac(&self, credential: &SecretBytes<32>) -> [u8; 32] {
        let key = hmac::Key::new(hmac::HMAC_SHA256, self.identity.key_bytes());
        let mut context = hmac::Context::with_key(&key);
        context.update(CLIENT_HMAC_DOMAIN);
        context.update(credential.expose());
        context
            .sign()
            .as_ref()
            .try_into()
            .expect("HMAC-SHA-256 always produces 32 bytes")
    }
}

impl fmt::Debug for FileClientStore {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("FileClientStore")
            .field("path", &self.path)
            .field("server_id", self.identity.server_id())
            .finish()
    }
}

impl ClientStore for FileClientStore {
    fn approve(
        &self,
        label: &str,
        credential: &SecretBytes<32>,
        binding: &ClientBinding,
        now: DateTime<Utc>,
    ) -> Result<ApprovedClient, PairingStoreError> {
        validate_label(label)?;
        let binding = canonicalize_binding(binding)?;
        let _transaction = acquire_transaction_lock(&self.path)?;
        let mut records = read_records(&self.path, self.sync_parent)?;
        if records.clients.len() >= MAX_CLIENTS {
            return Err(PairingStoreError::ClientLimit);
        }

        let credential_hmac = self.credential_hmac(credential);
        if records
            .clients
            .iter()
            .try_fold(false, |duplicate, client| {
                Ok::<_, PairingStoreError>(
                    duplicate
                        | bool::from(credential_hmac.ct_eq(&decode_hmac(&client.credential_hmac)?)),
                )
            })?
        {
            return Err(PairingStoreError::Duplicate);
        }

        let id = ClientId::random();
        if records.clients.iter().any(|client| client.id == id) {
            return Err(PairingStoreError::Duplicate);
        }
        let approved = ApprovedClient {
            id: id.clone(),
            label: label.to_owned(),
            binding: binding.clone(),
            created_at: now,
            last_used_at: now,
        };
        records.clients.push(StoredClient {
            id,
            label: label.to_owned(),
            credential_hmac: BASE64URL_NOPAD.encode(&credential_hmac),
            binding,
            created_at: now,
            last_used_at: now,
        });
        write_records(&self.path, &records, self.sync_parent)?;
        Ok(approved)
    }

    fn authenticate(
        &self,
        credential: &SecretBytes<32>,
        binding: &ClientBinding,
        now: DateTime<Utc>,
    ) -> Result<Option<ApprovedClient>, PairingStoreError> {
        let binding = canonicalize_binding(binding)?;
        let _transaction = acquire_transaction_lock(&self.path)?;
        let mut records = read_records(&self.path, self.sync_parent)?;
        let candidate = self.credential_hmac(credential);
        let mut matching_index = None;
        for (index, client) in records.clients.iter().enumerate() {
            let stored = decode_hmac(&client.credential_hmac)?;
            if bool::from(candidate.ct_eq(&stored)) {
                matching_index = Some(index);
            }
        }
        let Some(index) = matching_index else {
            return Ok(None);
        };
        if records.clients[index].binding != binding {
            return Err(PairingStoreError::BindingMismatch);
        }

        let should_persist = now.signed_duration_since(records.clients[index].last_used_at)
            >= LAST_USE_WRITE_INTERVAL;
        if should_persist {
            records.clients[index].last_used_at = now;
            write_records(&self.path, &records, self.sync_parent)?;
        }
        Ok(Some(to_approved(&records.clients[index])))
    }

    fn list(&self) -> Result<Vec<ApprovedClient>, PairingStoreError> {
        let _transaction = acquire_transaction_lock(&self.path)?;
        Ok(read_records(&self.path, self.sync_parent)?
            .clients
            .iter()
            .map(to_approved)
            .collect())
    }

    fn revoke(&self, id: &ClientId) -> Result<bool, PairingStoreError> {
        let _transaction = acquire_transaction_lock(&self.path)?;
        let mut records = read_records(&self.path, self.sync_parent)?;
        let original_len = records.clients.len();
        records.clients.retain(|client| client.id != *id);
        if records.clients.len() == original_len {
            return Ok(false);
        }
        write_records(&self.path, &records, self.sync_parent)?;
        Ok(true)
    }
}

fn to_approved(client: &StoredClient) -> ApprovedClient {
    ApprovedClient {
        id: client.id.clone(),
        label: client.label.clone(),
        binding: client.binding.clone(),
        created_at: client.created_at,
        last_used_at: client.last_used_at,
    }
}

fn validate_label(label: &str) -> Result<(), PairingStoreError> {
    if !(1..=80).contains(&label.chars().count()) {
        return Err(PairingStoreError::InvalidLabel);
    }
    Ok(())
}

fn canonicalize_binding(binding: &ClientBinding) -> Result<ClientBinding, PairingStoreError> {
    binding
        .canonicalized()
        .ok_or(PairingStoreError::InvalidBinding)
}

struct TransactionLock {
    _file: fs::File,
}

fn acquire_transaction_lock(path: &Path) -> Result<TransactionLock, PairingStoreError> {
    let parent = ensure_private_parent(path)?;
    let lock_path = transaction_lock_path(path)?;
    for _ in 0..LOCK_OPEN_ATTEMPTS {
        let (file, created) = match fs::symlink_metadata(&lock_path) {
            Ok(metadata) => {
                validate_lock_metadata(&metadata)?;
                (open_existing_lock(&lock_path)?, false)
            }
            Err(error) if error.kind() == io::ErrorKind::NotFound => {
                match create_lock(&lock_path) {
                    Ok(file) => (file, true),
                    Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
                    Err(_) => return Err(PairingStoreError::Storage),
                }
            }
            Err(_) => return Err(PairingStoreError::Storage),
        };

        if created {
            file.set_permissions(fs::Permissions::from_mode(PRIVATE_FILE_MODE))
                .map_err(|_| PairingStoreError::Storage)?;
        }
        let opened = file.metadata().map_err(|_| PairingStoreError::Storage)?;
        validate_lock_metadata(&opened)?;
        if created {
            file.sync_all().map_err(|_| PairingStoreError::Storage)?;
            sync_directory(parent).map_err(|_| PairingStoreError::Storage)?;
        }
        lock_exclusive(&file)?;

        let current = fs::symlink_metadata(&lock_path).map_err(|_| PairingStoreError::Storage)?;
        validate_lock_metadata(&current)?;
        if !same_file(&opened, &current) {
            return Err(PairingStoreError::Storage);
        }
        return Ok(TransactionLock { _file: file });
    }
    Err(PairingStoreError::Storage)
}

fn transaction_lock_path(path: &Path) -> Result<PathBuf, PairingStoreError> {
    let parent = path.parent().ok_or(PairingStoreError::Storage)?;
    let mut filename = path
        .file_name()
        .ok_or(PairingStoreError::Storage)?
        .to_os_string();
    filename.push(".lock");
    Ok(parent.join(filename))
}

fn create_lock(path: &Path) -> io::Result<fs::File> {
    let mut options = fs::OpenOptions::new();
    options
        .read(true)
        .write(true)
        .create_new(true)
        .mode(PRIVATE_FILE_MODE)
        .custom_flags(rustix::fs::OFlags::NOFOLLOW.bits() as i32);
    options.open(path)
}

fn open_existing_lock(path: &Path) -> Result<fs::File, PairingStoreError> {
    let mut options = fs::OpenOptions::new();
    options
        .read(true)
        .write(true)
        .custom_flags(rustix::fs::OFlags::NOFOLLOW.bits() as i32);
    options.open(path).map_err(|_| PairingStoreError::Storage)
}

fn validate_lock_metadata(metadata: &fs::Metadata) -> Result<(), PairingStoreError> {
    if metadata.file_type().is_symlink()
        || !metadata.is_file()
        || metadata.len() != 0
        || metadata.permissions().mode() & 0o777 != PRIVATE_FILE_MODE
        || metadata.nlink() != 1
    {
        return Err(PairingStoreError::Storage);
    }
    Ok(())
}

fn same_file(left: &fs::Metadata, right: &fs::Metadata) -> bool {
    left.dev() == right.dev() && left.ino() == right.ino()
}

fn lock_exclusive(file: &fs::File) -> Result<(), PairingStoreError> {
    loop {
        match rustix::fs::flock(file, rustix::fs::FlockOperation::LockExclusive) {
            Ok(()) => return Ok(()),
            Err(error) if error == rustix::io::Errno::INTR => continue,
            Err(_) => return Err(PairingStoreError::Storage),
        }
    }
}

fn read_records(
    path: &Path,
    sync_parent: fn(&Path) -> io::Result<()>,
) -> Result<ClientRecords, PairingStoreError> {
    ensure_private_parent(path)?;
    let metadata = match fs::symlink_metadata(path) {
        Ok(metadata) => metadata,
        Err(error) if error.kind() == io::ErrorKind::NotFound => {
            return Ok(ClientRecords::default());
        }
        Err(_) => return Err(PairingStoreError::Storage),
    };
    validate_file_metadata(&metadata)?;

    let mut options = fs::OpenOptions::new();
    options
        .read(true)
        .custom_flags(rustix::fs::OFlags::NOFOLLOW.bits() as i32);
    let file = options.open(path).map_err(|_| PairingStoreError::Storage)?;
    validate_file_metadata(&file.metadata().map_err(|_| PairingStoreError::Storage)?)?;
    let mut bytes = Vec::with_capacity(metadata.len() as usize);
    file.take(MAX_CLIENTS_FILE_BYTES + 1)
        .read_to_end(&mut bytes)
        .map_err(|_| PairingStoreError::Storage)?;
    if bytes.len() as u64 > MAX_CLIENTS_FILE_BYTES {
        return Err(PairingStoreError::Corrupt);
    }
    let header: SchemaHeader =
        serde_json::from_slice(&bytes).map_err(|_| PairingStoreError::Corrupt)?;
    let (records, migrated) = match header.schema_version {
        1 => {
            let legacy: LegacyClientRecords =
                serde_json::from_slice(&bytes).map_err(|_| PairingStoreError::Corrupt)?;
            if legacy.schema_version != 1 {
                return Err(PairingStoreError::Corrupt);
            }
            validate_legacy_clients(&legacy.clients)?;
            (ClientRecords::default(), true)
        }
        2 => {
            let legacy: SchemaTwoClientRecords =
                serde_json::from_slice(&bytes).map_err(|_| PairingStoreError::Corrupt)?;
            if legacy.schema_version != 2 {
                return Err(PairingStoreError::Corrupt);
            }
            validate_schema_two_clients(&legacy.clients)?;
            (migrate_schema_two(legacy), true)
        }
        SCHEMA_VERSION => (
            serde_json::from_slice(&bytes).map_err(|_| PairingStoreError::Corrupt)?,
            false,
        ),
        _ => return Err(PairingStoreError::Corrupt),
    };
    validate_records(&records)?;
    if migrated {
        write_records(path, &records, sync_parent)?;
    }
    Ok(records)
}

fn migrate_schema_two(records: SchemaTwoClientRecords) -> ClientRecords {
    ClientRecords {
        schema_version: SCHEMA_VERSION,
        clients: records
            .clients
            .into_iter()
            .filter_map(|client| {
                let SchemaTwoClientBinding::IrohEndpoint { endpoint_id } = client.binding else {
                    return None;
                };
                let parsed_endpoint = iroh::EndpointId::from_str(&endpoint_id).ok()?;
                if parsed_endpoint.to_string() != endpoint_id {
                    return None;
                }
                Some(StoredClient {
                    id: client.id,
                    label: client.label,
                    credential_hmac: client.credential_hmac,
                    binding: ClientBinding::IrohEndpoint { endpoint_id },
                    created_at: client.created_at,
                    last_used_at: client.last_used_at,
                })
            })
            .collect(),
    }
}

fn validate_file_metadata(metadata: &fs::Metadata) -> Result<(), PairingStoreError> {
    if metadata.file_type().is_symlink()
        || !metadata.is_file()
        || metadata.len() > MAX_CLIENTS_FILE_BYTES
        || metadata.permissions().mode() & 0o777 != PRIVATE_FILE_MODE
    {
        return Err(PairingStoreError::Corrupt);
    }
    Ok(())
}

fn validate_records(records: &ClientRecords) -> Result<(), PairingStoreError> {
    if records.schema_version != SCHEMA_VERSION || records.clients.len() > MAX_CLIENTS {
        return Err(PairingStoreError::Corrupt);
    }
    let mut ids = HashSet::with_capacity(records.clients.len());
    let mut hmacs = Vec::with_capacity(records.clients.len());
    for client in &records.clients {
        validate_label(&client.label).map_err(|_| PairingStoreError::Corrupt)?;
        if !ids.insert(client.id.as_str()) {
            return Err(PairingStoreError::Duplicate);
        }
        let decoded = decode_hmac(&client.credential_hmac)?;
        if canonicalize_binding(&client.binding).map_err(|_| PairingStoreError::Corrupt)?
            != client.binding
        {
            return Err(PairingStoreError::Corrupt);
        }
        if hmacs
            .iter()
            .any(|existing: &[u8; 32]| bool::from(existing.ct_eq(&decoded)))
        {
            return Err(PairingStoreError::Duplicate);
        }
        hmacs.push(decoded);
    }
    Ok(())
}

fn validate_legacy_clients(clients: &[LegacyStoredClient]) -> Result<(), PairingStoreError> {
    validate_client_fields(clients.iter().map(|client| {
        (
            &client.id,
            client.label.as_str(),
            client.credential_hmac.as_str(),
        )
    }))
}

fn validate_schema_two_clients(clients: &[SchemaTwoStoredClient]) -> Result<(), PairingStoreError> {
    validate_client_fields(clients.iter().map(|client| {
        (
            &client.id,
            client.label.as_str(),
            client.credential_hmac.as_str(),
        )
    }))
}

fn validate_client_fields<'a>(
    clients: impl Iterator<Item = (&'a ClientId, &'a str, &'a str)>,
) -> Result<(), PairingStoreError> {
    let mut ids = HashSet::new();
    let mut hmacs = Vec::new();
    for (id, label, credential_hmac) in clients {
        validate_label(label).map_err(|_| PairingStoreError::Corrupt)?;
        if !ids.insert(id.as_str()) {
            return Err(PairingStoreError::Duplicate);
        }
        let decoded = decode_hmac(credential_hmac)?;
        if hmacs
            .iter()
            .any(|existing: &[u8; 32]| bool::from(existing.ct_eq(&decoded)))
        {
            return Err(PairingStoreError::Duplicate);
        }
        hmacs.push(decoded);
        if ids.len() > MAX_CLIENTS {
            return Err(PairingStoreError::Corrupt);
        }
    }
    Ok(())
}

fn decode_hmac(encoded: &str) -> Result<[u8; 32], PairingStoreError> {
    let decoded = BASE64URL_NOPAD
        .decode(encoded.as_bytes())
        .map_err(|_| PairingStoreError::Corrupt)?;
    if BASE64URL_NOPAD.encode(&decoded) != encoded {
        return Err(PairingStoreError::Corrupt);
    }
    decoded.try_into().map_err(|_| PairingStoreError::Corrupt)
}

fn write_records(
    path: &Path,
    records: &ClientRecords,
    sync_parent: fn(&Path) -> io::Result<()>,
) -> Result<(), PairingStoreError> {
    validate_records(records)?;
    let serialized = serde_json::to_vec(records).map_err(|_| PairingStoreError::Storage)?;
    if serialized.len() as u64 > MAX_CLIENTS_FILE_BYTES {
        return Err(PairingStoreError::ClientLimit);
    }
    let parent = ensure_private_parent(path)?;
    match fs::symlink_metadata(path) {
        Ok(metadata) => validate_file_metadata(&metadata)?,
        Err(error) if error.kind() == io::ErrorKind::NotFound => {}
        Err(_) => return Err(PairingStoreError::Storage),
    }

    let temporary = create_temporary_path(parent)?;
    let mut options = fs::OpenOptions::new();
    options.write(true).create_new(true).mode(PRIVATE_FILE_MODE);
    let mut file = options
        .open(&temporary)
        .map_err(|_| PairingStoreError::Storage)?;
    let prepare_result = (|| -> Result<(), PairingStoreError> {
        file.set_permissions(fs::Permissions::from_mode(PRIVATE_FILE_MODE))
            .map_err(|_| PairingStoreError::Storage)?;
        file.write_all(&serialized)
            .map_err(|_| PairingStoreError::Storage)?;
        file.sync_all().map_err(|_| PairingStoreError::Storage)?;
        match fs::symlink_metadata(path) {
            Ok(metadata) => validate_file_metadata(&metadata)?,
            Err(error) if error.kind() == io::ErrorKind::NotFound => {}
            Err(_) => return Err(PairingStoreError::Storage),
        }
        Ok(())
    })();
    if let Err(error) = prepare_result {
        let _ = fs::remove_file(&temporary);
        return Err(error);
    }
    if fs::rename(&temporary, path).is_err() {
        let _ = fs::remove_file(&temporary);
        return Err(PairingStoreError::Storage);
    }
    sync_parent(parent).map_err(|_| PairingStoreError::CommitUncertain)
}

fn create_temporary_path(parent: &Path) -> Result<PathBuf, PairingStoreError> {
    for _ in 0..8 {
        let path = parent.join(format!(".clients-{}.tmp", Uuid::new_v4()));
        if !path.exists() {
            return Ok(path);
        }
    }
    Err(PairingStoreError::Storage)
}