telepath_server/

lib.rs

//! Target-side Telepath library.
//!
//! Runs on the MCU in `no_std` mode. Provides:
//! - [`TelepathServer`]: receive loop, COBS decode, dispatch, rzCOBS encode
//! - [`transport::Transport`]: non-blocking byte-stream I/O trait
//! - Re-export of `#[command]` attribute macro
//!
//! # Architecture
//!
//! ```text
//! Transport → FrameAccumulator → cobs_decode    → postcard::from_bytes → Dispatcher
//!                                                                       → postcard::to_slice → rzcobs_encode → Transport
//! ```
//!
//! # Usage
//!
//! ```rust,ignore
//! use telepath_server::{TelepathServer, command};
//!
//! #[command]
//! fn set_led(id: u8, brightness: u16) -> Result<(), ()> {
//!     Ok(())
//! }
//!
//! let mut server = TelepathServer::<_, 512>::new(transport, telepath_server::commands());
//! loop { server.poll(); }
//! ```
#![no_std]

pub mod transport;

mod resource;
pub use resource::ResourceRegistry;

#[cfg(feature = "profile")]
pub mod profile;
#[cfg(feature = "profile")]
pub use profile::init_dwt;

pub use telepath_macros::command;
use telepath_wire::{
    framing::{cobs_decode, rzcobs_encode, FrameAccumulator},
    Request, Response,
};
pub use telepath_wire::{
    PacketType, ResponseStatus, WireError, CMD_ID_DISCOVERY, MAX_PAYLOAD_SIZE,
};

// Re-exported for use in code generated by the #[command] macro so that callers
// only need `telepath-firmware` and `postcard` as direct dependencies.
#[doc(hidden)]
pub use linkme as __linkme;
#[doc(hidden)]
pub use postcard_schema as __postcard_schema;
#[doc(hidden)]
pub use telepath_wire::cmd_id::derive_cmd_id as __derive_cmd_id;

// ---------------------------------------------------------------------------
// CommandMetadata
// ---------------------------------------------------------------------------

/// Type-erased shim function signature.
///
/// Receives a postcard-serialized argument slice, writes a postcard-serialized
/// result into `output`, and returns the number of bytes written. The
/// `resources` parameter provides access to injected `#[resource]` values.
pub type ShimFn = fn(
    input: &[u8],
    output: &mut [u8],
    resources: &ResourceRegistry,
) -> Result<usize, DispatchError>;

/// Type alias for schema-writer function pointers.
///
/// Writes a postcard-serialized `postcard_schema::schema::NamedType` into `out`
/// and returns the number of bytes written.
pub type SchemaFn = fn(out: &mut [u8]) -> Result<usize, ()>;

/// Static metadata for a single registered RPC command.
///
/// Populated by the `#[command]` macro at compile time and collected into
/// [`TELEPATH_COMMANDS`] via `linkme` distributed slices.
#[derive(Clone, Copy)]
pub struct CommandMetadata {
    /// Human-readable function name (used for discovery).
    pub name: &'static str,
    /// Command ID: hash of (name + input schema + output schema).
    /// Computed at firmware build time for deterministic matching.
    pub id: u16,
    /// Type-erased shim that deserializes args, calls the function, and
    /// serializes the result.
    pub invoke: ShimFn,
    /// Writes the postcard-encoded args-tuple `NamedType` schema into the
    /// provided buffer. Returns the byte count written.
    pub args_schema: SchemaFn,
    /// Writes the postcard-encoded return-type `NamedType` schema into the
    /// provided buffer. Returns the byte count written.
    pub ret_schema: SchemaFn,
    /// Comma-separated argument names, e.g. `"a,b"` for `fn foo(a: i32, b: i32)`.
    /// Empty string for zero-argument commands.
    pub arg_names: &'static str,
}

/// All commands registered via `#[command]`, collected at link time.
///
/// Use [`commands()`] to access this as a `&'static [CommandMetadata]`.
#[linkme::distributed_slice]
pub static TELEPATH_COMMANDS: [CommandMetadata] = [..];

