bitwarden_server_communication_config/
client.rs

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

bitwarden_server_communication_config/
client.rs
