use std::path::{Path, PathBuf};
use async_trait::async_trait;
use crate::capability::{CheckpointImage, DiskSource, Staging};
use crate::error::{Error, Result};
fn lexically_resolved(path: &Path) -> PathBuf {
let mut out = PathBuf::new();
for component in path.components() {
match component {
std::path::Component::CurDir => {}
std::path::Component::ParentDir => {
if !out.pop() {
out.push(component);
}
}
other => out.push(other),
}
}
out
}
#[derive(Debug, Clone, Copy)]
enum Bring {
Device,
Copy,
Move,
}
#[derive(Debug, Clone)]
pub(super) struct StagingArea {
root: PathBuf,
}
impl StagingArea {
pub(super) fn new(runtime_dir: &Path) -> Self {
Self {
root: runtime_dir.join("staged"),
}
}
pub(super) async fn remove(&self) -> Result<()> {
match tokio::fs::remove_dir_all(&self.root).await {
Ok(()) => Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(Error::Io(e)),
}
}
fn staged_disk(&self, id: &str) -> Result<PathBuf> {
let mut components = Path::new(id).components();
if !matches!(components.next(), Some(std::path::Component::Normal(_)))
|| components.next().is_some()
{
return Err(Error::InvalidSpec(format!(
"disk id `{id}` must be a plain name"
)));
}
Ok(self.root.join(format!("{id}.ext4")))
}
async fn bring_in(&self, src: &Path, dst: &Path, how: Bring) -> Result<PathBuf> {
let resolved = lexically_resolved(src);
if resolved.starts_with(&self.root) {
return Ok(resolved);
}
if let Some(parent) = dst.parent() {
tokio::fs::create_dir_all(parent).await?;
}
match tokio::fs::remove_file(dst).await {
Err(e) if e.kind() != std::io::ErrorKind::NotFound => return Err(Error::Io(e)),
_ => {}
}
match how {
Bring::Device => tokio::fs::symlink(src, dst).await?,
Bring::Copy => {
tokio::fs::copy(src, dst).await?;
}
Bring::Move => match tokio::fs::rename(src, dst).await {
Err(e) if e.kind() == std::io::ErrorKind::CrossesDevices => {
tokio::fs::copy(src, dst).await?;
tokio::fs::remove_file(src).await?;
}
other => other?,
},
}
Ok(dst.to_path_buf())
}
}
#[async_trait]
impl Staging for StagingArea {
async fn stage_kernel(&self, src: &Path) -> Result<PathBuf> {
self.bring_in(src, &self.root.join("vmlinux"), Bring::Copy)
.await
}
async fn stage_disk(&self, id: &str, source: DiskSource<'_>) -> Result<PathBuf> {
let how = match source {
DiskSource::Device(_) => Bring::Device,
DiskSource::Image(_) => Bring::Copy,
DiskSource::Handover(_) => Bring::Move,
};
let into = self.staged_disk(id)?;
self.bring_in(source.path(), &into, how).await
}
async fn unstage_disk(&self, id: &str, dst: &Path) -> Result<bool> {
let staged = self.staged_disk(id)?;
if !tokio::fs::try_exists(&staged).await.unwrap_or(false) {
return Ok(false);
}
match tokio::fs::rename(&staged, dst).await {
Err(e) if e.kind() == std::io::ErrorKind::CrossesDevices => {
tokio::fs::copy(&staged, dst).await?;
tokio::fs::remove_file(&staged).await?;
}
other => other?,
}
Ok(true)
}
async fn stage_checkpoint(&self, image: &CheckpointImage) -> Result<CheckpointImage> {
let resolved = lexically_resolved(&image.dir);
if resolved.starts_with(&self.root) {
return Ok(CheckpointImage {
dir: resolved,
..image.clone()
});
}
let name = image.dir.file_name().ok_or_else(|| {
Error::InvalidSpec(format!(
"checkpoint dir {} has no usable name",
image.dir.display()
))
})?;
let dir = self.root.join("snapshots").join(name);
tokio::fs::create_dir_all(&dir).await?;
let mut entries = tokio::fs::read_dir(&image.dir).await?;
while let Some(entry) = entries.next_entry().await? {
if entry.file_type().await?.is_file() {
tokio::fs::copy(entry.path(), dir.join(entry.file_name())).await?;
}
}
Ok(CheckpointImage {
dir,
..image.clone()
})
}
}