RedBear-OS/local/recipes/system/diskd/source/src/main.rs

use std::collections::BTreeMap;
use std::io::{self, Write};
use std::process;
use std::time::Duration;

use libredox::Fd;
use log::{LevelFilter, Metadata, Record, error, info, warn};
use redox_scheme::scheme::{SchemeState, SchemeSync};
use redox_scheme::{CallerCtx, OpenResult, RequestKind, SignalBehavior, Socket};
use syscall::dirent::{DirEntry, DirentBuf, DirentKind};
use syscall::error::{EACCES, EBADF, EINTR, EINVAL, ENOENT, ENOTDIR, Error, Result};
use syscall::flag::{O_ACCMODE, O_DIRECTORY, O_STAT};
use syscall::schemev2::NewFdFlags;
use syscall::{MODE_DIR, Stat};

const SCHEME_NAME: &str = "diskd";
const SCHEME_ROOT_ID: usize = 0;
const DISK_PREFIX: &str = "disk.";
const PROBE_RETRY: usize = 5;
const PROBE_DELAY: Duration = Duration::from_millis(150);

struct StderrLogger;

impl log::Log for StderrLogger {
    fn enabled(&self, metadata: &Metadata<'_>) -> bool {
        metadata.level() <= LevelFilter::Info
    }

    fn log(&self, record: &Record<'_>) {
        if self.enabled(record.metadata()) {
            let _ = writeln!(
                io::stderr().lock(),
                "[{}] diskd: {}",
                record.level(),
                record.args()
            );
        }
    }

    fn flush(&self) {}
}

static LOGGER: StderrLogger = StderrLogger;

#[derive(Clone, Debug)]
struct BlockDevice {
    scheme: String,
    name: String,
    underlying_path: String,
}

impl BlockDevice {
    fn new(scheme: String, disk_index: u32, partition: Option<u32>) -> Self {
        let name = match partition {
            Some(p) => format!("{disk_index}p{p}"),
            None => format!("{disk_index}"),
        };
        let underlying_path = match partition {
            Some(p) => format!("/scheme/{scheme}/{disk_index}p{p}"),
            None => format!("/scheme/{scheme}/{disk_index}"),
        };
        Self {
            scheme,
            name,
            underlying_path,
        }
    }
}

enum Handle {
    SchemeRoot,
    List,
}

struct DiskdScheme {
    devices: Vec<BlockDevice>,
    handles: BTreeMap<usize, Handle>,
    next_id: usize,
}

impl DiskdScheme {
    fn new(devices: Vec<BlockDevice>) -> Self {
        let mut handles = BTreeMap::new();
        handles.insert(SCHEME_ROOT_ID, Handle::SchemeRoot);
        Self {
            devices,
            handles,
            next_id: 1,
        }
    }

    fn alloc_list(&mut self) -> usize {
        let id = self.next_id;
        self.next_id = self
            .next_id
            .checked_add(1)
            .expect("diskd: handle id overflow");
        self.handles.insert(id, Handle::List);
        id
    }

    fn handle(&self, id: usize) -> Result<&Handle> {
        self.handles.get(&id).ok_or_else(|| Error::new(EBADF))
    }

    fn device(&self, name: &str) -> Option<&BlockDevice> {
        self.devices.iter().find(|d| d.name == name)
    }
}

fn parse_partition_entry(line: &str) -> Option<(u32, Option<u32>)> {
    let trimmed = line.trim();
    if trimmed.is_empty() {
        return None;
    }
    if let Some(idx) = trimmed.find('p') {
        let (disk_part, part_part) = trimmed.split_at(idx);
        let part_str = &part_part[1..];
        let disk = disk_part.parse::<u32>().ok()?;
        let part = part_str.parse::<u32>().ok()?;
        Some((disk, Some(part)))
    } else {
        let disk = trimmed.parse::<u32>().ok()?;
        Some((disk, None))
    }
}

fn probe_scheme_listing(scheme_name: &str) -> Vec<(u32, Option<u32>)> {
    let path = format!("/scheme/{scheme_name}");
    let mut result = Vec::new();
    for _ in 0..PROBE_RETRY {
        match Fd::open(&path, O_DIRECTORY as i32, 0) {
            Ok(fd) => {
                let mut buffer = [0u8; 4096];
                match fd.read(&mut buffer) {
                    Ok(0) => return result,
                    Ok(n) => {
                        let text = String::from_utf8_lossy(&buffer[..n]);
                        for line in text.lines() {
                            if let Some(entry) = parse_partition_entry(line)
                                && !result.contains(&entry)
                            {
                                result.push(entry);
                            }
                        }
                        return result;
                    }
                    Err(err) => {
                        warn!("diskd: read {path} failed: {err}");
                        return result;
                    }
                }
            }
            Err(err) if err.errno() == ENOENT => {
                std::thread::sleep(PROBE_DELAY);
                continue;
            }
            Err(err) => {
                warn!("diskd: open {path} failed: {err}");
                return result;
            }
        }
    }
    result
}

