affine_doc_loader 0.1.9

use std::collections::{BTreeSet, HashSet, VecDeque};

use thiserror::Error;
use y_octo::{Any, Doc, Map, ReadDoc, ReadError, Value};

use crate::{BlobRef, get_blob_refs_from_binary};

pub const MAX_SOURCE_BINARY_BYTES: usize = 16 * 1024 * 1024;
pub const MAX_SOURCE_UPDATES: i64 = 10_000;
const MAX_SOURCE_BLOCKS: usize = 100_000;
const MAX_CHILDREN_PER_BLOCK: u64 = 10_000;
const MAX_SOURCE_REFS: usize = 10_000;
const MAX_BLOB_KEY_BYTES: usize = 1_024;
const MAX_WORKSPACE_DOCS: u64 = 100_000;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BlobRefParser {
  Immutable,
  MutableFallback,
}

impl BlobRefParser {
  pub const fn label(self) -> &'static str {
    match self {
      Self::Immutable => "immutable",
      Self::MutableFallback => "mutable_fallback",
    }
  }
}

#[derive(Debug, Eq, PartialEq)]
pub struct BlobRefExtraction {
  pub refs: Vec<BlobRef>,
  pub parser: BlobRefParser,
}

#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
pub enum BlobRefProjectionError {
  #[error("invalid binary")]
  InvalidBinary,
  #[error("client clock gap")]
  ClientClockGap,
  #[error("pending dependency")]
  PendingDependency,
  #[error("document parser does not support the input")]
  Unsupported,
  #[error("blob ref source is too large")]
  SourceTooLarge,
  #[error("workspace root is too large")]
  WorkspaceRootTooLarge,
  #[error("workspace root cannot be projected")]
  WorkspaceRootInvalid,
  #[error("workspace root contains too many documents")]
  WorkspaceRootDocCountTooLarge,
  #[error("blob ref tree is too large")]
  TreeTooLarge,
  #[error("blob ref tree root is invalid")]
  TreeRootInvalid,
  #[error("blob ref tree fanout is too large")]
  TreeFanoutTooLarge,
  #[error("blob ref tree child is invalid")]
  TreeChildInvalid,
  #[error("blob ref key is too large")]
  KeyTooLarge,
  #[error("blob ref count is too large")]
  RefCountTooLarge,
}

pub fn extract_blob_refs(blob: Vec<u8>) -> Result<BlobRefExtraction, BlobRefProjectionError> {
  if blob.len() > MAX_SOURCE_BINARY_BYTES {
    return Err(BlobRefProjectionError::SourceTooLarge);
  }
  let (doc, parser) = decode_blob_doc(&blob)?;
  if let Ok(refs) = get_blob_refs_from_binary(blob) {
    validate_extracted_refs(&doc, &refs)?;
    return Ok(BlobRefExtraction { refs, parser });
  }
  let refs = mutable_blob_refs(&doc).map_err(|()| BlobRefProjectionError::Unsupported)?;
  validate_extracted_refs(&doc, &refs)?;
  Ok(BlobRefExtraction {
    refs,
    parser: BlobRefParser::MutableFallback,
  })
}

pub fn page_workspace_doc_ids(
  blob: Vec<u8>,
  workspace_id: &str,
  cursor: Option<&str>,
  include_cursor: bool,
  limit: usize,
) -> Result<Vec<String>, BlobRefProjectionError> {
  if blob.len() > MAX_SOURCE_BINARY_BYTES {
    return Err(BlobRefProjectionError::WorkspaceRootTooLarge);
  }
  let (doc, _) = decode_blob_doc(&blob)?;
  let meta = doc
    .get_map("meta")
    .map_err(|_| BlobRefProjectionError::WorkspaceRootInvalid)?;
  let pages = meta
    .get("pages")
    .and_then(|value| value.to_array())
    .ok_or(BlobRefProjectionError::WorkspaceRootInvalid)?;
  validate_workspace_doc_count(pages.len())?;
  let cursor = cursor.unwrap_or("");
  let mut ids = BTreeSet::new();
  let mut admit = |id: String| {
    if (include_cursor && id.as_str() >= cursor) || (!include_cursor && id.as_str() > cursor) {
      ids.insert(id);
      if ids.len() > limit {
        ids.pop_last();
      }
    }
  };
  admit(workspace_id.to_string());
  for page in pages.iter() {
    let Some(page) = page.to_map() else {
      continue;
    };
    if let Some(id) = string_value(&page, "id") {
      admit(id);
    }
  }
  Ok(ids.into_iter().collect())
}

