memona-plugin-aindrive/src/lib.rs

//! An aindrive filesystem provider for Memona. It speaks to the drive's
//! remote MCP endpoint through the host's HTTP, with the drive's access token.

mod address;
mod codec;
mod mcp;
mod sdk;
#[cfg(test)]
mod session_tests;
#[cfg(test)]
mod test_host;

use sdk::{
  Capabilities, ChangePage, Entry, EntryKind, ErrorCode, FileInfo, Guest,
  GuestReader, GuestWriter, ListPage, Locator, Reader, Setting, WriteCondition,
  Writer,
};
use serde::de::DeserializeOwned;
use serde_json::json;
use std::cell::{Cell, RefCell};
use std::ops::Range;

#[cfg(not(test))]
use sdk::host;
#[cfg(test)]
use test_host as host;

const MAX_CHUNK: u32 = 1024 * 1024;
/// The drive's device reads at most 8 MiB in one call and cuts the rest, so
/// a longer file is refused before any of it is read.
const MAX_READ: u64 = 8 * 1024 * 1024;
/// A whole 8 MiB file in base64, with room for the reply around it.
const MAX_REPLY: usize = 12 * 1024 * 1024;
/// The largest file a save sends. `write_file` takes a file whole, and its
/// base64 has to be encoded within one call's fuel, which a single transfer,
/// the most Memona saves a remote page in, leaves room for.
const MAX_WRITE: usize = MAX_CHUNK as usize;
/// Without these the folder cannot even be browsed.
const REQUIRED_TOOLS: [&str; 3] = ["list_files", "stat", "read_file"];
/// A token offers these only when it may write.
const WRITE_TOOLS: [&str; 2] = ["write_file", "delete_path"];

struct Session {
  endpoint: String,
  authorization: String,
}

/// Whether a request may change the drive. Once one has gone out whole, a
/// lost answer leaves the change unknown rather than failed.
#[derive(Clone, Copy, PartialEq)]
enum Effect {
  Read,
  Change,
}

impl Session {
  fn from_config(config: &[(String, Setting)]) -> Result<Self, ErrorCode> {
    let lookup = |wanted: &str| {
      config
        .iter()
        .find(|(key, _)| key == wanted)
        .and_then(|(_, value)| match value {
          Setting::Text(text) | Setting::Secret(text) => {
            let text = text.trim();
            (!text.is_empty()).then(|| text.to_owned())
          }
          Setting::Toggle(_) => None,
        })
    };
    let endpoint =
      address::endpoint(&lookup("url").ok_or(ErrorCode::InvalidInput)?)?;
    let token = lookup("token").ok_or(ErrorCode::InvalidInput)?;
    Ok(Self {
      endpoint,
      authorization: format!("Bearer {token}"),
    })
  }
}

thread_local! {
  static SESSION: RefCell<Option<Session>> = const { RefCell::new(None) };
}

struct Aindrive;

fn transport(error: host::HostError) -> ErrorCode {
  match error {
    host::HostError::Denied => ErrorCode::Denied,
    host::HostError::Cancelled => ErrorCode::Cancelled,
    host::HostError::InvalidInput => ErrorCode::InvalidInput,
    host::HostError::Internal => ErrorCode::Internal,
    host::HostError::Unavailable => ErrorCode::Unavailable,
  }
}

/// The drive-relative path a locator names; the drive root is empty.
fn route(target: &Locator) -> String {
  let Locator::Path(path) = target;
  address::drive_path(path).to_owned()
}

fn header(name: &str, value: &str) -> host::Header {
  host::Header {
    name: name.to_owned(),
    value: value.to_owned(),
  }
}

