k-os/crates/servers/storage-server/src/server.rs

use libkos::syscall::Give;
use server_framework::{Decode, Encode, MALFORMED, Server};
use types::bounded::CapSlot;

use crate::filesystem::{DirEntry, FileHandle, FileSystem};
use crate::grants::Grants;
use crate::protocol::{DIR_ENTRY_SIZE, DIR_NAME_SIZE, MAX_PATH, StorageRequest, StorageResponse};
use crate::scope::Scope;

mod session;

use types::storage::MAX_CLIENTS;
/// Handles at or past this are never valid (the file systems hand out
/// fewer), so ownership fits a small table.
const MAX_HANDLES: usize = 16;
const MAX_LISTED: usize = 64;

/// What the server needs from the kernel: regions for client buffers
/// and badged copies of its endpoint for grants. Syscalls in the real
/// server, plain memory in tests.
pub trait Host {
    /// A buffer for a new client and the capability that lets the client
    /// map it, if there is one to send.
    fn create(&mut self, client: usize) -> Option<(&'static mut [u8], Option<CapSlot>)>;
    /// Unmaps a buffer from whoever it was sent to, so it can be reused.
    fn reclaim(&mut self, buffer: CapSlot);
    /// A copy of the server's endpoint carrying `badge`.
    fn mint(&mut self, badge: usize) -> Option<CapSlot>;
    /// Deletes `grant` and every copy made from it, in every process.
    fn revoke(&mut self, grant: CapSlot);
}

struct Client {
    /// Who the buffer belongs to now; `None` once its grant is released.
    badge: Option<usize>,
    buf: &'static mut [u8],
    /// What maps `buf`, sent again when the same client attaches again.
    cap: Option<CapSlot>,
}

#[derive(Debug, Clone, Copy)]
struct Owner {
    badge: usize,
    /// Opened through a scope that may write.
    write: bool,
}

pub struct StorageServer<F: FileSystem, H: Host> {
    fs: F,
    host: H,
    grants: Grants,
    clients: [Option<Client>; MAX_CLIENTS],
    owners: [Option<Owner>; MAX_HANDLES],
    log: Option<CapSlot>,
    /// The master key's and the seed's halves, from the gatekeeper through
    /// init; handed to the file system (and forgotten here) once all in.
    key: [Option<[usize; 2]>; 2],
    seed: [Option<[usize; 2]>; 2],
    /// Whose areas they open, said first.
    user: Option<types::user::UserFolder>,
    unlocked: bool,
}

impl<F: FileSystem, H: Host> StorageServer<F, H> {
    pub fn new(fs: F, host: H) -> Self {
        Self {
            fs,
            host,
            grants: Grants::new(),
            clients: [const { None }; MAX_CLIENTS],
            owners: [None; MAX_HANDLES],
            log: None,
            key: [None; 2],
            seed: [None; 2],
            user: None,
            unlocked: false,
        }
    }

    /// Where this server prints, once init has said.
    pub fn log(&self) -> Option<CapSlot> {
        self.log
    }

    /// Clients hold badged copies of the endpoint, so the unbadged sender is
    /// whoever was given the endpoint itself: init.
    fn set_log(&mut self, badge: usize, cap: Option<CapSlot>) -> StorageResponse {
        match (badge, cap, self.log) {
            (0, Some(cap), None) => {
                self.log = Some(cap);
                libkos::start::set_log(Some(cap));
                libkos::println!("storage: printing through the log server");
                StorageResponse::Ok
            }
            _ => StorageResponse::Error,
        }
    }

    /// Whether the encrypted areas are open: the key and seed are in.
    pub fn unlocked(&self) -> bool {
        self.unlocked
    }

    pub fn filesystem(&self) -> &F {
        &self.fs
    }

    #[cfg(test)]
    pub fn host(&self) -> &H {
        &self.host
    }

    #[cfg(test)]
    pub fn client_buffer_for_test(&mut self, badge: usize) -> Option<&mut [u8]> {
        self.client(badge).map(|c| &mut *c.buf)
    }

    fn client(&mut self, badge: usize) -> Option<&mut Client> {
        self.clients
            .iter_mut()
            .flatten()
            .find(|c| c.badge == Some(badge))
    }

