wavehunter: add QMDL storage

Instead of reading/writing to a single QMDL file, we now can manage
a directory of several files, and have the ability to start/stop writing
to them on the fly.

This commit also adds graceful exiting to the server, so we can perform
cleanup steps when the server's exiting.

wavehunter/src/config.rs

@@ -5,36 +5,40 @@ use serde::Deserialize;
 #[derive(Deserialize)]
 struct ConfigFile {
     qmdl_path: Option<String>,
+    qmdl_store_path: Option<String>,
     port: Option<u16>,
-    debug_mode: Option<bool>,
+    readonly_mode: Option<bool>,
 }
 
 #[derive(Debug)]
 pub struct Config {
     pub qmdl_path: String,
+    pub qmdl_store_path: String,
     pub port: u16,
-    pub debug_mode: bool,
+    pub readonly_mode: bool,
 }
 
 impl Default for Config {
     fn default() -> Self {
         Config {
             qmdl_path: "./wavehunter.qmdl".to_string(),
+            qmdl_store_path: "/data/wavehunter".to_string(),
             port: 8080,
-            debug_mode: false,
+            readonly_mode: false,
         }
     }
 }
 
 pub fn parse_config<P>(path: P) -> Result<Config, WavehunterError> where P: AsRef<std::path::Path> {
-    let config_file = std::fs::read_to_string(&path)
-        .map_err(|_| WavehunterError::MissingConfigFile(format!("{:?}", path.as_ref())))?;
-    let parsed_config: ConfigFile = toml::from_str(&config_file)
-        .map_err(WavehunterError::ConfigFileParsingError)?;
     let mut config = Config::default();
-    if let Some(path) = parsed_config.qmdl_path { config.qmdl_path = path }
-    if let Some(port) = parsed_config.port { config.port = port }
-    if let Some(debug_mode) = parsed_config.debug_mode { config.debug_mode = debug_mode }
+    if let Ok(config_file) = std::fs::read_to_string(&path) {
+        let parsed_config: ConfigFile = toml::from_str(&config_file)
+            .map_err(WavehunterError::ConfigFileParsingError)?;
+        if let Some(path) = parsed_config.qmdl_path { config.qmdl_path = path }
+        if let Some(path) = parsed_config.qmdl_store_path { config.qmdl_store_path = path }
+        if let Some(port) = parsed_config.port { config.port = port }
+        if let Some(debug_mode) = parsed_config.readonly_mode { config.readonly_mode = debug_mode }
+    }
     Ok(config)
 }
 

wavehunter/src/diag.rs

