// Copyright 2025 The ChromiumOS Authors

// Use of this source code is governed by a BSD-style license that can be

// found in the LICENSE file.

use std::fmt::Debug;

use std::fmt::Formatter;

use std::fs::File;

use std::io::BufReader;

use std::io::Read;

use std::io::Write;

use std::path::Path;

use std::path::PathBuf;

use anyhow::Context;

use anyhow::Result;

use crypto::CryptKey;

mod any_snapshot;

pub use any_snapshot::AnySnapshot;

// Use 4kB encrypted chunks by default (if encryption is used).

const DEFAULT_ENCRYPTED_CHUNK_SIZE_BYTES: usize = 1024 * 4;

/// Writer of serialized VM snapshots.

///

/// Each fragment is an opaque byte blob. Namespaces can be used to avoid fragment naming

/// collisions between devices.

///

/// In the current implementation, fragments are files and namespaces are directories, but the API

/// is kept abstract so that we can potentially support something like a single file archive

/// output.

#[derive(Clone, serde::Serialize, serde::Deserialize)]

pub struct SnapshotWriter {

    dir: PathBuf,

    /// If encryption is used, the plaintext key will be stored here.

    key: Option<CryptKey>,

}

impl Debug for SnapshotWriter {

    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {

        f.debug_struct("SnapshotWriter")

            .field("dir", &format!("{:?}", self.dir))

            .field("key", if self.key.is_some() { &"Some" } else { &"None" })

            .finish()

    }

}

impl SnapshotWriter {

    /// Creates a new `SnapshotWriter` that will writes its data to a dir at `root`. The path must

    /// not exist yet. If encryption is desired, set encrypt (Note: only supported downstream on

    /// Windows).

    // TODO(b/268094487): If the snapshot fails, we leave incomplete snapshot files at the

    // requested path. Consider building up the snapshot dir somewhere else and moving it into

    // place at the end.

    pub fn new(root: PathBuf, encrypt: bool) -> Result<Self> {

        std::fs::create_dir(&root)

            .with_context(|| format!("failed to create snapshot root dir: {}", root.display()))?;

        if encrypt {

            let key = crypto::generate_random_key();

            // Creating an empty CryptWriter will still write header information

            // to the file, and that header information is what we need. This

            // ensures we use a single key for *all* snapshot files.

            let mut writer = crypto::CryptWriter::new_from_key(

                File::create(root.join("enc_metadata")).context("failed to create enc_metadata")?,

                1024,

                &key,

            )

            .context("failed to create enc_metadata writer")?;

            writer.flush().context("flush of enc_metadata failed")?;

            return Ok(Self {

                dir: root,

                key: Some(key),

            });

        }

        Ok(Self {

            dir: root,

            key: None,

        })

    }

    /// Creates a snapshot fragment and get access to the `Write` impl representing it.

    pub fn raw_fragment(&self, name: &str) -> Result<Box<dyn Write>> {

        self.raw_fragment_with_chunk_size(name, DEFAULT_ENCRYPTED_CHUNK_SIZE_BYTES)

    }

    /// When encryption is used, allows direct control of the encrypted chunk size.

    pub fn raw_fragment_with_chunk_size(

        &self,

        name: &str,

        chunk_size_bytes: usize,

    ) -> Result<Box<dyn Write>> {

        let path = self.dir.join(name);

        let file = File::options()

            .write(true)

            .create_new(true)

            .open(&path)

            .with_context(|| {

                format!(

                    "failed to create snapshot fragment {name:?} at {}",

                    path.display()

                )

            })?;

        if let Some(key) = self.key.as_ref() {

            return Ok(Box::new(crypto::CryptWriter::new_from_key(

                file,

                chunk_size_bytes,

                key,

            )?));

        }

        Ok(Box::new(file))

    }

    /// Creates a snapshot fragment from a serialized representation of `v`.

