[pbs-devel] [RFC proxmox-backup 2/2 (nbd)] client: implement map/unmap commands for .img backups

[Stefan Reiter <s.reiter@proxmox.com>]

Allows mapping fixed-index .img files (usually from VM backups) to be
mapped to a local NBD device.

This uses the nbd-async crate, which implements an NBD server with an
async-trait based interface. We can very simply forward read requests to
an AsyncIndexReader, write requests are ignored.

Since unmapping requires some cleanup, a special 'unmap' command is
added, which uses a PID file to send SIGINT to the backup-client
instance started with 'map', which will handle the cleanup itself.

The client code is placed in the 'mount' module, which, while
admittedly a loose fit, allows reuse of the daemonizing code.

Signed-off-by: Stefan Reiter <s.reiter@proxmox.com>
---
 Cargo.toml                             |   2 +
 src/bin/proxmox-backup-client.rs       |   2 +
 src/bin/proxmox_backup_client/mount.rs | 159 +++++++++++++++++++++----
 src/tools.rs                           |   1 +
 4 files changed, 140 insertions(+), 24 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 74707f24..99ebd5ef 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ path = "src/lib.rs"
 
 [dependencies]
 apt-pkg-native = "0.3.1" # custom patched version
+async-trait = "0.1"
 base64 = "0.12"
 bitflags = "1.2.1"
 bytes = "0.5"
@@ -30,6 +31,7 @@ hyper = "0.13"
 lazy_static = "1.4"
 libc = "0.2"
 log = "0.4"
+nbd-async = "0.2"
 nix = "0.16"
 num-traits = "0.2"
 once_cell = "1.3.1"
diff --git a/src/bin/proxmox-backup-client.rs b/src/bin/proxmox-backup-client.rs
index 9a6f309d..88f9dba5 100644
--- a/src/bin/proxmox-backup-client.rs
+++ b/src/bin/proxmox-backup-client.rs
@@ -1960,6 +1960,8 @@ fn main() {
         .insert("status", status_cmd_def)
         .insert("key", key::cli())
         .insert("mount", mount_cmd_def())
+        .insert("map", map_cmd_def())
+        .insert("unmap", unmap_cmd_def())
         .insert("catalog", catalog_mgmt_cli())
         .insert("task", task_mgmt_cli())
         .insert("version", version_cmd_def)
diff --git a/src/bin/proxmox_backup_client/mount.rs b/src/bin/proxmox_backup_client/mount.rs
index 7646e98c..9e6aa1a0 100644
--- a/src/bin/proxmox_backup_client/mount.rs
+++ b/src/bin/proxmox_backup_client/mount.rs
@@ -3,9 +3,11 @@ use std::sync::Arc;
 use std::os::unix::io::RawFd;
 use std::path::Path;
 use std::ffi::OsStr;
+use std::collections::HashMap;
 
 use anyhow::{bail, format_err, Error};
 use serde_json::Value;
+use tokio::io::AsyncReadExt;
 use tokio::signal::unix::{signal, SignalKind};
 use nix::unistd::{fork, ForkResult, pipe};
 use futures::select;
@@ -23,6 +25,7 @@ use proxmox_backup::backup::{
     BackupDir,
     BackupGroup,
     BufferedDynamicReader,
+    AsyncIndexReader,
 };
 
 use proxmox_backup::client::*;
@@ -50,7 +53,35 @@ const API_METHOD_MOUNT: ApiMethod = ApiMethod::new(
             ("target", false, &StringSchema::new("Target directory path.").schema()),
             ("repository", true, &REPO_URL_SCHEMA),
             ("keyfile", true, &StringSchema::new("Path to encryption key.").schema()),
-            ("verbose", true, &BooleanSchema::new("Verbose output.").default(false).schema()),
+            ("verbose", true, &BooleanSchema::new("Verbose output and stay in foreground.").default(false).schema()),
+        ]),
+    )
+);
+
+#[sortable]
+const API_METHOD_MAP: ApiMethod = ApiMethod::new(
+    &ApiHandler::Sync(&mount),
+    &ObjectSchema::new(
+        "Map a drive image from a VM backup to a local NBD device. Use 'unmap' to undo.
+WARNING: Only do this with *trusted* backups!",
+        &sorted!([
+            ("snapshot", false, &StringSchema::new("Group/Snapshot path.").schema()),
+            ("archive-name", false, &StringSchema::new("Backup archive name.").schema()),
+            ("nbd-path", false, &StringSchema::new("Target NBD path (/dev/nbdX) or device number X.").schema()),
+            ("repository", true, &REPO_URL_SCHEMA),
+            ("keyfile", true, &StringSchema::new("Path to encryption key.").schema()),
+            ("verbose", true, &BooleanSchema::new("Verbose output and stay in foreground.").default(false).schema()),
+        ]),
+    )
+);
+
+#[sortable]
+const API_METHOD_UNMAP: ApiMethod = ApiMethod::new(
+    &ApiHandler::Sync(&unmap),
+    &ObjectSchema::new(
+        "Unmap a NBD device mapped with 'map' and release all resources.",
+        &sorted!([
+            ("nbd-path", false, &StringSchema::new("Path to NBD device (/dev/nbdX) or device number X.").schema()),
         ]),
     )
 );