@@ -0,0 +1,99 @@
+use std::sync::Arc;
+
+use axum::extract::State;
+use axum::http::StatusCode;
+use orca::diag_device::DiagDevice;
+use orca::diag_reader::DiagReader;
+use tokio::sync::RwLock;
+use tokio::sync::mpsc::{Receiver, self};
+use orca::qmdl::QmdlWriter;
+use log::{debug, info};
+use tokio::sync::mpsc::error::TryRecvError;
+use tokio::task::JoinHandle;
+use tokio_util::task::TaskTracker;
+
+use crate::error::WavehunterError;
+use crate::qmdl_store::QmdlStore;
+use crate::server::ServerState;
+
+pub enum DiagDeviceCtrlMessage {
+    StopRecording,
+    StartRecording(QmdlWriter<std::fs::File>),
+    Exit,
+}
+
+pub fn run_diag_read_thread(task_tracker: &TaskTracker, mut dev: DiagDevice, mut qmdl_file_rx: Receiver<DiagDeviceCtrlMessage>, qmdl_store_lock: Arc<RwLock<QmdlStore>>) -> JoinHandle<Result<(), WavehunterError>> {
+    let (tx, mut rx) = mpsc::channel::<(usize, usize)>(1);
+    let qmdl_store_lock_clone = qmdl_store_lock.clone();
+    task_tracker.spawn(async move {
+        while let Some((entry_idx, new_size)) = rx.recv().await {
+            let mut qmdl_store = qmdl_store_lock_clone.write().await;
+            qmdl_store.update_entry_size(entry_idx, new_size).await
+                .expect("failed to update qmdl file size");
+        }
+        info!("QMDL store size updater thread exiting...");
+    });
+
+    task_tracker.spawn_blocking(move || {
+        loop {
+            match qmdl_file_rx.try_recv() {
+                Ok(DiagDeviceCtrlMessage::StartRecording(qmdl_writer)) => {
+                    dev.qmdl_writer = Some(qmdl_writer);
+                },
+                Ok(DiagDeviceCtrlMessage::StopRecording) => dev.qmdl_writer = None,
+                Ok(DiagDeviceCtrlMessage::Exit) | Err(TryRecvError::Disconnected) => {
+                    info!("Diag reader thread exiting...");
+                    return Ok(())
+                },
+                // empty just means there's no message for us, so continue as normal
+                Err(TryRecvError::Empty) => {},
+            }
+
+            // remember the QmdlStore current entry index so we can update its size later
+            let qmdl_store_index = qmdl_store_lock.blocking_read().current_entry;
+
+            // TODO: once we're actually doing analysis, we'll wanna use the messages
+            // returned here. Until then, the DiagDevice has already written those messages
+            // to the QMDL file, so we can just ignore them.
+            debug!("reading response from diag device...");
+            let _messages = dev.read_response().map_err(WavehunterError::DiagReadError)?;
+            debug!("got diag response ({} messages)", _messages.len());
+
+            // keep track of how many bytes were written to the QMDL file so we can read
+            // a valid block of data from it in the HTTP server
+            if let Some(qmdl_writer) = dev.qmdl_writer.as_ref() {
+                debug!("total QMDL bytes written: {}, sending update...", qmdl_writer.total_written);
+                let index = qmdl_store_index.expect("DiagDevice had qmdl_writer, but QmdlStore didn't have current entry???");
+                tx.blocking_send((index, qmdl_writer.total_written)).unwrap();
+                debug!("done!");
+            } else {
+                debug!("no qmdl_writer set, continuing...");
+            }
+        }
+    })
+}
+
+pub async fn start_recording(State(state): State<Arc<ServerState>>) -> Result<(StatusCode, String), (StatusCode, String)> {
+    if state.readonly_mode {
+        return Err((StatusCode::FORBIDDEN, format!("server is in readonly mode")));
+    }
+    let mut qmdl_store = state.qmdl_store_lock.write().await;
+    let qmdl_file = qmdl_store.new_entry().await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("couldn't create new qmdl entry: {}", e)))?;
+    let qmdl_writer = QmdlWriter::new(qmdl_file.into_std().await);
+    state.diag_device_ctrl_sender.send(DiagDeviceCtrlMessage::StartRecording(qmdl_writer)).await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("couldn't send stop recording message: {}", e)))?;
+    Ok((StatusCode::ACCEPTED, format!("ok")))
+}
+
+pub async fn stop_recording(State(state): State<Arc<ServerState>>) -> Result<(StatusCode, String), (StatusCode, String)> {
+    if state.readonly_mode {
+        return Err((StatusCode::FORBIDDEN, format!("server is in readonly mode")));
+    }
+    let mut qmdl_store = state.qmdl_store_lock.write().await;
+    qmdl_store.close_current_entry().await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("couldn't close current qmdl entry: {}", e)))?;
+    state.diag_device_ctrl_sender.send(DiagDeviceCtrlMessage::StopRecording).await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("couldn't send stop recording message: {}", e)))?;
+    Ok((StatusCode::ACCEPTED, format!("ok")))
+}

wavehunter/src/error.rs

@@ -1,10 +1,10 @@
 use thiserror::Error;
 use orca::diag_device::DiagDeviceError;
 
