use crate::litebox::CopySourceKind;
use boxlite_shared::{
BoxByteStream, BoxliteError, BoxliteResult, DownloadRequest, FilesClient, UploadChunk,
};
use futures::StreamExt;
use tokio::fs::File;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tonic::transport::Channel;
const CHUNK_SIZE: usize = 1 << 20;
pub struct FilesInterface {
client: FilesClient<Channel>,
}
impl FilesInterface {
pub fn new(channel: Channel) -> Self {
Self {
client: FilesClient::new(channel),
}
}
pub async fn upload_tar(
&mut self,
tar_path: &std::path::Path,
dest_path: &str,
container_id: Option<&str>,
mkdir_parents: bool,
overwrite: bool,
) -> BoxliteResult<()> {
let dest = dest_path.to_string();
let cid = container_id.unwrap_or_default().to_string();
let mut file = File::open(tar_path)
.await
.map_err(|e| BoxliteError::Storage(format!("Failed to open tar file: {}", e)))?;
let mut chunks = Vec::new();
let mut buf = vec![0u8; CHUNK_SIZE];
let mut first = true;
loop {
match file.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
let chunk = UploadChunk {
dest_path: if first { dest.clone() } else { String::new() },
container_id: cid.clone(),
data: buf[..n].to_vec(),
mkdir_parents,
overwrite,
source_is_dir: None,
};
first = false;
chunks.push(chunk);
}
Err(e) => {
return Err(BoxliteError::Storage(format!(
"Failed to read tar file: {}",
e
)));
}
}
}
let stream = futures::stream::iter(chunks);
let response = self
.client
.upload(stream)
.await
.map_err(map_tonic_err)?
.into_inner();
if response.success {
Ok(())
} else {
Err(BoxliteError::Internal(
response.error.unwrap_or_else(|| "Upload failed".into()),
))
}
}
pub async fn download_tar(
&mut self,
container_src: &str,
container_id: Option<&str>,
include_parent: bool,
follow_symlinks: bool,
tar_dest: &std::path::Path,
) -> BoxliteResult<()> {
let request = DownloadRequest {
src_path: container_src.to_string(),
container_id: container_id.unwrap_or_default().to_string(),
include_parent,
follow_symlinks,
};
let mut stream = self
.client
.download(request)
.await
.map_err(map_tonic_err)?
.into_inner();
let mut file = File::create(tar_dest)
.await
.map_err(|e| BoxliteError::Storage(format!("Failed to create tar file: {}", e)))?;
loop {
match stream.message().await {
Ok(Some(chunk)) => {
file.write_all(&chunk.data).await.map_err(|e| {
BoxliteError::Storage(format!("Failed to write tar file: {}", e))
})?;
}
Ok(None) => break, Err(e) => return Err(map_tonic_err(e)),
}
}
file.flush()
.await
.map_err(|e| BoxliteError::Storage(format!("Failed to flush tar file: {}", e)))?;
Ok(())
}
pub async fn upload_stream<S>(
&mut self,
stream: S,
dest_path: &str,
container_id: Option<&str>,
mkdir_parents: bool,
overwrite: bool,
source: CopySourceKind,
) -> BoxliteResult<()>
where
S: futures::Stream<Item = std::io::Result<Vec<u8>>> + Send + 'static,
{
let dest = dest_path.to_string();
let cid = container_id.unwrap_or_default().to_string();
let stream_err: std::sync::Arc<tokio::sync::Mutex<Option<std::io::Error>>> =
std::sync::Arc::new(tokio::sync::Mutex::new(None));
let stream_err_slot = stream_err.clone();
let source_is_dir = source.to_wire();
let chunks = async_stream::stream! {
futures::pin_mut!(stream);
let mut first = true;
while let Some(item) = stream.next().await {
match item {
Ok(data) => {
yield UploadChunk {
dest_path: if first { dest.clone() } else { String::new() },
container_id: cid.clone(),
data,
mkdir_parents,
overwrite,
source_is_dir: if first { source_is_dir } else { None },
};
first = false;
}
Err(e) => {
*stream_err_slot.lock().await = Some(e);
break;
}
}
}
};
let response = self
.client
.upload(chunks)
.await
.map_err(map_tonic_err)?
.into_inner();
if let Some(e) = stream_err.lock().await.take() {
return Err(BoxliteError::Internal(format!(
"source stream failed during upload: {e}"
)));
}
if response.success {
Ok(())
} else {
Err(BoxliteError::Internal(
response.error.unwrap_or_else(|| "Upload failed".into()),
))
}
}
pub async fn download_stream(
&mut self,
container_src: &str,
container_id: Option<&str>,
include_parent: bool,
follow_symlinks: bool,
) -> BoxliteResult<(BoxByteStream, CopySourceKind)> {
let request = DownloadRequest {
src_path: container_src.to_string(),
container_id: container_id.unwrap_or_default().to_string(),
include_parent,
follow_symlinks,
};
let mut stream = self
.client
.download(request)
.await
.map_err(map_tonic_err)?
.into_inner();
let first = stream.message().await.map_err(map_tonic_err)?;
let source = CopySourceKind::from_wire(first.as_ref().and_then(|c| c.source_is_dir));
let mut first_data = first.map(|c| c.data);
let out = async_stream::stream! {
if let Some(data) = first_data.take().filter(|d| !d.is_empty()) {
yield Ok(data);
}
loop {
match stream.message().await {
Ok(Some(chunk)) => {
if !chunk.data.is_empty() {
yield Ok(chunk.data);
}
}
Ok(None) => break,
Err(e) => {
yield Err(std::io::Error::other(e));
break;
}
}
}
};
Ok((Box::pin(out), source))
}
}
fn map_tonic_err(err: tonic::Status) -> BoxliteError {
let message = err.message().to_owned();
match err.code() {
tonic::Code::FailedPrecondition => BoxliteError::Unsupported(message),
tonic::Code::InvalidArgument => BoxliteError::InvalidArgument(message),
tonic::Code::NotFound => BoxliteError::NotFound(message),
tonic::Code::ResourceExhausted => BoxliteError::ResourceExhausted(message),
_ => BoxliteError::Internal(err.to_string()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use tonic::Status;
#[test]
fn unreachable_destination_maps_to_unsupported() {
let error = map_tonic_err(Status::failed_precondition(
"/tmp/x is under the container's '/tmp' mount",
));
assert!(matches!(error, BoxliteError::Unsupported(_)));
assert_eq!(error.http().0, 400);
}
#[test]
fn missing_source_maps_to_not_found() {
let error = map_tonic_err(Status::not_found("source path does not exist"));
assert!(matches!(error, BoxliteError::NotFound(_)));
assert_eq!(error.http().0, 404);
}
#[test]
fn malformed_path_maps_to_invalid_argument() {
let error = map_tonic_err(Status::invalid_argument("path must not contain .."));
assert!(matches!(error, BoxliteError::InvalidArgument(_)));
assert_eq!(error.http().0, 400);
}
#[test]
fn oversized_upload_maps_to_resource_exhausted() {
let error = map_tonic_err(Status::resource_exhausted("upload too large"));
assert!(
matches!(error, BoxliteError::ResourceExhausted(_)),
"{error:?}"
);
assert_eq!(error.http().0, 429);
}
#[test]
fn unclassified_status_stays_internal() {
let error = map_tonic_err(Status::internal("failed to create temp file"));
assert!(matches!(error, BoxliteError::Internal(_)));
assert_eq!(error.http().0, 500);
}
#[test]
fn refusal_message_survives_the_remap() {
let error = map_tonic_err(Status::failed_precondition(
"/tmp/x is under the container's '/tmp' mount",
));
assert!(error.to_string().contains("'/tmp' mount"), "{error}");
}
}