firebase_rs_sdk/storage/
service.rs

1use std::sync::{Arc, Mutex};
2use std::time::Duration;
3
4use crate::app::FirebaseApp;
5use crate::app_check::FirebaseAppCheckInternal;
6use crate::auth::Auth;
7use crate::component::Provider;
8use crate::storage::constants::{
9    DEFAULT_HOST, DEFAULT_MAX_OPERATION_RETRY_TIME_MS, DEFAULT_MAX_UPLOAD_RETRY_TIME_MS,
10    DEFAULT_PROTOCOL,
11};
12use crate::storage::error::{internal_error, no_default_bucket, StorageResult};
13use crate::storage::location::Location;
14use crate::storage::reference::StorageReference;
15#[cfg(not(target_arch = "wasm32"))]
16use crate::storage::request::StreamingResponse;
17use crate::storage::request::{BackoffConfig, HttpClient, RequestInfo};
18use crate::storage::util::is_url;
19
20#[derive(Clone)]
21pub struct FirebaseStorageImpl {
22    app: FirebaseApp,
23    auth_provider: Provider,
24    app_check_provider: Provider,
25    firebase_version: Option<String>,
26    url_override: Option<String>,
27    state: Arc<Mutex<FirebaseStorageState>>,
28}
29
30struct FirebaseStorageState {
31    bucket: Option<Location>,
32    host: String,
33    protocol: String,
34    max_operation_retry_time_ms: u64,
35    max_upload_retry_time_ms: u64,
36    override_auth_token: Option<String>,
37    is_using_emulator: bool,
38}
39
40impl FirebaseStorageImpl {
41    #[allow(clippy::too_many_arguments)]
42    pub fn new(
43        app: FirebaseApp,
44        auth_provider: Provider,
45        app_check_provider: Provider,
46        url_override: Option<String>,
47        firebase_version: Option<String>,
48    ) -> StorageResult<Self> {
49        let host = DEFAULT_HOST.to_string();
50        let bucket = if let Some(url) = url_override.as_ref() {
51            Some(Location::from_bucket_spec(url, &host)?)
52        } else {
53            extract_bucket(&host, &app)?
54        };
55
56        let state = FirebaseStorageState {
57            bucket,
58            host,
59            protocol: DEFAULT_PROTOCOL.to_string(),
60            max_operation_retry_time_ms: DEFAULT_MAX_OPERATION_RETRY_TIME_MS,
61            max_upload_retry_time_ms: DEFAULT_MAX_UPLOAD_RETRY_TIME_MS,
62            override_auth_token: None,
63            is_using_emulator: false,
64        };
65
66        Ok(Self {
67            app,
68            auth_provider,
69            app_check_provider,
70            firebase_version,
71            url_override,
72            state: Arc::new(Mutex::new(state)),
73        })
74    }
75
76    pub fn app(&self) -> &FirebaseApp {
77        &self.app
78    }
79
80    pub fn host(&self) -> String {
81        self.state.lock().unwrap().host.clone()
82    }
83
84    pub fn protocol(&self) -> String {
85        self.state.lock().unwrap().protocol.clone()
86    }
87
88    pub fn auth_provider(&self) -> Provider {
89        self.auth_provider.clone()
90    }
91
92    pub fn app_check_provider(&self) -> Provider {
93        self.app_check_provider.clone()
94    }
95
96    pub fn firebase_version(&self) -> Option<&str> {
97        self.firebase_version.as_deref()
98    }
99
100    pub fn bucket(&self) -> Option<Location> {
101        self.state.lock().unwrap().bucket.clone()
102    }
103
104    pub fn max_upload_retry_time(&self) -> u64 {
105        self.state.lock().unwrap().max_upload_retry_time_ms
106    }
107
108    pub fn max_operation_retry_time(&self) -> u64 {
109        self.state.lock().unwrap().max_operation_retry_time_ms
110    }
111
112    pub fn set_max_upload_retry_time(&self, millis: u64) {
113        self.state.lock().unwrap().max_upload_retry_time_ms = millis;
114    }
115
116    pub fn set_max_operation_retry_time(&self, millis: u64) {
117        self.state.lock().unwrap().max_operation_retry_time_ms = millis;
118    }
119
120    pub fn is_using_emulator(&self) -> bool {
121        self.state.lock().unwrap().is_using_emulator
122    }
123
124    pub fn connect_emulator(
125        &self,
126        host: &str,
127        port: u16,
128        mock_user_token: Option<String>,
129    ) -> StorageResult<()> {
130        let host_string = format!("{host}:{port}");
131        let bucket = self.compute_bucket_for_host(&host_string)?;
132        let mut state = self.state.lock().unwrap();
133        state.host = host_string;
134        state.bucket = bucket;
135        state.protocol = "http".to_string();
136        state.is_using_emulator = true;
137        state.override_auth_token = mock_user_token;
138        Ok(())
139    }
140
141    pub fn set_host(&self, host: &str) -> StorageResult<()> {
142        let bucket = self.compute_bucket_for_host(host)?;
143        let mut state = self.state.lock().unwrap();
144        state.host = host.to_string();
145        state.bucket = bucket;
146        Ok(())
147    }
148
149    fn compute_bucket_for_host(&self, host: &str) -> StorageResult<Option<Location>> {
150        if let Some(url) = self.url_override.as_ref() {
151            Ok(Some(Location::from_bucket_spec(url, host)?))
152        } else {
153            extract_bucket(host, &self.app)
154        }
155    }
156
157    pub fn make_storage_reference(&self, location: Location) -> StorageReference {
158        StorageReference::new(self.clone(), location)
159    }
160
161    pub fn root_reference(&self) -> StorageResult<StorageReference> {
162        let state = self.state.lock().unwrap();
163        let bucket = state.bucket.clone().ok_or_else(no_default_bucket)?;
164        Ok(StorageReference::new(self.clone(), bucket))
165    }
166
167    pub fn reference_from_path(&self, path: Option<&str>) -> StorageResult<StorageReference> {
168        let location = match path {
169            Some(path) if is_url(path) => Location::from_url(path, &self.host())?,
170            Some(path) => {
171                let base = self.bucket().ok_or_else(no_default_bucket)?;
172                let child_path = crate::storage::path::child(base.path(), path);
173                Location::new(base.bucket(), child_path)
174            }
175            None => self.bucket().ok_or_else(no_default_bucket)?,
176        };
177        Ok(StorageReference::new(self.clone(), location))
178    }
179
180    pub fn http_client(&self) -> StorageResult<HttpClient> {
181        let timeout = Duration::from_millis(self.max_operation_retry_time());
182        let config = BackoffConfig::standard_operation().with_total_timeout(timeout);
183        HttpClient::new(self.is_using_emulator(), config)
184    }
185
186    pub fn upload_http_client(&self) -> StorageResult<HttpClient> {
187        let timeout = Duration::from_millis(self.max_upload_retry_time());
188        let config = BackoffConfig::upload_operation(timeout);
189        HttpClient::new(self.is_using_emulator(), config)
190    }
191
192    pub async fn run_request<O>(&self, info: RequestInfo<O>) -> StorageResult<O> {
193        let client = self.http_client()?;
194        let info = self.prepare_request(info).await?;
195        client.execute(info).await
196    }
197
198    pub async fn run_upload_request<O>(&self, info: RequestInfo<O>) -> StorageResult<O> {
199        let client = self.upload_http_client()?;
200        let info = self.prepare_request(info).await?;
201        client.execute(info).await
202    }
203
204    #[cfg(not(target_arch = "wasm32"))]
205    pub async fn run_streaming_request<O>(
206        &self,
207        info: RequestInfo<O>,
208    ) -> StorageResult<StreamingResponse> {
209        let client = self.http_client()?;
210        let info = self.prepare_request(info).await?;
211        client.execute_streaming(info).await
212    }
213
214    async fn prepare_request<O>(&self, mut info: RequestInfo<O>) -> StorageResult<RequestInfo<O>> {
215        if let Some(token) = self.auth_token().await? {
216            if !token.is_empty() {
217                info.headers
218                    .insert("Authorization".to_string(), format!("Firebase {token}"));
219            }
220        }
221
222        if let Some(headers) = self.app_check_headers().await? {
223            if !headers.token.is_empty() {
224                info.headers
225                    .insert("X-Firebase-AppCheck".to_string(), headers.token);
226            }
227            if let Some(heartbeat) = headers.heartbeat {
228                if !heartbeat.is_empty() {
229                    info.headers
230                        .insert("X-Firebase-Client".to_string(), heartbeat);
231                }
232            }
233        }
234
235        if !info.headers.contains_key("X-Firebase-Storage-Version") {
236            let version = format!(
237                "webjs/{}",
238                self.firebase_version.as_deref().unwrap_or("AppManager")
239            );
240            info.headers
241                .insert("X-Firebase-Storage-Version".to_string(), version);
242        }
243
244        if let Some(app_id) = self.app.options().app_id {
245            if !app_id.is_empty() {
246                info.headers
247                    .entry("X-Firebase-GMPID".to_string())
248                    .or_insert(app_id);
249            }
250        }
251
252        Ok(info)
253    }
254
255    async fn auth_token(&self) -> StorageResult<Option<String>> {
256        if let Some(token) = {
257            let state = self.state.lock().unwrap();
258            state.override_auth_token.clone()
259        } {
260            return Ok(Some(token));
261        }
262
263        let auth = match self
264            .auth_provider
265            .get_immediate_with_options::<Auth>(None, true)
266        {
267            Ok(Some(auth)) => auth,
268            Ok(None) => return Ok(None),
269            Err(err) => {
270                return Err(internal_error(format!(
271                    "failed to resolve auth provider: {err}"
272                )))
273            }
274        };
275
276        match auth.get_token(false).await {
277            Ok(Some(token)) if token.is_empty() => Ok(None),
278            Ok(Some(token)) => Ok(Some(token)),
279            Ok(None) => Ok(None),
280            Err(err) => Err(internal_error(format!(
281                "failed to obtain auth token: {err}"
282            ))),
283        }
284    }
285
286    async fn app_check_headers(&self) -> StorageResult<Option<AppCheckHeaders>> {
287        let app_check = match self
288            .app_check_provider
289            .get_immediate_with_options::<FirebaseAppCheckInternal>(None, true)
290        {
291            Ok(Some(app_check)) => app_check,
292            Ok(None) => return Ok(None),
293            Err(err) => {
294                return Err(internal_error(format!(
295                    "failed to resolve app check provider: {err}"
296                )))
297            }
298        };
299
300        let token = match app_check.get_token(false).await {
301            Ok(result) => result.token,
302            Err(err) => {
303                if let Some(cached) = err.cached_token() {
304                    cached.token.clone()
305                } else {
306                    return Err(internal_error(format!(
307                        "failed to obtain App Check token: {err}"
308                    )));
309                }
310            }
311        };
312
313        if token.is_empty() {
314            Ok(None)
315        } else {
316            let heartbeat = app_check.heartbeat_header().await.map_err(|err| {
317                internal_error(format!(
318                    "failed to obtain App Check heartbeat header: {err}"
319                ))
320            })?;
321
322            Ok(Some(AppCheckHeaders { token, heartbeat }))
323        }
324    }
325}
326
327struct AppCheckHeaders {
328    token: String,
329    heartbeat: Option<String>,
330}
331
332fn extract_bucket(host: &str, app: &FirebaseApp) -> StorageResult<Option<Location>> {
333    let options = app.options();
334    match options.storage_bucket {
335        Some(bucket) => Ok(Some(Location::from_bucket_spec(&bucket, host)?)),
336        None => Ok(None),
337    }
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343    use crate::app::initialize_app;
344    use crate::app::{FirebaseAppSettings, FirebaseOptions};
345    use crate::app_check::{
346        box_app_check_future, AppCheckOptions, AppCheckProvider, AppCheckProviderFuture,
347        AppCheckToken,
348    };
349    use crate::app_check::{
350        clear_registry, clear_state_for_tests, initialize_app_check, test_guard, token_with_ttl,
351    };
352    use crate::component::types::{ComponentError, DynService, InstanceFactoryOptions};
353    use crate::component::{Component, ComponentType};
354    use crate::storage::request::{RequestInfo, ResponseHandler};
355    use reqwest::Method;
356    use std::future::Future;
357    use std::sync::atomic::{AtomicUsize, Ordering};
358    use std::sync::Arc;
359    use std::time::Duration;
360
361    fn unique_settings(prefix: &str) -> FirebaseAppSettings {
362        static COUNTER: AtomicUsize = AtomicUsize::new(0);
363        FirebaseAppSettings {
364            name: Some(format!(
365                "{prefix}-{}",
366                COUNTER.fetch_add(1, Ordering::SeqCst)
367            )),
368            ..Default::default()
369        }
370    }
371
372    fn base_options() -> FirebaseOptions {
373        FirebaseOptions {
374            storage_bucket: Some("my-bucket".into()),
375            app_id: Some("1:123:web:abc".into()),
376            ..Default::default()
377        }
378    }
379
380    fn test_request() -> RequestInfo<()> {
381        let handler: ResponseHandler<()> = Arc::new(|_| Ok(()));
382        RequestInfo::new(
383            "https://example.com",
384            Method::GET,
385            Duration::from_secs(5),
386            handler,
387        )
388    }
389
390    async fn build_storage_with<F, Fut>(configure: F) -> FirebaseStorageImpl
391    where
392        F: Fn(&FirebaseApp) -> Fut,
393        Fut: Future<Output = ()>,
394    {
395        let app = initialize_app(base_options(), Some(unique_settings("storage-service")))
396            .await
397            .expect("failed to initialize app");
398        configure(&app).await;
399
400        let container = app.container();
401        let auth_provider = container.get_provider("auth-internal");
402        let app_check_provider = container.get_provider("app-check-internal");
403        FirebaseStorageImpl::new(
404            app,
405            auth_provider,
406            app_check_provider,
407            None,
408            Some("test-sdk".into()),
409        )
410        .expect("storage construction should succeed")
411    }
412
413    #[tokio::test]
414    async fn prepare_request_adds_headers_for_emulator_override() {
415        let storage = build_storage_with(|_| async {}).await;
416        storage
417            .connect_emulator("localhost", 9199, Some("mock-token".into()))
418            .unwrap();
419
420        let prepared = storage.prepare_request(test_request()).await.unwrap();
421
422        assert_eq!(
423            prepared.headers.get("Authorization"),
424            Some(&"Firebase mock-token".to_string())
425        );
426
427        let expected_version = format!(
428            "webjs/{}",
429            storage.firebase_version().unwrap_or("AppManager")
430        );
431        assert_eq!(
432            prepared.headers.get("X-Firebase-Storage-Version"),
433            Some(&expected_version)
434        );
435
436        assert_eq!(
437            prepared.headers.get("X-Firebase-GMPID"),
438            Some(&"1:123:web:abc".to_string())
439        );
440
441        assert!(prepared.headers.get("X-Firebase-AppCheck").is_none());
442    }
443
444    #[derive(Clone)]
445    struct StaticAppCheckProvider;
446
447    impl AppCheckProvider for StaticAppCheckProvider {
448        fn get_token(
449            &self,
450        ) -> AppCheckProviderFuture<'_, crate::app_check::AppCheckResult<AppCheckToken>> {
451            box_app_check_future(async {
452                token_with_ttl("app-check-token", Duration::from_secs(60))
453            })
454        }
455    }
456
457    async fn register_app_check(app: &FirebaseApp) {
458        let provider = Arc::new(StaticAppCheckProvider);
459        let options = AppCheckOptions::new(provider);
460        let app_check = initialize_app_check(Some(app.clone()), options)
461            .await
462            .expect("initialize app check");
463        let internal = Arc::new(FirebaseAppCheckInternal::new(app_check));
464
465        let factory = {
466            let internal = internal.clone();
467            Arc::new(
468                move |_: &crate::component::ComponentContainer,
469                      _: InstanceFactoryOptions|
470                      -> Result<DynService, ComponentError> {
471                    Ok(internal.clone() as DynService)
472                },
473            )
474        };
475
476        let component = Component::new("app-check-internal", factory, ComponentType::Private);
477        app.container().add_or_overwrite_component(component);
478    }
479
480    #[tokio::test]
481    async fn prepare_request_includes_app_check_header_when_available() {
482        let _guard = test_guard();
483        clear_state_for_tests();
484        clear_registry();
485        let storage = build_storage_with(|app| {
486            let app = app.clone();
487            async move { register_app_check(&app).await }
488        })
489        .await;
490        let prepared = storage.prepare_request(test_request()).await.unwrap();
491
492        assert_eq!(
493            prepared.headers.get("X-Firebase-AppCheck"),
494            Some(&"app-check-token".to_string())
495        );
496
497        clear_state_for_tests();
498        clear_registry();
499    }
500}