atrium_api/agent/
atp_agent.rs

1//! Implementation of [`AtpAgent`] and definitions of [`AtpSessionStore`] for it.
2
3mod inner;
4pub mod store;
5
6use self::store::AtpSessionStore;
7use super::{
8    Agent, CloneWithProxy, Configure, SessionManager, inner::Wrapper,
9    utils::SessionWithEndpointStore,
10};
11use crate::{
12    client::com::atproto::Service,
13    did_doc::DidDocument,
14    types::{TryFromUnknown, string::Did},
15};
16use atrium_xrpc::{Error, HttpClient, OutputDataOrBytes, XrpcClient, XrpcRequest};
17use http::{Request, Response};
18use serde::{Serialize, de::DeserializeOwned};
19use std::{convert, fmt::Debug, ops::Deref, sync::Arc};
20
21/// Type alias for the [com::atproto::server::create_session::Output](crate::com::atproto::server::create_session::Output)
22pub type AtpSession = crate::com::atproto::server::create_session::Output;
23
24pub struct CredentialSession<S, T>
25where
26    S: AtpSessionStore + Send + Sync,
27    T: XrpcClient + Send + Sync,
28    S::Error: std::error::Error + Send + Sync + 'static,
29{
30    store: Arc<SessionWithEndpointStore<S, AtpSession>>,
31    inner: Arc<inner::Client<S, T>>,
32    atproto_service: Service<inner::Client<S, T>>,
33}
34
35impl<S, T> CredentialSession<S, T>
36where
37    S: AtpSessionStore + Send + Sync,
38    T: XrpcClient + Send + Sync,
39    S::Error: std::error::Error + Send + Sync + 'static,
40{
41    pub fn new(xrpc: T, store: S) -> Self {
42        let store = Arc::new(SessionWithEndpointStore::new(store, xrpc.base_uri()));
43        let inner = Arc::new(inner::Client::new(Arc::clone(&store), xrpc));
44        let atproto_service = Service::new(Arc::clone(&inner));
45        Self { store, inner, atproto_service }
46    }
47    /// Start a new session with this agent.
48    pub async fn login(
49        &self,
50        identifier: impl AsRef<str>,
51        password: impl AsRef<str>,
52    ) -> Result<AtpSession, Error<crate::com::atproto::server::create_session::Error>> {
53        let result = self
54            .atproto_service
55            .server
56            .create_session(
57                crate::com::atproto::server::create_session::InputData {
58                    allow_takendown: None,
59                    auth_factor_token: None,
60                    identifier: identifier.as_ref().into(),
61                    password: password.as_ref().into(),
62                }
63                .into(),
64            )
65            .await?;
66        self.store.set(result.clone()).await.ok();
67        if let Some(did_doc) = result
68            .did_doc
69            .as_ref()
70            .and_then(|value| DidDocument::try_from_unknown(value.clone()).ok())
71        {
72            self.store.update_endpoint(&did_doc);
73        }
74        Ok(result)
75    }
76    /// Resume a pre-existing session with this agent.
77    pub async fn resume_session(
78        &self,
79        session: AtpSession,
80    ) -> Result<(), Error<crate::com::atproto::server::get_session::Error>> {
81        self.store.set(session.clone()).await.ok();
82        let result = self.atproto_service.server.get_session().await;
83        match result {
84            Ok(output) => {
85                assert_eq!(output.data.did, session.data.did);
86                if let Ok(Some(mut session)) = self.store.get().await {
87                    session.did_doc = output.data.did_doc.clone();
88                    session.email = output.data.email;
89                    session.email_confirmed = output.data.email_confirmed;
90                    session.handle = output.data.handle;
91                    self.store.set(session).await.ok();
92                }
93                if let Some(did_doc) = output
94                    .data
95                    .did_doc
96                    .as_ref()
97                    .and_then(|value| DidDocument::try_from_unknown(value.clone()).ok())
98                {
99                    self.store.update_endpoint(&did_doc);
100                }
101                Ok(())
102            }
103            Err(err) => {
104                self.store.clear().await.ok();
105                Err(err)
106            }
107        }
108    }
109    /// Get the current session.
110    pub async fn get_session(&self) -> Option<AtpSession> {
111        self.store.get().await.ok().and_then(convert::identity)
112    }
113    /// Get the current endpoint.
114    pub async fn get_endpoint(&self) -> String {
115        self.store.get_endpoint()
116    }
117    /// Get the current labelers header.
118    pub async fn get_labelers_header(&self) -> Option<Vec<String>> {
119        self.inner.get_labelers_header().await
120    }
121    /// Get the current proxy header.
122    pub async fn get_proxy_header(&self) -> Option<String> {
123        self.inner.get_proxy_header().await
124    }
125}
126
127impl<S, T> HttpClient for CredentialSession<S, T>
128where
129    S: AtpSessionStore + Send + Sync,
130    T: XrpcClient + Send + Sync,
131    S::Error: std::error::Error + Send + Sync + 'static,
132{
133    async fn send_http(
134        &self,
135        request: Request<Vec<u8>>,
136    ) -> Result<Response<Vec<u8>>, Box<dyn std::error::Error + Send + Sync + 'static>> {
137        self.inner.send_http(request).await
138    }
139}
140
141impl<S, T> XrpcClient for CredentialSession<S, T>
142where
143    S: AtpSessionStore + Send + Sync,
144    T: XrpcClient + Send + Sync,
145    S::Error: std::error::Error + Send + Sync + 'static,
146{
147    fn base_uri(&self) -> String {
148        self.inner.base_uri()
149    }
150    async fn send_xrpc<P, I, O, E>(
151        &self,
152        request: &XrpcRequest<P, I>,
153    ) -> Result<OutputDataOrBytes<O>, Error<E>>
154    where
155        P: Serialize + Send + Sync,
156        I: Serialize + Send + Sync,
157        O: DeserializeOwned + Send + Sync,
158        E: DeserializeOwned + Send + Sync + Debug,
159    {
160        self.inner.send_xrpc(request).await
161    }
162}
163
164impl<S, T> SessionManager for CredentialSession<S, T>
165where
166    S: AtpSessionStore + Send + Sync,
167    T: XrpcClient + Send + Sync,
168    S::Error: std::error::Error + Send + Sync + 'static,
169{
170    async fn did(&self) -> Option<Did> {
171        self.store.get().await.ok().and_then(|session| session.map(|session| session.data.did))
172    }
173}
174
175impl<S, T> Configure for CredentialSession<S, T>
176where
177    S: AtpSessionStore + Send + Sync,
178    T: XrpcClient + Send + Sync,
179    S::Error: std::error::Error + Send + Sync + 'static,
180{
181    fn configure_endpoint(&self, endpoint: String) {
182        self.inner.configure_endpoint(endpoint);
183    }
184    fn configure_labelers_header(&self, labeler_dids: Option<Vec<(Did, bool)>>) {
185        self.inner.configure_labelers_header(labeler_dids);
186    }
187    fn configure_proxy_header(&self, did: Did, service_type: impl AsRef<str>) {
188        self.inner.configure_proxy_header(did, service_type);
189    }
190}
191
192impl<S, T> CloneWithProxy for CredentialSession<S, T>
193where
194    S: AtpSessionStore + Send + Sync,
195    S::Error: std::error::Error + Send + Sync + 'static,
196    T: XrpcClient + Send + Sync,
197{
198    fn clone_with_proxy(&self, did: Did, service_type: impl AsRef<str>) -> Self {
199        let inner = Arc::new(self.inner.clone_with_proxy(did, service_type));
200        let atproto_service = Service::new(Arc::clone(&inner));
201        Self { store: Arc::clone(&self.store), inner, atproto_service }
202    }
203}
204
205/// An ATP "Agent".
206/// Manages session token lifecycles and provides convenience methods.
207///
208/// This will be deprecated in the near future. Use [`Agent`] directly
209/// with a [`CredentialSession`] instead:
210///
211/// # Example
212///
213/// ```
214/// use atrium_api::agent::atp_agent::{store::MemorySessionStore, CredentialSession};
215/// use atrium_api::agent::Agent;
216/// use atrium_xrpc_client::reqwest::ReqwestClient;
217///
218/// let session = CredentialSession::new(
219///     ReqwestClient::new("https://bsky.social"),
220///     MemorySessionStore::default(),
221/// );
222/// let agent = Agent::new(session);
223/// ```
224pub struct AtpAgent<S, T>
225where
226    S: AtpSessionStore + Send + Sync,
227    T: XrpcClient + Send + Sync,
228    S::Error: std::error::Error + Send + Sync + 'static,
229{
230    session_manager: Wrapper<CredentialSession<S, T>>,
231    inner: Agent<Wrapper<CredentialSession<S, T>>>,
232}
233
234impl<S, T> AtpAgent<S, T>
235where
236    S: AtpSessionStore + Send + Sync,
237    T: XrpcClient + Send + Sync,
238    S::Error: std::error::Error + Send + Sync + 'static,
239{
240    /// Create a new agent.
241    pub fn new(xrpc: T, store: S) -> Self {
242        let session_manager = Wrapper::new(CredentialSession::new(xrpc, store));
243        let inner = Agent::new(session_manager.clone());
244        Self { session_manager, inner }
245    }
246    /// Start a new session with this agent.
247    pub async fn login(
248        &self,
249        identifier: impl AsRef<str>,
250        password: impl AsRef<str>,
251    ) -> Result<AtpSession, Error<crate::com::atproto::server::create_session::Error>> {
252        self.session_manager.login(identifier, password).await
253    }
254    // /// Resume a pre-existing session with this agent.
255    pub async fn resume_session(
256        &self,
257        session: AtpSession,
258    ) -> Result<(), Error<crate::com::atproto::server::get_session::Error>> {
259        self.session_manager.resume_session(session).await
260    }
261    /// Get the current session.
262    pub async fn get_session(&self) -> Option<AtpSession> {
263        self.session_manager.get_session().await
264    }
265    /// Get the current endpoint.
266    pub async fn get_endpoint(&self) -> String {
267        self.session_manager.get_endpoint().await
268    }
269    /// Get the current labelers header.
270    pub async fn get_labelers_header(&self) -> Option<Vec<String>> {
271        self.session_manager.get_labelers_header().await
272    }
273    /// Get the current proxy header.
274    pub async fn get_proxy_header(&self) -> Option<String> {
275        self.session_manager.get_proxy_header().await
276    }
277}
278
279impl<S, T> Deref for AtpAgent<S, T>
280where
281    S: AtpSessionStore + Send + Sync,
282    T: XrpcClient + Send + Sync,
283    S::Error: std::error::Error + Send + Sync + 'static,
284{
285    type Target = Agent<Wrapper<CredentialSession<S, T>>>;
286
287    fn deref(&self) -> &Self::Target {
288        &self.inner
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::store::MemorySessionStore;
295    use super::*;
296    use crate::{
297        agent::AtprotoServiceType,
298        com::atproto::server::create_session::OutputData,
299        did_doc::{DidDocument, Service, VerificationMethod},
300        types::TryIntoUnknown,
301    };
302    use atrium_xrpc::HttpClient;
303    use http::{HeaderMap, HeaderName, HeaderValue, Request, Response};
304    use std::collections::HashMap;
305    use tokio::sync::RwLock;
306    #[cfg(target_arch = "wasm32")]
307    use wasm_bindgen_test::wasm_bindgen_test;
308
309    #[derive(Default)]
310    struct MockResponses {
311        create_session: Option<crate::com::atproto::server::create_session::OutputData>,
312        get_session: Option<crate::com::atproto::server::get_session::OutputData>,
313    }
314
315    #[derive(Default)]
316    struct MockClient {
317        responses: MockResponses,
318        counts: Arc<RwLock<HashMap<String, usize>>>,
319        headers: Arc<RwLock<Vec<HeaderMap<HeaderValue>>>>,
320    }
321
322    impl HttpClient for MockClient {
323        async fn send_http(
324            &self,
325            request: Request<Vec<u8>>,
326        ) -> Result<Response<Vec<u8>>, Box<dyn std::error::Error + Send + Sync + 'static>> {
327            // tick tokio time
328            #[cfg(not(target_arch = "wasm32"))]
329            tokio::time::sleep(std::time::Duration::from_micros(10)).await;
330
331            self.headers.write().await.push(request.headers().clone());
332            let builder =
333                Response::builder().header(http::header::CONTENT_TYPE, "application/json");
334            let token = request
335                .headers()
336                .get(http::header::AUTHORIZATION)
337                .and_then(|value| value.to_str().ok())
338                .and_then(|value| value.split(' ').last());
339            if token == Some("expired") {
340                return Ok(builder.status(http::StatusCode::BAD_REQUEST).body(
341                    serde_json::to_vec(&atrium_xrpc::error::ErrorResponseBody {
342                        error: Some(String::from("ExpiredToken")),
343                        message: Some(String::from("Token has expired")),
344                    })?,
345                )?);
346            }
347            let mut body = Vec::new();
348            if let Some(nsid) = request.uri().path().strip_prefix("/xrpc/") {
349                *self.counts.write().await.entry(nsid.into()).or_default() += 1;
350                match nsid {
351                    crate::com::atproto::server::create_session::NSID => {
352                        if let Some(output) = &self.responses.create_session {
353                            body.extend(serde_json::to_vec(output)?);
354                        }
355                    }
356                    crate::com::atproto::server::get_session::NSID => {
357                        if token == Some("access") {
358                            if let Some(output) = &self.responses.get_session {
359                                body.extend(serde_json::to_vec(output)?);
360                            }
361                        }
362                    }
363                    crate::com::atproto::server::refresh_session::NSID => {
364                        if token == Some("refresh") {
365                            body.extend(serde_json::to_vec(
366                                &crate::com::atproto::server::refresh_session::OutputData {
367                                    access_jwt: String::from("access"),
368                                    active: None,
369                                    did: "did:web:example.com".parse().expect("valid"),
370                                    did_doc: None,
371                                    email: None,
372                                    email_auth_factor: None,
373                                    email_confirmed: None,
374                                    handle: "example.com".parse().expect("valid"),
375                                    refresh_jwt: String::from("refresh"),
376                                    status: None,
377                                },
378                            )?);
379                        }
380                    }
381                    crate::com::atproto::server::describe_server::NSID => {
382                        body.extend(serde_json::to_vec(
383                            &crate::com::atproto::server::describe_server::OutputData {
384                                available_user_domains: Vec::new(),
385                                blob_upload_limit: None,
386                                contact: None,
387                                did: "did:web:example.com".parse().expect("valid"),
388                                invite_code_required: None,
389                                links: None,
390                                phone_verification_required: None,
391                            },
392                        )?);
393                    }
394                    _ => {}
395                }
396            }
397            if body.is_empty() {
398                Ok(builder.status(http::StatusCode::UNAUTHORIZED).body(serde_json::to_vec(
399                    &atrium_xrpc::error::ErrorResponseBody {
400                        error: Some(String::from("AuthenticationRequired")),
401                        message: Some(String::from("Invalid identifier or password")),
402                    },
403                )?)?)
404            } else {
405                Ok(builder.status(http::StatusCode::OK).body(body)?)
406            }
407        }
408    }
409
410    impl XrpcClient for MockClient {
411        fn base_uri(&self) -> String {
412            "http://localhost:8080".into()
413        }
414    }
415
416    fn session_data() -> OutputData {
417        OutputData {
418            access_jwt: String::from("access"),
419            active: None,
420            did: "did:web:example.com".parse().expect("valid"),
421            did_doc: None,
422            email: None,
423            email_auth_factor: None,
424            email_confirmed: None,
425            handle: "example.com".parse().expect("valid"),
426            refresh_jwt: String::from("refresh"),
427            status: None,
428        }
429    }
430
431    #[tokio::test]
432    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
433    async fn test_new() {
434        let agent = AtpAgent::new(MockClient::default(), MemorySessionStore::default());
435        assert_eq!(agent.get_session().await, None);
436    }
437
438    #[tokio::test]
439    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
440    async fn test_login() {
441        let session_data = session_data();
442        // success
443        {
444            let client = MockClient {
445                responses: MockResponses {
446                    create_session: Some(crate::com::atproto::server::create_session::OutputData {
447                        ..session_data.clone()
448                    }),
449                    ..Default::default()
450                },
451                ..Default::default()
452            };
453            let agent = AtpAgent::new(client, MemorySessionStore::default());
454            agent.login("test", "pass").await.expect("login should be succeeded");
455            assert_eq!(agent.get_session().await, Some(session_data.into()));
456        }
457        // failure with `createSession` error
458        {
459            let client = MockClient {
460                responses: MockResponses { ..Default::default() },
461                ..Default::default()
462            };
463            let agent = AtpAgent::new(client, MemorySessionStore::default());
464            agent.login("test", "bad").await.expect_err("login should be failed");
465            assert_eq!(agent.get_session().await, None);
466        }
467    }
468
469    #[tokio::test]
470    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
471    async fn test_xrpc_get_session() {
472        let session_data = session_data();
473        let client = MockClient {
474            responses: MockResponses {
475                get_session: Some(crate::com::atproto::server::get_session::OutputData {
476                    active: session_data.active,
477                    did: session_data.did.clone(),
478                    did_doc: session_data.did_doc.clone(),
479                    email: session_data.email.clone(),
480                    email_auth_factor: session_data.email_auth_factor,
481                    email_confirmed: session_data.email_confirmed,
482                    handle: session_data.handle.clone(),
483                    status: session_data.status.clone(),
484                }),
485                ..Default::default()
486            },
487            ..Default::default()
488        };
489        let agent = AtpAgent::new(client, MemorySessionStore::default());
490        agent
491            .session_manager
492            .store
493            .set(session_data.clone().into())
494            .await
495            .expect("set session should be succeeded");
496        let output = agent
497            .api
498            .com
499            .atproto
500            .server
501            .get_session()
502            .await
503            .expect("get session should be succeeded");
504        assert_eq!(output.did.as_str(), "did:web:example.com");
505    }
506
507    #[tokio::test]
508    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
509    async fn test_xrpc_get_session_with_refresh() {
510        let mut session_data = session_data();
511        session_data.access_jwt = String::from("expired");
512        let client = MockClient {
513            responses: MockResponses {
514                get_session: Some(crate::com::atproto::server::get_session::OutputData {
515                    active: session_data.active,
516                    did: session_data.did.clone(),
517                    did_doc: session_data.did_doc.clone(),
518                    email: session_data.email.clone(),
519                    email_auth_factor: session_data.email_auth_factor,
520                    email_confirmed: session_data.email_confirmed,
521                    handle: session_data.handle.clone(),
522                    status: session_data.status.clone(),
523                }),
524                ..Default::default()
525            },
526            ..Default::default()
527        };
528        let agent = AtpAgent::new(client, MemorySessionStore::default());
529        agent
530            .session_manager
531            .store
532            .set(session_data.clone().into())
533            .await
534            .expect("set session should be succeeded");
535        let output = agent
536            .api
537            .com
538            .atproto
539            .server
540            .get_session()
541            .await
542            .expect("get session should be succeeded");
543        assert_eq!(output.did.as_str(), "did:web:example.com");
544        assert_eq!(
545            agent
546                .session_manager
547                .store
548                .get()
549                .await
550                .expect("session should be stored")
551                .map(|session| session.data.access_jwt),
552            Some("access".into())
553        );
554    }
555
556    #[cfg(not(target_arch = "wasm32"))]
557    #[tokio::test]
558    async fn test_xrpc_get_session_with_duplicated_refresh() {
559        let mut session_data = session_data();
560        session_data.access_jwt = String::from("expired");
561        let client = MockClient {
562            responses: MockResponses {
563                get_session: Some(crate::com::atproto::server::get_session::OutputData {
564                    active: session_data.active,
565                    did: session_data.did.clone(),
566                    did_doc: session_data.did_doc.clone(),
567                    email: session_data.email.clone(),
568                    email_auth_factor: session_data.email_auth_factor,
569                    email_confirmed: session_data.email_confirmed,
570                    handle: session_data.handle.clone(),
571                    status: session_data.status.clone(),
572                }),
573                ..Default::default()
574            },
575            ..Default::default()
576        };
577        let counts = Arc::clone(&client.counts);
578        let agent = Arc::new(AtpAgent::new(client, MemorySessionStore::default()));
579        agent
580            .session_manager
581            .store
582            .set(session_data.clone().into())
583            .await
584            .expect("set session should be succeeded");
585        let handles = (0..3).map(|_| {
586            let agent = Arc::clone(&agent);
587            tokio::spawn(async move { agent.api.com.atproto.server.get_session().await })
588        });
589        let results = futures::future::join_all(handles).await;
590        for result in &results {
591            let output = result
592                .as_ref()
593                .expect("task should be successfully executed")
594                .as_ref()
595                .expect("get session should be succeeded");
596            assert_eq!(output.did.as_str(), "did:web:example.com");
597        }
598        assert_eq!(
599            agent
600                .session_manager
601                .store
602                .get()
603                .await
604                .expect("session should be stored")
605                .map(|session| session.data.access_jwt),
606            Some("access".into())
607        );
608        assert_eq!(
609            counts.read().await.clone(),
610            HashMap::from_iter([
611                ("com.atproto.server.refreshSession".into(), 1),
612                ("com.atproto.server.getSession".into(), 3)
613            ])
614        );
615    }
616
617    #[tokio::test]
618    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
619    async fn test_resume_session() {
620        let session_data = session_data();
621        // success
622        {
623            let client = MockClient {
624                responses: MockResponses {
625                    get_session: Some(crate::com::atproto::server::get_session::OutputData {
626                        active: session_data.active,
627                        did: session_data.did.clone(),
628                        did_doc: session_data.did_doc.clone(),
629                        email: session_data.email.clone(),
630                        email_auth_factor: session_data.email_auth_factor,
631                        email_confirmed: session_data.email_confirmed,
632                        handle: session_data.handle.clone(),
633                        status: session_data.status.clone(),
634                    }),
635                    ..Default::default()
636                },
637                ..Default::default()
638            };
639            let agent = AtpAgent::new(client, MemorySessionStore::default());
640            assert_eq!(agent.get_session().await, None);
641            agent
642                .resume_session(
643                    OutputData {
644                        email: Some(String::from("test@example.com")),
645                        ..session_data.clone()
646                    }
647                    .into(),
648                )
649                .await
650                .expect("resume_session should be succeeded");
651            assert_eq!(agent.get_session().await, Some(session_data.clone().into()));
652        }
653        // failure with `getSession` error
654        {
655            let client = MockClient {
656                responses: MockResponses { ..Default::default() },
657                ..Default::default()
658            };
659            let agent = AtpAgent::new(client, MemorySessionStore::default());
660            assert_eq!(agent.get_session().await, None);
661            agent
662                .resume_session(session_data.clone().into())
663                .await
664                .expect_err("resume_session should be failed");
665            assert_eq!(agent.get_session().await, None);
666        }
667    }
668
669    #[tokio::test]
670    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
671    async fn test_resume_session_with_refresh() {
672        let session_data = session_data();
673        let client = MockClient {
674            responses: MockResponses {
675                get_session: Some(crate::com::atproto::server::get_session::OutputData {
676                    active: session_data.active,
677                    did: session_data.did.clone(),
678                    did_doc: session_data.did_doc.clone(),
679                    email: session_data.email.clone(),
680                    email_auth_factor: session_data.email_auth_factor,
681                    email_confirmed: session_data.email_confirmed,
682                    handle: session_data.handle.clone(),
683                    status: session_data.status.clone(),
684                }),
685                ..Default::default()
686            },
687            ..Default::default()
688        };
689        let agent = AtpAgent::new(client, MemorySessionStore::default());
690        agent
691            .resume_session(
692                OutputData { access_jwt: "expired".into(), ..session_data.clone() }.into(),
693            )
694            .await
695            .expect("resume_session should be succeeded");
696        assert_eq!(agent.get_session().await, Some(session_data.clone().into()));
697    }
698
699    #[tokio::test]
700    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
701    async fn test_login_with_diddoc() {
702        let session_data = session_data();
703        let did_doc = DidDocument {
704            context: None,
705            id: "did:plc:ewvi7nxzyoun6zhxrhs64oiz".into(),
706            also_known_as: Some(vec!["at://atproto.com".into()]),
707            verification_method: Some(vec![VerificationMethod {
708                id: "did:plc:ewvi7nxzyoun6zhxrhs64oiz#atproto".into(),
709                r#type: "Multikey".into(),
710                controller: "did:plc:ewvi7nxzyoun6zhxrhs64oiz".into(),
711                public_key_multibase: Some(
712                    "zQ3shXjHeiBuRCKmM36cuYnm7YEMzhGnCmCyW92sRJ9pribSF".into(),
713                ),
714            }]),
715            service: Some(vec![Service {
716                id: "#atproto_pds".into(),
717                r#type: "AtprotoPersonalDataServer".into(),
718                service_endpoint: "https://bsky.social".into(),
719            }]),
720        };
721        // success
722        {
723            let client = MockClient {
724                responses: MockResponses {
725                    create_session: Some(crate::com::atproto::server::create_session::OutputData {
726                        did_doc: Some(
727                            did_doc
728                                .clone()
729                                .try_into_unknown()
730                                .expect("failed to convert to unknown"),
731                        ),
732                        ..session_data.clone()
733                    }),
734                    ..Default::default()
735                },
736                ..Default::default()
737            };
738            let agent = AtpAgent::new(client, MemorySessionStore::default());
739            agent.login("test", "pass").await.expect("login should be succeeded");
740            assert_eq!(agent.get_endpoint().await, "https://bsky.social");
741            assert_eq!(agent.api.com.atproto.server.xrpc.base_uri(), "https://bsky.social");
742        }
743        // invalid services
744        {
745            let client = MockClient {
746                responses: MockResponses {
747                    create_session: Some(crate::com::atproto::server::create_session::OutputData {
748                        did_doc: Some(
749                            DidDocument {
750                                service: Some(vec![
751                                    Service {
752                                        id: "#pds".into(), // not `#atproto_pds`
753                                        r#type: "AtprotoPersonalDataServer".into(),
754                                        service_endpoint: "https://bsky.social".into(),
755                                    },
756                                    Service {
757                                        id: "#atproto_pds".into(),
758                                        r#type: "AtprotoPersonalDataServer".into(),
759                                        service_endpoint: "htps://bsky.social".into(), // invalid url (not `https`)
760                                    },
761                                ]),
762                                ..did_doc.clone()
763                            }
764                            .try_into_unknown()
765                            .expect("failed to convert to unknown"),
766                        ),
767                        ..session_data.clone()
768                    }),
769                    ..Default::default()
770                },
771                ..Default::default()
772            };
773            let agent = AtpAgent::new(client, MemorySessionStore::default());
774            agent.login("test", "pass").await.expect("login should be succeeded");
775            // not updated
776            assert_eq!(agent.get_endpoint().await, "http://localhost:8080");
777            assert_eq!(agent.api.com.atproto.server.xrpc.base_uri(), "http://localhost:8080");
778        }
779    }
780
781    #[tokio::test]
782    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
783    async fn test_configure_labelers_header() {
784        let client = MockClient::default();
785        let headers = Arc::clone(&client.headers);
786        let agent = AtpAgent::new(client, MemorySessionStore::default());
787
788        agent
789            .api
790            .com
791            .atproto
792            .server
793            .describe_server()
794            .await
795            .expect("describe_server should be succeeded");
796        assert_eq!(headers.read().await.last(), Some(&HeaderMap::new()));
797
798        agent.configure_labelers_header(Some(vec![(
799            "did:plc:test1".parse().expect("did should be valid"),
800            false,
801        )]));
802        agent
803            .api
804            .com
805            .atproto
806            .server
807            .describe_server()
808            .await
809            .expect("describe_server should be succeeded");
810        assert_eq!(
811            headers.read().await.last(),
812            Some(&HeaderMap::from_iter([(
813                HeaderName::from_static("atproto-accept-labelers"),
814                HeaderValue::from_static("did:plc:test1"),
815            )]))
816        );
817
818        agent.configure_labelers_header(Some(vec![
819            ("did:plc:test1".parse().expect("did should be valid"), true),
820            ("did:plc:test2".parse().expect("did should be valid"), false),
821        ]));
822        agent
823            .api
824            .com
825            .atproto
826            .server
827            .describe_server()
828            .await
829            .expect("describe_server should be succeeded");
830        assert_eq!(
831            headers.read().await.last(),
832            Some(&HeaderMap::from_iter([(
833                HeaderName::from_static("atproto-accept-labelers"),
834                HeaderValue::from_static("did:plc:test1;redact, did:plc:test2"),
835            )]))
836        );
837
838        assert_eq!(
839            agent.get_labelers_header().await,
840            Some(vec![String::from("did:plc:test1;redact"), String::from("did:plc:test2")])
841        );
842    }
843
844    #[tokio::test]
845    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
846    async fn test_configure_proxy_header() {
847        let client = MockClient::default();
848        let headers = Arc::clone(&client.headers);
849        let agent = AtpAgent::new(client, MemorySessionStore::default());
850
851        agent
852            .api
853            .com
854            .atproto
855            .server
856            .describe_server()
857            .await
858            .expect("describe_server should be succeeded");
859        assert_eq!(headers.read().await.last(), Some(&HeaderMap::new()));
860
861        agent.configure_proxy_header(
862            "did:plc:test1".parse().expect("did should be valid"),
863            AtprotoServiceType::AtprotoLabeler,
864        );
865        agent
866            .api
867            .com
868            .atproto
869            .server
870            .describe_server()
871            .await
872            .expect("describe_server should be succeeded");
873        assert_eq!(
874            headers.read().await.last(),
875            Some(&HeaderMap::from_iter([(
876                HeaderName::from_static("atproto-proxy"),
877                HeaderValue::from_static("did:plc:test1#atproto_labeler"),
878            ),]))
879        );
880
881        agent.configure_proxy_header(
882            "did:plc:test1".parse().expect("did should be valid"),
883            "atproto_labeler",
884        );
885        agent
886            .api
887            .com
888            .atproto
889            .server
890            .describe_server()
891            .await
892            .expect("describe_server should be succeeded");
893        assert_eq!(
894            headers.read().await.last(),
895            Some(&HeaderMap::from_iter([(
896                HeaderName::from_static("atproto-proxy"),
897                HeaderValue::from_static("did:plc:test1#atproto_labeler"),
898            ),]))
899        );
900
901        agent
902            .api_with_proxy(
903                "did:plc:test2".parse().expect("did should be valid"),
904                "atproto_labeler",
905            )
906            .com
907            .atproto
908            .server
909            .describe_server()
910            .await
911            .expect("describe_server should be succeeded");
912        assert_eq!(
913            headers.read().await.last(),
914            Some(&HeaderMap::from_iter([(
915                HeaderName::from_static("atproto-proxy"),
916                HeaderValue::from_static("did:plc:test2#atproto_labeler"),
917            ),]))
918        );
919
920        agent
921            .api
922            .com
923            .atproto
924            .server
925            .describe_server()
926            .await
927            .expect("describe_server should be succeeded");
928        assert_eq!(
929            headers.read().await.last(),
930            Some(&HeaderMap::from_iter([(
931                HeaderName::from_static("atproto-proxy"),
932                HeaderValue::from_static("did:plc:test1#atproto_labeler"),
933            ),]))
934        );
935
936        assert_eq!(
937            agent.get_proxy_header().await,
938            Some(String::from("did:plc:test1#atproto_labeler"))
939        );
940    }
941
942    #[tokio::test]
943    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
944    async fn test_agent_did() {
945        let session_data = session_data();
946        let client = MockClient { responses: MockResponses::default(), ..Default::default() };
947        let agent = AtpAgent::new(client, MemorySessionStore::default());
948        assert_eq!(agent.did().await, None);
949        agent
950            .session_manager
951            .store
952            .set(session_data.clone().into())
953            .await
954            .expect("set session should be succeeded");
955        assert_eq!(agent.did().await, Some(session_data.did));
956    }
957}
atrium_api/agent/atp_agent.rs

atrium_api/agent/
atp_agent.rs
