mcp_utils/client/
oauth_handler.rs

1use crate::client::manager::{ElicitationRequest, McpClientEvent, OAuthHandlerContext, UrlElicitationCompleteParams};
2use aether_auth::{OAuthCallback, OAuthError, OAuthHandler, accept_oauth_callback};
3use futures::future::BoxFuture;
4use rmcp::model::{CreateElicitationRequestParams, ElicitationAction};
5use std::num::NonZeroU16;
6use tokio::net::TcpListener;
7use tokio::sync::{mpsc, oneshot};
8
9pub const AETHER_OAUTH_ELICITATION_ID: &str = "aether-oauth";
10
11/// `OAuthHandler` that dispatches the OAuth authorization URL to the host
12pub struct ElicitingOAuthHandler {
13    listener: TcpListener,
14    redirect_uri: String,
15    server_name: String,
16    event_sender: mpsc::Sender<McpClientEvent>,
17}
18
19impl ElicitingOAuthHandler {
20    pub fn new(ctx: OAuthHandlerContext) -> Result<Self, std::io::Error> {
21        let (port, listener) = {
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            (port, TcpListener::from_std(std_listener)?)
27        };
28
29        Ok(Self {
30            listener,
31            redirect_uri: format!("http://localhost:{port}/"),
32            server_name: ctx.server_name,
33            event_sender: ctx.tx,
34        })
35    }
36}
37
38impl OAuthHandler for ElicitingOAuthHandler {
39    fn redirect_uri(&self) -> &str {
40        &self.redirect_uri
41    }
42
43    fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result<OAuthCallback, OAuthError>> {
44        let auth_url = auth_url.to_string();
45        Box::pin(async move {
46            let (response_sender, response_rx) = oneshot::channel();
47            let request = ElicitationRequest {
48                server_name: self.server_name.clone(),
49                request: CreateElicitationRequestParams::UrlElicitationParams {
50                    meta: None,
51                    message: "Open this URL to authorize MCP server access.".to_string(),
52                    url: auth_url,
53                    elicitation_id: AETHER_OAUTH_ELICITATION_ID.to_string(),
54                },
55                response_sender,
56            };
57
58            self.event_sender
59                .send(McpClientEvent::Elicitation(request))
60                .await
61                .map_err(|_| OAuthError::Rmcp("OAuth prompt channel closed".to_string()))?;
62
63            let callback = tokio::select! {
64                callback = accept_oauth_callback(&self.listener) => callback,
65                response = response_rx => match response {
66                    Ok(result) if matches!(result.action, ElicitationAction::Decline | ElicitationAction::Cancel) => {
67                        Err(OAuthError::UserCancelled)
68                    }
69                    Ok(_) | Err(_) => accept_oauth_callback(&self.listener).await,
70                },
71            }?;
72
73            let complete = UrlElicitationCompleteParams {
74                server_name: self.server_name.clone(),
75                elicitation_id: AETHER_OAUTH_ELICITATION_ID.to_string(),
76            };
77
78            if self.event_sender.send(McpClientEvent::UrlElicitationComplete(complete)).await.is_err() {
79                tracing::warn!("Failed to send OAuth URL elicitation completion: receiver dropped");
80            }
81
82            Ok(callback)
83        })
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[tokio::test]
92    async fn configured_callback_port_uses_registered_localhost_redirect() {
93        let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
94        let port = probe.local_addr().unwrap().port();
95        drop(probe);
96        let (tx, _) = mpsc::channel(1);
97
98        let handler = ElicitingOAuthHandler::new(OAuthHandlerContext {
99            server_name: "slack".to_string(),
100            callback_port: NonZeroU16::new(port),
101            tx,
102        })
103        .unwrap();
104
105        assert_eq!(handler.redirect_uri(), format!("http://localhost:{port}/"));
106    }
107
108    #[tokio::test]
109    async fn configured_callback_port_fails_when_already_in_use() {
110        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
111        let port = listener.local_addr().unwrap().port();
112        let (tx, _) = mpsc::channel(1);
113
114        let error = ElicitingOAuthHandler::new(OAuthHandlerContext {
115            server_name: "slack".to_string(),
116            callback_port: NonZeroU16::new(port),
117            tx,
118        })
119        .err()
120        .expect("occupied callback port should fail");
121
122        assert_eq!(error.kind(), std::io::ErrorKind::AddrInUse);
123    }
124}
mcp_utils/client/oauth_handler.rs

mcp_utils/client/
oauth_handler.rs
