Skip to repository content494 lines · 18.2 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T04:27:50.994Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
oauth_callback_server.rs
1//! Loopback OAuth 2.0 callback server and shared HTML response page.
2//!
3//! Used by Zed's OAuth-based sign-in flows (e.g. MCP servers, ChatGPT
4//! Subscription) to receive the authorization code redirect from the user's
5//! browser. The HTML response page rendered to the browser is kept alongside
6//! the server so all OAuth callback presentation lives in one place.
7
8/// Generate a styled HTML page for OAuth callback responses.
9///
10/// Returns a complete HTML document (no HTTP headers) with a centered card
11/// layout styled to match Zed's dark theme. The `title` is rendered as a
12/// heading and `message` as body text below it.
13///
14/// When `is_error` is true, a red X icon is shown instead of the green
15/// checkmark.
16pub fn oauth_callback_page(title: &str, message: &str, is_error: bool) -> String {
17 let title = html_escape(title);
18 let message = html_escape(message);
19 let (icon_bg, icon_svg) = if is_error {
20 (
21 "#f38ba8",
22 r#"<svg viewBox="0 0 24 24"><line x1="18" y1="6" x2="6" y2="18"/><line x1="6" y1="6" x2="18" y2="18"/></svg>"#,
23 )
24 } else {
25 (
26 "#a6e3a1",
27 r#"<svg viewBox="0 0 24 24"><polyline points="20 6 9 17 4 12"/></svg>"#,
28 )
29 };
30 format!(
31 r#"<!DOCTYPE html>
32<html lang="en">
33<head>
34<meta charset="utf-8">
35<meta name="viewport" content="width=device-width, initial-scale=1">
36<title>{title} — Omega</title>
37<style>
38 * {{ margin: 0; padding: 0; box-sizing: border-box; }}
39 body {{
40 font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
41 background: #1e1e2e;
42 color: #cdd6f4;
43 display: flex;
44 align-items: center;
45 justify-content: center;
46 min-height: 100vh;
47 padding: 1rem;
48 }}
49 .card {{
50 background: #313244;
51 border-radius: 12px;
52 padding: 2.5rem;
53 max-width: 420px;
54 width: 100%;
55 text-align: center;
56 box-shadow: 0 4px 24px rgba(0, 0, 0, 0.3);
57 }}
58 .icon {{
59 width: 48px;
60 height: 48px;
61 margin: 0 auto 1.5rem;
62 background: {icon_bg};
63 border-radius: 50%;
64 display: flex;
65 align-items: center;
66 justify-content: center;
67 }}
68 .icon svg {{
69 width: 24px;
70 height: 24px;
71 stroke: #1e1e2e;
72 stroke-width: 3;
73 fill: none;
74 }}
75 h1 {{
76 font-size: 1.25rem;
77 font-weight: 600;
78 margin-bottom: 0.75rem;
79 color: #cdd6f4;
80 }}
81 p {{
82 font-size: 0.925rem;
83 line-height: 1.5;
84 color: #a6adc8;
85 }}
86 .brand {{
87 margin-top: 1.5rem;
88 font-size: 0.8rem;
89 color: #585b70;
90 letter-spacing: 0.05em;
91 }}
92</style>
93</head>
94<body>
95<div class="card">
96 <div class="icon">
97 {icon_svg}
98 </div>
99 <h1>{title}</h1>
100 <p>{message}</p>
101 <div class="brand">Omega</div>
102</div>
103</body>
104</html>"#,
105 title = title,
106 message = message,
107 icon_bg = icon_bg,
108 icon_svg = icon_svg,
109 )
110}
111
112fn html_escape(input: &str) -> String {
113 let mut output = String::with_capacity(input.len());
114 for ch in input.chars() {
115 match ch {
116 '&' => output.push_str("&"),
117 '<' => output.push_str("<"),
118 '>' => output.push_str(">"),
119 '"' => output.push_str("""),
120 '\'' => output.push_str("'"),
121 _ => output.push(ch),
122 }
123 }
124 output
125}
126
127#[cfg(not(target_family = "wasm"))]
128mod server {
129 use super::oauth_callback_page;
130 use anyhow::{Context as _, Result, anyhow};
131 use std::str::FromStr;
132 use std::time::Duration;
133 use url::Url;
134
135 /// Parsed OAuth callback parameters from the authorization server redirect.
136 pub struct OAuthCallbackParams {
137 pub code: String,
138 pub state: String,
139 }
140
141 /// Configuration for the loopback OAuth callback server.
142 ///
143 /// OAuth servers compare `redirect_uri` against a per-client allow-list using
144 /// exact string matching (RFC 6749 §3.1.2), so the `host`, `preferred_port`,
145 /// and `path` here must match what's registered for the OAuth client_id.
146 #[derive(Clone, Copy)]
147 pub struct OAuthCallbackServerConfig {
148 /// Host portion of the redirect URI (typically `127.0.0.1` or `localhost`).
149 pub host: &'static str,
150 /// Preferred port. Use `0` for an OS-assigned ephemeral port.
151 pub preferred_port: u16,
152 /// Optional fallback port if `preferred_port` is unavailable. Only used
153 /// when `preferred_port` is non-zero.
154 pub fallback_port: Option<u16>,
155 /// Callback path on the redirect URI (e.g. `/callback`, `/auth/callback`).
156 pub path: &'static str,
157 }
158
159 impl Default for OAuthCallbackServerConfig {
160 fn default() -> Self {
161 Self {
162 host: "127.0.0.1",
163 preferred_port: 0,
164 fallback_port: None,
165 path: "/callback",
166 }
167 }
168 }
169
170 impl OAuthCallbackParams {
171 /// Parse the query string from a callback URL like
172 /// `http://127.0.0.1:<port>/callback?code=...&state=...`.
173 pub fn parse_query(query: &str) -> Result<Self> {
174 let mut code: Option<String> = None;
175 let mut state: Option<String> = None;
176 let mut error: Option<String> = None;
177 let mut error_description: Option<String> = None;
178
179 for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
180 match key.as_ref() {
181 "code" => {
182 if !value.is_empty() {
183 code = Some(value.into_owned());
184 }
185 }
186 "state" => {
187 if !value.is_empty() {
188 state = Some(value.into_owned());
189 }
190 }
191 "error" => {
192 if !value.is_empty() {
193 error = Some(value.into_owned());
194 }
195 }
196 "error_description" => {
197 if !value.is_empty() {
198 error_description = Some(value.into_owned());
199 }
200 }
201 _ => {}
202 }
203 }
204
205 if let Some(error_code) = error {
206 anyhow::bail!(
207 "OAuth authorization failed: {} ({})",
208 error_code,
209 error_description.as_deref().unwrap_or("no description")
210 );
211 }
212
213 let code = code.ok_or_else(|| anyhow!("missing 'code' parameter in OAuth callback"))?;
214 let state =
215 state.ok_or_else(|| anyhow!("missing 'state' parameter in OAuth callback"))?;
216
217 Ok(Self { code, state })
218 }
219 }
220
221 /// How long to wait for the browser to complete the OAuth flow before giving
222 /// up and releasing the loopback port.
223 const OAUTH_CALLBACK_TIMEOUT: Duration = Duration::from_secs(2 * 60);
224
225 /// Start a loopback HTTP server to receive the OAuth authorization callback.
226 ///
227 /// Binds to an ephemeral loopback port. Returns `(redirect_uri, callback_future)`.
228 /// The caller should use the redirect URI in the authorization request, open
229 /// the browser, then await the future to receive the callback.
230 pub fn start_oauth_callback_server() -> Result<(
231 String,
232 futures::channel::oneshot::Receiver<Result<OAuthCallbackParams>>,
233 )> {
234 start_oauth_callback_server_with_config(OAuthCallbackServerConfig::default())
235 }
236
237 /// Start a loopback HTTP server with custom host/port/path.
238 ///
239 /// Use this when the OAuth client requires a specific redirect URI that the
240 /// default ephemeral-port `http://127.0.0.1:<port>/callback` doesn't match.
241 pub fn start_oauth_callback_server_with_config(
242 config: OAuthCallbackServerConfig,
243 ) -> Result<(
244 String,
245 futures::channel::oneshot::Receiver<Result<OAuthCallbackParams>>,
246 )> {
247 let server = bind_callback_server(&config)?;
248 let port = server
249 .server_addr()
250 .to_ip()
251 .ok_or_else(|| anyhow!("server not bound to a TCP address"))?
252 .port();
253
254 let redirect_uri = format!("http://{}:{}{}", config.host, port, config.path);
255 let expected_path = config.path;
256
257 let (tx, rx) = futures::channel::oneshot::channel();
258
259 std::thread::spawn(move || {
260 let deadline = std::time::Instant::now() + OAUTH_CALLBACK_TIMEOUT;
261
262 loop {
263 if tx.is_canceled() {
264 return;
265 }
266 let remaining = deadline.saturating_duration_since(std::time::Instant::now());
267 if remaining.is_zero() {
268 return;
269 }
270
271 let timeout = remaining.min(Duration::from_millis(500));
272 let Some(request) = (match server.recv_timeout(timeout) {
273 Ok(req) => req,
274 Err(_) => {
275 let _ = tx.send(Err(anyhow!("OAuth callback server I/O error")));
276 return;
277 }
278 }) else {
279 continue;
280 };
281
282 let raw_url = request.url().to_string();
283 let raw_path = raw_url.split('?').next().unwrap_or(&raw_url);
284 if raw_path == CANCEL_PATH {
285 let response = tiny_http::Response::from_string("Cancelled")
286 .with_status_code(200)
287 .with_header(
288 tiny_http::Header::from_str("Content-Type: text/plain")
289 .expect("failed to construct response header"),
290 )
291 .with_header(
292 tiny_http::Header::from_str("Connection: close")
293 .expect("failed to construct response header"),
294 );
295 if let Err(err) = request.respond(response) {
296 log::error!("Failed to send OAuth cancel response: {}", err);
297 }
298 let _ = tx.send(Err(anyhow!(
299 "OAuth callback server was cancelled by another sign-in attempt"
300 )));
301 return;
302 }
303
304 let result = handle_oauth_callback_request(&request, expected_path);
305
306 let (status_code, body) = match &result {
307 Ok(_) => (
308 200,
309 oauth_callback_page(
310 "Authorization Successful",
311 "You can close this tab and return to Omega.",
312 false,
313 ),
314 ),
315 Err(err) => {
316 log::error!("OAuth callback error: {}", err);
317 (
318 400,
319 oauth_callback_page(
320 "Authorization Failed",
321 "Something went wrong. Please try again from Omega.",
322 true,
323 ),
324 )
325 }
326 };
327
328 let response = tiny_http::Response::from_string(body)
329 .with_status_code(status_code)
330 .with_header(
331 tiny_http::Header::from_str("Content-Type: text/html")
332 .expect("failed to construct response header"),
333 )
334 .with_header(
335 tiny_http::Header::from_str("Keep-Alive: timeout=0,max=0")
336 .expect("failed to construct response header"),
337 );
338 if let Err(err) = request.respond(response) {
339 log::error!("Failed to send OAuth callback response: {}", err);
340 }
341
342 let _ = tx.send(result);
343 return;
344 }
345 });
346
347 Ok((redirect_uri, rx))
348 }
349
350 fn handle_oauth_callback_request(
351 request: &tiny_http::Request,
352 expected_path: &str,
353 ) -> Result<OAuthCallbackParams> {
354 let url = Url::parse(&format!("http://localhost{}", request.url()))
355 .context("malformed callback request URL")?;
356
357 if url.path() != expected_path {
358 anyhow::bail!("unexpected path in OAuth callback: {}", url.path());
359 }
360
361 let query = url
362 .query()
363 .ok_or_else(|| anyhow!("OAuth callback has no query string"))?;
364 OAuthCallbackParams::parse_query(query)
365 }
366
367 /// Callback path reserved for evicting a previously-running OAuth callback
368 /// server bound to the same port. Always handled, regardless of `config.path`.
369 const CANCEL_PATH: &str = "/cancel";
370
371 const BIND_MAX_ATTEMPTS: u32 = 10;
372 const BIND_RETRY_DELAY: Duration = Duration::from_millis(200);
373 const CANCEL_REQUEST_TIMEOUT: Duration = Duration::from_secs(2);
374
375 fn bind_callback_server(config: &OAuthCallbackServerConfig) -> Result<tiny_http::Server> {
376 // Ephemeral ports always succeed; skip the cancel-retry dance entirely.
377 if config.preferred_port == 0 {
378 let addr = format!("{}:0", config.host);
379 return tiny_http::Server::http(&addr).map_err(|err| {
380 anyhow!(err).context(format!(
381 "Failed to bind loopback listener for OAuth callback on {addr}"
382 ))
383 });
384 }
385
386 match try_bind_with_cancel(config.host, config.preferred_port) {
387 Ok(server) => Ok(server),
388 Err(primary_err) => {
389 let Some(fallback_port) = config.fallback_port else {
390 return Err(primary_err.context(format!(
391 "Failed to bind loopback listener for OAuth callback on {}:{}",
392 config.host, config.preferred_port,
393 )));
394 };
395 log::warn!(
396 "OAuth callback port {}:{} unavailable; falling back to port {}",
397 config.host,
398 config.preferred_port,
399 fallback_port,
400 );
401 try_bind_with_cancel(config.host, fallback_port).map_err(|fallback_err| {
402 fallback_err.context(format!(
403 "Failed to bind loopback listener for OAuth callback on {}:{} or {}:{}",
404 config.host, config.preferred_port, config.host, fallback_port,
405 ))
406 })
407 }
408 }
409 }
410
411 /// Attempts to bind to a fixed `host:port`. On `AddrInUse`, sends a single
412 /// `GET /cancel` to the existing listener (to evict a previous OAuth flow
413 /// from this or a compatible client) and retries.
414 fn try_bind_with_cancel(host: &'static str, port: u16) -> Result<tiny_http::Server> {
415 let addr = format!("{host}:{port}");
416 let mut cancel_attempted = false;
417 let mut last_err: Option<anyhow::Error> = None;
418
419 for _ in 0..BIND_MAX_ATTEMPTS {
420 match tiny_http::Server::http(&addr) {
421 Ok(server) => return Ok(server),
422 Err(err) => {
423 let is_addr_in_use = err
424 .downcast_ref::<std::io::Error>()
425 .map(|io_err| io_err.kind() == std::io::ErrorKind::AddrInUse)
426 .unwrap_or(false);
427
428 if !is_addr_in_use {
429 return Err(anyhow!(err).context(format!(
430 "Failed to bind loopback listener for OAuth callback on {addr}"
431 )));
432 }
433
434 if !cancel_attempted {
435 cancel_attempted = true;
436 if let Err(cancel_err) = send_cancel_request(host, port) {
437 log::warn!(
438 "Failed to cancel previous OAuth callback server on {addr}: {cancel_err}"
439 );
440 }
441 }
442
443 last_err = Some(anyhow!(err));
444 std::thread::sleep(BIND_RETRY_DELAY);
445 }
446 }
447 }
448
449 Err(last_err
450 .unwrap_or_else(|| anyhow!("unknown bind error"))
451 .context(format!(
452 "OAuth callback port {addr} remained in use after {BIND_MAX_ATTEMPTS} attempts"
453 )))
454 }
455
456 /// Sends `GET /cancel` to a listener on `host:port`, asking it to shut down.
457 ///
458 /// Best-effort: errors here are surfaced to the caller for logging but do
459 /// not block the subsequent rebind attempt.
460 fn send_cancel_request(host: &str, port: u16) -> std::io::Result<()> {
461 use std::io::{Read as _, Write as _};
462 use std::net::{TcpStream, ToSocketAddrs as _};
463
464 let addr = format!("{host}:{port}")
465 .to_socket_addrs()?
466 .next()
467 .ok_or_else(|| {
468 std::io::Error::new(
469 std::io::ErrorKind::InvalidInput,
470 format!("could not resolve {host}:{port}"),
471 )
472 })?;
473 let mut stream = TcpStream::connect_timeout(&addr, CANCEL_REQUEST_TIMEOUT)?;
474 stream.set_read_timeout(Some(CANCEL_REQUEST_TIMEOUT))?;
475 stream.set_write_timeout(Some(CANCEL_REQUEST_TIMEOUT))?;
476
477 stream.write_all(b"GET /cancel HTTP/1.1\r\n")?;
478 stream.write_all(format!("Host: {host}:{port}\r\n").as_bytes())?;
479 stream.write_all(b"Connection: close\r\n\r\n")?;
480
481 // Drain the response so the server can close cleanly. We don't care
482 // about the body; errors here are harmless.
483 let mut buf = [0u8; 64];
484 let _ = stream.read(&mut buf);
485 Ok(())
486 }
487}
488
489#[cfg(not(target_family = "wasm"))]
490pub use server::{
491 OAuthCallbackParams, OAuthCallbackServerConfig, start_oauth_callback_server,
492 start_oauth_callback_server_with_config,
493};
494