IronStorage/crates/storage/src/document.rs

//! Lossless, presentation-neutral documents for arbitrary pass entries.

use std::{error::Error, fmt, ops::Range};

use sha2::{Digest as _, Sha256};

use crate::{
    command::EditRequest,
    crypto::{KeyStore, SecretProvider},
    otp::{OtpAlgorithm, OtpKind, OtpUri},
    recipient::SigningPolicy,
    repository::{EntryPath, Repository, SecretBytes},
    write::{EditSession, EntryCommitter, VaultWriter, WriteError, WriteOutcome},
};

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct EntryFieldId(u64);

impl EntryFieldId {
    pub fn value(self) -> u64 {
        self.0
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EntryFieldKind {
    Password,
    Username,
    Email,
    Url,
    OtpUri,
    Field,
    Note,
    Blank,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EntrySensitivity {
    Sensitive,
    Ordinary,
    Empty,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct EntryFieldMetadata {
    kind: EntryFieldKind,
    sensitivity: EntrySensitivity,
    name: Option<String>,
    otp: Option<EntryOtpMetadata>,
    value: Range<usize>,
}

impl EntryFieldMetadata {
    pub fn kind(&self) -> EntryFieldKind {
        self.kind
    }

    pub fn sensitivity(&self) -> EntrySensitivity {
        self.sensitivity
    }

    pub fn name(&self) -> Option<&str> {
        self.name.as_deref()
    }

    pub fn otp(&self) -> Option<&EntryOtpMetadata> {
        self.otp.as_ref()
    }
}

/// Non-secret presentation metadata parsed from a validated `otpauth` URI.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct EntryOtpMetadata {
    kind: OtpKind,
    issuer: Option<String>,
    account: String,
    algorithm: OtpAlgorithm,
    digits: u32,
    period: Option<u64>,
    counter: Option<u64>,
}

impl EntryOtpMetadata {
    pub fn kind(&self) -> OtpKind {
        self.kind
    }

    pub fn issuer(&self) -> Option<&str> {
        self.issuer.as_deref()
    }

    pub fn account(&self) -> &str {
        &self.account
    }

    pub fn algorithm(&self) -> OtpAlgorithm {
        self.algorithm
    }

    pub fn digits(&self) -> u32 {
        self.digits
    }

    pub fn period(&self) -> Option<u64> {
        self.period
    }

    pub fn counter(&self) -> Option<u64> {
        self.counter
    }
}

pub struct EntryField {
    id: EntryFieldId,
    contents: SecretBytes,
    ending: Vec<u8>,
    metadata: EntryFieldMetadata,
}

impl EntryField {
    pub fn id(&self) -> EntryFieldId {
        self.id
    }

    pub fn contents(&self) -> &SecretBytes {
        &self.contents
    }

    pub fn value(&self) -> &[u8] {
        &self.contents.expose()[self.metadata.value.clone()]
    }

    pub fn metadata(&self) -> &EntryFieldMetadata {
        &self.metadata
    }
}

impl fmt::Debug for EntryField {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("EntryField")
            .field("id", &self.id)
            .field("metadata", &self.metadata)
            .field("contents", &"[REDACTED]")
            .finish()
    }
}

pub struct EntryFieldDraft {
    contents: SecretBytes,
}

impl EntryFieldDraft {
    pub fn line(contents: Vec<u8>) -> Result<Self, DocumentError> {
        validate_line(&contents)?;
        Ok(Self {
            contents: SecretBytes::new(contents),
        })
    }

    pub fn field(name: impl Into<String>, value: Vec<u8>) -> Result<Self, DocumentError> {
        let name = name.into();
        validate_name(&name)?;
        validate_line(&value)?;
        let value = SecretBytes::new(value);
        let mut contents = Vec::with_capacity(name.len() + value.expose().len() + 2);
        contents.extend_from_slice(name.as_bytes());
        contents.extend_from_slice(b": ");
        contents.extend_from_slice(value.expose());
        Ok(Self {
            contents: SecretBytes::new(contents),
        })
    }

    pub fn otp_uri(uri: Vec<u8>) -> Result<Self, DocumentError> {
        validate_line(&uri)?;
        let encoded = SecretBytes::new(uri);
        OtpUri::parse(SecretBytes::new(encoded.expose().to_vec()))
            .map_err(|_| DocumentError::InvalidOtpUri)?;
        Ok(Self { contents: encoded })
    }

    pub fn blank() -> Self {
        Self {
            contents: SecretBytes::new(Vec::new()),
        }
    }

    fn render(self) -> SecretBytes {
        self.contents
    }
}

impl fmt::Debug for EntryFieldDraft {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("EntryFieldDraft([REDACTED])")
    }
}

#[derive(Clone, Copy, Eq, Hash, PartialEq)]
pub struct DocumentConflictToken([u8; 32]);

impl DocumentConflictToken {
    pub fn as_bytes(&self) -> &[u8; 32] {
        &self.0
    }
}

impl fmt::Debug for DocumentConflictToken {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("DocumentConflictToken([REDACTED])")
    }
}