/// Read the list of scheme names from /scheme/ that start with DISK_PREFIX.
///
/// This is the idiomatic Redox approach: enumerate the namespace manager's
/// getdents output (which lists all registered schemes) and filter by prefix.
/// Works with any disk.* scheme regardless of naming convention — no hardcoded
/// prefix list needed.
fn discover_disk_schemes() -> Vec<String> {
    let mut schemes = Vec::new();

    for _ in 0..PROBE_RETRY {
        match Fd::open("/scheme/", O_DIRECTORY as i32, 0) {
            Ok(fd) => {
                // The /scheme/ directory is served by initnsmgr. Its read()
                // returns newline-separated scheme names (the text encoding of
                // getdents from the namespace manager).
                let mut buf = [0u8; 8192];
                match fd.read(&mut buf) {
                    Ok(0) | Err(_) => break,
                    Ok(n) => {
                        let text = String::from_utf8_lossy(&buf[..n]);
                        for line in text.lines() {
                            let name = line.trim();
                            if name.starts_with(DISK_PREFIX) && !name.is_empty() {
                                schemes.push(name.to_string());
                            }
                        }
                        break;
                    }
                }
            }
            Err(_) => {
                std::thread::sleep(PROBE_DELAY);
            }
        }
    }

    schemes.sort();
    schemes.dedup();
    schemes
}

fn scan_devices() -> Vec<BlockDevice> {
    let disk_schemes = discover_disk_schemes();
    if disk_schemes.is_empty() {
        warn!("diskd: no disk.* schemes found in /scheme/");
        return Vec::new();
    }

    info!(
        "diskd: found {} disk scheme(s): {:?}",
        disk_schemes.len(),
        disk_schemes
    );

    let mut devices = Vec::new();
    for scheme_name in &disk_schemes {
        let entries = probe_scheme_listing(scheme_name);
        if entries.is_empty() {
            continue;
        }
        info!(
            "diskd: discovered scheme {scheme_name} with {} entries",
            entries.len()
        );
        for (disk_index, partition) in entries {
            devices.push(BlockDevice::new(scheme_name.clone(), disk_index, partition));
        }
    }

    devices.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.scheme.cmp(&b.scheme)));
    let mut deduped: Vec<BlockDevice> = Vec::with_capacity(devices.len());
    for d in devices {
        if !deduped
            .iter()
            .any(|x| x.name == d.name && x.scheme == d.scheme)
        {
            deduped.push(d);
        }
    }
    deduped
}

impl SchemeSync for DiskdScheme {
    fn scheme_root(&mut self) -> Result<usize> {
        Ok(SCHEME_ROOT_ID)
    }

    fn openat(
        &mut self,
        dirfd: usize,
        path: &str,
        flags: usize,
        _fcntl_flags: u32,
        ctx: &CallerCtx,
    ) -> Result<OpenResult> {
        if !matches!(self.handle(dirfd)?, Handle::SchemeRoot) {
            return Err(Error::new(EACCES));
        }

        let trimmed = path.trim_matches('/');

        if trimmed.is_empty() {
            if flags & (O_DIRECTORY | O_STAT) == 0 {
                return Err(Error::new(EINVAL));
            }
            if ctx.uid != 0 {
                return Err(Error::new(EACCES));
            }
            let id = self.alloc_list();
            return Ok(OpenResult::ThisScheme {
                number: id,
                flags: NewFdFlags::empty(),
            });
        }

        if ctx.uid != 0 {
            return Err(Error::new(EACCES));
        }

        let device = self.device(trimmed).ok_or(Error::new(ENOENT))?;
        let underlying = device.underlying_path.clone();
        let fd = Fd::open(&device.underlying_path, (flags & O_ACCMODE) as i32, 0).inspect_err(
            |err| {
                warn!("diskd: failed to open {underlying} for caller: {err}");
            },
        )?;
        Ok(OpenResult::OtherScheme { fd: fd.into_raw() })
    }

