diff --git a/src/agent/onefuzz/src/uploader.rs b/src/agent/onefuzz/src/uploader.rs index dcbfc0d2ed..cb2d90abf3 100644 --- a/src/agent/onefuzz/src/uploader.rs +++ b/src/agent/onefuzz/src/uploader.rs @@ -7,7 +7,7 @@ use anyhow::{Context, Result}; use futures::stream::TryStreamExt; use reqwest::{Body, Client, Response, StatusCode, Url}; use reqwest_retry::{ - send_retry_reqwest_default, RetryCheck, SendRetry, DEFAULT_RETRY_PERIOD, MAX_RETRY_ATTEMPTS, + send_retry_reqwest, RetryCheck, SendRetry, DEFAULT_RETRY_PERIOD, MAX_RETRY_ATTEMPTS, }; use serde::Serialize; use tokio::{fs, io}; @@ -42,44 +42,39 @@ impl BlobUploader { url }; - // Check if the file already exists before uploading - if let Ok(head) = self - .client - .head(url.clone()) - .send_retry( - |code| match code { - StatusCode::NOT_FOUND => RetryCheck::Fail, - _ => RetryCheck::Retry, - }, - DEFAULT_RETRY_PERIOD, - MAX_RETRY_ATTEMPTS, - ) - .await - { - if head.status() == StatusCode::OK { - return Ok(head); - } - } - let content_length = format!("{}", file_len); - let resp = send_retry_reqwest_default(|| { - let file = fs::File::from_std(std::fs::File::open(file_path)?); - let reader = io::BufReader::new(file); - let codec = codec::BytesCodec::new(); - let file_stream = codec::FramedRead::new(reader, codec) - .map_ok(bytes::BytesMut::freeze) - .into_stream(); - - let request_builder = self - .client - .put(url.clone()) - .header("Content-Length", &content_length) - .header("x-ms-blob-type", "BlockBlob") - .body(Body::wrap_stream(file_stream)); - - Ok(request_builder) - }) + let resp = send_retry_reqwest( + || { + let file = fs::File::from_std(std::fs::File::open(file_path)?); + let reader = io::BufReader::new(file); + let codec = codec::BytesCodec::new(); + let file_stream = codec::FramedRead::new(reader, codec) + .map_ok(bytes::BytesMut::freeze) + .into_stream(); + + let request_builder = self + .client + .put(url.clone()) + // https://learn.microsoft.com/en-us/rest/api/storageservices/put-blob-from-url#request-headers + .header("Content-Length", &content_length) + .header("x-ms-blob-type", "BlockBlob") + // upload only if the the destination blob does not exist + .header("If-None-Match", "*") + .body(Body::wrap_stream(file_stream)); + + Ok(request_builder) + }, + |status| { + if status == StatusCode::PRECONDITION_FAILED || status == StatusCode::CONFLICT { + RetryCheck::Succeed + } else { + RetryCheck::Retry + } + }, + DEFAULT_RETRY_PERIOD, + MAX_RETRY_ATTEMPTS, + ) .await .context("BlobUploader.upload")?;