pub struct EntryDocument {
    path: EntryPath,
    fields: Vec<EntryField>,
    next_id: u64,
    conflict_token: DocumentConflictToken,
    session: EditSession,
}

impl EntryDocument {
    fn from_session(session: EditSession) -> Self {
        let path = session.path().clone();
        let conflict_token = conflict_token(&session);
        let mut fields = parse_lines(session.plaintext().expose());
        classify_all(&mut fields);
        let next_id = fields.len() as u64;
        Self {
            path,
            fields,
            next_id,
            conflict_token,
            session,
        }
    }

    pub fn path(&self) -> &EntryPath {
        &self.path
    }

    pub fn fields(&self) -> &[EntryField] {
        &self.fields
    }

    pub fn field(&self, id: EntryFieldId) -> Option<&EntryField> {
        self.fields.iter().find(|field| field.id == id)
    }

    pub fn password(&self) -> Option<&EntryField> {
        self.fields.first()
    }

    /// Copy one structured field value into an independently zeroizing buffer.
    ///
    /// Frontends use this for explicit presentation actions without reparsing
    /// the pass entry or copying label syntax alongside the selected value.
    pub fn copy_field_value(&self, id: EntryFieldId) -> Result<SecretBytes, DocumentError> {
        self.field(id)
            .map(|field| SecretBytes::new(field.value().to_vec()))
            .ok_or(DocumentError::UnknownField { id })
    }

    /// Replace only a structured field's value while preserving its storage-
    /// supplied name and syntax.
    pub fn replace_field_value(
        &mut self,
        id: EntryFieldId,
        value: Vec<u8>,
    ) -> Result<(), DocumentError> {
        let field = self.field(id).ok_or(DocumentError::UnknownField { id })?;
        let draft = match field.metadata().kind() {
            EntryFieldKind::OtpUri => EntryFieldDraft::otp_uri(value)?,
            EntryFieldKind::Password | EntryFieldKind::Note | EntryFieldKind::Blank => {
                EntryFieldDraft::line(value)?
            }
            EntryFieldKind::Username
            | EntryFieldKind::Email
            | EntryFieldKind::Url
            | EntryFieldKind::Field => EntryFieldDraft::field(
                field
                    .metadata()
                    .name()
                    .ok_or(DocumentError::InvalidFieldName)?,
                value,
            )?,
        };
        self.update(id, draft)
    }

    pub fn conflict_token(&self) -> DocumentConflictToken {
        self.conflict_token
    }