/// Returns the complete set of commands registered by `#[command]`.
pub fn commands() -> &'static [CommandMetadata] {
    &TELEPATH_COMMANDS
}

// ---------------------------------------------------------------------------
// DispatchError
// ---------------------------------------------------------------------------

/// Errors that can occur during command dispatch.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DispatchError {
    /// No command with the given ID was found in the registry.
    UnknownCommand,
    /// Argument deserialization failed (malformed or truncated payload).
    DeserializeError,
    /// Result serialization failed (output buffer too small).
    SerializeError,
    /// The request payload exceeded [`MAX_PAYLOAD_SIZE`].
    PayloadTooLarge,
    /// A `#[resource]`-annotated argument could not be resolved from the
    /// server's [`ResourceRegistry`].
    ResourceUnavailable,
}

// ---------------------------------------------------------------------------
// TelepathServer
// ---------------------------------------------------------------------------

/// RPC server that runs on the target MCU.
///
/// `T` is the transport type implementing [`transport::Transport`].
/// `N` is the size of the internal receive accumulator and transmit buffers.
///
/// # Type parameter guidance
///
/// Choose `N` ≥ 512 to accommodate a max-payload frame with COBS overhead.
pub struct TelepathServer<T, const N: usize> {
    transport: T,
    rx_accum: FrameAccumulator<N>,
    tx_buf: [u8; N],
    /// Command registry slice. Pass `telepath_server::commands()` for the
    /// full linkme-populated registry, or a manual slice for testing.
    commands: &'static [CommandMetadata],
    /// Type-keyed resource registry for `#[resource]` injection.
    resources: ResourceRegistry,
}

impl<T, const N: usize> TelepathServer<T, N> {
    /// Create a new server with the given transport and command registry.
    pub fn new(transport: T, commands: &'static [CommandMetadata]) -> Self {
        #[cfg(feature = "profile")]
        profile::init_dwt();
        Self {
            transport,
            rx_accum: FrameAccumulator::new(),
            tx_buf: [0u8; N],
            commands,
            resources: ResourceRegistry::new(),
        }
    }

