use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use futures::stream::{self, TryStreamExt};
use sha2::{Digest, Sha256};
use std::sync::Arc;
use crate::error::{SailError, TransportKind};
use crate::image::{
AddLocalDirFile, BaseImage, ImageArchitecture, ImageBuildStep, ImageSpec, PackageInstall,
RunCommand,
};
use crate::pb::image::v1 as pbimage;
use crate::pb::imagebuilder::v1 as pbimg;
use crate::Client;
pub(crate) const MAX_LOCAL_FILE_BYTES: u64 = 5 * 1024 * 1024 * 1024;
pub(crate) const MAX_LOCAL_DIR_FILES: usize = 50_000;
pub(crate) const MAX_LOCAL_DIR_RELATIVE_PATH_BYTES: usize = 1024;
const UPLOAD_CONCURRENCY: usize = 16;
const BUILD_POLL_INTERVAL: Duration = Duration::from_secs(1);
const UPLOAD_BASE_TIMEOUT: Duration = Duration::from_mins(5);
const MIN_UPLOAD_BYTES_PER_SEC: u64 = 1 << 20;
const UNBOUNDED_BUILD_RPC_BUDGET: Duration = Duration::from_mins(1);
fn invalid(message: String) -> SailError {
SailError::InvalidArgument { message }
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImageBuildStatus {
Unspecified,
Queued,
Building,
Ready,
Failed,
}
impl ImageBuildStatus {
pub fn as_str(self) -> &'static str {
match self {
ImageBuildStatus::Unspecified => "unspecified",
ImageBuildStatus::Queued => "queued",
ImageBuildStatus::Building => "building",
ImageBuildStatus::Ready => "ready",
ImageBuildStatus::Failed => "failed",
}
}
fn from_pb(status: i32) -> ImageBuildStatus {
match pbimage::ImageBuildStatus::try_from(status) {
Ok(pbimage::ImageBuildStatus::Queued) => ImageBuildStatus::Queued,
Ok(pbimage::ImageBuildStatus::Building) => ImageBuildStatus::Building,
Ok(pbimage::ImageBuildStatus::Ready) => ImageBuildStatus::Ready,
Ok(pbimage::ImageBuildStatus::Failed) => ImageBuildStatus::Failed,
_ => ImageBuildStatus::Unspecified,
}
}
}
#[derive(Debug, Clone)]
pub struct ImageBuild {
pub image_id: String,
pub status: ImageBuildStatus,
pub error_message: String,
}
#[derive(Debug, Clone)]
pub(crate) enum LocalFileUploadPlan {
AlreadyExists,
SinglePart {
upload_url: String,
headers: HashMap<String, String>,
},
}
#[derive(Debug, Clone)]
pub enum ImageDefinitionStep {
AptInstall(Vec<String>),
PipInstall(Vec<String>),
RunCommand(String),
AddLocalFile {
local_path: PathBuf,
remote_path: String,
mode: Option<u32>,
},
AddLocalDir {
local_path: PathBuf,
remote_path: String,
ignore: Vec<String>,
ignore_file: Option<PathBuf>,
},
}
#[derive(Debug, Clone, Default)]
pub struct ImageDefinition {
pub base: Option<BaseImage>,
pub architecture: ImageArchitecture,
pub env: HashMap<String, String>,
pub python_version: String,
pub steps: Vec<ImageDefinitionStep>,
}
#[doc(hidden)]
pub fn is_builtin_base_spec(spec: &ImageSpec) -> bool {
matches!(spec.base, Some(BaseImage::Debian | BaseImage::Devbox))
&& spec.build_steps.is_empty()
&& spec.env.is_empty()
&& spec.python_version.is_empty()
}
fn validate_remote_path(target: &str) -> Result<(), SailError> {
if !target.starts_with('/') {
return Err(invalid(format!("remotePath {target:?} must be absolute")));
}
if target.len() > 1 && target.ends_with('/') {
return Err(invalid(format!(
"remotePath {target:?} must not end with '/'"
)));
}
for ch in target.chars() {
let code = ch as u32;
if code < 0x20 || code == 0x7f || matches!(ch, '"' | '\\' | '$' | ' ') {
return Err(invalid(format!(
"remotePath {target:?} contains an unsupported character"
)));
}
}
if target.split('/').any(|segment| segment == "..") {
return Err(invalid(format!(
"remotePath {target:?} must not contain '..'"
)));
}
Ok(())
}
fn validate_mode(mode: Option<u32>) -> Result<u32, SailError> {
match mode {
None | Some(0) => Ok(0),
Some(mode) if mode <= 0o777 => Ok(mode),
Some(mode) => Err(invalid(format!(
"mode 0o{mode:o} must fit in the low 9 bits"
))),
}
}
async fn hash_file(path: &Path) -> Result<(String, u64), SailError> {
let path = path.to_path_buf();
tokio::task::spawn_blocking(move || {
use std::io::Read;
let file = std::fs::File::open(&path)
.map_err(|err| invalid(format!("cannot read {}: {err}", path.display())))?;
let mut reader = std::io::BufReader::new(file);
let mut hasher = Sha256::new();
let mut buf = vec![0u8; 64 * 1024];
let mut size: u64 = 0;
loop {
let n = reader
.read(&mut buf)
.map_err(|err| invalid(format!("cannot read {}: {err}", path.display())))?;
if n == 0 {
break;
}
hasher.update(&buf[..n]);
size += n as u64;
}
Ok((format!("{:x}", hasher.finalize()), size))
})
.await
.map_err(|err| SailError::Internal {
message: format!("hashing task failed: {err}"),
})?
}
struct WalkedFile {
abs_path: PathBuf,
relative_path: String,
mode: u32,
}
fn walk_dir(
root: &Path,
matcher: &ignore::gitignore::Gitignore,
) -> Result<Vec<WalkedFile>, SailError> {
fn recurse(
root: &Path,
dir: &Path,
rel: &str,
matcher: &ignore::gitignore::Gitignore,
out: &mut Vec<WalkedFile>,
) -> Result<(), SailError> {
let mut entries: Vec<_> = std::fs::read_dir(dir)
.map_err(|err| invalid(format!("cannot read {}: {err}", dir.display())))?
.collect::<Result<_, _>>()
.map_err(|err| invalid(format!("cannot read {}: {err}", dir.display())))?;
entries.sort_by_key(std::fs::DirEntry::file_name);
for entry in entries {
let name = entry
.file_name()
.to_str()
.ok_or_else(|| {
invalid(format!(
"addLocalDir: {} has a non-UTF-8 file name",
entry.path().display()
))
})?
.to_string();
let rel_path = if rel.is_empty() {
name.clone()
} else {
format!("{rel}/{name}")
};
let file_type = entry
.file_type()
.map_err(|err| invalid(format!("cannot stat {}: {err}", entry.path().display())))?;
if file_type.is_symlink() {
continue;
}
let is_dir = file_type.is_dir();
if matcher
.matched_path_or_any_parents(&rel_path, is_dir)
.is_ignore()
{
continue;
}
if is_dir {
recurse(root, &entry.path(), &rel_path, matcher, out)?;
continue;
}
if !file_type.is_file() {
continue;
}
if rel_path.len() > MAX_LOCAL_DIR_RELATIVE_PATH_BYTES {
return Err(invalid(format!(
"relative path {rel_path} exceeds {MAX_LOCAL_DIR_RELATIVE_PATH_BYTES} bytes"
)));
}
let metadata = entry
.metadata()
.map_err(|err| invalid(format!("cannot stat {}: {err}", entry.path().display())))?;
if metadata.len() > MAX_LOCAL_FILE_BYTES {
return Err(invalid(format!(
"{} ({} bytes) exceeds the {MAX_LOCAL_FILE_BYTES}-byte per-file limit",
entry.path().display(),
metadata.len()
)));
}
out.push(WalkedFile {
abs_path: entry.path(),
relative_path: rel_path,
mode: unix_mode(&metadata),
});
if out.len() > MAX_LOCAL_DIR_FILES {
return Err(invalid(format!(
"{} has more than {MAX_LOCAL_DIR_FILES} files (max {MAX_LOCAL_DIR_FILES})",
root.display()
)));
}
}
Ok(())
}
let mut out = Vec::new();
recurse(root, root, "", matcher, &mut out)?;
Ok(out)
}
#[cfg(unix)]
fn unix_mode(metadata: &std::fs::Metadata) -> u32 {
use std::os::unix::fs::PermissionsExt;
metadata.permissions().mode() & 0o777
}
#[cfg(not(unix))]
fn unix_mode(_metadata: &std::fs::Metadata) -> u32 {
0o644
}
fn ignore_matcher(
root: &Path,
patterns: &[String],
ignore_file: Option<&Path>,
) -> Result<ignore::gitignore::Gitignore, SailError> {
let mut builder = ignore::gitignore::GitignoreBuilder::new(root);
if let Some(file) = ignore_file {
if let Some(err) = builder.add(file) {
return Err(invalid(format!(
"cannot read ignore file {}: {err}",
file.display()
)));
}
}
for pattern in patterns {
builder
.add_line( None, pattern)
.map_err(|err| invalid(format!("invalid ignore pattern {pattern:?}: {err}")))?;
}
builder
.build()
.map_err(|err| invalid(format!("invalid ignore patterns: {err}")))
}
fn base_image_to_pb(base: BaseImage) -> pbimage::BaseImage {
match base {
BaseImage::Debian => pbimage::BaseImage::Debian,
BaseImage::Devbox => pbimage::BaseImage::Devbox,
BaseImage::Unspecified => pbimage::BaseImage::Unspecified,
}
}
fn architecture_to_pb(arch: ImageArchitecture) -> pbimage::ImageArchitecture {
match arch {
ImageArchitecture::Amd64 => pbimage::ImageArchitecture::Amd64,
ImageArchitecture::Arm64 => pbimage::ImageArchitecture::Arm64,
ImageArchitecture::Unspecified => pbimage::ImageArchitecture::Unspecified,
}
}
fn build_step_to_pb(step: &ImageBuildStep) -> pbimage::ImageBuildStep {
use pbimage::image_build_step::Step;
let packages = |p: &PackageInstall| pbimage::PackageInstall {
packages: p.packages.clone(),
};
let inner = match step {
ImageBuildStep::AptInstall(p) => Step::AptInstall(packages(p)),
ImageBuildStep::PipInstall(p) => Step::PipInstall(packages(p)),
ImageBuildStep::RunCommand(c) => Step::RunCommand(pbimage::RunCommand {
command: c.command.clone(),
}),
ImageBuildStep::AddLocalFile(f) => Step::AddLocalFile(pbimage::AddLocalFile {
content_sha256: f.content_sha256.clone(),
remote_path: f.remote_path.clone(),
mode: f.mode,
}),
ImageBuildStep::AddLocalDir(d) => Step::AddLocalDir(pbimage::AddLocalDir {
remote_path: d.remote_path.clone(),
files: d
.files
.iter()
.map(|file| pbimage::AddLocalDirFile {
relative_path: file.relative_path.clone(),
content_sha256: file.content_sha256.clone(),
mode: file.mode,
})
.collect(),
}),
};
pbimage::ImageBuildStep { step: Some(inner) }
}
pub(crate) fn image_spec_to_pb(spec: &ImageSpec) -> pbimage::ImageSpec {
pbimage::ImageSpec {
source: spec
.base
.map(|base| pbimage::image_spec::Source::Base(base_image_to_pb(base) as i32)),
build_steps: spec.build_steps.iter().map(build_step_to_pb).collect(),
env: spec.env.clone(),
architecture: architecture_to_pb(spec.architecture) as i32,
python_version: spec.python_version.clone(),
}
}
impl Client {
pub(crate) async fn prepare_local_file_upload(
&self,
content_sha256: &str,
content_length: u64,
) -> Result<LocalFileUploadPlan, SailError> {
let request = pbimg::PrepareLocalFileUploadRequest {
content_sha256: content_sha256.to_string(),
content_length,
};
let response = self
.imagebuilder()
.prepare_local_file_upload(request)
.await?;
use pbimg::prepare_local_file_upload_response::Outcome;
match response.outcome {
Some(Outcome::AlreadyExists(_)) => Ok(LocalFileUploadPlan::AlreadyExists),
Some(Outcome::SinglePart(plan)) => Ok(LocalFileUploadPlan::SinglePart {
upload_url: plan.upload_url,
headers: plan.required_headers,
}),
None => Err(SailError::Internal {
message: "prepare_local_file_upload returned no outcome".to_string(),
}),
}
}
pub async fn build_image(
&self,
spec: &ImageSpec,
retry_timeout_secs: f64,
) -> Result<ImageBuild, SailError> {
let request = pbimg::BuildImageRequest {
image: Some(image_spec_to_pb(spec)),
};
let response = self
.imagebuilder()
.build_image(request, retry_timeout_secs)
.await?;
Ok(ImageBuild {
image_id: response.image_id,
status: ImageBuildStatus::from_pb(response.status),
error_message: response.error_message,
})
}
pub async fn get_image_build_status(
&self,
image_id: &str,
retry_timeout_secs: f64,
) -> Result<ImageBuild, SailError> {
let request = pbimg::GetImageBuildStatusRequest {
image_id: image_id.to_string(),
};
let response = self
.imagebuilder()
.get_image_build_status(request, retry_timeout_secs)
.await?;
Ok(ImageBuild {
image_id: response.image_id,
status: ImageBuildStatus::from_pb(response.status),
error_message: response.error_message,
})
}
#[doc(hidden)]
pub async fn resolve_local_file_step(
&self,
local_path: &Path,
remote_path: &str,
mode: Option<u32>,
) -> Result<crate::image::AddLocalFile, SailError> {
let metadata = std::fs::metadata(local_path).map_err(|_| {
invalid(format!(
"addLocalFile: {} does not exist or is not a file",
local_path.display()
))
})?;
if !metadata.is_file() {
return Err(invalid(format!(
"addLocalFile: {} is not a file",
local_path.display()
)));
}
if metadata.len() > MAX_LOCAL_FILE_BYTES {
return Err(invalid(format!(
"addLocalFile: {} ({} bytes) exceeds the {MAX_LOCAL_FILE_BYTES}-byte limit",
local_path.display(),
metadata.len()
)));
}
let mode = validate_mode(mode)?;
let mut target = remote_path.to_string();
if target.ends_with('/') {
let basename = local_path
.file_name()
.map(|name| name.to_string_lossy().into_owned())
.unwrap_or_default();
target = format!("{target}{basename}");
}
validate_remote_path(&target)?;
let (digest, size) = hash_file(local_path).await?;
if size > MAX_LOCAL_FILE_BYTES {
return Err(invalid(format!(
"addLocalFile: {} ({size} bytes) exceeds the {MAX_LOCAL_FILE_BYTES}-byte limit",
local_path.display()
)));
}
let http = reqwest::Client::new();
self.upload_local_content(&http, &digest, local_path, size)
.await?;
Ok(crate::image::AddLocalFile {
content_sha256: digest,
remote_path: target,
mode,
})
}
#[doc(hidden)]
pub async fn resolve_local_dir_step(
&self,
local_path: &Path,
remote_path: &str,
ignore: &[String],
ignore_file: Option<&Path>,
) -> Result<crate::image::AddLocalDir, SailError> {
let target = remote_path.trim_end_matches('/').to_string();
if target.is_empty() {
return Err(invalid(
"addLocalDir: remotePath must not be '/'".to_string(),
));
}
validate_remote_path(&target)?;
let walk_root = local_path.to_path_buf();
let ignore_owned = ignore.to_vec();
let ignore_file_owned = ignore_file.map(Path::to_path_buf);
let has_ignore = !ignore.is_empty() || ignore_file.is_some();
let walked = tokio::task::spawn_blocking(move || {
let metadata = std::fs::metadata(&walk_root).map_err(|_| {
invalid(format!(
"addLocalDir: {} does not exist or is not a directory",
walk_root.display()
))
})?;
if !metadata.is_dir() {
return Err(invalid(format!(
"addLocalDir: {} is not a directory",
walk_root.display()
)));
}
let matcher = ignore_matcher(&walk_root, &ignore_owned, ignore_file_owned.as_deref())?;
let walked = walk_dir(&walk_root, &matcher)?;
if walked.is_empty() {
let qualifier = if has_ignore {
" after applying ignore patterns"
} else {
""
};
return Err(invalid(format!(
"addLocalDir: {} contains no files{qualifier}",
walk_root.display()
)));
}
Ok(walked)
})
.await
.map_err(|err| SailError::Internal {
message: format!("directory walk task failed: {err}"),
})??;
let mut uploads: HashMap<String, (PathBuf, u64)> = HashMap::new();
let mut files = Vec::with_capacity(walked.len());
for file in walked {
let (digest, size) = hash_file(&file.abs_path).await?;
if size > MAX_LOCAL_FILE_BYTES {
return Err(invalid(format!(
"addLocalDir: {} ({size} bytes) exceeds the {MAX_LOCAL_FILE_BYTES}-byte \
per-file limit",
file.abs_path.display()
)));
}
uploads
.entry(digest.clone())
.or_insert_with(|| (file.abs_path.clone(), size));
files.push(AddLocalDirFile {
relative_path: file.relative_path,
content_sha256: digest,
mode: file.mode,
});
}
files.sort_by(|a, b| a.relative_path.cmp(&b.relative_path));
let http = reqwest::Client::new();
stream::iter(uploads.into_iter().map(Ok::<_, SailError>))
.try_for_each_concurrent(UPLOAD_CONCURRENCY, |(digest, (source, size))| {
let http = http.clone();
async move {
self.upload_local_content(&http, &digest, &source, size)
.await
}
})
.await?;
Ok(crate::image::AddLocalDir {
remote_path: target,
files,
})
}
pub async fn resolve_image(&self, def: &ImageDefinition) -> Result<ImageSpec, SailError> {
let mut steps = Vec::with_capacity(def.steps.len());
for step in &def.steps {
steps.push(match step {
ImageDefinitionStep::AptInstall(packages) => {
ImageBuildStep::AptInstall(PackageInstall {
packages: packages.clone(),
})
}
ImageDefinitionStep::PipInstall(packages) => {
ImageBuildStep::PipInstall(PackageInstall {
packages: packages.clone(),
})
}
ImageDefinitionStep::RunCommand(command) => {
ImageBuildStep::RunCommand(RunCommand {
command: command.clone(),
})
}
ImageDefinitionStep::AddLocalFile {
local_path,
remote_path,
mode,
} => ImageBuildStep::AddLocalFile(
self.resolve_local_file_step(local_path, remote_path, *mode)
.await?,
),
ImageDefinitionStep::AddLocalDir {
local_path,
remote_path,
ignore,
ignore_file,
} => ImageBuildStep::AddLocalDir(
self.resolve_local_dir_step(
local_path,
remote_path,
ignore,
ignore_file.as_deref(),
)
.await?,
),
});
}
Ok(ImageSpec {
base: def.base,
build_steps: steps,
env: def.env.clone(),
architecture: def.architecture,
python_version: def.python_version.clone(),
})
}
async fn upload_local_content(
&self,
http: &reqwest::Client,
digest: &str,
source: &Path,
size: u64,
) -> Result<(), SailError> {
let plan = self.prepare_local_file_upload(digest, size).await?;
let LocalFileUploadPlan::SinglePart {
upload_url,
headers,
} = plan
else {
return Ok(());
};
let file = tokio::fs::File::open(source)
.await
.map_err(|err| invalid(format!("cannot read {}: {err}", source.display())))?;
let (request, streamed_digest) = sized_put_request(http, &upload_url, file, size, &headers);
let response = tokio::time::timeout(upload_timeout(size), request.send())
.await
.map_err(|_| SailError::Transport {
kind: TransportKind::Timeout,
message: format!("local file upload stalled ({size} bytes not delivered in time)"),
source: None,
})?
.map_err(|err| SailError::Transport {
kind: TransportKind::Connection,
message: format!("local file upload failed: {err}"),
source: None,
})?;
if !response.status().is_success() {
return Err(SailError::Api {
message: format!(
"local file upload failed: HTTP {} {}",
response.status().as_u16(),
response.status().canonical_reason().unwrap_or("")
),
status: response.status().as_u16(),
body: serde_json::Value::Null,
});
}
let streamed = streamed_digest.lock().unwrap().take();
if streamed.as_deref() != Some(digest) {
return Err(invalid(format!(
"{} changed while it was being uploaded; retry the build",
source.display()
)));
}
Ok(())
}
#[doc(hidden)]
pub async fn build_spec_with_timeout(
&self,
spec: &ImageSpec,
timeout: Duration,
) -> Result<ImageBuild, SailError> {
let deadline = Instant::now().checked_add(timeout);
match deadline {
None => self.build_spec_to_ready(spec, None).await,
Some(_) => tokio::time::timeout(timeout, self.build_spec_to_ready(spec, deadline))
.await
.unwrap_or_else(|_| {
Err(SailError::Transport {
kind: TransportKind::Timeout,
message: "timed out building the image".to_string(),
source: None,
})
}),
}
}
pub async fn build_image_definition(
&self,
def: &ImageDefinition,
timeout: Duration,
) -> Result<ImageSpec, SailError> {
let deadline = Instant::now().checked_add(timeout);
let work = async {
let spec = self.resolve_image(def).await?;
if is_builtin_base_spec(&spec) {
return Ok(spec);
}
self.build_spec_to_ready(&spec, deadline).await?;
Ok(spec)
};
match deadline {
None => work.await,
Some(_) => tokio::time::timeout(timeout, work)
.await
.unwrap_or_else(|_| {
Err(SailError::Transport {
kind: TransportKind::Timeout,
message: "timed out building the image".to_string(),
source: None,
})
}),
}
}
#[doc(hidden)]
pub async fn build_spec_to_ready(
&self,
spec: &ImageSpec,
deadline: Option<Instant>,
) -> Result<ImageBuild, SailError> {
let rpc_budget = || {
deadline.map_or(UNBOUNDED_BUILD_RPC_BUDGET.as_secs_f64(), |deadline| {
deadline
.saturating_duration_since(Instant::now())
.as_secs_f64()
})
};
let mut build = self.build_image(spec, rpc_budget()).await?;
loop {
match build.status {
ImageBuildStatus::Ready => return Ok(build),
ImageBuildStatus::Failed => {
let message = if build.error_message.is_empty() {
"image build failed".to_string()
} else {
build.error_message.clone()
};
return Err(SailError::ImageBuild { message });
}
_ => {}
}
let nap = match deadline {
None => BUILD_POLL_INTERVAL,
Some(deadline) => {
let left = deadline.saturating_duration_since(Instant::now());
if left.is_zero() {
return Err(SailError::Transport {
kind: TransportKind::Timeout,
message: format!(
"timed out waiting for image build {}",
build.image_id
),
source: None,
});
}
left.min(BUILD_POLL_INTERVAL)
}
};
tokio::time::sleep(nap).await;
build = self
.get_image_build_status(&build.image_id, rpc_budget())
.await?;
}
}
}
fn sized_put_request(
http: &reqwest::Client,
upload_url: &str,
file: tokio::fs::File,
size: u64,
headers: &HashMap<String, String>,
) -> (
reqwest::RequestBuilder,
Arc<std::sync::Mutex<Option<String>>>,
) {
let (body, streamed_digest) = SizedFileBody::new(file, size);
let mut request = http.put(upload_url).body(reqwest::Body::wrap(body));
for (name, value) in headers {
request = request.header(name, value);
}
(request, streamed_digest)
}
fn upload_timeout(size: u64) -> Duration {
UPLOAD_BASE_TIMEOUT + Duration::from_secs(size / MIN_UPLOAD_BYTES_PER_SEC)
}
struct SizedFileBody {
reader: tokio_util::io::ReaderStream<tokio::fs::File>,
remaining: u64,
hasher: Option<sha2::Sha256>,
streamed_digest: Arc<std::sync::Mutex<Option<String>>>,
}
impl SizedFileBody {
fn new(file: tokio::fs::File, size: u64) -> (Self, Arc<std::sync::Mutex<Option<String>>>) {
let streamed_digest = Arc::new(std::sync::Mutex::new(None));
let mut hasher = Some(sha2::Sha256::new());
if size == 0 {
*streamed_digest.lock().unwrap() =
Some(format!("{:x}", hasher.take().unwrap().finalize()));
}
(
SizedFileBody {
reader: tokio_util::io::ReaderStream::new(file),
remaining: size,
hasher,
streamed_digest: Arc::clone(&streamed_digest),
},
streamed_digest,
)
}
}
impl http_body::Body for SizedFileBody {
type Data = bytes::Bytes;
type Error = std::io::Error;
fn poll_frame(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
use futures::Stream;
match std::pin::Pin::new(&mut self.reader).poll_next(cx) {
std::task::Poll::Ready(Some(Ok(chunk))) => {
self.remaining = self.remaining.saturating_sub(chunk.len() as u64);
if let Some(hasher) = self.hasher.as_mut() {
hasher.update(&chunk);
}
if self.remaining == 0 {
if let Some(hasher) = self.hasher.take() {
*self.streamed_digest.lock().unwrap() =
Some(format!("{:x}", hasher.finalize()));
}
}
std::task::Poll::Ready(Some(Ok(http_body::Frame::data(chunk))))
}
std::task::Poll::Ready(Some(Err(err))) => std::task::Poll::Ready(Some(Err(err))),
std::task::Poll::Ready(None) => {
if let Some(hasher) = self.hasher.take() {
*self.streamed_digest.lock().unwrap() =
Some(format!("{:x}", hasher.finalize()));
}
std::task::Poll::Ready(None)
}
std::task::Poll::Pending => std::task::Poll::Pending,
}
}
fn is_end_stream(&self) -> bool {
self.remaining == 0
}
fn size_hint(&self) -> http_body::SizeHint {
http_body::SizeHint::with_exact(self.remaining)
}
}
#[cfg(test)]
mod tests {
#[test]
fn upload_budget_scales_with_content_size() {
assert_eq!(upload_timeout(0), Duration::from_mins(5));
assert_eq!(
upload_timeout(1 << 30),
Duration::from_mins(5) + Duration::from_secs(1024)
);
}
#[tokio::test]
async fn upload_body_advertises_its_exact_size() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("payload.bin");
std::fs::write(&path, b"0123456789").expect("write");
let file = tokio::fs::File::open(&path).await.expect("open");
let (body, _digest) = SizedFileBody::new(file, 10);
assert_eq!(http_body::Body::size_hint(&body).exact(), Some(10));
assert!(!http_body::Body::is_end_stream(&body));
}
#[tokio::test]
async fn presigned_put_uses_content_length_framing() {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("payload.bin");
std::fs::write(&path, b"0123456789").expect("write");
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut sock, _) = listener.accept().await.expect("accept");
let mut raw = Vec::new();
let mut buf = [0u8; 4096];
loop {
let n = sock.read(&mut buf).await.expect("read");
raw.extend_from_slice(&buf[..n]);
if let Some(head_end) = raw.windows(4).position(|w| w == b"\r\n\r\n") {
let head = String::from_utf8_lossy(&raw[..head_end]).to_lowercase();
let body_len = raw.len() - (head_end + 4);
if body_len >= 10 {
sock.write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n")
.await
.expect("respond");
return head;
}
}
}
});
let file = tokio::fs::File::open(&path).await.expect("open");
let headers = HashMap::from([(
"Content-Type".to_string(),
"application/octet-stream".to_string(),
)]);
let (request, streamed_digest) = sized_put_request(
&reqwest::Client::new(),
&format!("http://{addr}/upload"),
file,
10,
&headers,
);
let response = request.send().await.expect("send");
assert!(response.status().is_success());
assert_eq!(
streamed_digest.lock().unwrap().as_deref(),
Some("84d89877f0d4041efb6bf91a16f0248f2fd573e6af05c19f96bedb9f882f7882")
);
let head = server.await.expect("server");
assert!(
head.contains("content-length: 10"),
"missing sized framing in request head: {head}"
);
assert!(
!head.contains("transfer-encoding"),
"request must not be chunked: {head}"
);
}
use super::*;
#[test]
fn remote_path_rules_match_the_wrappers() {
assert!(validate_remote_path("/app/config.json").is_ok());
assert!(validate_remote_path("relative").is_err());
assert!(validate_remote_path("/app/").is_err());
assert!(validate_remote_path("/app/../etc").is_err());
assert!(validate_remote_path("/app/with space").is_err());
assert!(validate_remote_path("/app/$HOME").is_err());
assert!(validate_mode(Some(0o600)).is_ok());
assert!(validate_mode(Some(0o1777)).is_err());
}
#[tokio::test]
async fn resolve_walks_hashes_and_respects_gitignore() {
let dir = tempfile::tempdir().expect("tempdir");
std::fs::create_dir_all(dir.path().join("src/generated")).unwrap();
std::fs::write(dir.path().join("src/keep.py"), b"keep").unwrap();
std::fs::write(dir.path().join("src/skip.pyc"), b"skip").unwrap();
std::fs::write(dir.path().join("src/generated/gen.py"), b"gen").unwrap();
std::fs::write(dir.path().join("top.txt"), b"top").unwrap();
let matcher = ignore_matcher(
dir.path(),
&["*.pyc".to_string(), "src/generated/".to_string()],
None,
)
.expect("matcher");
let walked = walk_dir(dir.path(), &matcher).expect("walk");
let mut paths: Vec<_> = walked.iter().map(|f| f.relative_path.clone()).collect();
paths.sort();
assert_eq!(paths, ["src/keep.py", "top.txt"]);
let (digest, size) = hash_file(&dir.path().join("top.txt")).await.expect("hash");
assert_eq!(size, 3);
assert_eq!(
digest,
"28720365c5e7476a011e4f43ac003ee5f16247a263b9d623aa85ed311d73bf39"
);
}
}