/// Sends one JSON-RPC body to the drive and returns the reply in the pieces
/// the host handed over. They are not joined here: Memona's host charges
/// fuel for every byte a plugin copies, and a file's reply joined whole would
/// spend most of a call's budget before any of it was decoded.
fn post(body: &[u8], effect: Effect) -> Result<Vec<Vec<u8>>, ErrorCode> {
  let (url, authorization) = SESSION.with_borrow(|session| {
    session
      .as_ref()
      .map(|session| (session.endpoint.clone(), session.authorization.clone()))
      .ok_or(ErrorCode::Internal)
  })?;
  let call = host::http(&host::Request {
    method: host::Method::Post,
    url,
    headers: vec![
      header("content-type", "application/json"),
      header("accept", "application/json, text/event-stream"),
      header("authorization", &authorization),
    ],
  })
  .map_err(transport)?;
  for piece in body.chunks(MAX_CHUNK as usize) {
    call.push(piece).map_err(transport)?;
  }
  let unsettled = |error: ErrorCode| match (effect, error) {
    (Effect::Change, ErrorCode::Unavailable | ErrorCode::Cancelled) => {
      ErrorCode::OutcomeUnknown
    }
    _ => error,
  };
  let head = call
    .respond()
    .map_err(|error| unsettled(transport(error)))?;
  // A server error can come after the drive's server already acted, while a
  // refused token, permission, drive or rate comes before it does anything.
  mcp::status(head.status).map_err(|error| match head.status {
    500.. => unsettled(error),
    _ => error,
  })?;
  let mut pieces = Vec::new();
  let mut total = 0;
  loop {
    let piece = call
      .pull(MAX_CHUNK)
      .map_err(|error| unsettled(transport(error)))?;
    if piece.is_empty() {
      return Ok(pieces);
    }
    total += piece.len();
    if total > MAX_REPLY {
      return Err(ErrorCode::Unavailable);
    }
    pieces.push(piece);
  }
}

/// A reply small enough to read as one buffer, which it usually already is.
fn joined(mut pieces: Vec<Vec<u8>>) -> Vec<u8> {
  match pieces.len() {
    1 => pieces.pop().unwrap_or_default(),
    _ => pieces.concat(),
  }
}

/// Calls one of the drive's tools that only reads, and reads its structured
/// answer. A call that timed out on the drive's device changed nothing.
fn tool<T: DeserializeOwned>(
  name: &str,
  arguments: &serde_json::Value,
) -> Result<T, ErrorCode> {
  let pieces = post(&mcp::call(name, arguments), Effect::Read)?;
  mcp::reply(&joined(pieces)).map_err(|error| match error {
    ErrorCode::OutcomeUnknown => ErrorCode::Unavailable,
    error => error,
  })
}

/// Sends a call that changes the drive and reads its structured answer.
fn change<T: DeserializeOwned>(body: &[u8]) -> Result<T, ErrorCode> {
  mcp::reply(&joined(post(body, Effect::Change)?))
}

fn info_of(entry: &mcp::Entry) -> FileInfo {
  FileInfo {
    kind: if entry.is_dir {
      EntryKind::Directory
    } else {
      EntryKind::File
    },
    // A folder's size is the device's directory record, not its contents.
    size: if entry.is_dir { 0 } else { entry.size },
    revision: entry.revision(),
  }
}

/// One file read whole from the drive and handed over in pieces, decoded
/// from the reply as each piece is asked for.
struct Download {
  reply: Vec<Vec<u8>>,
  /// Where the base64 contents sit in the reply, read as one run.
  encoded: Range<usize>,
  size: u64,
  position: Cell<u64>,
  revision: Option<String>,
}

impl Download {
  /// Where the characters of group `index`, three bytes each, start.
  fn chars(&self, index: usize) -> usize {
    self.encoded.start + index * 4
  }

  /// The bytes of one group, the last of which may be padded.
  fn group(&self, index: usize) -> Result<Vec<u8>, ErrorCode> {
    let mut chars = [0_u8; 4];
    let mut filled = 0;
    let at = self.chars(index);
    for slice in mcp::slices(&self.reply, at..at + 4) {
      chars[filled..filled + slice.len()].copy_from_slice(slice);
      filled += slice.len();
    }
    let bytes = codec::last(chars)?;
    let final_group = self.encoded.len() / 4 - 1;
    if filled != 4 || (index != final_group && bytes.len() != 3) {
      return Err(ErrorCode::Unavailable);
    }
    Ok(bytes)
  }

  /// Decodes at most `wanted` bytes from byte `start`. A start inside a
  /// group of three answers only the rest of that group, a short piece the
  /// host reads past, so every later piece starts where a group does and no
  /// decoded bytes are shifted into place.
  fn decode(&self, start: u64, wanted: usize) -> Result<Vec<u8>, ErrorCode> {
    let first = usize::try_from(start / 3).map_err(|_| ErrorCode::Internal)?;
    let skip = (start % 3) as usize;
    if skip != 0 {
      let bytes = self.group(first)?;
      let rest = bytes.get(skip..).unwrap_or_default();
      return Ok(rest[..rest.len().min(wanted)].to_vec());
    }
    let final_group = self.encoded.len() / 4 - 1;
    let end = first + wanted.div_ceil(3);
    let whole = end.min(final_group);
    let mut bytes = Vec::with_capacity((end - first) * 3);
    bytes.resize((whole - first) * 3, 0);
    let chars = self.chars(first)..self.chars(whole);
    codec::decode(mcp::slices(&self.reply, chars), &mut bytes)?;
    if end > final_group {
      bytes.extend(self.group(final_group)?);
    }
    bytes.truncate(wanted);
    Ok(bytes)
  }
}