    /// Register a resource for `#[resource]` injection.
    ///
    /// The value is moved into the server's internal registry and made
    /// available to command shims that declare a matching `#[resource]`
    /// parameter.
    pub fn resource<R: 'static>(mut self, val: R) -> Self {
        self.resources.insert(val);
        self
    }

    /// Look up a command by its ID using linear scan.
    ///
    /// Linear scan is intentional: embedded command counts are typically ≤ 64,
    /// making hash-map overhead unjustified.
    pub fn find_command(&self, id: u16) -> Option<&CommandMetadata> {
        self.commands.iter().find(|cmd| cmd.id == id)
    }

    /// Dispatch a pre-decoded payload slice to the matching command handler.
    ///
    /// Returns the number of bytes written to `output` on success.
    pub fn dispatch(
        &mut self,
        cmd_id: u16,
        input: &[u8],
        output: &mut [u8],
    ) -> Result<usize, DispatchError> {
        if cmd_id == telepath_wire::CMD_ID_DISCOVERY {
            return self.handle_discovery(input, output);
        }
        let cmd = self
            .find_command(cmd_id)
            .ok_or(DispatchError::UnknownCommand)?;
        (cmd.invoke)(input, output, &self.resources)
    }

    /// Handle a Discovery request (CmdID 0x0000) with offset-based pagination.
    ///
    /// Builds a [`telepath_wire::DiscoveryPage`] containing entries starting at
    /// `request.offset`, limited by `MAX_PAYLOAD_SIZE`. Each entry includes
    /// postcard-serialized schema bytes for the argument tuple and return type.
    ///
    /// Empty `input` is treated as `offset=0` for backward compatibility with
    /// hosts that send raw discovery requests without a `DiscoveryRequest` payload.
    fn handle_discovery(&self, input: &[u8], output: &mut [u8]) -> Result<usize, DispatchError> {
        use telepath_wire::{DiscoveryPage, DiscoveryRequest};

        let offset = if input.is_empty() {
            0u16
        } else {
            postcard::from_bytes::<DiscoveryRequest>(input)
                .map_err(|_| DispatchError::DeserializeError)?
                .offset
        };

        let total = self
            .commands
            .iter()
            .filter(|c| c.id != telepath_wire::CMD_ID_DISCOVERY)
            .count() as u16;

        // DiscoveryPage header overhead: ≤3 B for total (u16 varint) +
        // ≤3 B for offset (u16 varint) + ≤5 B for the entries-slice length
        // varint. 16 B is a conservative bound; derived from MAX_PAYLOAD_SIZE
        // so the entries budget updates automatically if the limit changes.
        const PAGE_HEADER_BUDGET: usize = 16;
        const ENTRIES_RAW_MAX: usize = MAX_PAYLOAD_SIZE - PAGE_HEADER_BUDGET;

        let mut raw_entries = [0u8; ENTRIES_RAW_MAX];
        let mut raw_cursor = 0usize;
        let mut page_count = 0u32;

        // Upper bound on postcard_schema::schema::NamedType bytes for a single
        // command schema. Measured empirically: typical primitive schemas are
        // 20–60 bytes. 128 bytes gives ~2× headroom for deeply nested types.
        // Exceeding this limit returns SerializeError at discovery time.
        const SCHEMA_SCRATCH_LEN: usize = 128;

        // Per-entry scratch buffers; reused each iteration.
        let mut args_scratch = [0u8; SCHEMA_SCRATCH_LEN];
        let mut ret_scratch = [0u8; SCHEMA_SCRATCH_LEN];

        let iter = self
            .commands
            .iter()
            .filter(|c| c.id != telepath_wire::CMD_ID_DISCOVERY)
            .skip(offset as usize);

        for cmd in iter {
            let n_args =
                (cmd.args_schema)(&mut args_scratch).map_err(|_| DispatchError::SerializeError)?;
            let n_ret =
                (cmd.ret_schema)(&mut ret_scratch).map_err(|_| DispatchError::SerializeError)?;
            let entry = telepath_wire::DiscoveryEntry {
                id: cmd.id,
                name: cmd.name,
                args_schema: &args_scratch[..n_args],
                ret_schema: &ret_scratch[..n_ret],
                arg_names: cmd.arg_names,
            };
            // Pre-measure the entry by serializing into a temp scratch.
            let mut entry_tmp = [0u8; 300];
            let entry_bytes = postcard::to_slice(&entry, &mut entry_tmp)
                .map_err(|_| DispatchError::SerializeError)?;
            let entry_size = entry_bytes.len();

            if raw_cursor + entry_size > ENTRIES_RAW_MAX {
                if raw_cursor == 0 {
                    // This entry alone exceeds the page budget; it can never
                    // fit regardless of paging. Signal a hard error so the host
                    // receives a SystemError instead of an infinite stall.
                    return Err(DispatchError::SerializeError);
                }
                break; // page is full — more entries on next page
            }
            raw_entries[raw_cursor..raw_cursor + entry_size].copy_from_slice(entry_bytes);
            raw_cursor += entry_size;
            page_count += 1;
        }

        // Build the entries field: varint(count) ++ raw_entries[..raw_cursor].
        let mut entries_combined = [0u8; ENTRIES_RAW_MAX + 5];
        let cnt_bytes = postcard::to_slice(&page_count, &mut entries_combined)
            .map_err(|_| DispatchError::SerializeError)?;
        let cnt_len = cnt_bytes.len();
        entries_combined[cnt_len..cnt_len + raw_cursor].copy_from_slice(&raw_entries[..raw_cursor]);
        let entries_len = cnt_len + raw_cursor;

        let page = DiscoveryPage {
            total,
            offset,
            entries: &entries_combined[..entries_len],
        };
        let written =
            postcard::to_slice(&page, output).map_err(|_| DispatchError::SerializeError)?;
        Ok(written.len())
    }
}