    pub fn add(
        &mut self,
        index: usize,
        draft: EntryFieldDraft,
    ) -> Result<EntryFieldId, DocumentError> {
        if index > self.fields.len() {
            return Err(DocumentError::InvalidIndex { index });
        }
        let final_newline = self.has_final_newline();
        let default_ending = self.default_ending();
        let id = EntryFieldId(self.next_id);
        self.next_id = self
            .next_id
            .checked_add(1)
            .ok_or(DocumentError::FieldIdExhausted)?;
        self.fields.insert(
            index,
            EntryField {
                id,
                contents: draft.render(),
                ending: default_ending.clone(),
                metadata: blank_metadata(),
            },
        );
        normalize_endings(&mut self.fields, &default_ending, final_newline);
        classify_all(&mut self.fields);
        Ok(id)
    }

    pub fn update(
        &mut self,
        id: EntryFieldId,
        draft: EntryFieldDraft,
    ) -> Result<(), DocumentError> {
        let field = self
            .fields
            .iter_mut()
            .find(|field| field.id == id)
            .ok_or(DocumentError::UnknownField { id })?;
        field.contents = draft.render();
        classify_all(&mut self.fields);
        Ok(())
    }

    pub fn remove(&mut self, id: EntryFieldId) -> Result<EntryField, DocumentError> {
        let index = self
            .fields
            .iter()
            .position(|field| field.id == id)
            .ok_or(DocumentError::UnknownField { id })?;
        let final_newline = self.has_final_newline();
        let default_ending = self.default_ending();
        let removed = self.fields.remove(index);
        normalize_endings(&mut self.fields, &default_ending, final_newline);
        classify_all(&mut self.fields);
        Ok(removed)
    }

    pub fn reorder(&mut self, id: EntryFieldId, index: usize) -> Result<(), DocumentError> {
        if index >= self.fields.len() {
            return Err(DocumentError::InvalidIndex { index });
        }
        let current = self
            .fields
            .iter()
            .position(|field| field.id == id)
            .ok_or(DocumentError::UnknownField { id })?;
        if current == index {
            return Ok(());
        }
        let final_newline = self.has_final_newline();
        let default_ending = self.default_ending();
        let field = self.fields.remove(current);
        self.fields.insert(index, field);
        normalize_endings(&mut self.fields, &default_ending, final_newline);
        classify_all(&mut self.fields);
        Ok(())
    }

    pub fn serialize(&self) -> SecretBytes {
        let capacity = self
            .fields
            .iter()
            .map(|field| field.contents.expose().len() + field.ending.len())
            .sum();
        let mut output = Vec::with_capacity(capacity);
        for field in &self.fields {
            output.extend_from_slice(field.contents.expose());
            output.extend_from_slice(&field.ending);
        }
        SecretBytes::new(output)
    }

    fn has_final_newline(&self) -> bool {
        self.fields
            .last()
            .is_some_and(|field| !field.ending.is_empty())
    }

    fn default_ending(&self) -> Vec<u8> {
        self.fields
            .iter()
            .find(|field| !field.ending.is_empty())
            .map(|field| field.ending.clone())
            .unwrap_or_else(|| b"\n".to_vec())
    }
}

impl fmt::Debug for EntryDocument {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("EntryDocument")
            .field("path", &self.path)
            .field("field_count", &self.fields.len())
            .field("conflict_token", &self.conflict_token)
            .finish_non_exhaustive()
    }
}

pub struct EntryDocumentService<'a> {
    repository: &'a Repository,
    keys: &'a KeyStore,
}

impl<'a> EntryDocumentService<'a> {
    pub fn new(repository: &'a Repository, keys: &'a KeyStore) -> Self {
        Self { repository, keys }
    }

    pub fn open(
        &self,
        entry: &str,
        secrets: &mut impl SecretProvider,
    ) -> Result<EntryDocument, DocumentError> {
        let session = VaultWriter::new(self.repository, self.keys).begin_edit(
            &EditRequest {
                entry: entry.to_owned(),
            },
            secrets,
        )?;
        Ok(EntryDocument::from_session(session))
    }

