Skip to main content

helix_driver_host/
upload.rs

1//! Shared local-file uploader for native and FFI host drivers.
2
3use 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        // reqwest 错误默认可能携带完整 URL;先剥离,避免签名查询串进入 PortError/trace/log。
45        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}