    fn getdents<'buf>(
        &mut self,
        id: usize,
        mut buf: DirentBuf<&'buf mut [u8]>,
        opaque_offset: u64,
    ) -> Result<DirentBuf<&'buf mut [u8]>> {
        if !matches!(self.handle(id)?, Handle::List) {
            return Err(Error::new(ENOTDIR));
        }
        let offset = usize::try_from(opaque_offset).unwrap_or(usize::MAX);
        for (i, device) in self.devices.iter().enumerate().skip(offset) {
            if let Err(err) = buf.entry(DirEntry {
                inode: 0,
                next_opaque_id: (i as u64) + 1,
                name: &device.name,
                kind: DirentKind::BlockDev,
            }) {
                if err.errno == EINVAL {
                    break;
                }
                return Err(err);
            }
        }
        Ok(buf)
    }

    fn fstat(&mut self, id: usize, stat: &mut Stat, _ctx: &CallerCtx) -> Result<()> {
        match self.handle(id)? {
            Handle::SchemeRoot => {
                stat.st_mode = MODE_DIR;
                stat.st_size = 0;
                Ok(())
            }
            Handle::List => {
                stat.st_mode = MODE_DIR;
                stat.st_size = 0;
                Ok(())
            }
        }
    }

    fn fpath(&mut self, id: usize, buf: &mut [u8], _ctx: &CallerCtx) -> Result<usize> {
        let mut writer = FpathBuf::new(buf);
        write!(&mut writer, "{SCHEME_NAME}:").map_err(|_| Error::new(EINVAL))?;
        match self.handle(id)? {
            Handle::SchemeRoot => {}
            Handle::List => {
                write!(&mut writer, "/").map_err(|_| Error::new(EINVAL))?;
            }
        }
        Ok(writer.written())
    }

    fn read(
        &mut self,
        _id: usize,
        _buf: &mut [u8],
        _offset: u64,
        _fcntl_flags: u32,
        _ctx: &CallerCtx,
    ) -> Result<usize> {
        Err(Error::new(EBADF))
    }

    fn write(
        &mut self,
        _id: usize,
        _buf: &[u8],
        _offset: u64,
        _fcntl_flags: u32,
        _ctx: &CallerCtx,
    ) -> Result<usize> {
        Err(Error::new(EBADF))
    }

    fn fsize(&mut self, id: usize, _ctx: &CallerCtx) -> Result<u64> {
        match self.handle(id)? {
            Handle::SchemeRoot | Handle::List => Ok(0),
        }
    }

    fn on_close(&mut self, id: usize) {
        if id == SCHEME_ROOT_ID {
            return;
        }
        self.handles.remove(&id);
    }
}

struct FpathBuf<'a> {
    buf: &'a mut [u8],
    written: usize,
}

impl<'a> FpathBuf<'a> {
    fn new(buf: &'a mut [u8]) -> Self {
        Self { buf, written: 0 }
    }

    fn written(&self) -> usize {
        self.written
    }
}

impl Write for FpathBuf<'_> {
    fn write(&mut self, src: &[u8]) -> io::Result<usize> {
        if self.written >= self.buf.len() {
            return Ok(0);
        }
        let avail = self.buf.len() - self.written;
        let count = src.len().min(avail);
        self.buf[self.written..self.written + count].copy_from_slice(&src[..count]);
        self.written += count;
        Ok(count)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(target_os = "redox")]
fn enter_null_namespace() {
    if let Err(err) = libredox::call::setrens(0, 0) {
        error!("diskd: setrens(0, 0) failed: {err}");
    }
}

#[cfg(not(target_os = "redox"))]
fn enter_null_namespace() {}

#[cfg(target_os = "redox")]
fn run_daemon() -> io::Result<()> {
    enter_null_namespace();

    let devices = scan_devices();
    info!("diskd: discovered {} block device entries", devices.len());
    for d in &devices {
        info!("diskd:   {} -> {}", d.name, d.underlying_path);
    }

    let socket = Socket::create()
        .map_err(|err| io::Error::other(format!("diskd: failed to create scheme socket: {err}")))?;
    let mut state = SchemeState::new();
    let mut scheme = DiskdScheme::new(devices);

    info!("diskd: scheme {SCHEME_NAME} ready");

    loop {
        let request = match socket.next_request(SignalBehavior::Restart) {
            Ok(Some(req)) => req,
            Ok(None) => {
                info!("diskd: scheme socket closed; exiting");
                return Ok(());
            }
            Err(err) if err.errno == EINTR => continue,
            Err(err) => {
                error!("diskd: next_request failed: {err}");
                return Err(io::Error::other(format!("diskd: {err}")));
            }
        };

        if let RequestKind::Call(call_request) = request.kind() {
            let response = call_request.handle_sync(&mut scheme, &mut state);
            if let Err(err) = socket.write_response(response, SignalBehavior::Restart) {
                error!("diskd: write_response failed: {err}");
                return Err(io::Error::other(format!("diskd: {err}")));
            }
        }
    }
}

#[cfg(not(target_os = "redox"))]
fn run_daemon() -> io::Result<()> {
    info!("diskd: host build: scheme serving disabled outside Redox");
    Ok(())
}

fn main() {
    let _ = log::set_logger(&LOGGER);
    log::set_max_level(LevelFilter::Info);

    match run_daemon() {
        Ok(()) => process::exit(0),
        Err(err) => {
            error!("diskd: fatal: {err}");
            process::exit(1);
        }
    }
}