+use crate::qmdl_store::QmdlStoreError;
+
 #[derive(Error, Debug)]
 pub enum WavehunterError {
-    #[error("Missing config file: {0}")]
-    MissingConfigFile(String),
     #[error("Config file parsing error: {0}")]
     ConfigFileParsingError(#[from] toml::de::Error),
     #[error("Diag intialization error: {0}")]
@@ -13,4 +13,8 @@ pub enum WavehunterError {
     DiagReadError(DiagDeviceError),
     #[error("Tokio error: {0}")]
     TokioError(#[from] tokio::io::Error),
+    #[error("QmdlStore error: {0}")]
+    QmdlStoreError(#[from] QmdlStoreError),
+    #[error("No QMDL store found at path {0}, but can't create a new one due to readonly mode")]
+    NoStoreReadonlyMode(String),
 }

wavehunter/src/main.rs

@@ -3,64 +3,106 @@ mod error;
 mod pcap;
 mod server;
 mod stats;
+mod qmdl_store;
+mod diag;
 
 use crate::config::{parse_config, parse_args};
+use crate::diag::run_diag_read_thread;
+use crate::qmdl_store::QmdlStore;
 use crate::server::{ServerState, get_qmdl, serve_static};
 use crate::pcap::get_pcap;
-use crate::stats::{get_system_stats, get_diag_stats};
+use crate::stats::get_system_stats;
 use crate::error::WavehunterError;
 
 use axum::response::Redirect;
+use diag::{DiagDeviceCtrlMessage, start_recording, stop_recording};
+use log::{info, error};
 use orca::diag_device::DiagDevice;
-use orca::diag_reader::DiagReader;
-use axum::routing::get;
+use axum::routing::{get, post};
 use axum::Router;
-use tokio::fs::File;
-use log::debug;
+use orca::qmdl::QmdlWriter;
+use stats::get_qmdl_manifest;
+use tokio::sync::mpsc::{self, Sender};
+use tokio::task::JoinHandle;
+use tokio_util::task::TaskTracker;
 use std::net::SocketAddr;
 use tokio::net::TcpListener;
-use tokio::sync::RwLock;
-use tokio::task::JoinHandle;
+use tokio::sync::{RwLock, oneshot};
 use std::sync::Arc;
 
-fn run_diag_read_thread(mut dev: DiagDevice, bytes_read_lock: Arc<RwLock<usize>>) -> JoinHandle<Result<(), WavehunterError>> {
-    tokio::task::spawn_blocking(move || {
-        loop {
-            // TODO: once we're actually doing analysis, we'll wanna use the messages
-            // returned here. Until then, the DiagDevice has already written those messages
-            // to the QMDL file, so we can just ignore them.
-            debug!("reading response from diag device...");
-            let _messages = dev.read_response().map_err(WavehunterError::DiagReadError)?;
-            debug!("got diag response ({} messages)", _messages.len());
-
-            // keep track of how many bytes were written to the QMDL file so we can read
-            // a valid block of data from it in the HTTP server
-            debug!("total QMDL bytes written: {}, updating state...", dev.qmdl_writer.total_written);
-            let mut bytes_read = bytes_read_lock.blocking_write();
-            *bytes_read = dev.qmdl_writer.total_written;
-            debug!("done!");
-        }
-    })
-}
-
-async fn run_server(config: &config::Config, qmdl_bytes_written: Arc<RwLock<usize>>) -> Result<(), WavehunterError> {
+async fn run_server(
+    task_tracker: &TaskTracker,
+    config: &config::Config,
+    qmdl_store_lock: Arc<RwLock<QmdlStore>>,
+    server_shutdown_rx: oneshot::Receiver<()>,
+    diag_device_sender: Sender<DiagDeviceCtrlMessage>
+) -> JoinHandle<()> {
     let state = Arc::new(ServerState {
-        qmdl_bytes_written,
-        qmdl_path: config.qmdl_path.clone(),
+        qmdl_store_lock,
+        diag_device_ctrl_sender: diag_device_sender,
+        readonly_mode: config.readonly_mode,
     });
 
     let app = Router::new()
-        .route("/api/pcap/latest.pcap", get(get_pcap))
-        .route("/api/qmdl/latest.qmdl", get(get_qmdl))
+        .route("/api/pcap/*name", get(get_pcap))
+        .route("/api/qmdl/*name", get(get_qmdl))
         .route("/api/system-stats", get(get_system_stats))
-        .route("/api/diag-stats", get(get_diag_stats))
+        .route("/api/qmdl-manifest", get(get_qmdl_manifest))
+        .route("/api/start-recording", post(start_recording))
+        .route("/api/stop-recording", post(stop_recording))
         .route("/", get(|| async { Redirect::permanent("/index.html") }))
         .route("/*path", get(serve_static))
         .with_state(state);
     let addr = SocketAddr::from(([127, 0, 0, 1], config.port));
     let listener = TcpListener::bind(&addr).await.unwrap();
-    axum::serve(listener, app).await.unwrap();
-    Ok(())
+    task_tracker.spawn(async move {
+        info!("The orca is hunting for stingrays...");
+        axum::serve(listener, app)
+            .with_graceful_shutdown(server_shutdown_signal(server_shutdown_rx))
+            .await.unwrap();
+    })
+}
+
+async fn server_shutdown_signal(server_shutdown_rx: oneshot::Receiver<()>) {
+    server_shutdown_rx.await.unwrap();
+    info!("Server received shutdown signal, exiting...");
+}
+
+async fn init_qmdl_store(config: &config::Config) -> Result<QmdlStore, WavehunterError> {
+    match (QmdlStore::exists(&config.qmdl_store_path).await?, config.readonly_mode) {
+        (true, _) => Ok(QmdlStore::load(&config.qmdl_store_path).await?),
+        (false, false) => Ok(QmdlStore::create(&config.qmdl_store_path).await?),
+        (false, true) => Err(WavehunterError::NoStoreReadonlyMode(config.qmdl_store_path.clone())),
+    }
+}
+
+fn run_ctrl_c_thread(
+    task_tracker: &TaskTracker,
+    diag_device_sender: Sender<DiagDeviceCtrlMessage>,
+    server_shutdown_tx: oneshot::Sender<()>,
+    qmdl_store_lock: Arc<RwLock<QmdlStore>>
+) -> JoinHandle<Result<(), WavehunterError>> {
+    task_tracker.spawn(async move {
+        match tokio::signal::ctrl_c().await {
+            Ok(()) => {
+                let mut qmdl_store = qmdl_store_lock.write().await;
+                if qmdl_store.current_entry.is_some() {
+                    info!("Closing current QMDL entry...");
+                    qmdl_store.close_current_entry().await?;
+                    info!("Done!");
+                }
+
+                server_shutdown_tx.send(())
+                    .expect("couldn't send server shutdown signal");
+                diag_device_sender.send(DiagDeviceCtrlMessage::Exit).await
+                    .expect("couldn't send Exit message to diag thread");
+            },
+            Err(err) => {
+                error!("Unable to listen for shutdown signal: {}", err);
+            }
+        }
+        Ok(())
+    })
 }
 
 #[tokio::main]
@@ -70,23 +112,27 @@ async fn main() -> Result<(), WavehunterError> {
     let args = parse_args();
     let config = parse_config(&args.config_path)?;
 
-    let qmdl_bytes_lock: Arc<RwLock<usize>>;
-    if !config.debug_mode {
-        let mut dev = DiagDevice::new(&config.qmdl_path)
+    let task_tracker = TaskTracker::new();
+
+    let qmdl_store_lock = Arc::new(RwLock::new(init_qmdl_store(&config).await?));
+    let (tx, rx) = mpsc::channel::<DiagDeviceCtrlMessage>(1);
+    if !config.readonly_mode {
+        let qmdl_file = qmdl_store_lock.write().await.new_entry().await?;
+        let qmdl_writer = QmdlWriter::new(qmdl_file.into_std().await);
+        let mut dev = DiagDevice::new(Some(qmdl_writer))
             .map_err(WavehunterError::DiagInitError)?;
         dev.config_logs()
             .map_err(WavehunterError::DiagInitError)?;
-        qmdl_bytes_lock = Arc::new(RwLock::new(dev.qmdl_writer.total_written));
 
-        // TODO: handle exiting gracefully
-        let _read_thread_handle = run_diag_read_thread(dev, qmdl_bytes_lock.clone());
-    } else {
-        let qmdl_file = File::open(&config.qmdl_path).await.expect("couldn't open QMDL file");
-        let qmdl_file_size = qmdl_file.metadata().await.expect("couldn't get QMDL file metadata")
-            .len() as usize;
-        qmdl_bytes_lock = Arc::new(RwLock::new(qmdl_file_size));
+        run_diag_read_thread(&task_tracker, dev, rx, qmdl_store_lock.clone());
     }
 
-    println!("The orca is hunting for stingrays...");
-    run_server(&config, qmdl_bytes_lock).await
+    let (server_shutdown_tx, server_shutdown_rx) = oneshot::channel::<()>();
+    run_ctrl_c_thread(&task_tracker, tx.clone(), server_shutdown_tx, qmdl_store_lock.clone());
+    run_server(&task_tracker, &config, qmdl_store_lock.clone(), server_shutdown_rx, tx).await;
+
+    task_tracker.close();
+    task_tracker.wait().await;
+
+    Ok(())
 }

wavehunter/src/pcap.rs

@@ -6,10 +6,9 @@ use orca::qmdl::{QmdlReader, QmdlReaderError};
 use orca::diag_reader::DiagReader;
 use axum::body::Body;
 use axum::http::header::CONTENT_TYPE;
-use axum::extract::State;
+use axum::extract::{State, Path};
 use axum::http::StatusCode;
 use axum::response::{Response, IntoResponse};
-use std::fs::File;
 use std::io::Write;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -22,14 +21,19 @@ use tokio::sync::mpsc;
 // written so far. This is done by spawning a blocking thread (a tokio thread
 // capable of handling blocking operations) which streams chunks of pcap data to
 // a channel that's piped to the client.
-pub async fn get_pcap(State(state): State<Arc<ServerState>>) -> Result<Response, (StatusCode, String)> {
-    let qmdl_bytes_written = *state.qmdl_bytes_written.read().await;
-    if qmdl_bytes_written == 0 {
+pub async fn get_pcap(State(state): State<Arc<ServerState>>, Path(qmdl_name): Path<String>) -> Result<Response, (StatusCode, String)> {
+    let qmdl_store = state.qmdl_store_lock.read().await;
+    let entry = qmdl_store.entry_for_name(&qmdl_name)
+        .ok_or((StatusCode::NOT_FOUND, format!("couldn't find qmdl file with name {}", qmdl_name)))?;
+    if entry.size_bytes == 0 {
         return Err((
             StatusCode::SERVICE_UNAVAILABLE,
             "QMDL file is empty, try again in a bit!".to_string()
         ));
     }
+    let qmdl_file = qmdl_store.open_entry(&entry).await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("{:?}", e)))?
+        .into_std().await;
 
     let (tx, rx) = mpsc::channel(1);
     let channel_reader = ChannelReader { rx };
@@ -37,8 +41,7 @@ pub async fn get_pcap(State(state): State<Arc<ServerState>>) -> Result<Response,
     tokio::task::spawn_blocking(move || {
         // the QMDL reader should stop at the last successfully written data
         // chunk (qmdl_bytes_written)
-        let qmdl_file = File::open(&state.qmdl_path).unwrap();
-        let mut qmdl_reader = QmdlReader::new(qmdl_file, Some(qmdl_bytes_written));
+        let mut qmdl_reader = QmdlReader::new(qmdl_file, Some(entry.size_bytes));
 
         let mut gsmtap_parser = GsmtapParser::new();
         let mut pcap_writer = GsmtapPcapWriter::new(channel_writer).unwrap();

wavehunter/src/qmdl_store.rs

@@ -0,0 +1,212 @@
+use std::path::{PathBuf, Path};
+use thiserror::Error;
+use tokio::{fs::{self, File, try_exists}, io::AsyncWriteExt};
+use serde::{Deserialize, Serialize};
+use chrono::{DateTime, Local};
+
+#[derive(Debug, Error)]
+pub enum QmdlStoreError {
+    #[error("Can't close an entry when there's no current entry")]
+    NoCurrentEntry,
+    #[error("Couldn't create file: {0}")]
+    CreateFileError(tokio::io::Error),
+    #[error("Couldn't read file: {0}")]
+    ReadFileError(tokio::io::Error),
+    #[error("Couldn't open directory at path: {0}")]
+    OpenDirError(tokio::io::Error),
+    #[error("Couldn't read manifest file: {0}")]
+    ReadManifestError(tokio::io::Error),
+    #[error("Couldn't write manifest file: {0}")]
+    WriteManifestError(tokio::io::Error),
+    #[error("Couldn't parse QMDL store manifest file: {0}")]
+    ParseManifestError(toml::de::Error)
+}
+
+pub struct QmdlStore {
+    pub path: PathBuf,
+    pub manifest: Manifest,
+    pub current_entry: Option<usize>, // index into manifest
+}
+
+#[derive(Deserialize, Serialize, Clone, PartialEq, Debug)]
+pub struct Manifest {
+    pub entries: Vec<ManifestEntry>,
+}
+
+#[derive(Deserialize, Serialize, Clone, PartialEq, Debug)]
+pub struct ManifestEntry {
+    pub name: String,
+    pub start_time: DateTime<Local>,
+    pub end_time: Option<DateTime<Local>>,
+    pub size_bytes: usize,
+}
+
+impl ManifestEntry {
+    fn new() -> Self {
+        let now = Local::now();
+        ManifestEntry {
+            name: format!("{}", now.timestamp()),
+            start_time: now,
+            end_time: None,
+            size_bytes: 0,
+        }
+    }
+}
+
+impl QmdlStore {
+    // Returns whether a directory with a "manifest.toml" exists at the given
+    // path (though doesn't check if that manifest is valid)
+    pub async fn exists<P>(path: P) -> Result<bool, QmdlStoreError> where P: AsRef<Path> {
+        let manifest_path = path.as_ref().join("manifest.toml");
+        let dir_exists = try_exists(path).await.map_err(QmdlStoreError::OpenDirError)?;
+        let manifest_exists = try_exists(manifest_path).await.map_err(QmdlStoreError::ReadManifestError)?;
+        Ok(dir_exists && manifest_exists)
+    }
+
+    // Loads an existing QmdlStore at the given path. Errors if no store exists,
+    // or if it's malformed.
+    pub async fn load<P>(path: P) -> Result<Self, QmdlStoreError> where P: AsRef<Path> {
+        let path: PathBuf = path.as_ref().to_path_buf();
+        let manifest = QmdlStore::read_manifest(&path).await?;
+        Ok(QmdlStore {
+            path,
+            manifest,
+            current_entry: None,
+        })
+    }
+
+    // Creates a new QmdlStore at the given path. This involves creating a dir
+    // and writing an empty manifest.
+    pub async fn create<P>(path: P) -> Result<Self, QmdlStoreError> where P: AsRef<Path> {
+        let manifest_path = path.as_ref().join("manifest.toml");
+        fs::create_dir_all(&path).await
+            .map_err(QmdlStoreError::OpenDirError)?;
+        let mut manifest_file = File::create(&manifest_path).await
+            .map_err(QmdlStoreError::WriteManifestError)?;
+        let empty_manifest = Manifest { entries: Vec::new() };
+        let empty_manifest_contents = toml::to_string_pretty(&empty_manifest)
+            .expect("failed to serialize manifest");
+        manifest_file.write_all(empty_manifest_contents.as_bytes()).await
+            .map_err(QmdlStoreError::WriteManifestError)?;
+        QmdlStore::load(path).await
+    }
+
+    async fn read_manifest<P>(path: P) -> Result<Manifest, QmdlStoreError> where P: AsRef<Path> {
+        let manifest_path = path.as_ref().join("manifest.toml");
+        let file_contents = fs::read_to_string(&manifest_path).await
+            .map_err(QmdlStoreError::ReadManifestError)?;
+        toml::from_str(&file_contents)
+            .map_err(QmdlStoreError::ParseManifestError)
+    }
+
+    // Closes the current entry (if needed), creates a new entry based on the
+    // current time, and updates the manifest
+    pub async fn new_entry(&mut self) -> Result<File, QmdlStoreError> {
+        // if we've already got an entry open, close it
+        if self.current_entry.is_some() {
+            self.close_current_entry().await?;
+        }
+        let new_entry = ManifestEntry::new();
+        let mut file_path = self.path.join(&new_entry.name);
+        file_path.set_extension("qmdl");
+        let file = File::options()
+            .create(true)
+            .write(true)
+            .open(&file_path).await
+            .map_err(QmdlStoreError::CreateFileError)?;
+        self.manifest.entries.push(new_entry);
+        self.current_entry = Some(self.manifest.entries.len() - 1);
+        self.write_manifest().await?;
+        Ok(file)
+    }
+
+    pub async fn open_entry(&self, entry: &ManifestEntry) -> Result<File, QmdlStoreError> {
+        let mut file_path = self.path.join(&entry.name);
+        file_path.set_extension("qmdl");
+        File::open(file_path).await
+            .map_err(QmdlStoreError::ReadFileError)
+    }
+
+    // Sets the current entry's end_time, updates the manifest, and unsets the
+    // current entry
+    pub async fn close_current_entry(&mut self) -> Result<(), QmdlStoreError> {
+        let entry_index = self.current_entry.take()
+            .ok_or(QmdlStoreError::NoCurrentEntry)?;
+        self.manifest.entries[entry_index].end_time = Some(Local::now());
+        self.write_manifest().await
+    }
+
+    // Sets the given entry's size, updating the manifest
+    pub async fn update_entry_size(&mut self, entry_index: usize, size_bytes: usize) -> Result<(), QmdlStoreError> {
+        self.manifest.entries[entry_index].size_bytes = size_bytes;
+        self.write_manifest().await
+    }
+
+    async fn write_manifest(&mut self) -> Result<(), QmdlStoreError> {
+        let mut manifest_file = File::options()
+            .write(true)
+            .open(self.path.join("manifest.toml")).await
+            .map_err(QmdlStoreError::WriteManifestError)?;
+        let manifest_contents = toml::to_string_pretty(&self.manifest)
+            .expect("failed to serialize manifest");
+        manifest_file.write_all(manifest_contents.as_bytes()).await
+            .map_err(QmdlStoreError::WriteManifestError)?;
+        Ok(())
+    }
+
+    // Finds an entry by filename
+    pub fn entry_for_name(&self, name: &str) -> Option<ManifestEntry> {
+        self.manifest.entries.iter()
+            .find(|entry| entry.name == name)
+            .cloned()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use tempdir::TempDir;
+    use super::*;
+
+    #[tokio::test]
+    async fn test_load_from_empty_dir() {
+        let dir = TempDir::new("qmdl_store_test").unwrap();
+        assert!(!QmdlStore::exists(dir.path()).await.unwrap());
+        let _created_store = QmdlStore::create(dir.path()).await.unwrap();
+        assert!(QmdlStore::exists(dir.path()).await.unwrap());
+        let loaded_store = QmdlStore::load(dir.path()).await.unwrap();
+        assert_eq!(loaded_store.manifest.entries.len(), 0);
+    }
+
+    #[tokio::test]
+    async fn test_creating_updating_and_closing_entries() {
+        let dir = TempDir::new("qmdl_store_test").unwrap();
+        let mut store = QmdlStore::create(dir.path()).await.unwrap();
+        let _ = store.new_entry().await.unwrap();
+        let entry_index = store.current_entry.unwrap();
+        assert_eq!(QmdlStore::read_manifest(dir.path()).await.unwrap(), store.manifest);
+
+        store.update_entry_size(entry_index, 1000).await.unwrap();
+        assert_eq!(store.manifest.entries[entry_index].size_bytes, 1000);
+        assert_eq!(QmdlStore::read_manifest(dir.path()).await.unwrap(), store.manifest);
+
+        assert!(store.manifest.entries[entry_index].end_time.is_none());
+        store.close_current_entry().await.unwrap();
+        let entry = store.entry_for_name(&store.manifest.entries[entry_index].name).unwrap();
+        assert!(entry.end_time.is_some());
+        assert_eq!(QmdlStore::read_manifest(dir.path()).await.unwrap(), store.manifest);
+
+        assert!(matches!(store.close_current_entry().await, Err(QmdlStoreError::NoCurrentEntry)));
+    }
+
+    #[tokio::test]
+    async fn test_repeated_new_entries() {
+        let dir = TempDir::new("qmdl_store_test").unwrap();
+        let mut store = QmdlStore::create(dir.path()).await.unwrap();
+        let _ = store.new_entry().await.unwrap();
+        let entry_index = store.current_entry.unwrap();
+        let _ = store.new_entry().await.unwrap();
+        let new_entry_index = store.current_entry.unwrap();
+        assert_ne!(entry_index, new_entry_index);
+        assert_eq!(store.manifest.entries.len(), 2);
+    }
+}

wavehunter/src/server.rs

@@ -5,22 +5,28 @@ use axum::http::{StatusCode, HeaderValue};
 use axum::response::{Response, IntoResponse};
 use axum::extract::Path;
 use tokio::io::AsyncReadExt;
+use tokio::sync::mpsc::Sender;
 use std::sync::Arc;
 use tokio::sync::RwLock;
-use tokio::fs::File as AsyncFile;
 use tokio_util::io::ReaderStream;
 use include_dir::{include_dir, Dir};
 
+use crate::DiagDeviceCtrlMessage;
+use crate::qmdl_store::QmdlStore;
+
 pub struct ServerState {
-    pub qmdl_bytes_written: Arc<RwLock<usize>>,
-    pub qmdl_path: String,
+    pub qmdl_store_lock: Arc<RwLock<QmdlStore>>,
+    pub diag_device_ctrl_sender: Sender<DiagDeviceCtrlMessage>,
+    pub readonly_mode: bool,
 }
 
-pub async fn get_qmdl(State(state): State<Arc<ServerState>>) -> Result<Response, (StatusCode, String)> {
-    let qmdl_file = AsyncFile::open(&state.qmdl_path).await
+pub async fn get_qmdl(State(state): State<Arc<ServerState>>, Path(qmdl_name): Path<String>) -> Result<Response, (StatusCode, String)> {
+    let qmdl_store = state.qmdl_store_lock.read().await;
+    let entry = qmdl_store.entry_for_name(&qmdl_name)
+        .ok_or((StatusCode::NOT_FOUND, format!("couldn't find qmdl file with name {}", qmdl_name)))?;
+    let qmdl_file = qmdl_store.open_entry(&entry).await
         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("error opening QMDL file: {}", e)))?;
-    let qmdl_bytes_written = *state.qmdl_bytes_written.read().await;
-    let limited_qmdl_file = qmdl_file.take(qmdl_bytes_written as u64);
+    let limited_qmdl_file = qmdl_file.take(entry.size_bytes as u64);
     let qmdl_stream = ReaderStream::new(limited_qmdl_file);
 
     let headers = [(CONTENT_TYPE, "application/octet-stream")];

wavehunter/src/stats.rs

@@ -1,10 +1,11 @@
+use std::sync::Arc;
+
+use crate::qmdl_store::ManifestEntry;
 use crate::server::ServerState;
 
 use axum::Json;
 use axum::extract::State;
 use axum::http::StatusCode;
-use axum::response::IntoResponse;
-use std::sync::Arc;
 use log::error;
 use serde::Serialize;
 use tokio::process::Command;
@@ -97,7 +98,8 @@ fn humanize_kb(kb: usize) -> String {
 }
 
 pub async fn get_system_stats(State(state): State<Arc<ServerState>>) -> Result<Json<SystemStats>, (StatusCode, String)> {
-    match SystemStats::new(&state.qmdl_path).await {
+    let qmdl_store = state.qmdl_store_lock.read().await;
+    match SystemStats::new(qmdl_store.path.to_str().unwrap()).await {
         Ok(stats) => Ok(Json(stats)),
         Err(err) => {
             error!("error getting system stats: {}", err);
@@ -109,13 +111,18 @@ pub async fn get_system_stats(State(state): State<Arc<ServerState>>) -> Result<J
     }
 }
 
-#[derive(Debug, Serialize)]
-pub struct DiagStats {
-    bytes_written: usize,
+#[derive(Serialize)]
+pub struct ManifestStats {
+    pub entries: Vec<ManifestEntry>,
+    pub current_entry: Option<ManifestEntry>,
 }
 
-pub async fn get_diag_stats(State(state): State<Arc<ServerState>>) -> impl IntoResponse {
-    Json(DiagStats {
-        bytes_written: *state.qmdl_bytes_written.read().await,
-    })
+pub async fn get_qmdl_manifest(State(state): State<Arc<ServerState>>) -> Result<Json<ManifestStats>, (StatusCode, String)> {
+    let qmdl_store = state.qmdl_store_lock.read().await;
+    let mut entries = qmdl_store.manifest.entries.clone();
+    let current_entry = qmdl_store.current_entry.map(|index| entries.remove(index));
+    Ok(Json(ManifestStats {
+        entries,
+        current_entry,
+    }))
 }