pub fn live_blob_keys(
  blob: Vec<u8>,
  workspace_root: bool,
) -> Result<(Vec<Box<str>>, BlobRefParser), BlobRefProjectionError> {
  if blob.len() > MAX_SOURCE_BINARY_BYTES {
    return Err(BlobRefProjectionError::SourceTooLarge);
  }
  let (doc, parser) = decode_blob_doc(&blob)?;

  let mut keys = BTreeSet::new();
  if workspace_root
    && let Ok(meta) = doc.get_map("meta")
    && let Some(avatar) = string_value(&meta, "avatar")
    && !avatar.is_empty()
  {
    insert_key(&mut keys, avatar)?;
  }

  let Ok(blocks) = doc.get_map("blocks") else {
    return Ok((keys.into_iter().map(String::into_boxed_str).collect(), parser));
  };
  validate_tree_size(blocks.len())?;

  let mut roots = Vec::new();
  for value in blocks.values() {
    let Some(block) = value.to_map() else {
      continue;
    };
    if string_value(&block, "sys:flavour").as_deref() == Some("affine:page") {
      roots.push(string_value(&block, "sys:id").ok_or(BlobRefProjectionError::TreeRootInvalid)?);
    }
  }
  if roots.len() != 1 {
    return Err(BlobRefProjectionError::TreeRootInvalid);
  }

  let root = roots.pop().expect("one page root");
  let mut queue = VecDeque::from([root.clone()]);
  let mut discovered = HashSet::from([root]);
  while let Some(block_id) = queue.pop_front() {
    let Some(block) = blocks.get(&block_id).and_then(|value| value.to_map()) else {
      continue;
    };
    if matches!(
      string_value(&block, "sys:flavour").as_deref(),
      Some("affine:attachment" | "affine:image")
    ) && let Some(key) = string_value(&block, "prop:sourceId")
    {
      insert_key(&mut keys, key)?;
    }
    if let Some(children) = block.get("sys:children").and_then(|value| value.to_array()) {
      validate_tree_fanout(children.len())?;
      for child in children.iter() {
        let Some(Any::String(child_id)) = child.to_any() else {
          return Err(BlobRefProjectionError::TreeChildInvalid);
        };
        if discovered.insert(child_id.clone()) {
          if discovered.len() > MAX_SOURCE_BLOCKS {
            return Err(BlobRefProjectionError::TreeTooLarge);
          }
          queue.push_back(child_id);
        }
      }
    }
  }
  Ok((keys.into_iter().map(String::into_boxed_str).collect(), parser))
}

fn decode_blob_doc(blob: &[u8]) -> Result<(Doc, BlobRefParser), BlobRefProjectionError> {
  match ReadDoc::from_full_update_v1(blob.to_vec()) {
    Ok(_) => {}
    Err(ReadError::IncompleteSnapshot("client clock gap")) => {
      return Err(BlobRefProjectionError::ClientClockGap);
    }
    Err(ReadError::IncompleteSnapshot(_)) => return Err(BlobRefProjectionError::PendingDependency),
    Err(ReadError::InvalidUpdate(_)) => return Err(BlobRefProjectionError::InvalidBinary),
    Err(ReadError::ResourceLimit(_)) => return Err(BlobRefProjectionError::Unsupported),
  }
  let mut doc = Doc::default();
  doc
    .apply_update_from_binary_v1(blob)
    .map_err(|_| BlobRefProjectionError::InvalidBinary)?;
  if doc.has_pending_updates() {
    return Err(BlobRefProjectionError::ClientClockGap);
  }
  Ok((doc, BlobRefParser::Immutable))
}

fn mutable_blob_refs(doc: &Doc) -> Result<Vec<BlobRef>, ()> {
  let blocks = doc.get_map("blocks").map_err(|_| ())?;
  Ok(blocks.values().filter_map(blob_ref).collect())
}

fn validate_extracted_refs(doc: &Doc, refs: &[BlobRef]) -> Result<(), BlobRefProjectionError> {
  let blocks = doc.get_map("blocks").map_err(|_| BlobRefProjectionError::Unsupported)?;
  validate_tree_size(blocks.len())?;
  if refs.len() > MAX_SOURCE_REFS {
    return Err(BlobRefProjectionError::RefCountTooLarge);
  }
  if refs
    .iter()
    .any(|reference| reference.blob_key.len() > MAX_BLOB_KEY_BYTES)
  {
    return Err(BlobRefProjectionError::KeyTooLarge);
  }
  Ok(())
}

fn validate_workspace_doc_count(count: u64) -> Result<(), BlobRefProjectionError> {
  if count > MAX_WORKSPACE_DOCS {
    return Err(BlobRefProjectionError::WorkspaceRootDocCountTooLarge);
  }
  Ok(())
}

fn validate_tree_size(count: u64) -> Result<(), BlobRefProjectionError> {
  if count > MAX_SOURCE_BLOCKS as u64 {
    return Err(BlobRefProjectionError::TreeTooLarge);
  }
  Ok(())
}

fn validate_tree_fanout(count: u64) -> Result<(), BlobRefProjectionError> {
  if count > MAX_CHILDREN_PER_BLOCK {
    return Err(BlobRefProjectionError::TreeFanoutTooLarge);
  }
  Ok(())
}

fn blob_ref(value: Value) -> Option<BlobRef> {
  let block = value.to_map()?;
  let flavour = string_value(&block, "sys:flavour")?;
  if !matches!(flavour.as_str(), "affine:attachment" | "affine:image") {
    return None;
  }
  Some(BlobRef {
    blob_key: string_value(&block, "prop:sourceId")?,
    block_id: string_value(&block, "sys:id")?,
    flavour,
  })
}

fn insert_key(keys: &mut BTreeSet<String>, key: String) -> Result<(), BlobRefProjectionError> {
  if key.len() > MAX_BLOB_KEY_BYTES {
    return Err(BlobRefProjectionError::KeyTooLarge);
  }
  keys.insert(key);
  if keys.len() > MAX_SOURCE_REFS {
    return Err(BlobRefProjectionError::RefCountTooLarge);
  }
  Ok(())
}

fn string_value(map: &Map, key: &str) -> Option<String> {
  match map.get(key)? {
    Value::Any(Any::String(value)) => Some(value),
    Value::Text(value) => Some(value.to_string()),
    _ => None,
  }
}

#[cfg(test)]
#[path = "tests/blob_refs/tests.rs"]
mod tests;