bitwarden_server_communication_config/
client.rs

1#[cfg(test)]
2use crate::AcquiredCookie;
3use crate::{
4    AcquireCookieError, BootstrapConfig, ServerCommunicationConfig,
5    ServerCommunicationConfigPlatformApi, ServerCommunicationConfigRepository,
6};
7
8/// Server communication configuration client
9pub struct ServerCommunicationConfigClient<R, P>
10where
11    R: ServerCommunicationConfigRepository,
12    P: ServerCommunicationConfigPlatformApi,
13{
14    repository: R,
15    platform_api: P,
16}
17
18impl<R, P> ServerCommunicationConfigClient<R, P>
19where
20    R: ServerCommunicationConfigRepository,
21    P: ServerCommunicationConfigPlatformApi,
22{
23    /// Creates a new server communication configuration client
24    ///
25    /// # Arguments
26    ///
27    /// * `repository` - Cookie storage implementation (e.g a StateProvider hook)
28    /// * `platform_api` - Cookie acquistion implementation
29    pub fn new(repository: R, platform_api: P) -> Self {
30        Self {
31            repository,
32            platform_api,
33        }
34    }
35
36    /// Retrieves the server communication configuration for a hostname
37    pub async fn get_config(
38        &self,
39        hostname: String,
40    ) -> Result<ServerCommunicationConfig, R::GetError> {
41        Ok(self
42            .repository
43            .get(hostname)
44            .await?
45            .unwrap_or(ServerCommunicationConfig {
46                bootstrap: BootstrapConfig::Direct,
47            }))
48    }
49
50    /// Determines if cookie bootstrapping is needed for this hostname
51    pub async fn needs_bootstrap(&self, hostname: String) -> bool {
52        if let Ok(Some(config)) = self.repository.get(hostname).await
53            && let BootstrapConfig::SsoCookieVendor(vendor_config) = config.bootstrap
54        {
55            return vendor_config.cookie_value.is_none();
56        }
57        false
58    }
59
60    /// Returns cookies to include in HTTP requests
61    ///
62    /// Returns the stored cookies as-is. For sharded cookies, each entry includes
63    /// the full cookie name with its `-{N}` suffix (e.g., `AWSELBAuthSessionCookie-0`).
64    pub async fn cookies(&self, hostname: String) -> Vec<(String, String)> {
65        if let Ok(Some(config)) = self.repository.get(hostname).await
66            && let BootstrapConfig::SsoCookieVendor(vendor_config) = config.bootstrap
67            && let Some(acquired_cookies) = vendor_config.cookie_value
68        {
69            return acquired_cookies
70                .into_iter()
71                .map(|cookie| (cookie.name, cookie.value))
72                .collect();
73        }
74        Vec::new()
75    }
76
77    /// Sets the server communication configuration for a hostname
78    ///
79    /// This method saves the provided communication configuration to the repository.
80    /// Typically called when receiving the `/api/config` response from the server.
81    ///
82    /// # Arguments
83    ///
84    /// * `hostname` - The server hostname (e.g., "vault.acme.com")
85    /// * `config` - The server communication configuration to store
86    ///
87    /// # Errors
88    ///
89    /// Returns an error if the repository save operation fails
90    pub async fn set_communication_type(
91        &self,
92        hostname: String,
93        config: ServerCommunicationConfig,
94    ) -> Result<(), R::SaveError> {
95        self.repository.save(hostname, config).await
96    }
97
98    /// Acquires a cookie from the platform and saves it to the repository
99    ///
100    /// This method calls the platform API to trigger cookie acquisition (e.g., browser
101    /// redirect to IdP), then validates and stores the acquired cookie in the repository.
102    ///
103    /// # Arguments
104    ///
105    /// * `hostname` - The server hostname (e.g., "vault.acme.com")
106    ///
107    /// # Errors
108    ///
109    /// Returns an error if:
110    /// - Cookie acquisition was cancelled by the user ([`AcquireCookieError::Cancelled`])
111    /// - Server configuration doesn't support SSO cookies
112    ///   ([`AcquireCookieError::UnsupportedConfiguration`])
113    /// - Acquired cookie name doesn't match expected name
114    ///   ([`AcquireCookieError::CookieNameMismatch`])
115    /// - Repository operations fail ([`AcquireCookieError::RepositoryGetError`] or
116    ///   [`AcquireCookieError::RepositorySaveError`])
117    pub async fn acquire_cookie(&self, hostname: &str) -> Result<(), AcquireCookieError> {
118        // Get existing configuration - we need this to know what cookie to expect
119        let mut config = self
120            .repository
121            .get(hostname.to_string())
122            .await
123            .map_err(|e| AcquireCookieError::RepositoryGetError(format!("{:?}", e)))?
124            .ok_or(AcquireCookieError::UnsupportedConfiguration)?;
125
126        // Verify this is an SSO cookie vendor configuration and get mutable reference
127        let BootstrapConfig::SsoCookieVendor(ref mut vendor_config) = config.bootstrap else {
128            return Err(AcquireCookieError::UnsupportedConfiguration);
129        };
130
131        let expected_cookie_name = vendor_config
132            .cookie_name
133            .as_ref()
134            .ok_or(AcquireCookieError::UnsupportedConfiguration)?;
135
136        // Call platform API to acquire cookies
137        let cookies = self
138            .platform_api
139            .acquire_cookies(hostname.to_string())
140            .await
141            .ok_or(AcquireCookieError::Cancelled)?;
142
143        // Validate that all cookies match the expected base name
144        // Cookie names should either:
145        // 1. Exactly match the expected name (unsharded cookie)
146        // 2. Match the pattern {expected_name}-{N} where N is a digit (sharded cookies)
147        //
148        // AWS ALB shards cookies > 4KB with naming pattern: {base_name}-{N}
149        // where N starts at 0 (e.g., AWSELBAuthSessionCookie-0, AWSELBAuthSessionCookie-1)
150        let all_cookies_match = cookies.iter().all(|cookie| {
151            cookie.name == *expected_cookie_name
152                || cookie
153                    .name
154                    .strip_prefix(&format!("{}-", expected_cookie_name))
155                    .is_some_and(|suffix| suffix.chars().all(|c| c.is_ascii_digit()))
156        });
157
158        if !all_cookies_match {
159            // Find the first mismatched cookie for error reporting
160            let mismatched = cookies
161                .iter()
162                .find(|cookie| {
163                    cookie.name != *expected_cookie_name
164                        && !cookie
165                            .name
166                            .strip_prefix(&format!("{}-", expected_cookie_name))
167                            .is_some_and(|suffix| suffix.chars().all(|c| c.is_ascii_digit()))
168                })
169                .expect("all_cookies_match is false, so at least one cookie must not match");
170
171            return Err(AcquireCookieError::CookieNameMismatch {
172                expected: expected_cookie_name.clone(),
173                actual: mismatched.name.clone(),
174            });
175        }
176
177        // Update the cookie values using the mutable reference we already have
178        vendor_config.cookie_value = Some(cookies);
179
180        // Save the updated config
181        self.repository
182            .save(hostname.to_string(), config)
183            .await
184            .map_err(|e| AcquireCookieError::RepositorySaveError(format!("{:?}", e)))?;
185
186        Ok(())
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use std::collections::HashMap;
193
194    use tokio::sync::RwLock;
195
196    use super::*;
197    use crate::SsoCookieVendorConfig;
198
199    /// Mock in-memory repository for testing
200    #[derive(Default, Clone)]
201    struct MockRepository {
202        storage: std::sync::Arc<RwLock<HashMap<String, ServerCommunicationConfig>>>,
203    }
204
205    impl ServerCommunicationConfigRepository for MockRepository {
206        type GetError = ();
207        type SaveError = ();
208
209        async fn get(&self, hostname: String) -> Result<Option<ServerCommunicationConfig>, ()> {
210            Ok(self.storage.read().await.get(&hostname).cloned())
211        }
212
213        async fn save(
214            &self,
215            hostname: String,
216            config: ServerCommunicationConfig,
217        ) -> Result<(), ()> {
218            self.storage.write().await.insert(hostname, config);
219            Ok(())
220        }
221    }
222
223    /// Mock platform API for testing
224    #[derive(Clone)]
225    struct MockPlatformApi {
226        cookies_to_return: std::sync::Arc<RwLock<Option<Vec<AcquiredCookie>>>>,
227    }
228
229    impl MockPlatformApi {
230        fn new() -> Self {
231            Self {
232                cookies_to_return: std::sync::Arc::new(RwLock::new(None)),
233            }
234        }
235
236        async fn set_cookies(&self, cookies: Option<Vec<AcquiredCookie>>) {
237            *self.cookies_to_return.write().await = cookies;
238        }
239    }
240
241    #[async_trait::async_trait]
242    impl ServerCommunicationConfigPlatformApi for MockPlatformApi {
243        async fn acquire_cookies(&self, _hostname: String) -> Option<Vec<AcquiredCookie>> {
244            self.cookies_to_return.read().await.clone()
245        }
246    }
247
248    #[tokio::test]
249    async fn get_config_returns_direct_when_not_found() {
250        let repo = MockRepository::default();
251        let platform_api = MockPlatformApi::new();
252        let client = ServerCommunicationConfigClient::new(repo, platform_api);
253
254        let config = client
255            .get_config("vault.example.com".to_string())
256            .await
257            .unwrap();
258
259        assert!(matches!(config.bootstrap, BootstrapConfig::Direct));
260    }
261
262    #[tokio::test]
263    async fn get_config_returns_saved_config() {
264        let repo = MockRepository::default();
265        let config = ServerCommunicationConfig {
266            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
267                idp_login_url: Some("https://example.com".to_string()),
268                cookie_name: Some("TestCookie".to_string()),
269                cookie_domain: Some("example.com".to_string()),
270                cookie_value: Some(vec![AcquiredCookie {
271                    name: "TestCookie".to_string(),
272                    value: "value123".to_string(),
273                }]),
274            }),
275        };
276
277        repo.save("vault.example.com".to_string(), config.clone())
278            .await
279            .unwrap();
280
281        let platform_api = MockPlatformApi::new();
282        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
283        let retrieved = client
284            .get_config("vault.example.com".to_string())
285            .await
286            .unwrap();
287
288        assert!(matches!(
289            retrieved.bootstrap,
290            BootstrapConfig::SsoCookieVendor(_)
291        ));
292    }
293
294    #[tokio::test]
295    async fn needs_bootstrap_true_when_cookie_missing() {
296        let repo = MockRepository::default();
297        let config = ServerCommunicationConfig {
298            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
299                idp_login_url: Some("https://example.com".to_string()),
300                cookie_name: Some("TestCookie".to_string()),
301                cookie_domain: Some("example.com".to_string()),
302                cookie_value: None,
303            }),
304        };
305
306        repo.save("vault.example.com".to_string(), config)
307            .await
308            .unwrap();
309
310        let platform_api = MockPlatformApi::new();
311        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
312        assert!(
313            client
314                .needs_bootstrap("vault.example.com".to_string())
315                .await
316        );
317    }
318
319    #[tokio::test]
320    async fn needs_bootstrap_false_when_cookie_present() {
321        let repo = MockRepository::default();
322        let config = ServerCommunicationConfig {
323            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
324                idp_login_url: Some("https://example.com".to_string()),
325                cookie_name: Some("TestCookie".to_string()),
326                cookie_domain: Some("example.com".to_string()),
327                cookie_value: Some(vec![AcquiredCookie {
328                    name: "TestCookie".to_string(),
329                    value: "value123".to_string(),
330                }]),
331            }),
332        };
333
334        repo.save("vault.example.com".to_string(), config)
335            .await
336            .unwrap();
337
338        let platform_api = MockPlatformApi::new();
339        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
340        assert!(
341            !client
342                .needs_bootstrap("vault.example.com".to_string())
343                .await
344        );
345    }
346
347    #[tokio::test]
348    async fn needs_bootstrap_false_for_direct() {
349        let repo = MockRepository::default();
350        let config = ServerCommunicationConfig {
351            bootstrap: BootstrapConfig::Direct,
352        };
353
354        repo.save("vault.example.com".to_string(), config)
355            .await
356            .unwrap();
357
358        let platform_api = MockPlatformApi::new();
359        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
360        assert!(
361            !client
362                .needs_bootstrap("vault.example.com".to_string())
363                .await
364        );
365    }
366
367    #[tokio::test]
368    async fn cookies_returns_empty_for_direct() {
369        let repo = MockRepository::default();
370        let config = ServerCommunicationConfig {
371            bootstrap: BootstrapConfig::Direct,
372        };
373
374        repo.save("vault.example.com".to_string(), config)
375            .await
376            .unwrap();
377
378        let platform_api = MockPlatformApi::new();
379        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
380        let cookies = client.cookies("vault.example.com".to_string()).await;
381
382        assert!(cookies.is_empty());
383    }
384
385    #[tokio::test]
386    async fn cookies_returns_empty_when_value_none() {
387        let repo = MockRepository::default();
388        let config = ServerCommunicationConfig {
389            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
390                idp_login_url: Some("https://example.com".to_string()),
391                cookie_name: Some("TestCookie".to_string()),
392                cookie_domain: Some("example.com".to_string()),
393                cookie_value: None,
394            }),
395        };
396
397        repo.save("vault.example.com".to_string(), config)
398            .await
399            .unwrap();
400
401        let platform_api = MockPlatformApi::new();
402        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
403        let cookies = client.cookies("vault.example.com".to_string()).await;
404
405        assert!(cookies.is_empty());
406    }
407
408    #[tokio::test]
409    async fn cookies_returns_unsharded_cookie_without_suffix() {
410        let repo = MockRepository::default();
411        let config = ServerCommunicationConfig {
412            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
413                idp_login_url: Some("https://example.com".to_string()),
414                cookie_name: Some("AWSELBAuthSessionCookie".to_string()),
415                cookie_domain: Some("example.com".to_string()),
416                cookie_value: Some(vec![AcquiredCookie {
417                    name: "AWSELBAuthSessionCookie".to_string(),
418                    value: "eyJhbGciOiJFUzI1NiIsImtpZCI6Im...".to_string(),
419                }]),
420            }),
421        };
422
423        repo.save("vault.example.com".to_string(), config)
424            .await
425            .unwrap();
426
427        let platform_api = MockPlatformApi::new();
428        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
429        let cookies = client.cookies("vault.example.com".to_string()).await;
430
431        // Single cookie without suffix
432        assert_eq!(cookies.len(), 1);
433        assert_eq!(cookies[0].0, "AWSELBAuthSessionCookie");
434        assert_eq!(cookies[0].1, "eyJhbGciOiJFUzI1NiIsImtpZCI6Im...");
435    }
436
437    #[tokio::test]
438    async fn cookies_returns_empty_when_no_config() {
439        let repo = MockRepository::default();
440        let platform_api = MockPlatformApi::new();
441        let client = ServerCommunicationConfigClient::new(repo, platform_api);
442        let cookies = client.cookies("vault.example.com".to_string()).await;
443
444        assert!(cookies.is_empty());
445    }
446
447    #[tokio::test]
448    async fn cookies_returns_shards_with_numbered_suffixes() {
449        let repo = MockRepository::default();
450        let config = ServerCommunicationConfig {
451            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
452                idp_login_url: Some("https://example.com".to_string()),
453                cookie_name: Some("AWSELBAuthSessionCookie".to_string()),
454                cookie_domain: Some("example.com".to_string()),
455                cookie_value: Some(vec![
456                    AcquiredCookie {
457                        name: "AWSELBAuthSessionCookie-0".to_string(),
458                        value: "shard0value".to_string(),
459                    },
460                    AcquiredCookie {
461                        name: "AWSELBAuthSessionCookie-1".to_string(),
462                        value: "shard1value".to_string(),
463                    },
464                    AcquiredCookie {
465                        name: "AWSELBAuthSessionCookie-2".to_string(),
466                        value: "shard2value".to_string(),
467                    },
468                ]),
469            }),
470        };
471
472        repo.save("vault.example.com".to_string(), config)
473            .await
474            .unwrap();
475
476        let platform_api = MockPlatformApi::new();
477        let client = ServerCommunicationConfigClient::new(repo, platform_api);
478        let cookies = client.cookies("vault.example.com".to_string()).await;
479
480        // Each shard is returned as stored with -N suffix
481        assert_eq!(cookies.len(), 3);
482        assert_eq!(
483            cookies[0],
484            (
485                "AWSELBAuthSessionCookie-0".to_string(),
486                "shard0value".to_string()
487            )
488        );
489        assert_eq!(
490            cookies[1],
491            (
492                "AWSELBAuthSessionCookie-1".to_string(),
493                "shard1value".to_string()
494            )
495        );
496        assert_eq!(
497            cookies[2],
498            (
499                "AWSELBAuthSessionCookie-2".to_string(),
500                "shard2value".to_string()
501            )
502        );
503    }
504
505    #[tokio::test]
506    async fn acquire_cookie_saves_when_cookie_returned() {
507        let repo = MockRepository::default();
508        let platform_api = MockPlatformApi::new();
509
510        // Setup existing config with SsoCookieVendor
511        let config = ServerCommunicationConfig {
512            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
513                idp_login_url: Some("https://example.com".to_string()),
514                cookie_name: Some("TestCookie".to_string()),
515                cookie_domain: Some("example.com".to_string()),
516                cookie_value: None,
517            }),
518        };
519        repo.save("vault.example.com".to_string(), config)
520            .await
521            .unwrap();
522
523        // Configure platform API to return a cookie with correct name
524        platform_api
525            .set_cookies(Some(vec![AcquiredCookie {
526                name: "TestCookie".to_string(),
527                value: "acquired-cookie-value".to_string(),
528            }]))
529            .await;
530
531        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
532
533        // Call acquire_cookie - should succeed
534        client.acquire_cookie("vault.example.com").await.unwrap();
535
536        // Verify cookie was saved
537        let saved_config = repo
538            .get("vault.example.com".to_string())
539            .await
540            .unwrap()
541            .unwrap();
542
543        if let BootstrapConfig::SsoCookieVendor(vendor_config) = saved_config.bootstrap {
544            assert_eq!(vendor_config.cookie_value.as_ref().unwrap().len(), 1);
545            assert_eq!(
546                vendor_config.cookie_value.as_ref().unwrap()[0].name,
547                "TestCookie"
548            );
549            assert_eq!(
550                vendor_config.cookie_value.as_ref().unwrap()[0].value,
551                "acquired-cookie-value"
552            );
553        } else {
554            panic!("Expected SsoCookieVendor config");
555        }
556    }
557
558    #[tokio::test]
559    async fn acquire_cookie_returns_cancelled_when_none() {
560        let repo = MockRepository::default();
561        let platform_api = MockPlatformApi::new();
562
563        // Setup existing config
564        let config = ServerCommunicationConfig {
565            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
566                idp_login_url: Some("https://example.com".to_string()),
567                cookie_name: Some("TestCookie".to_string()),
568                cookie_domain: Some("example.com".to_string()),
569                cookie_value: None,
570            }),
571        };
572        repo.save("vault.example.com".to_string(), config)
573            .await
574            .unwrap();
575
576        // Platform API returns None (user cancelled)
577        platform_api.set_cookies(None).await;
578
579        let client = ServerCommunicationConfigClient::new(repo, platform_api);
580
581        let result = client.acquire_cookie("vault.example.com").await;
582
583        assert!(matches!(result, Err(AcquireCookieError::Cancelled)));
584    }
585
586    #[tokio::test]
587    async fn acquire_cookie_returns_unsupported_for_direct_config() {
588        let repo = MockRepository::default();
589        let platform_api = MockPlatformApi::new();
590
591        // Setup Direct config
592        let config = ServerCommunicationConfig {
593            bootstrap: BootstrapConfig::Direct,
594        };
595        repo.save("vault.example.com".to_string(), config)
596            .await
597            .unwrap();
598
599        // Platform API returns a cookie
600        platform_api
601            .set_cookies(Some(vec![AcquiredCookie {
602                name: "TestCookie".to_string(),
603                value: "cookie-value".to_string(),
604            }]))
605            .await;
606
607        let client = ServerCommunicationConfigClient::new(repo, platform_api);
608
609        let result = client.acquire_cookie("vault.example.com").await;
610
611        // Should return UnsupportedConfiguration because config is Direct
612        assert!(matches!(
613            result,
614            Err(AcquireCookieError::UnsupportedConfiguration)
615        ));
616    }
617
618    #[tokio::test]
619    async fn acquire_cookie_validates_cookie_name() {
620        let repo = MockRepository::default();
621        let platform_api = MockPlatformApi::new();
622
623        // Setup config expecting "ExpectedCookie"
624        let config = ServerCommunicationConfig {
625            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
626                idp_login_url: Some("https://example.com".to_string()),
627                cookie_name: Some("ExpectedCookie".to_string()),
628                cookie_domain: Some("example.com".to_string()),
629                cookie_value: None,
630            }),
631        };
632        repo.save("vault.example.com".to_string(), config)
633            .await
634            .unwrap();
635
636        // Platform API returns wrong cookie name
637        platform_api
638            .set_cookies(Some(vec![AcquiredCookie {
639                name: "WrongCookie".to_string(),
640                value: "some-value".to_string(),
641            }]))
642            .await;
643
644        let client = ServerCommunicationConfigClient::new(repo, platform_api);
645
646        let result = client.acquire_cookie("vault.example.com").await;
647
648        // Should return CookieNameMismatch
649        match result {
650            Err(AcquireCookieError::CookieNameMismatch { expected, actual }) => {
651                assert_eq!(expected, "ExpectedCookie");
652                assert_eq!(actual, "WrongCookie");
653            }
654            _ => panic!("Expected CookieNameMismatch error"),
655        }
656    }
657
658    #[tokio::test]
659    async fn acquire_cookie_returns_unsupported_when_no_config() {
660        let repo = MockRepository::default();
661        let platform_api = MockPlatformApi::new();
662
663        // No config saved for this hostname
664
665        let client = ServerCommunicationConfigClient::new(repo, platform_api);
666
667        let result = client.acquire_cookie("vault.example.com").await;
668
669        // Should return UnsupportedConfiguration because no config exists
670        assert!(matches!(
671            result,
672            Err(AcquireCookieError::UnsupportedConfiguration)
673        ));
674    }
675
676    #[tokio::test]
677    async fn acquire_cookie_accepts_sharded_cookies_with_numbered_suffixes() {
678        let repo = MockRepository::default();
679        let platform_api = MockPlatformApi::new();
680
681        // Setup config expecting "AWSELBAuthSessionCookie"
682        let config = ServerCommunicationConfig {
683            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
684                idp_login_url: Some("https://example.com".to_string()),
685                cookie_name: Some("AWSELBAuthSessionCookie".to_string()),
686                cookie_domain: Some("example.com".to_string()),
687                cookie_value: None,
688            }),
689        };
690        repo.save("vault.example.com".to_string(), config)
691            .await
692            .unwrap();
693
694        // Platform API returns multiple shards with -N suffixes
695        platform_api
696            .set_cookies(Some(vec![
697                AcquiredCookie {
698                    name: "AWSELBAuthSessionCookie-0".to_string(),
699                    value: "shard0value".to_string(),
700                },
701                AcquiredCookie {
702                    name: "AWSELBAuthSessionCookie-1".to_string(),
703                    value: "shard1value".to_string(),
704                },
705                AcquiredCookie {
706                    name: "AWSELBAuthSessionCookie-2".to_string(),
707                    value: "shard2value".to_string(),
708                },
709            ]))
710            .await;
711
712        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
713
714        // Should succeed
715        client.acquire_cookie("vault.example.com").await.unwrap();
716
717        // Verify all shards were saved
718        let saved_config = repo
719            .get("vault.example.com".to_string())
720            .await
721            .unwrap()
722            .unwrap();
723
724        if let BootstrapConfig::SsoCookieVendor(vendor_config) = saved_config.bootstrap {
725            assert_eq!(vendor_config.cookie_value.as_ref().unwrap().len(), 3);
726            assert_eq!(
727                vendor_config.cookie_value.as_ref().unwrap()[0].name,
728                "AWSELBAuthSessionCookie-0"
729            );
730            assert_eq!(
731                vendor_config.cookie_value.as_ref().unwrap()[1].name,
732                "AWSELBAuthSessionCookie-1"
733            );
734            assert_eq!(
735                vendor_config.cookie_value.as_ref().unwrap()[2].name,
736                "AWSELBAuthSessionCookie-2"
737            );
738        } else {
739            panic!("Expected SsoCookieVendor config");
740        }
741    }
742
743    #[tokio::test]
744    async fn acquire_cookie_accepts_unsharded_cookie_without_suffix() {
745        let repo = MockRepository::default();
746        let platform_api = MockPlatformApi::new();
747
748        // Setup config
749        let config = ServerCommunicationConfig {
750            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
751                idp_login_url: Some("https://example.com".to_string()),
752                cookie_name: Some("SessionCookie".to_string()),
753                cookie_domain: Some("example.com".to_string()),
754                cookie_value: None,
755            }),
756        };
757        repo.save("vault.example.com".to_string(), config)
758            .await
759            .unwrap();
760
761        // Platform API returns single cookie without suffix
762        platform_api
763            .set_cookies(Some(vec![AcquiredCookie {
764                name: "SessionCookie".to_string(),
765                value: "single-cookie-value".to_string(),
766            }]))
767            .await;
768
769        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
770
771        // Should succeed
772        client.acquire_cookie("vault.example.com").await.unwrap();
773
774        // Verify value was saved
775        let saved_config = repo
776            .get("vault.example.com".to_string())
777            .await
778            .unwrap()
779            .unwrap();
780
781        if let BootstrapConfig::SsoCookieVendor(vendor_config) = saved_config.bootstrap {
782            assert_eq!(vendor_config.cookie_value.as_ref().unwrap().len(), 1);
783            assert_eq!(
784                vendor_config.cookie_value.as_ref().unwrap()[0].name,
785                "SessionCookie"
786            );
787            assert_eq!(
788                vendor_config.cookie_value.as_ref().unwrap()[0].value,
789                "single-cookie-value"
790            );
791        } else {
792            panic!("Expected SsoCookieVendor config");
793        }
794    }
795
796    #[tokio::test]
797    async fn set_communication_type_saves_direct_config() {
798        let repo = MockRepository::default();
799        let platform_api = MockPlatformApi::new();
800        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
801
802        let config = ServerCommunicationConfig {
803            bootstrap: BootstrapConfig::Direct,
804        };
805
806        // Call set_communication_type
807        client
808            .set_communication_type("vault.example.com".to_string(), config.clone())
809            .await
810            .unwrap();
811
812        // Verify config was saved
813        let saved_config = repo
814            .get("vault.example.com".to_string())
815            .await
816            .unwrap()
817            .unwrap();
818
819        assert!(matches!(saved_config.bootstrap, BootstrapConfig::Direct));
820    }
821
822    #[tokio::test]
823    async fn set_communication_type_saves_sso_cookie_vendor_config() {
824        let repo = MockRepository::default();
825        let platform_api = MockPlatformApi::new();
826        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
827
828        let config = ServerCommunicationConfig {
829            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
830                idp_login_url: Some("https://idp.example.com/login".to_string()),
831                cookie_name: Some("SessionCookie".to_string()),
832                cookie_domain: Some("vault.example.com".to_string()),
833                cookie_value: None,
834            }),
835        };
836
837        // Call set_communication_type
838        client
839            .set_communication_type("vault.example.com".to_string(), config.clone())
840            .await
841            .unwrap();
842
843        // Verify config was saved
844        let saved_config = repo
845            .get("vault.example.com".to_string())
846            .await
847            .unwrap()
848            .unwrap();
849
850        if let BootstrapConfig::SsoCookieVendor(vendor_config) = saved_config.bootstrap {
851            assert_eq!(
852                vendor_config.idp_login_url,
853                Some("https://idp.example.com/login".to_string())
854            );
855            assert_eq!(vendor_config.cookie_name, Some("SessionCookie".to_string()));
856            assert_eq!(
857                vendor_config.cookie_domain,
858                Some("vault.example.com".to_string())
859            );
860            assert!(vendor_config.cookie_value.is_none());
861        } else {
862            panic!("Expected SsoCookieVendor config");
863        }
864    }
865
866    #[tokio::test]
867    async fn set_communication_type_overwrites_existing_config() {
868        let repo = MockRepository::default();
869        let platform_api = MockPlatformApi::new();
870
871        // Setup existing Direct config
872        let old_config = ServerCommunicationConfig {
873            bootstrap: BootstrapConfig::Direct,
874        };
875        repo.save("vault.example.com".to_string(), old_config)
876            .await
877            .unwrap();
878
879        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
880
881        // Overwrite with SsoCookieVendor config
882        let new_config = ServerCommunicationConfig {
883            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
884                idp_login_url: Some("https://new-idp.example.com/login".to_string()),
885                cookie_name: Some("NewCookie".to_string()),
886                cookie_domain: Some("vault.example.com".to_string()),
887                cookie_value: None,
888            }),
889        };
890
891        client
892            .set_communication_type("vault.example.com".to_string(), new_config)
893            .await
894            .unwrap();
895
896        // Verify new config replaced old config
897        let saved_config = repo
898            .get("vault.example.com".to_string())
899            .await
900            .unwrap()
901            .unwrap();
902
903        if let BootstrapConfig::SsoCookieVendor(vendor_config) = saved_config.bootstrap {
904            assert_eq!(
905                vendor_config.idp_login_url,
906                Some("https://new-idp.example.com/login".to_string())
907            );
908            assert_eq!(vendor_config.cookie_name, Some("NewCookie".to_string()));
909        } else {
910            panic!("Expected SsoCookieVendor config");
911        }
912    }
913
914    #[tokio::test]
915    async fn set_communication_type_preserves_per_hostname_isolation() {
916        let repo = MockRepository::default();
917        let platform_api = MockPlatformApi::new();
918        let client = ServerCommunicationConfigClient::new(repo.clone(), platform_api);
919
920        // Save config for first hostname
921        let config1 = ServerCommunicationConfig {
922            bootstrap: BootstrapConfig::Direct,
923        };
924        client
925            .set_communication_type("vault1.example.com".to_string(), config1)
926            .await
927            .unwrap();
928
929        // Save different config for second hostname
930        let config2 = ServerCommunicationConfig {
931            bootstrap: BootstrapConfig::SsoCookieVendor(SsoCookieVendorConfig {
932                idp_login_url: Some("https://idp.example.com/login".to_string()),
933                cookie_name: Some("TestCookie".to_string()),
934                cookie_domain: Some("vault2.example.com".to_string()),
935                cookie_value: None,
936            }),
937        };
938        client
939            .set_communication_type("vault2.example.com".to_string(), config2)
940            .await
941            .unwrap();
942
943        // Verify both configs are stored independently
944        let saved_config1 = repo
945            .get("vault1.example.com".to_string())
946            .await
947            .unwrap()
948            .unwrap();
949        assert!(matches!(saved_config1.bootstrap, BootstrapConfig::Direct));
950
951        let saved_config2 = repo
952            .get("vault2.example.com".to_string())
953            .await
954            .unwrap()
955            .unwrap();
956        assert!(matches!(
957            saved_config2.bootstrap,
958            BootstrapConfig::SsoCookieVendor(_)
959        ));
960    }
961}
bitwarden_server_communication_config/client.rs

bitwarden_server_communication_config/
client.rs
