helix_driver_host/
upload.rs1use std::sync::Arc;
4use std::time::Duration;
5
6use bytes::Bytes;
7use futures_util::stream;
8use reqwest::header;
9use reqwest::{Client, Method};
10use tokio::io::AsyncReadExt;
11
12use helix_core::effect::{FileUploadProgress, FileUploadRequest, FileUploadResponse};
13use helix_core::ports::{FileUploadProgressReporter, FileUploader};
14use helix_core::PortError;
15
16#[derive(Clone)]
17pub struct SharedFileUploader {
18 client: Client,
19 timeout: Duration,
20}
21
22impl Default for SharedFileUploader {
23 fn default() -> Self {
24 Self::new(Client::new())
25 }
26}
27
28impl SharedFileUploader {
29 pub fn new(client: Client) -> Self {
30 Self {
31 client,
32 timeout: Duration::from_secs(30),
33 }
34 }
35
36 pub fn with_timeout(mut self, timeout: Duration) -> Self {
37 self.timeout = timeout;
38 self
39 }
40
41 fn classify_error(e: reqwest::Error) -> PortError {
42 let is_timeout = e.is_timeout();
43 let is_network = e.is_connect() || e.is_request();
44 let sanitized = e.without_url();
46 if is_timeout {
47 PortError::Transport(format!("timeout: {sanitized}"))
48 } else if is_network {
49 PortError::Transport(format!("network: {sanitized}"))
50 } else {
51 PortError::Http(sanitized.to_string())
52 }
53 }
54}
55
56#[async_trait::async_trait]
57impl FileUploader for SharedFileUploader {
58 async fn upload(&self, req: FileUploadRequest) -> Result<FileUploadResponse, PortError> {
59 self.upload_with_progress(req, None).await
60 }
61
62 async fn upload_with_progress(
63 &self,
64 req: FileUploadRequest,
65 progress: Option<Arc<dyn FileUploadProgressReporter>>,
66 ) -> Result<FileUploadResponse, PortError> {
67 let method = Method::from_bytes(req.method.as_bytes())
68 .map_err(|e| PortError::Other(format!("invalid upload method: {e}")))?;
69 if method != Method::PUT {
70 return Err(PortError::Other(format!(
71 "unsupported upload method: {}",
72 req.method
73 )));
74 }
75
76 let FileUploadRequest {
77 local_path,
78 object_key,
79 urls,
80 headers,
81 content_type,
82 size,
83 ..
84 } = req;
85 let (upload_url, public_url) = urls.into_parts();
86
87 let metadata = tokio::fs::metadata(&local_path)
88 .await
89 .map_err(|e| PortError::Other(format!("stat upload file failed: {e}")))?;
90 let actual_size = metadata.len();
91 if let Some(expected) = size {
92 if actual_size != expected {
93 return Err(PortError::Other(format!(
94 "upload file size mismatch: expected {expected}, got {actual_size}"
95 )));
96 }
97 }
98 let file = tokio::fs::File::open(&local_path)
99 .await
100 .map_err(|e| PortError::Other(format!("open upload file failed: {e}")))?;
101 let body = reqwest::Body::wrap_stream(stream::unfold(
102 (file, 0_u64, progress),
103 move |(mut file, completed_bytes, progress)| async move {
104 let mut chunk = vec![0_u8; 64 * 1024];
105 match file.read(&mut chunk).await {
106 Ok(0) => None,
107 Ok(read) => {
108 chunk.truncate(read);
109 let completed_bytes = completed_bytes.saturating_add(read as u64);
110 if let Some(reporter) = progress.as_ref() {
111 reporter.report(FileUploadProgress {
112 completed_bytes,
113 total_bytes: actual_size,
114 });
115 }
116 Some((
117 Ok::<Bytes, std::io::Error>(Bytes::from(chunk)),
118 (file, completed_bytes, progress),
119 ))
120 }
121 Err(err) => Some((
122 Err::<Bytes, std::io::Error>(err),
123 (file, completed_bytes, progress),
124 )),
125 }
126 },
127 ));
128
129 let has_ticket_content_type = headers
130 .iter()
131 .any(|(name, _)| name.eq_ignore_ascii_case(header::CONTENT_TYPE.as_str()));
132 let mut builder = self.client.request(method, &upload_url);
133 for (name, value) in headers {
134 builder = builder.header(&name, value);
135 }
136 builder = builder.header(header::CONTENT_LENGTH, actual_size);
137 if let Some(content_type) = content_type.filter(|_| !has_ticket_content_type) {
138 builder = builder.header(header::CONTENT_TYPE, content_type);
139 }
140
141 let response = tokio::time::timeout(self.timeout, builder.body(body).send())
142 .await
143 .map_err(|_| {
144 PortError::Transport(format!("timeout after {} ms", self.timeout.as_millis()))
145 })?
146 .map_err(Self::classify_error)?;
147 let status = response.status();
148 if !status.is_success() {
149 return Err(PortError::Http(format!(
150 "upload status {}",
151 status.as_u16()
152 )));
153 }
154
155 let etag = response
156 .headers()
157 .get(header::ETAG)
158 .and_then(|value| value.to_str().ok())
159 .map(str::to_owned);
160
161 Ok(FileUploadResponse {
162 object_key,
163 public_url,
164 etag,
165 })
166 }
167}