impl GuestReader for Download {
  fn pull(&self, limit: u32) -> Result<Vec<u8>, ErrorCode> {
    let start = self.position.get();
    let wanted = u64::from(limit.min(MAX_CHUNK)).min(self.size - start);
    if wanted == 0 {
      return Ok(Vec::new());
    }
    let bytes = self.decode(start, wanted as usize)?;
    self.position.set(start + bytes.len() as u64);
    Ok(bytes)
  }

  fn revision(&self) -> Option<String> {
    self.revision.clone()
  }
}

/// Starts reading one file. `read_file` has no offset and no revision, so
/// the file is statted first and read whole. If its length then disagrees
/// with the stat, the file changed in between, and the read is refused as a
/// conflict rather than paired with the wrong revision.
fn download(target: &Locator, offset: u64) -> Result<Download, ErrorCode> {
  let path = route(target);
  let entry: mcp::Entry = tool("stat", &json!({ "path": path }))?;
  if entry.is_dir {
    return Err(ErrorCode::InvalidInput);
  }
  if entry.size > MAX_READ {
    return Err(ErrorCode::Unsupported);
  }
  let (reply, encoded) = if offset >= entry.size {
    (Vec::new(), 0..0)
  } else {
    let length = usize::try_from(entry.size.div_ceil(3) * 4)
      .map_err(|_| ErrorCode::Unsupported)?;
    let reply = post(
      &mcp::call("read_file", &json!({ "path": path, "encoding": "base64" })),
      Effect::Read,
    )?;
    let encoded = mcp::contents(&reply, length)?;
    (reply, encoded)
  };
  Ok(Download {
    reply,
    encoded,
    size: entry.size,
    position: Cell::new(offset.min(entry.size)),
    revision: entry.revision(),
  })
}

/// One file being saved. `write_file` takes a file whole, so its bytes are
/// held until `finish` and nothing reaches the drive before then: a save
/// abandoned on the way leaves no partial file.
struct Upload {
  path: String,
  /// What `write_file` compares before writing, as its `base_revision`.
  expected: Option<String>,
  bytes: RefCell<Vec<u8>>,
  finished: Cell<bool>,
}

impl GuestWriter for Upload {
  fn push(&self, data: Vec<u8>) -> Result<(), ErrorCode> {
    if self.finished.get() || data.len() > MAX_CHUNK as usize {
      return Err(ErrorCode::InvalidInput);
    }
    let mut bytes = self.bytes.borrow_mut();
    if bytes.len() + data.len() > MAX_WRITE {
      return Err(ErrorCode::Unsupported);
    }
    // A save of one transfer is kept as it arrived rather than copied.
    if bytes.is_empty() {
      *bytes = data;
    } else {
      bytes.extend_from_slice(&data);
    }
    Ok(())
  }

  /// One call is enough: the file goes out whole, and aindrive compares the
  /// expected revision and writes under one lock, then answers with the new
  /// revision a later conditional save compares against.
  fn finish(&self) -> Result<Option<FileInfo>, ErrorCode> {
    if self.finished.replace(true) {
      return Err(ErrorCode::InvalidInput);
    }
    let bytes = self.bytes.take();
    let body = mcp::write(&self.path, self.expected.as_deref(), &bytes);
    let written: mcp::Written = change(&body)?;
    Ok(Some(FileInfo {
      kind: EntryKind::File,
      size: bytes.len() as u64,
      revision: written.revision,
    }))
  }
}

/// Starts saving one file under `condition`. aindrive spells create-only as
/// the revision `none`, and an unconditional save as no revision at all.
fn upload(
  target: &Locator,
  condition: WriteCondition,
  size_hint: Option<u64>,
) -> Result<Upload, ErrorCode> {
  let path = route(target);
  if path.is_empty() {
    return Err(ErrorCode::InvalidInput);
  }
  if size_hint.is_some_and(|size| size > MAX_WRITE as u64) {
    return Err(ErrorCode::Unsupported);
  }
  let expected = match condition {
    WriteCondition::CreateOnly => Some("none".to_owned()),
    WriteCondition::IfMatch(revision) => Some(revision),
    WriteCondition::Overwrite => None,
  };
  Ok(Upload {
    path,
    expected,
    bytes: RefCell::new(Vec::new()),
    finished: Cell::new(false),
  })
}

