initial version

.gitignore

@@ -0,0 +1,2 @@
+target
+Cargo.lock

Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "nbs"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+nbd = { version = "0.2.3", git = "https://github.com/icewind1991/rust-nbd" }
+tracing = "0.1.37"
+tracing-subscriber = "0.3.16"
+thiserror = "1.0.38"
+miette = "5.5.0"
+serde = { version = "1.0.152", features = ["derive"] }
+toml = "0.7.2"
+clap = { version = "4.1.6", features = ["derive"] }
+signal-hook = "0.3.15"

README.md

@@ -0,0 +1,39 @@
+# NBS
+
+NBD Block Server
+
+## What
+
+A basic NBD block server with a single gimmick.
+
+## Why
+
+The main reason for using this over any other NBD server is its ability to reload the config without affecting any existing connection.
+
+This allows for booting a device off an NBD device and changing the export configuration to point to a new root image. Without affecting the booted devices.
+Then, when the device is rebooted, it will connect to the new root image.
+
+## How
+
+Create a config file `config.toml`
+
+```toml
+[listen]
+port = 10809
+
+[exports]
+main = { path = "./src/main.rs", readonly = true }
+block = "/tmp/block.bin"
+```
+
+Run the server with
+
+```bash
+nbs -c config.toml
+```
+
+When the configuration is changed, it can be reloaded by sending `SIGHUP` to the server.
+
+```bash
+pkill -sighup nbs
+```

config.sample.toml

@@ -0,0 +1,7 @@
+[listen]
+port = 10809 # optional, defaults to 10809
+address = "127.0.0.1" # optional, defaults to 0.0.0.
+
+[exports]
+main = { path = "./src/main.rs", readonly = true }
+block = "/tmp/block.bin"

src/config.rs