    /// A client attaching again gets its buffer again; a new one gets a
    /// buffer a released grant left behind, or a fresh one.
    fn attach(&mut self, badge: usize) -> (StorageResponse, Option<Give>) {
        if let Some(c) = self.client(badge) {
            return (
                StorageResponse::Attached { size: c.buf.len() },
                c.cap.map(Give::from),
            );
        }
        if let Some(c) = self
            .clients
            .iter_mut()
            .flatten()
            .find(|c| c.badge.is_none())
        {
            c.badge = Some(badge);
            return (
                StorageResponse::Attached { size: c.buf.len() },
                c.cap.map(Give::from),
            );
        }
        let Some(index) = self.clients.iter().position(Option::is_none) else {
            return (StorageResponse::Error, None);
        };
        let Some((buf, cap)) = self.host.create(index) else {
            return (StorageResponse::Error, None);
        };
        let size = buf.len();
        if let Some(slot) = self.clients.get_mut(index) {
            *slot = Some(Client {
                badge: Some(badge),
                buf,
                cap,
            });
        }
        (StorageResponse::Attached { size }, cap.map(Give::from))
    }

    /// Takes `badge`'s grant back from everyone, with everything granted
    /// from it: their capabilities, open files and buffers.
    fn release(&mut self, badge: usize) -> StorageResponse {
        if badge == 0 {
            return StorageResponse::Error;
        }
        let Self {
            fs,
            host,
            grants,
            clients,
            owners,
            ..
        } = self;
        grants.release(badge, |gone, cap| {
            host.revoke(cap);
            for (raw, owner) in owners.iter_mut().enumerate() {
                if owner.is_some_and(|o| o.badge == gone) {
                    fs.close(FileHandle::new(raw as u16));
                    *owner = None;
                }
            }
            for c in clients
                .iter_mut()
                .flatten()
                .filter(|c| c.badge == Some(gone))
            {
                if let Some(buffer) = c.cap {
                    host.reclaim(buffer);
                }
                c.buf.fill(0);
                c.badge = None;
            }
        });
        StorageResponse::Ok
    }

    fn owner(&self, handle: FileHandle) -> Option<Owner> {
        self.owners.get(handle.raw() as usize).copied().flatten()
    }

