mcp_utils/client/
oauth_handler.rs

1use crate::client::config::loopback_redirect_uri;
2use crate::client::manager::{ElicitationRequest, McpClientEvent, OAuthHandlerContext};
3use aether_auth::{OAuthError, OAuthHandler, accept_oauth_callback};
4use futures::future::BoxFuture;
5use rmcp::model::{ElicitRequestParams, ElicitationAction};
6use std::num::NonZeroU16;
7use tokio::net::TcpListener;
8use tokio::sync::{mpsc, oneshot};
9
10const AETHER_OAUTH_ELICITATION_ID: &str = "aether-oauth";
11
12/// `OAuthHandler` that dispatches the OAuth authorization URL to the host.
13pub struct ElicitingOAuthHandler {
14    listener: TcpListener,
15    redirect_uri: String,
16    server_name: String,
17    event_sender: mpsc::Sender<McpClientEvent>,
18}
19
20impl ElicitingOAuthHandler {
21    pub fn new(ctx: OAuthHandlerContext) -> Result<Self, std::io::Error> {
22        let port = ctx.callback_port.map_or(0, NonZeroU16::get);
23        let std_listener = std::net::TcpListener::bind(("127.0.0.1", port))?;
24        let port = std_listener.local_addr()?.port();
25        std_listener.set_nonblocking(true)?;
26        Ok(Self {
27            listener: TcpListener::from_std(std_listener)?,
28            redirect_uri: loopback_redirect_uri(port),
29            server_name: ctx.server_name,
30            event_sender: ctx.tx,
31        })
32    }
33}
34
35impl OAuthHandler for ElicitingOAuthHandler {
36    fn redirect_uri(&self) -> &str {
37        &self.redirect_uri
38    }
39
40    fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result<String, OAuthError>> {
41        let auth_url = auth_url.to_string();
42        Box::pin(async move {
43            let (response_sender, response_rx) = oneshot::channel();
44            self.event_sender
45                .send(McpClientEvent::Elicitation(Box::new(ElicitationRequest {
46                    server_name: self.server_name.clone(),
47                    request: ElicitRequestParams::UrlElicitationParams {
48                        meta: None,
49                        message: "Open this URL to authorize MCP server access.".to_string(),
50                        url: auth_url,
51                        elicitation_id: AETHER_OAUTH_ELICITATION_ID.to_string(),
52                    },
53                    response_sender,
54                })))
55                .await
56                .map_err(|_| OAuthError::Rmcp("OAuth prompt channel closed".to_string()))?;
57
58            let result = tokio::select! {
59                callback = accept_oauth_callback(&self.listener) => callback,
60                response = response_rx => match response {
61                    Ok(result) if matches!(result.action, ElicitationAction::Decline | ElicitationAction::Cancel) => {
62                        Err(OAuthError::UserCancelled)
63                    }
64                    Ok(_) | Err(_) => accept_oauth_callback(&self.listener).await,
65                },
66            };
67            if !matches!(result, Err(OAuthError::UserCancelled)) {
68                let _ = self
69                    .event_sender
70                    .send(McpClientEvent::ElicitationComplete {
71                        server_name: self.server_name.clone(),
72                        elicitation_id: AETHER_OAUTH_ELICITATION_ID.to_string(),
73                    })
74                    .await;
75            }
76            result
77        })
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use rmcp::model::ElicitResult;
85    use std::sync::Arc;
86    use tokio::{io::AsyncWriteExt, task::yield_now};
87
88    #[tokio::test]
89    async fn accepting_browser_prompt_keeps_waiting_for_callback() {
90        let (tx, mut rx) = mpsc::channel(1);
91        let handler = Arc::new(
92            ElicitingOAuthHandler::new(OAuthHandlerContext {
93                server_name: "slack".to_string(),
94                callback_port: None,
95                tx,
96            })
97            .unwrap(),
98        );
99        let port = handler
100            .redirect_uri()
101            .strip_prefix("http://localhost:")
102            .and_then(|value| value.strip_suffix('/'))
103            .unwrap()
104            .parse::<u16>()
105            .unwrap();
106        let authorize = {
107            let handler = Arc::clone(&handler);
108            tokio::spawn(async move { handler.authorize("https://example.com/oauth").await })
109        };
110        let McpClientEvent::Elicitation(request) = rx.recv().await.unwrap() else {
111            panic!("expected OAuth elicitation");
112        };
113        request.response_sender.send(ElicitResult::new(ElicitationAction::Accept)).unwrap();
114        yield_now().await;
115        assert!(!authorize.is_finished());
116
117        let mut callback = tokio::net::TcpStream::connect(("127.0.0.1", port)).await.unwrap();
118        callback.write_all(b"GET /?code=test-code&state=test-state HTTP/1.1\r\nHost: localhost\r\n\r\n").await.unwrap();
119
120        assert!(authorize.await.unwrap().unwrap().contains("code=test-code&state=test-state"));
121    }
122
123    #[tokio::test]
124    async fn configured_callback_port_uses_registered_redirect() {
125        let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
126        let port = probe.local_addr().unwrap().port();
127        drop(probe);
128        let (tx, _) = mpsc::channel(1);
129        let handler = ElicitingOAuthHandler::new(OAuthHandlerContext {
130            server_name: "slack".to_string(),
131            callback_port: NonZeroU16::new(port),
132            tx,
133        })
134        .unwrap();
135        assert_eq!(handler.redirect_uri(), format!("http://localhost:{port}/"));
136    }
137
138    #[tokio::test]
139    async fn configured_callback_port_fails_when_in_use() {
140        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
141        let port = listener.local_addr().unwrap().port();
142        let (tx, _) = mpsc::channel(1);
143        let error = ElicitingOAuthHandler::new(OAuthHandlerContext {
144            server_name: "slack".to_string(),
145            callback_port: NonZeroU16::new(port),
146            tx,
147        })
148        .err()
149        .unwrap();
150        assert_eq!(error.kind(), std::io::ErrorKind::AddrInUse);
151    }
152}
mcp_utils/client/oauth_handler.rs

mcp_utils/client/
oauth_handler.rs