impl<T: transport::Transport, const N: usize> TelepathServer<T, N> {
    /// Process any pending bytes from the transport.
    ///
    /// Call this in a tight loop. Reads all available bytes, accumulates them
    /// into COBS frames, and sends a response for each complete request.
    pub fn poll(&mut self) {
        let mut byte = [0u8; 1];
        loop {
            let n = self.transport.read(&mut byte);
            if n == 0 {
                break;
            }
            if self.rx_accum.feed(byte[0]) {
                self.process_frame();
                self.rx_accum.reset();
            }
        }
    }

    /// Decode and dispatch a complete COBS frame, then encode and send the response.
    fn process_frame(&mut self) {
        let frame = match self.rx_accum.frame() {
            Some(f) => f,
            None => return,
        };

        // COBS decode into a stack buffer.
        let mut decoded = [0u8; N];
        #[cfg(feature = "profile")]
        let t0 = profile::cycles_now();
        let decoded_len = match cobs_decode(frame, &mut decoded) {
            Ok(n) => n,
            Err(_) => return,
        };
        #[cfg(feature = "profile")]
        {
            use core::sync::atomic::Ordering;
            let dt = profile::cycles_now().wrapping_sub(t0) as u64;
            profile::DECODE_CYCLES.fetch_add(dt, Ordering::Relaxed);
            profile::DECODED_BYTES.fetch_add(decoded_len as u32, Ordering::Relaxed);
        }

        // Deserialize Request (args borrows from decoded[]).
        let req: Request<'_> = match postcard::from_bytes(&decoded[..decoded_len]) {
            Ok(r) => r,
            Err(_) => return,
        };

        // Reject packets that are not properly typed as Request.
        if req.kind != PacketType::Request {
            return;
        }

        // Reject oversized argument payloads before dispatch.
        if req.args.len() > MAX_PAYLOAD_SIZE {
            return;
        }

        let seq_no = req.seq_no;
        let cmd_id = req.cmd_id;
        let args = req.args;

        // Dispatch; clamp oversized return payloads to SystemError.
        let mut payload_buf = [0u8; N];
        let (status, payload_len) = match self.dispatch(cmd_id, args, &mut payload_buf) {
            Ok(n) if n > MAX_PAYLOAD_SIZE => (ResponseStatus::SystemError, 0),
            Ok(n) => (ResponseStatus::Ok, n),
            Err(_) => (ResponseStatus::SystemError, 0),
        };

        // Build and serialize Response.
        let resp = Response {
            kind: PacketType::Response,
            seq_no,
            status,
            payload: &payload_buf[..payload_len],
        };
        let mut serialized = [0u8; N];
        let serialized_len = match postcard::to_slice(&resp, &mut serialized) {
            Ok(s) => s.len(),
            Err(_) => return,
        };

        // rzCOBS encode into tx_buf and write (upstream framing).
        #[cfg(feature = "profile")]
        let t1 = profile::cycles_now();
        let n = match rzcobs_encode(&serialized[..serialized_len], &mut self.tx_buf) {
            Ok(n) => n,
            Err(_) => return,
        };
        #[cfg(feature = "profile")]
        {
            use core::sync::atomic::Ordering;
            let dt = profile::cycles_now().wrapping_sub(t1) as u64;
            profile::ENCODE_CYCLES.fetch_add(dt, Ordering::Relaxed);
            profile::ENCODED_BYTES.fetch_add(serialized_len as u32, Ordering::Relaxed);
            profile::SAMPLE_COUNT.fetch_add(1, Ordering::Relaxed);
        }
        self.transport.write(&self.tx_buf[..n]);
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    extern crate std;
    use super::*;

    fn noop_shim(
        _input: &[u8],
        _output: &mut [u8],
        _resources: &ResourceRegistry,
    ) -> Result<usize, DispatchError> {
        Ok(0)
    }

    fn noop_schema(_out: &mut [u8]) -> Result<usize, ()> {
        Ok(0)
    }

    static TEST_COMMANDS: [CommandMetadata; 1] = [CommandMetadata {
        name: "ping",
        id: 0x0001,
        invoke: noop_shim,
        args_schema: noop_schema,
        ret_schema: noop_schema,
        arg_names: "",
    }];

    struct FakeTransport;

