Open an aindrive drive as a read-only folder in Memona (#1)

src/lib.rs

@@ -0,0 +1,411 @@
+//! 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 decode;
+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;
+/// Without these the folder cannot even be browsed.
+const REQUIRED_TOOLS: [&str; 3] = ["list_files", "stat", "read_file"];
+
+struct Session {
+  endpoint: String,
+  authorization: String,
+}
+
+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]) -> 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 head = call.respond().map_err(transport)?;
+  mcp::status(head.status)?;
+  let mut pieces = Vec::new();
+  let mut total = 0;
+  loop {
+    let piece = call.pull(MAX_CHUNK).map_err(transport)?;
+    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 and reads its structured answer.
+fn tool<T: DeserializeOwned>(
+  name: &str,
+  arguments: &serde_json::Value,
+) -> Result<T, ErrorCode> {
+  mcp::reply(&joined(post(&mcp::call(name, arguments))?))
+}
+
+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 = decode::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);
+    decode::run(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" }),
+    ))?;
+    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(),
+  })
+}
+
+/// Writing is not offered yet, so no transfer is ever started.
+enum Upload {}
+
+impl GuestWriter for Upload {
+  fn push(&self, _data: Vec<u8>) -> Result<(), ErrorCode> {
+    match *self {}
+  }
+
+  fn finish(&self) -> Result<Option<FileInfo>, ErrorCode> {
+    match *self {}
+  }
+}
+
+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).and_then(|pieces| mcp::tools(&joined(pieces)));
+    let tools = match probe {
+      Ok(tools) => tools,
+      Err(error) => {
+        SESSION.set(None);
+        return Err(error);
+      }
+    };
+    if !REQUIRED_TOOLS
+      .iter()
+      .all(|required| tools.iter().any(|tool| tool == required))
+    {
+      SESSION.set(None);
+      return Err(ErrorCode::Denied);
+    }
+    Ok(Capabilities {
+      write: false,
+      mkdir: false,
+      rename: false,
+      remove: false,
+      conditional_write: false,
+      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> {
+    Err(ErrorCode::Unsupported)
+  }
+
+  fn mkdir(_target: Locator) -> Result<FileInfo, ErrorCode> {
+    Err(ErrorCode::Unsupported)
+  }
+
+  fn rename(
+    _target: Locator,
+    _destination: String,
+  ) -> Result<FileInfo, ErrorCode> {
+    Err(ErrorCode::Unsupported)
+  }
+
+  fn remove(_target: Locator) -> Result<(), ErrorCode> {
+    Err(ErrorCode::Unsupported)
+  }
+
+  fn poll(_cursor: Option<String>) -> Result<ChangePage, ErrorCode> {
+    Err(ErrorCode::Unsupported)
+  }
+}
+
+#[cfg(not(test))]
+sdk::export!(Aindrive with_types_in sdk);