    pub fn save(
        &self,
        document: EntryDocument,
        signing: Option<&SigningPolicy>,
        committer: &mut impl EntryCommitter,
    ) -> Result<WriteOutcome, DocumentError> {
        let replacement = document.serialize();
        Ok(VaultWriter::new(self.repository, self.keys).finish_edit(
            document.session,
            replacement,
            "IronStorage structured editor",
            signing,
            committer,
        )?)
    }

    /// Save without consuming the document, allowing an in-process editor to
    /// present a conflict or encryption failure and retain the user's draft.
    pub fn save_recoverable(
        &self,
        document: &EntryDocument,
        signing: Option<&SigningPolicy>,
        committer: &mut impl EntryCommitter,
    ) -> Result<WriteOutcome, DocumentError> {
        let replacement = document.serialize();
        Ok(
            VaultWriter::new(self.repository, self.keys).finish_edit_recoverable(
                &document.session,
                replacement,
                "IronStorage structured editor",
                signing,
                committer,
            )?,
        )
    }
}

#[derive(Debug)]
pub enum DocumentError {
    Write(WriteError),
    InvalidIndex { index: usize },
    UnknownField { id: EntryFieldId },
    LineBreak,
    InvalidFieldName,
    InvalidOtpUri,
    FieldIdExhausted,
}

impl fmt::Display for DocumentError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Write(error) => error.fmt(formatter),
            Self::InvalidIndex { index } => write!(formatter, "invalid entry field index: {index}"),
            Self::UnknownField { id } => write!(formatter, "unknown entry field: {}", id.value()),
            Self::LineBreak => formatter.write_str("an entry field may not contain a line break"),
            Self::InvalidFieldName => formatter.write_str("entry field name is invalid"),
            Self::InvalidOtpUri => formatter.write_str("entry OTP URI is invalid"),
            Self::FieldIdExhausted => formatter.write_str("entry field identifiers are exhausted"),
        }
    }
}

impl Error for DocumentError {}

impl From<WriteError> for DocumentError {
    fn from(error: WriteError) -> Self {
        Self::Write(error)
    }
}

fn parse_lines(contents: &[u8]) -> Vec<EntryField> {
    let mut fields = Vec::new();
    let mut start = 0;
    while start < contents.len() {
        let newline = contents[start..]
            .iter()
            .position(|byte| *byte == b'\n')
            .map(|offset| start + offset);
        let (content_end, end, ending) = match newline {
            Some(newline) if newline > start && contents[newline - 1] == b'\r' => {
                (newline - 1, newline + 1, b"\r\n".to_vec())
            }
            Some(newline) => (newline, newline + 1, b"\n".to_vec()),
            None => (contents.len(), contents.len(), Vec::new()),
        };
        fields.push(EntryField {
            id: EntryFieldId(fields.len() as u64),
            contents: SecretBytes::new(contents[start..content_end].to_vec()),
            ending,
            metadata: blank_metadata(),
        });
        start = end;
    }
    fields
}

fn classify_all(fields: &mut [EntryField]) {
    for (index, field) in fields.iter_mut().enumerate() {
        field.metadata = classify(index, field.contents.expose());
    }
}