    pub fn write_fragment<T: serde::Serialize>(&self, name: &str, v: &T) -> Result<()> {

        let mut w = std::io::BufWriter::new(self.raw_fragment(name)?);

        ciborium::into_writer(v, &mut w)?;

        w.flush()?;

        Ok(())

    }

    /// Creates new namespace and returns a `SnapshotWriter` that writes to it. Namespaces can be

    /// nested.

    pub fn add_namespace(&self, name: &str) -> Result<Self> {

        let dir = self.dir.join(name);

        std::fs::create_dir(&dir).with_context(|| {

            format!(

                "failed to create nested snapshot writer {name:?} at {}",

                dir.display()

            )

        })?;

        Ok(Self {

            dir,

            key: self.key.clone(),

        })

    }

}

/// Reads snapshots created by `SnapshotWriter`.

#[derive(Clone, serde::Serialize, serde::Deserialize)]

pub struct SnapshotReader {

    dir: PathBuf,

    /// If encryption is used, the plaintext key will be stored here.

    key: Option<CryptKey>,

}

impl Debug for SnapshotReader {

    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {

        f.debug_struct("SnapshotReader")

            .field("dir", &format!("{:?}", self.dir))

            .field("key", if self.key.is_some() { &"Some" } else { &"None" })

            .finish()

    }

}

impl SnapshotReader {

    /// Reads a snapshot at `root`. Set require_encrypted to require an encrypted snapshot.

    pub fn new(root: &Path, require_encrypted: bool) -> Result<Self> {

        let enc_metadata_path = root.join("enc_metadata");

        if Path::exists(&enc_metadata_path) {

            let key = Some(

                crypto::CryptReader::extract_key(

                    File::open(&enc_metadata_path).context("failed to open encryption metadata")?,

                )

                .context("failed to load snapshot key")?,

            );

            return Ok(Self {

                dir: root.to_path_buf(),

                key,

            });

        } else if require_encrypted {

            return Err(anyhow::anyhow!("snapshot was not encrypted"));

        }

        Ok(Self {

            dir: root.to_path_buf(),

            key: None,

        })

    }

    /// Gets access to a `Read` impl that represents a fragment.

    pub fn raw_fragment(&self, name: &str) -> Result<Box<dyn Read>> {

        let path = self.dir.join(name);

        let file = File::open(&path).with_context(|| {

            format!(

                "failed to open snapshot fragment {name:?} at {}",

                path.display()

            )

        })?;

        if let Some(key) = self.key.as_ref() {

            return Ok(Box::new(crypto::CryptReader::from_file_and_key(file, key)?));

        }

        Ok(Box::new(file))

    }

    /// Reads a fragment.

    pub fn read_fragment<T: serde::de::DeserializeOwned>(&self, name: &str) -> Result<T> {

        Ok(ciborium::from_reader(BufReader::new(

            self.raw_fragment(name)?,

        ))?)

    }

    /// Reads the names of all fragments in this namespace.

    pub fn list_fragments(&self) -> Result<Vec<String>> {

        let mut result = Vec::new();

        for entry in std::fs::read_dir(&self.dir)? {

            let entry = entry?;

            if entry.file_type()?.is_file() {

                result.push(entry.file_name().to_string_lossy().into_owned());

            }

        }

        Ok(result)

    }

    /// Open a namespace.

    pub fn namespace(&self, name: &str) -> Result<Self> {

        let dir = self.dir.join(name);

        Ok(Self {

            dir,

            key: self.key.clone(),

        })

    }

    /// Reads the names of all child namespaces

    pub fn list_namespaces(&self) -> Result<Vec<String>> {

        let mut result = Vec::new();

        for entry in std::fs::read_dir(&self.dir)? {

            let entry = entry?;

            if entry.path().is_dir() {

                if let Some(file_name) = entry.path().file_name() {

                    result.push(file_name.to_string_lossy().into_owned());

                }

            }

        }

        Ok(result)

    }

}