@@ -0,0 +1,155 @@
+use crate::error::{ConfigError, HandshakeError};
+use nbd::Export;
+use serde::Deserialize;
+use std::collections::HashMap;
+use std::fmt::{Display, Formatter};
+use std::fs::{read_to_string, File, OpenOptions};
+use std::net::ToSocketAddrs;
+use std::path::{Path, PathBuf};
+use std::sync::{Arc, Mutex};
+
+#[derive(Debug, Deserialize)]
+pub struct Config {
+    #[serde(default)]
+    pub listen: ListenConfig,
+    pub exports: Exports,
+}
+
+impl Config {
+    pub fn load(path: impl AsRef<Path>) -> Result<Config, ConfigError> {
+        let content = read_to_string(path)?;
+        Ok(toml::from_str(&content)?)
+    }
+}
+
+#[derive(Debug, Deserialize)]
+pub struct ListenConfig {
+    #[serde(default = "default_address")]
+    pub address: String,
+    #[serde(default = "default_port")]
+    pub port: u16,
+}
+
+impl Display for ListenConfig {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}", self.address, self.port)
+    }
+}
+
+impl Default for ListenConfig {
+    fn default() -> Self {
+        ListenConfig {
+            address: default_address(),
+            port: default_port(),
+        }
+    }
+}
+
+fn default_address() -> String {
+    "0.0.0.0".into()
+}
+
+fn default_port() -> u16 {
+    10809
+}
+
+impl ToSocketAddrs for ListenConfig {
+    type Iter = <(String, u16) as ToSocketAddrs>::Iter;
+
+    fn to_socket_addrs(&self) -> std::io::Result<Self::Iter> {
+        (self.address.as_str(), self.port).to_socket_addrs()
+    }
+}
+
+#[derive(Deserialize)]
+#[serde(untagged)]
+enum ExportConfigDe {
+    Simple(PathBuf),
+    Options(ExportConfigRaw),
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct ExportConfigRaw {
+    #[serde(default)]
+    pub readonly: bool,
+    pub path: PathBuf,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+#[serde(from = "ExportConfigDe")]
+pub struct ExportConfig {
+    pub readonly: bool,
+    pub path: PathBuf,
+}
+
+impl From<ExportConfigDe> for ExportConfig {
+    fn from(value: ExportConfigDe) -> Self {
+        match value {
+            ExportConfigDe::Simple(path) => ExportConfig {
+                readonly: false,
+                path,
+            },
+            ExportConfigDe::Options(export) => ExportConfig {
+                readonly: export.readonly,
+                path: export.path,
+            },
+        }
+    }
+}
+
+impl ExportConfig {
+    pub fn export(&self) -> Result<Export<File>, HandshakeError> {
+        let meta = self.path.metadata().map_err(|e| HandshakeError::Open {
+            err: e,
+            path: self.path.clone(),
+        })?;
+        let file = OpenOptions::new()
+            .read(true)
+            .write(!self.readonly)
+            .open(&self.path)
+            .map_err(|e| HandshakeError::Open {
+                err: e,
+                path: self.path.clone(),
+            })?;
+        Ok(Export {
+            readonly: self.readonly || meta.permissions().readonly(),
+            size: meta.len(),
+            data: file,
+            resizeable: false,
+            rotational: false,
+            send_trim: false,
+            send_flush: false,
+        })
+    }
+}
+
+#[derive(Debug, Clone, Deserialize)]
+#[serde(from = "HashMap<String, ExportConfig>")]
+pub struct Exports {
+    exports: Arc<Mutex<HashMap<String, ExportConfig>>>,
+}
+
+impl From<HashMap<String, ExportConfig>> for Exports {
+    fn from(value: HashMap<String, ExportConfig>) -> Self {
+        Exports {
+            exports: Arc::new(Mutex::new(value)),
+        }
+    }
+}
+
+impl Exports {
+    pub fn get(&self, name: &str) -> Option<ExportConfig> {
+        self.exports.lock().unwrap().get(name).cloned()
+    }
+
+    #[allow(dead_code)]
+    pub fn update(&self, other: Exports) -> usize {
+        let other = match Arc::try_unwrap(other.exports) {
+            Ok(mutex) => mutex.into_inner().unwrap(),
+            Err(arc) => arc.lock().unwrap().clone(),
+        };
+        let count = other.len();
+        *self.exports.lock().unwrap() = other;
+        count
+    }
+}

src/error.rs

@@ -0,0 +1,75 @@
+use miette::Diagnostic;
+use std::io::{Error as IoError, ErrorKind};
+use std::path::PathBuf;
+use thiserror::Error;
+
+#[derive(Debug, Error, Diagnostic)]
+pub enum Error {
+    #[error("Failed to listen to port")]
+    Connection(IoError),
+    #[error("Failed to listen to port")]
+    Listen(IoError),
+    #[error(transparent)]
+    #[diagnostic(transparent)]
+    Nbd(#[from] NbdError),
+    #[error(transparent)]
+    #[diagnostic(transparent)]
+    Config(#[from] ConfigError),
+}
+
+#[derive(Debug, Error, Diagnostic)]
+pub enum NbdError {
+    #[error(transparent)]
+    Io(IoError),
+    #[error("Client disconnected unexpectedly")]
+    Disconnected,
+    #[error(transparent)]
+    #[diagnostic(transparent)]
+    Handshake(HandshakeError),
+}
+
+impl From<NbdError> for IoError {
+    fn from(value: NbdError) -> Self {
+        match value {
+            NbdError::Handshake(e) => e.into(),
+            NbdError::Io(e) => e,
+            NbdError::Disconnected => ErrorKind::UnexpectedEof.into(),
+        }
+    }
+}
+
+impl From<IoError> for NbdError {
+    fn from(value: IoError) -> Self {
+        match value.kind() {
+            ErrorKind::UnexpectedEof => NbdError::Disconnected,
+            _ => NbdError::Io(value),
+        }
+    }
+}
+
+#[derive(Debug, Error, Diagnostic)]
+pub enum HandshakeError {
+    #[error("Unknown export {0}")]
+    UnknownExport(String),
+    #[error("Failed to open {path}")]
+    Open { path: PathBuf, err: IoError },
+}
+
+impl From<HandshakeError> for IoError {
+    fn from(value: HandshakeError) -> Self {
+        match value {
+            HandshakeError::Open { err, .. } => err,
+            HandshakeError::UnknownExport(export) => {
+                IoError::new(ErrorKind::InvalidData, format!("Unknown export: {export}"))
+            }
+        }
+    }
+}
+
+#[derive(Debug, Error, Diagnostic)]
+pub enum ConfigError {
+    #[error("Failed to read config file")]
+    Read(#[from] IoError),
+    #[error("Failed to parse config file")]
+    Parse(#[from] toml::de::Error),
+}

src/main.rs

@@ -0,0 +1,80 @@
+mod config;
+mod error;
+
+use crate::config::{Config, Exports};
+use crate::error::{Error, HandshakeError, NbdError};
+use clap::Parser;
+use nbd::server::{handshake, transmission};
+use signal_hook::consts::SIGHUP;
+use signal_hook::iterator::exfiltrator::SignalOnly;
+use signal_hook::iterator::SignalsInfo;
+use std::net::{TcpListener, TcpStream};
+use std::path::PathBuf;
+use std::thread::spawn;
+use tracing::{error, info};
+
+#[derive(Debug, Parser)]
+struct Args {
+    #[arg(short, long)]
+    config: PathBuf,
+}
+
+pub type Result<T, E = Error> = std::result::Result<T, E>;
+
+#[tracing::instrument(skip(exports))]
+fn handle_client(mut stream: TcpStream, exports: Exports) -> Result<(), NbdError> {
+    let file = handshake(&mut stream, move |name| {
+        let export_cfg = exports
+            .get(name)
+            .ok_or_else(|| HandshakeError::UnknownExport(name.into()))?;
+        info!(name = name, export = ?export_cfg, "opening export");
+        Ok(export_cfg.export()?)
+    })?;
+    info!("connected");
+    transmission(&mut stream, file)?;
+    info!("disconnected");
+    Ok(())
+}
+
+fn main() -> Result<()> {
+    tracing_subscriber::fmt::init();
+    let args = Args::parse();
+    let config = Config::load(&args.config)?;
+    let listener = TcpListener::bind(&config.listen).map_err(Error::Listen)?;
+    info!("Listening on {}", config.listen);
+
+    let exports = config.exports.clone();
+    spawn(move || {
+        let mut reload_signals = SignalsInfo::<SignalOnly>::new([SIGHUP]).unwrap();
+        for _ in &mut reload_signals {
+            info!("Reloading config");
+            match Config::load(&args.config) {
+                Ok(updated) => {
+                    let count = exports.update(updated.exports);
+                    info!("Registered {count} exports");
+                }
+                Err(e) => {
+                    error!(error = ?e, "Failed to load updated config");
+                }
+            }
+        }
+    });
+
+    for stream in listener.incoming() {
+        match stream {
+            Ok(stream) => {
+                let exports = config.exports.clone();
+                spawn(move || {
+                    if let Err(e) = handle_client(stream, exports) {
+                        error!("{e}");
+                    }
+                });
+            }
+            Err(e) => {
+                let e = Error::Connection(e);
+                error!("{e}");
+            }
+        }
+    }
+    Ok(())
+}