impl Guest for Aindrive {
  type Reader = Download;
  type Writer = Upload;

  /// Asking which tools the token may call checks the address and the token
  /// without waking the drive's device, so a connection can be saved while
  /// that machine sleeps; its files answer once it is back.
  fn open(
    _provider_id: String,
    config: Vec<(String, Setting)>,
  ) -> Result<Capabilities, ErrorCode> {
    SESSION.set(None);
    SESSION.set(Some(Session::from_config(&config)?));
    let probe = post(mcp::TOOLS, Effect::Read)
      .and_then(|pieces| mcp::tools(&joined(pieces)));
    let tools = match probe {
      Ok(tools) => tools,
      Err(error) => {
        SESSION.set(None);
        return Err(error);
      }
    };
    let offers = |wanted: &[&str]| {
      wanted
        .iter()
        .all(|wanted| tools.iter().any(|tool| tool == wanted))
    };
    if !offers(&REQUIRED_TOOLS) {
      SESSION.set(None);
      return Err(ErrorCode::Denied);
    }
    // A read-only token is a read-only folder, not one whose saves fail.
    let writable = offers(&WRITE_TOOLS);
    Ok(Capabilities {
      write: writable,
      mkdir: false,
      rename: false,
      remove: writable,
      conditional_write: writable,
      watch: false,
      virtual_directories: false,
    })
  }

  fn close() {
    SESSION.set(None);
  }

  /// The stat tool refuses the drive root, which always exists, so the root
  /// is listed instead: that still tells whether the device can be reached.
  fn stat(target: Locator) -> Result<FileInfo, ErrorCode> {
    let path = route(&target);
    if path.is_empty() {
      let _: mcp::Listing = tool("list_files", &json!({ "path": "" }))?;
      return Ok(FileInfo {
        kind: EntryKind::Directory,
        size: 0,
        revision: None,
      });
    }
    let entry: mcp::Entry = tool("stat", &json!({ "path": path }))?;
    Ok(info_of(&entry))
  }

  /// aindrive returns a whole folder at once, so the page is taken from it.
  fn list(
    target: Locator,
    cursor: Option<String>,
    limit: u32,
  ) -> Result<ListPage, ErrorCode> {
    let start: usize = match cursor {
      Some(value) => value.parse().map_err(|_| ErrorCode::InvalidInput)?,
      None => 0,
    };
    let listing: mcp::Listing =
      tool("list_files", &json!({ "path": route(&target) }))?;
    let children: Vec<Entry> = listing
      .entries
      .iter()
      .filter(|entry| address::showable(&entry.name))
      .map(|entry| Entry {
        name: entry.name.clone(),
        info: info_of(entry),
      })
      .collect();
    let end = start.saturating_add(limit as usize).min(children.len());
    let entries = children.get(start..end).unwrap_or_default().to_vec();
    Ok(ListPage {
      entries,
      cursor: (end < children.len()).then(|| end.to_string()),
    })
  }

  fn open_read(target: Locator, offset: u64) -> Result<Reader, ErrorCode> {
    Ok(Reader::new(download(&target, offset)?))
  }

  fn open_write(
    target: Locator,
    condition: WriteCondition,
    size_hint: Option<u64>,
  ) -> Result<Writer, ErrorCode> {
    Ok(Writer::new(upload(&target, condition, size_hint)?))
  }

  fn mkdir(_target: Locator) -> Result<FileInfo, ErrorCode> {
    Err(ErrorCode::Unsupported)
  }

  fn rename(
    _target: Locator,
    _destination: String,
  ) -> Result<FileInfo, ErrorCode> {
    Err(ErrorCode::Unsupported)
  }

  /// `delete_path` removes a folder with everything in it, and aindrive
  /// leaves some entries, such as `.git`, out of every listing, so a folder
  /// cannot be shown to be empty. Only files are removed.
  fn remove(target: Locator) -> Result<(), ErrorCode> {
    let path = route(&target);
    if path.is_empty() {
      return Err(ErrorCode::InvalidInput);
    }
    let entry: mcp::Entry = tool("stat", &json!({ "path": path }))?;
    if entry.is_dir {
      return Err(ErrorCode::Unsupported);
    }
    let _: mcp::Deleted =
      change(&mcp::call("delete_path", &json!({ "path": path })))?;
    Ok(())
  }

  fn poll(_cursor: Option<String>) -> Result<ChangePage, ErrorCode> {
    Err(ErrorCode::Unsupported)
  }
}

#[cfg(not(test))]
sdk::export!(Aindrive with_types_in sdk);