@@ -65,6 +96,23 @@ pub fn mount_cmd_def() -> CliCommand {
         .completion_cb("target", tools::complete_file_name)
 }
 
+pub fn map_cmd_def() -> CliCommand {
+
+    CliCommand::new(&API_METHOD_MAP)
+        .arg_param(&["snapshot", "archive-name", "nbd-path"])
+        .completion_cb("repository", complete_repository)
+        .completion_cb("snapshot", complete_group_or_snapshot)
+        .completion_cb("archive-name", complete_pxar_archive_name)
+        .completion_cb("nbd-path", tools::complete_file_name)
+}
+
+pub fn unmap_cmd_def() -> CliCommand {
+
+    CliCommand::new(&API_METHOD_UNMAP)
+        .arg_param(&["nbd-path"])
+        .completion_cb("nbd-path", tools::complete_file_name)
+}
+
 fn mount(
     param: Value,
     _info: &ApiMethod,
@@ -100,9 +148,11 @@ fn mount(
 async fn mount_do(param: Value, pipe: Option<RawFd>) -> Result<Value, Error> {
     let repo = extract_repository_from_value(&param)?;
     let archive_name = tools::required_string_param(&param, "archive-name")?;
-    let target = tools::required_string_param(&param, "target")?;
     let client = connect(repo.host(), repo.user())?;
 
+    let target = param["target"].as_str();
+    let nbd_path = param["nbd-path"].as_str();
+
     record_repository(&repo);
 
     let path = tools::required_string_param(&param, "snapshot")?;
@@ -124,9 +174,17 @@ async fn mount_do(param: Value, pipe: Option<RawFd>) -> Result<Value, Error> {
     };
 
     let server_archive_name = if archive_name.ends_with(".pxar") {
+        if let None = target {
+            bail!("use the 'mount' command to mount pxar archives");
+        }
         format!("{}.didx", archive_name)
+    } else if archive_name.ends_with(".img") {
+        if let None = nbd_path {
+            bail!("use the 'map' command to map drive images");
+        }
+        format!("{}.fidx", archive_name)
     } else {
-        bail!("Can only mount pxar archives.");
+        bail!("Can only mount/map pxar archives and drive images.");
     };
 
     let client = BackupReader::start(
@@ -141,27 +199,9 @@ async fn mount_do(param: Value, pipe: Option<RawFd>) -> Result<Value, Error> {
 
     let (manifest, _) = client.download_manifest().await?;
 
-    let file_info = manifest.lookup_file_info(&archive_name)?;
-
-    if server_archive_name.ends_with(".didx") {
-        let index = client.download_dynamic_index(&manifest, &server_archive_name).await?;
-        let most_used = index.find_most_used_chunks(8);
-        let chunk_reader = RemoteChunkReader::new(client.clone(), crypt_config, file_info.chunk_crypt_mode(), most_used);
-        let reader = BufferedDynamicReader::new(index, chunk_reader);
-        let archive_size = reader.archive_size();
-        let reader: proxmox_backup::pxar::fuse::Reader =
-            Arc::new(BufferedDynamicReadAt::new(reader));
-        let decoder = proxmox_backup::pxar::fuse::Accessor::new(reader, archive_size).await?;
-        let options = OsStr::new("ro,default_permissions");
-
-        let session = proxmox_backup::pxar::fuse::Session::mount(
-            decoder,
-            &options,
-            false,
-            Path::new(target),
-        )
-        .map_err(|err| format_err!("pxar mount failed: {}", err))?;
+    let file_info = manifest.lookup_file_info(&server_archive_name)?;
 
+    let daemonize = || -> Result<(), Error> {
         if let Some(pipe) = pipe {
             nix::unistd::chdir(Path::new("/")).unwrap();
             // Finish creation of daemon by redirecting filedescriptors.
@@ -182,6 +222,31 @@ async fn mount_do(param: Value, pipe: Option<RawFd>) -> Result<Value, Error> {
             nix::unistd::close(pipe).unwrap();
         }
 
+        Ok(())
+    };
+
+    let options = OsStr::new("ro,default_permissions");
+
+    if server_archive_name.ends_with(".didx") {
+        let index = client.download_dynamic_index(&manifest, &server_archive_name).await?;
+        let most_used = index.find_most_used_chunks(8);
+        let chunk_reader = RemoteChunkReader::new(client.clone(), crypt_config, file_info.chunk_crypt_mode(), most_used);
+        let reader = BufferedDynamicReader::new(index, chunk_reader);
+        let archive_size = reader.archive_size();
+        let reader: proxmox_backup::pxar::fuse::Reader =
+            Arc::new(BufferedDynamicReadAt::new(reader));
+        let decoder = proxmox_backup::pxar::fuse::Accessor::new(reader, archive_size).await?;
+
+        let session = proxmox_backup::pxar::fuse::Session::mount(
+            decoder,
+            &options,
+            false,
+            Path::new(target.unwrap()),
+        )
+        .map_err(|err| format_err!("pxar mount failed: {}", err))?;
+
+        daemonize()?;
+
         let mut interrupt = signal(SignalKind::interrupt())?;
         select! {
             res = session.fuse() => res?,
@@ -189,9 +254,55 @@ async fn mount_do(param: Value, pipe: Option<RawFd>) -> Result<Value, Error> {
                 // exit on interrupted
             }
         }
+    } else if server_archive_name.ends_with(".fidx") {
+        tools::nbd::check_module_loaded()?;
+
+        // we can unwrap since we fail earlier on None
+        let nbd_path = if let Ok(num) = nbd_path.unwrap().parse::<u8>() {
+            format!("/dev/nbd{}", num)
+        } else {
+            nbd_path.unwrap().to_owned()
+        };
+
+        let index = client.download_fixed_index(&manifest, &server_archive_name).await?;
+        let size = index.index_bytes();
+        let chunk_reader = RemoteChunkReader::new(client.clone(), crypt_config, file_info.chunk_crypt_mode(), HashMap::new());
+        let mut reader = AsyncIndexReader::new(index, chunk_reader);
+
+        // attempt a single read to check if the reader is configured correctly
+        let _ = reader.read_u8().await?;
+
+        daemonize()?;
+
+        let mut interrupt = signal(SignalKind::interrupt())?;
+
+        // continue polling until complete or interrupted (which also happens on unmap)
+        select! {
+            res = tools::nbd::map(size, reader, &nbd_path).fuse() => res?,
+            _ = interrupt.recv().fuse() => {
+                // exit on interrupted
+            }
+        }
     } else {
-        bail!("unknown archive file extension (expected .pxar)");
+        bail!("unknown archive file extension (expected .pxar or .img)");
     }
 
     Ok(Value::Null)
 }
+
+fn unmap(
+    param: Value,
+    _info: &ApiMethod,
+    _rpcenv: &mut dyn RpcEnvironment,
+) -> Result<Value, Error> {
+
+    let mut path = tools::required_string_param(&param, "nbd-path")?.to_owned();
+
+    if let Ok(num) = path.parse::<u8>() {
+        path = format!("/dev/nbd{}", num);
+    }
+
+    tools::nbd::unmap(path)?;
+
+    Ok(Value::Null)
+}
diff --git a/src/tools.rs b/src/tools.rs
index 8792bf0c..bb514600 100644
--- a/src/tools.rs
+++ b/src/tools.rs
@@ -34,6 +34,7 @@ pub mod ticket;
 pub mod statistics;
 pub mod systemd;
 pub mod nom;
+pub mod nbd;
 
 mod wrapped_reader_stream;
 pub use wrapped_reader_stream::*;
-- 
2.20.1