    #[test]
    fn find_known_command() {
        let server = TelepathServer::<FakeTransport, 256>::new(FakeTransport, &TEST_COMMANDS);
        assert!(server.find_command(0x0001).is_some());
    }

    #[test]
    fn find_unknown_command_returns_none() {
        let server = TelepathServer::<FakeTransport, 256>::new(FakeTransport, &TEST_COMMANDS);
        assert!(server.find_command(0xFFFF).is_none());
    }

    #[test]
    fn dispatch_unknown_returns_error() {
        let mut server = TelepathServer::<FakeTransport, 256>::new(FakeTransport, &TEST_COMMANDS);
        let mut out = [0u8; 256];
        assert_eq!(
            server.dispatch(0xFFFF, &[], &mut out),
            Err(DispatchError::UnknownCommand)
        );
    }

    // ---------------------------------------------------------------------------
    // poll() integration test using a loopback transport
    // ---------------------------------------------------------------------------

    use telepath_wire::framing::{cobs_encode, rzcobs_decode};
    use telepath_wire::{PacketType, Request, Response, ResponseStatus};

    /// A ping shim that writes `0xDEADBEEFu32` as postcard to output.
    fn ping_shim(
        _input: &[u8],
        output: &mut [u8],
        _resources: &ResourceRegistry,
    ) -> Result<usize, DispatchError> {
        let val: u32 = 0xDEAD_BEEF;
        let s = postcard::to_slice(&val, output).map_err(|_| DispatchError::SerializeError)?;
        Ok(s.len())
    }

    static PING_COMMANDS: [CommandMetadata; 1] = [CommandMetadata {
        name: "ping",
        id: 0x0001,
        invoke: ping_shim,
        args_schema: noop_schema,
        ret_schema: noop_schema,
        arg_names: "",
    }];

    /// Loopback transport: bytes written are available for reading.
    struct LoopbackTransport {
        rx: std::vec::Vec<u8>,
        tx: std::vec::Vec<u8>,
    }

    impl LoopbackTransport {
        fn new(rx: std::vec::Vec<u8>) -> Self {
            Self {
                rx,
                tx: std::vec::Vec::new(),
            }
        }
    }

    impl transport::Transport for LoopbackTransport {
        fn read(&mut self, buf: &mut [u8]) -> usize {
            if self.rx.is_empty() {
                return 0;
            }
            let n = buf.len().min(self.rx.len());
            buf[..n].copy_from_slice(&self.rx[..n]);
            self.rx.drain(..n);
            n
        }

        fn write(&mut self, buf: &[u8]) -> usize {
            self.tx.extend_from_slice(buf);
            buf.len()
        }
    }

    #[test]
    fn poll_ping_roundtrip() {
        // Build a COBS-framed ping request.
        let req = Request {
            kind: PacketType::Request,
            seq_no: 42,
            cmd_id: 0x0001,
            args: &[],
        };
        let mut ser_buf = [0u8; 64];
        let serialized = postcard::to_slice(&req, &mut ser_buf).unwrap();
        let mut framed = [0u8; 64];
        let n = cobs_encode(serialized, &mut framed).unwrap();

        let transport = LoopbackTransport::new(framed[..n].to_vec());
        let mut server = TelepathServer::<LoopbackTransport, 512>::new(transport, &PING_COMMANDS);
        server.poll();

        // Decode the response from tx buffer.
        let tx = &server.transport.tx;
        assert!(!tx.is_empty(), "server must have written a response");

        // Find the 0x00 delimiter.
        let delim = tx
            .iter()
            .position(|&b| b == 0x00)
            .expect("no frame delimiter");
        let mut decoded = [0u8; 512];
        let m = rzcobs_decode(&tx[..delim], &mut decoded).unwrap();

        let resp: Response<'_> = postcard::from_bytes(&decoded[..m]).unwrap();
        assert_eq!(resp.seq_no, 42);
        assert_eq!(resp.status, ResponseStatus::Ok);

        let val: u32 = postcard::from_bytes(resp.payload).unwrap();
        assert_eq!(val, 0xDEAD_BEEF);
    }
}