    /// Serves one request from `badge`, which holds `scope`.
    fn serve(
        &mut self,
        badge: usize,
        scope: Scope,
        request: StorageRequest,
    ) -> (StorageResponse, Option<Give>) {
        use StorageRequest::*;
        use StorageResponse as R;
        match request {
            Attach => return self.attach(badge),
            Release => return (self.release(badge), None),
            _ => {}
        }
        let allowed = match request {
            Read { handle, .. } | Close { handle } | FileSize { handle } => {
                self.owner(handle).is_some_and(|o| o.badge == badge)
            }
            Write { handle, .. } => self
                .owner(handle)
                .is_some_and(|o| o.badge == badge && o.write),
            Create { .. } | Delete { .. } | Mkdir { .. } | Rmdir { .. } => scope.write,
            _ => true,
        };
        if !allowed {
            return (R::Denied, None);
        }
        let mutates = matches!(
            request,
            Write { .. } | Create { .. } | Delete { .. } | Mkdir { .. } | Rmdir { .. }
        );
        if mutates && !self.fs.writable() {
            return (R::ReadOnly, None);
        }
        // Split borrows: the client's buffer, the file system, the owners.
        let Self {
            fs,
            host,
            grants,
            clients,
            owners,
            ..
        } = self;
        let Some(client) = clients
            .iter_mut()
            .flatten()
            .find(|c| c.badge == Some(badge))
        else {
            return (R::NotAttached, None);
        };
        let buf = &mut *client.buf;
        let room = buf.len();
        let sent = |len: usize| ..len.min(MAX_PATH).min(room);
        // The client's path, as a path on disk inside its scope.
        let path_len = match request {
            Open { path_len }
            | Create { path_len }
            | Delete { path_len }
            | Mkdir { path_len }
            | Rmdir { path_len }
            | ListDir { path_len } => Some(path_len),
            _ => None,
        };
        let disk = match path_len.map(|len| scope.resolve(&buf[sent(len)])) {
            Some(None) => return (R::Denied, None),
            resolved => resolved.flatten(),
        };
        let path = disk.as_ref().map_or(&[][..], |p| p.as_bytes());
        let opened = |owners: &mut [Option<Owner>; MAX_HANDLES], h: Option<FileHandle>| match h {
            Some(h) => match owners.get_mut(h.raw() as usize) {
                Some(owner) => {
                    *owner = Some(Owner {
                        badge,
                        write: scope.write,
                    });
                    R::Handle(h)
                }
                // A handle the table cannot track is one we cannot police.
                None => R::Error,
            },
            None => R::NotFound,
        };
        let response = match request {
            Attach | SetLog | SetKey { .. } | SetSeed { .. } | SetUser { .. } | Lock | Purge { .. } | Release => R::Error,
            Grant { path_len, write } => {
                let narrower = scope.narrow(&buf[sent(path_len)], write);
                return match narrower.and_then(|s| grants.add(badge, s, |b| host.mint(b))) {
                    // The holder may pass it on: that is what it asked for.
                    Some(cap) => (R::Granted, Some(Give::keeping_delegate(cap))),
                    None => (R::Denied, None),
                };
            }
            Open { .. } => {
                let h = fs.open(path);
                opened(owners, h)
            }
            Create { .. } => match fs.create(path) {
                None => R::Error,
                h => opened(owners, h),
            },
            Read { handle, len } => match fs.read(handle, &mut buf[..len.min(room)]) {
                Some(bytes) => R::Data { bytes },
                None => R::Error,
            },
            Write { handle, len } => match fs.write(handle, &buf[..len.min(room)]) {
                Some(bytes) => R::Written { bytes },
                None => R::Error,
            },
            Close { handle } => {
                fs.close(handle);
                if let Some(owner) = owners.get_mut(handle.raw() as usize) {
                    *owner = None;
                }
                R::Ok
            }
            FileSize { handle } => match fs.file_size(handle) {
                Some(size) => R::Size(size),
                None => R::Error,
            },
            ListDir { .. } => {
                let mut entries = [DirEntry::empty(); MAX_LISTED];
                let room = (buf.len() / DIR_ENTRY_SIZE).min(MAX_LISTED);
                match fs.list_dir(path, &mut entries[..room]) {
                    Some(count) => {
                        for (entry, record) in entries
                            .iter()
                            .take(count)
                            .zip(buf.chunks_mut(DIR_ENTRY_SIZE))
                        {
                            write_entry(entry, record);
                        }
                        R::DirEntries { count }
                    }
                    None => R::NotFound,
                }
            }
            Delete { .. } => {
                if fs.delete(path) {
                    R::Ok
                } else {
                    R::NotFound
                }
            }
            Mkdir { .. } => {
                if fs.mkdir(path) {
                    R::Ok
                } else {
                    R::Error
                }
            }
            Rmdir { .. } => {
                if fs.rmdir(path) {
                    R::Ok
                } else {
                    R::Error
                }
            }
        };
        (response, None)
    }
}

fn write_entry(entry: &DirEntry, record: &mut [u8]) {
    record.fill(0);
    let name = entry.name.as_bytes();
    let n = name.len().min(DIR_NAME_SIZE);
    record[..n].copy_from_slice(&name[..n]);
    record[DIR_NAME_SIZE..DIR_NAME_SIZE + 8].copy_from_slice(&(entry.size as u64).to_le_bytes());
    let flags = if entry.is_dir {
        types::storage::DIR_FLAG_DIRECTORY
    } else {
        0
    };
    record[DIR_NAME_SIZE + 8..DIR_ENTRY_SIZE].copy_from_slice(&flags.to_le_bytes());
}

impl<F: FileSystem, H: Host> Server for StorageServer<F, H> {
    type Request = StorageRequest;
    type Response = StorageResponse;

    /// Without a badge: as the unbadged client.
    fn handle(&mut self, request: StorageRequest) -> StorageResponse {
        self.serve(0, Scope::EVERYTHING, request).0
    }

    fn handle_call(
        &mut self,
        badge: usize,
        words: [usize; 4],
        cap: Option<CapSlot>,
    ) -> ([usize; 4], Option<Give>) {
        match StorageRequest::decode(words) {
            Some(StorageRequest::SetLog) => (self.set_log(badge, cap).encode(), None),
            Some(StorageRequest::SetKey { half, words }) => (self.set_half(badge, false, half, words).encode(), None),
            Some(StorageRequest::SetSeed { half, words }) => (self.set_half(badge, true, half, words).encode(), None),
            Some(StorageRequest::SetUser { id, guest }) => (self.set_user(badge, id, guest).encode(), None),
            Some(StorageRequest::Lock) => (self.lock(badge).encode(), None),
            Some(StorageRequest::Purge { path_len }) => (self.purge(badge, path_len).encode(), None),
            // A badge storage never issued, or one it has taken back.
            Some(request) => match self.grants.scope(badge) {
                None => (StorageResponse::Denied.encode(), None),
                Some(scope) => {
                    let (response, cap) = self.serve(badge, scope, request);
                    (response.encode(), cap)
                }
            },
            None => ([MALFORMED, 0, 0, 0], None),
        }
    }
}