fn classify(index: usize, line: &[u8]) -> EntryFieldMetadata {
    if index == 0 {
        return EntryFieldMetadata {
            kind: EntryFieldKind::Password,
            sensitivity: EntrySensitivity::Sensitive,
            name: Some("password".to_owned()),
            otp: None,
            value: 0..line.len(),
        };
    }
    if line.is_empty() {
        return blank_metadata();
    }
    if line.starts_with(b"otpauth://")
        && let Ok(uri) = OtpUri::parse(SecretBytes::new(line.to_vec()))
    {
        let otp = EntryOtpMetadata {
            kind: uri.kind(),
            issuer: uri.issuer().map(str::to_owned),
            account: uri.account().to_owned(),
            algorithm: uri.algorithm(),
            digits: uri.digits(),
            period: uri.period(),
            counter: uri.counter(),
        };
        return EntryFieldMetadata {
            kind: EntryFieldKind::OtpUri,
            sensitivity: EntrySensitivity::Sensitive,
            name: Some("otp".to_owned()),
            otp: Some(otp),
            value: 0..line.len(),
        };
    }
    if let Some(colon) = line.iter().position(|byte| *byte == b':') {
        let raw_name = trim_ascii(&line[..colon]);
        if !raw_name.is_empty()
            && let Ok(name) = std::str::from_utf8(raw_name)
        {
            let value_start = colon + 1 + usize::from(line.get(colon + 1) == Some(&b' '));
            let kind = semantic_field_kind(name);
            return EntryFieldMetadata {
                kind,
                sensitivity: field_sensitivity(name, kind),
                name: Some(name.to_owned()),
                otp: None,
                value: value_start..line.len(),
            };
        }
    }
    EntryFieldMetadata {
        kind: EntryFieldKind::Note,
        sensitivity: EntrySensitivity::Sensitive,
        name: None,
        otp: None,
        value: 0..line.len(),
    }
}

fn semantic_field_kind(name: &str) -> EntryFieldKind {
    match name.trim().to_ascii_lowercase().as_str() {
        "user" | "username" | "login" => EntryFieldKind::Username,
        "email" | "e-mail" => EntryFieldKind::Email,
        "url" | "uri" | "website" => EntryFieldKind::Url,
        _ => EntryFieldKind::Field,
    }
}

fn field_sensitivity(name: &str, kind: EntryFieldKind) -> EntrySensitivity {
    if matches!(
        kind,
        EntryFieldKind::Username | EntryFieldKind::Email | EntryFieldKind::Url
    ) {
        return EntrySensitivity::Ordinary;
    }
    match name.trim().to_ascii_lowercase().as_str() {
        "title" | "site" | "host" => EntrySensitivity::Ordinary,
        _ => EntrySensitivity::Sensitive,
    }
}

fn blank_metadata() -> EntryFieldMetadata {
    EntryFieldMetadata {
        kind: EntryFieldKind::Blank,
        sensitivity: EntrySensitivity::Empty,
        name: None,
        otp: None,
        value: 0..0,
    }
}

fn trim_ascii(mut value: &[u8]) -> &[u8] {
    while value.first().is_some_and(u8::is_ascii_whitespace) {
        value = &value[1..];
    }
    while value.last().is_some_and(u8::is_ascii_whitespace) {
        value = &value[..value.len() - 1];
    }
    value
}

fn validate_line(value: &[u8]) -> Result<(), DocumentError> {
    if value.iter().any(|byte| matches!(byte, b'\r' | b'\n')) {
        Err(DocumentError::LineBreak)
    } else {
        Ok(())
    }
}

fn validate_name(name: &str) -> Result<(), DocumentError> {
    if name.trim().is_empty() || name.contains([':', '\r', '\n']) || name.trim() != name {
        Err(DocumentError::InvalidFieldName)
    } else {
        Ok(())
    }
}

fn normalize_endings(fields: &mut [EntryField], default: &[u8], final_newline: bool) {
    let last = fields.len().saturating_sub(1);
    for (index, field) in fields.iter_mut().enumerate() {
        if index < last && field.ending.is_empty() {
            field.ending = default.to_vec();
        } else if index == last {
            if final_newline && field.ending.is_empty() {
                field.ending = default.to_vec();
            } else if !final_newline {
                field.ending.clear();
            }
        }
    }
}

fn conflict_token(session: &EditSession) -> DocumentConflictToken {
    let mut digest = Sha256::new();
    digest.update(b"ironstorage-entry-document-v1\0");
    digest.update(session.path().to_string().as_bytes());
    digest.update(b"\0");
    match session.original_ciphertext() {
        Some(ciphertext) => {
            digest.update(b"present\0");
            digest.update(ciphertext.as_bytes());
        }
        None => digest.update(b"missing\0"),
    }
    DocumentConflictToken(digest.finalize().into())
}