1046 lines
45 KiB
Rust
1046 lines
45 KiB
Rust
//! WebSocket relay connection management.
|
||
//!
|
||
//! This module provides a shared `RelayConnection` that handles the WebSocket
|
||
//! connection to the grok.com relay server with automatic reconnection.
|
||
//! It is used by both `run_headless` and `run_leader` modes to avoid code duplication.
|
||
use super::proxy;
|
||
use crate::auth::{GrokAuth, GrokComConfig};
|
||
use crate::{teprintln, tprintln};
|
||
use futures_util::{SinkExt as _, StreamExt as _};
|
||
use std::sync::Arc;
|
||
use tokio::sync::mpsc;
|
||
use tokio::time::Duration;
|
||
use tokio_tungstenite::{
|
||
connect_async,
|
||
tungstenite::{Message, Utf8Bytes, client::IntoClientRequest},
|
||
};
|
||
use tokio_util::sync::CancellationToken;
|
||
use tracing::{debug, info, warn};
|
||
const KEEPALIVE_INTERVAL_SECS: u64 = 15;
|
||
/// Read-side liveness deadline. The write half pings every
|
||
/// `KEEPALIVE_INTERVAL_SECS`, and a healthy peer answers each ping with a
|
||
/// pong, so a live connection delivers an inbound frame at least that often.
|
||
/// If *nothing* arrives for this long the connection is treated as dead and
|
||
/// the session is torn down so the reconnect loop can take over.
|
||
///
|
||
/// Without this, a half-open TCP connection (e.g. the proxy/NAT leg still
|
||
/// ACKing our tiny pings while the upstream relay leg is gone) blocks
|
||
/// `ws_inbound.next()` forever and the agent never reconnects — sessions
|
||
/// stay bricked until the process is killed (server sees a 1006 close, the
|
||
/// client never notices).
|
||
const READ_LIVENESS_TIMEOUT_SECS: u64 = 4 * KEEPALIVE_INTERVAL_SECS;
|
||
/// Upper bound on a single auth-recovery attempt — a backstop against an
|
||
/// indefinitely wedged relay loop, NOT a bound on a healthy refresh. It must
|
||
/// stay comfortably above the refresh path's own internal worst case so it
|
||
/// only fires when something is truly stuck: `refresh_chain` waits up to 45s
|
||
/// for `auth.json.lock` (`REFRESH_LOCK_TIMEOUT`) before IdP IO, which has its
|
||
/// own timeouts (30s external refresher; 10–15s per OIDC request with short
|
||
/// retries). When this fires the recovery future is dropped (the file lock
|
||
/// releases on drop) and the loop falls through to reconnect backoff, which
|
||
/// retries recovery on the next 401.
|
||
const AUTH_RECOVERY_TIMEOUT_SECS: u64 = 180;
|
||
const BASE_DELAY_SECS: u64 = 1;
|
||
const MAX_DELAY_SECS: u64 = 60;
|
||
const CONNECT_TIMEOUT_SECS: u64 = 30;
|
||
/// JSON-RPC auth error code
|
||
const AUTH_ERROR_CODE: i64 = -32000;
|
||
use crate::auth::AuthManager;
|
||
/// Config for the grok.com WebSocket relay. Fields are private so the only
|
||
/// constructor is [`RelayConfig::for_session`] — "no relay without a session
|
||
/// bearer" is a compile-time guarantee.
|
||
#[derive(Clone)]
|
||
pub struct RelayConfig {
|
||
ws_url: String,
|
||
ws_origin: String,
|
||
token_header: String,
|
||
auth: GrokAuth,
|
||
auth_manager: Option<Arc<AuthManager>>,
|
||
}
|
||
impl RelayConfig {
|
||
/// Session gate: builds only for a grok.com first-party session
|
||
/// (`is_xai_auth`: x.ai-issuer OIDC or external credential) with a
|
||
/// non-empty bearer. BYOK/ApiKey, non-x.ai issuers (enterprise OIDC,
|
||
/// third-party external providers), and deprecated WebLogin get `None`
|
||
/// (relay-off; the leader still serves clients over IPC).
|
||
pub(crate) fn for_session(
|
||
session: &GrokAuth,
|
||
ctx: &GrokComConfig,
|
||
alpha_test_key: Option<String>,
|
||
auth_manager: Option<Arc<AuthManager>>,
|
||
) -> Option<Self> {
|
||
if !session.is_xai_auth() || session.key.is_empty() {
|
||
return None;
|
||
}
|
||
let _ = alpha_test_key;
|
||
Some(Self {
|
||
ws_url: ctx.grok_ws_url.clone(),
|
||
ws_origin: ctx.grok_ws_origin.clone(),
|
||
token_header: ctx.token_header.clone(),
|
||
auth: session.clone(),
|
||
auth_manager,
|
||
})
|
||
}
|
||
}
|
||
/// Callback type for first connection event.
|
||
pub type FirstConnectCallback = Box<dyn FnOnce() + Send + 'static>;
|
||
/// Handle to a running relay connection.
|
||
///
|
||
/// The relay maintains a persistent WebSocket connection to grok.com with
|
||
/// automatic reconnection on disconnection.
|
||
pub struct RelayHandle {
|
||
/// Cancel token to stop the relay connection loop
|
||
cancel: CancellationToken,
|
||
}
|
||
impl RelayHandle {
|
||
/// Stop the relay connection.
|
||
pub fn stop(&self) {
|
||
self.cancel.cancel();
|
||
}
|
||
/// Check if the relay is still running.
|
||
pub fn is_running(&self) -> bool {
|
||
!self.cancel.is_cancelled()
|
||
}
|
||
}
|
||
impl Drop for RelayHandle {
|
||
fn drop(&mut self) {
|
||
self.cancel.cancel();
|
||
}
|
||
}
|
||
/// Spawn a relay connection task that maintains a WebSocket connection.
|
||
///
|
||
/// The task runs in the background, automatically reconnecting on disconnection.
|
||
/// Messages from the relay are sent to `to_agent_tx`, and messages to send to
|
||
/// the relay should be sent via the returned sender.
|
||
///
|
||
/// # Arguments
|
||
/// * `config` - Relay connection configuration
|
||
/// * `to_agent_tx` - Channel to send messages received from the relay
|
||
/// * `parent_cancel` - Parent cancellation token (relay stops when parent is cancelled)
|
||
///
|
||
/// # Returns
|
||
/// A tuple of (sender for outbound messages, handle to control the relay)
|
||
pub fn spawn_relay_connection(
|
||
config: RelayConfig,
|
||
to_agent_tx: mpsc::UnboundedSender<String>,
|
||
parent_cancel: CancellationToken,
|
||
) -> (mpsc::UnboundedSender<String>, RelayHandle) {
|
||
spawn_relay_connection_with_callback(config, to_agent_tx, Some(parent_cancel), None)
|
||
}
|
||
/// Spawn a relay connection with an optional first-connection callback.
|
||
///
|
||
/// Same as `spawn_relay_connection` but allows providing a callback that will be
|
||
/// called once when the first successful connection is established.
|
||
pub fn spawn_relay_connection_with_callback(
|
||
config: RelayConfig,
|
||
to_agent_tx: mpsc::UnboundedSender<String>,
|
||
parent_cancel: Option<CancellationToken>,
|
||
on_first_connect: Option<FirstConnectCallback>,
|
||
) -> (mpsc::UnboundedSender<String>, RelayHandle) {
|
||
let cancel = parent_cancel.map_or(CancellationToken::new(), |c| c.child_token());
|
||
let cancel_clone = cancel.clone();
|
||
let (agent_to_ws_tx, agent_to_ws_rx) = mpsc::unbounded_channel::<String>();
|
||
tokio::spawn(async move {
|
||
run_relay_loop(
|
||
config,
|
||
to_agent_tx,
|
||
agent_to_ws_rx,
|
||
cancel_clone,
|
||
on_first_connect,
|
||
)
|
||
.await;
|
||
});
|
||
let handle = RelayHandle { cancel };
|
||
(agent_to_ws_tx, handle)
|
||
}
|
||
/// Check if a connection error is an HTTP 401 from the WebSocket handshake.
|
||
fn is_handshake_unauthorized(err: &anyhow::Error) -> bool {
|
||
use tokio_tungstenite::tungstenite::Error as WsError;
|
||
err.downcast_ref::<WsError>()
|
||
.map(|ws_err| {
|
||
matches!(
|
||
ws_err, WsError::Http(resp) if resp.status() ==
|
||
reqwest::StatusCode::UNAUTHORIZED
|
||
)
|
||
})
|
||
.unwrap_or(false)
|
||
}
|
||
/// Attempt auth recovery after a 401. Returns `true` to reconnect
|
||
/// immediately, `false` to exit or fall through to backoff.
|
||
async fn attempt_auth_recovery(
|
||
config: &mut RelayConfig,
|
||
cancel: &CancellationToken,
|
||
context: &str,
|
||
) -> bool {
|
||
let Some(ref am) = config.auth_manager else {
|
||
teprintln!("Authentication required. Run `grok login` to re-authenticate.");
|
||
cancel.cancel();
|
||
return false;
|
||
};
|
||
info!("auth recovery: relay {context}, attempting refresh");
|
||
let mut recovery = am.unauthorized_recovery(
|
||
Some(config.auth.clone()),
|
||
crate::auth::recovery::RecoverySource::Relay,
|
||
);
|
||
let recovered = match tokio::time::timeout(
|
||
Duration::from_secs(AUTH_RECOVERY_TIMEOUT_SECS),
|
||
recovery.next(),
|
||
)
|
||
.await
|
||
{
|
||
Ok(res) => res,
|
||
Err(_) => {
|
||
warn!(
|
||
timeout_secs = AUTH_RECOVERY_TIMEOUT_SECS,
|
||
"auth recovery: relay {context}, refresh timed out"
|
||
);
|
||
xai_grok_telemetry::unified_log::warn(
|
||
"auth recovery: relay refresh timed out",
|
||
None,
|
||
Some(serde_json::json!(
|
||
{ "context" : context, "timeout_secs" :
|
||
AUTH_RECOVERY_TIMEOUT_SECS, }
|
||
)),
|
||
);
|
||
return false;
|
||
}
|
||
};
|
||
match recovered {
|
||
Ok(new_auth) if new_auth.key == config.auth.key => {
|
||
info!("auth recovery: relay {context}, token unchanged, backing off");
|
||
xai_grok_telemetry::unified_log::info(
|
||
"auth recovery: relay token unchanged, backing off",
|
||
None,
|
||
Some(serde_json::json!(
|
||
{ "context" : context, "key_prefix" : crate
|
||
::auth::token_suffix(& new_auth.key), }
|
||
)),
|
||
);
|
||
false
|
||
}
|
||
Ok(new_auth) => {
|
||
info!("auth recovery: relay {context}, recovered, reconnecting");
|
||
xai_grok_telemetry::unified_log::info(
|
||
"auth recovery: relay recovered",
|
||
None,
|
||
Some(serde_json::json!(
|
||
{ "context" : context, "new_key_prefix" : crate
|
||
::auth::token_suffix(& new_auth.key), }
|
||
)),
|
||
);
|
||
config.auth = new_auth;
|
||
true
|
||
}
|
||
Err(e) if crate::auth::recovery::relay_should_cancel(&e) => {
|
||
teprintln!("{e}");
|
||
xai_grok_telemetry::unified_log::warn(
|
||
"auth recovery: relay giving up (terminal)",
|
||
None,
|
||
Some(serde_json::json!({ "context" : context, "error" : format!("{e}") })),
|
||
);
|
||
cancel.cancel();
|
||
false
|
||
}
|
||
Err(e) => {
|
||
warn!(error = % e, "auth recovery: relay {context}, refresh failed");
|
||
xai_grok_telemetry::unified_log::debug(
|
||
"auth recovery: relay refresh failed",
|
||
None,
|
||
Some(serde_json::json!({ "context" : context, "error" : format!("{e}") })),
|
||
);
|
||
false
|
||
}
|
||
}
|
||
}
|
||
/// Internal function that runs the reconnection loop.
|
||
async fn run_relay_loop(
|
||
mut config: RelayConfig,
|
||
to_agent_tx: mpsc::UnboundedSender<String>,
|
||
mut agent_to_ws_rx: mpsc::UnboundedReceiver<String>,
|
||
cancel: CancellationToken,
|
||
mut on_first_connect: Option<FirstConnectCallback>,
|
||
) {
|
||
let mut reconnect_attempts = 0u32;
|
||
let mut delay_secs = BASE_DELAY_SECS;
|
||
let mut first_connection = true;
|
||
let target_host = url::Url::parse(&config.ws_url)
|
||
.ok()
|
||
.and_then(|u| u.host_str().map(|h| h.to_string()));
|
||
let proxy_url = target_host
|
||
.as_deref()
|
||
.and_then(proxy::resolve_proxy_for_host);
|
||
if let Some(ref url) = proxy_url {
|
||
info!(
|
||
proxy = % url, target = target_host.as_deref().unwrap_or("unknown"),
|
||
"Using HTTP CONNECT proxy for relay connections"
|
||
);
|
||
}
|
||
loop {
|
||
if cancel.is_cancelled() {
|
||
info!("Relay connection cancelled, stopping");
|
||
break;
|
||
}
|
||
tracing::info!(
|
||
target : crate ::instrumentation::TARGET, event = "relay_connecting", ws_url
|
||
= % config.ws_url, attempt = reconnect_attempts,
|
||
);
|
||
match connect_to_relay(&config, proxy_url.as_deref(), &cancel).await {
|
||
Ok(ws) => {
|
||
tracing::info!(
|
||
target : crate ::instrumentation::TARGET, event = "relay_connected",
|
||
ws_url = % config.ws_url,
|
||
);
|
||
reconnect_attempts = 0;
|
||
delay_secs = BASE_DELAY_SECS;
|
||
if first_connection {
|
||
if let Some(callback) = on_first_connect.take() {
|
||
callback();
|
||
}
|
||
first_connection = false;
|
||
}
|
||
let result =
|
||
run_websocket_session(ws, &to_agent_tx, &mut agent_to_ws_rx, &cancel).await;
|
||
match result {
|
||
Ok(SessionEndReason::Normal) => {
|
||
info!("WebSocket session ended normally");
|
||
}
|
||
Ok(SessionEndReason::AuthError) => {
|
||
if attempt_auth_recovery(&mut config, &cancel, "Auth error").await {
|
||
continue;
|
||
}
|
||
}
|
||
Err(e) => {
|
||
warn!(error = ? e, "WebSocket session ended with error");
|
||
}
|
||
}
|
||
if cancel.is_cancelled() {
|
||
break;
|
||
}
|
||
tracing::info!(
|
||
target : crate ::instrumentation::TARGET, event =
|
||
"relay_disconnected", ws_url = % config.ws_url,
|
||
);
|
||
tprintln!("Disconnected from Grok WebSocket server");
|
||
info!("WebSocket disconnected, will reconnect");
|
||
}
|
||
Err(e) => {
|
||
let handshake_401 = is_handshake_unauthorized(&e);
|
||
tracing::info!(
|
||
target : crate ::instrumentation::TARGET, event =
|
||
"relay_connection_failed", ws_url = % config.ws_url, error = % e,
|
||
handshake_401,
|
||
);
|
||
if handshake_401 {
|
||
if attempt_auth_recovery(&mut config, &cancel, "Handshake 401").await {
|
||
continue;
|
||
}
|
||
} else {
|
||
warn!(error = % e, "Failed to connect to WebSocket server");
|
||
}
|
||
}
|
||
}
|
||
if cancel.is_cancelled() {
|
||
break;
|
||
}
|
||
reconnect_attempts += 1;
|
||
delay_secs = std::cmp::min(delay_secs * 2, MAX_DELAY_SECS);
|
||
info!(delay_secs, attempt = reconnect_attempts, "Reconnecting...");
|
||
tprintln!(
|
||
"Attempting to reconnect in {} seconds... (attempt #{})",
|
||
delay_secs,
|
||
reconnect_attempts
|
||
);
|
||
tokio::select! {
|
||
_ = cancel.cancelled() => break, _ =
|
||
tokio::time::sleep(Duration::from_secs(delay_secs)) => {}
|
||
}
|
||
}
|
||
}
|
||
/// Reason why a WebSocket session ended.
|
||
#[derive(Debug, PartialEq, Eq)]
|
||
pub(crate) enum SessionEndReason {
|
||
/// Normal disconnection (server closed, network error, etc.)
|
||
Normal,
|
||
/// Authentication error that may be recoverable with token refresh
|
||
AuthError,
|
||
}
|
||
/// Build an HTTP request with the relay authentication headers.
|
||
fn build_relay_request(config: &RelayConfig) -> anyhow::Result<axum::http::Request<()>> {
|
||
let mut req = config.ws_url.clone().into_client_request()?;
|
||
req.headers_mut().insert(
|
||
"Origin",
|
||
axum::http::header::HeaderValue::from_str(&config.ws_origin)?,
|
||
);
|
||
req.headers_mut().insert(
|
||
"Authorization",
|
||
axum::http::header::HeaderValue::from_str(&format!("Bearer {}", config.auth.key))?,
|
||
);
|
||
req.headers_mut().insert(
|
||
"X-XAI-Token-Auth",
|
||
axum::http::header::HeaderValue::from_str(&config.token_header)?,
|
||
);
|
||
req.headers_mut().insert(
|
||
"x-userid",
|
||
axum::http::header::HeaderValue::from_str(&config.auth.user_id)?,
|
||
);
|
||
req.headers_mut().insert(
|
||
"x-grok-client-version",
|
||
axum::http::header::HeaderValue::from_static(xai_grok_version::VERSION),
|
||
);
|
||
req.headers_mut().insert(
|
||
crate::http::CLIENT_MODE_HEADER,
|
||
axum::http::header::HeaderValue::from_static(crate::http::process_client_mode()),
|
||
);
|
||
Ok(req)
|
||
}
|
||
/// Attempt to connect to the relay WebSocket server.
|
||
///
|
||
/// If `proxy_url` is `Some`, the connection is established through an HTTP
|
||
/// CONNECT tunnel. Otherwise, a direct connection is used.
|
||
async fn connect_to_relay(
|
||
config: &RelayConfig,
|
||
proxy_url: Option<&str>,
|
||
cancel: &CancellationToken,
|
||
) -> anyhow::Result<
|
||
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
|
||
> {
|
||
let req = build_relay_request(config)?;
|
||
let connect_timeout = Duration::from_secs(CONNECT_TIMEOUT_SECS);
|
||
tokio::select! {
|
||
_ = cancel.cancelled() => { anyhow::bail!("Connection cancelled"); } result =
|
||
tokio::time::timeout(connect_timeout, async { if let Some(proxy_url) = proxy_url
|
||
{ let target_host = req.uri().host().ok_or_else(||
|
||
anyhow::anyhow!("WebSocket URL has no host")) ?; let target_port = req.uri()
|
||
.port_u16().unwrap_or(443); let tunneled_stream =
|
||
proxy::connect_via_proxy(proxy_url, target_host, target_port,). await ?; let (ws,
|
||
resp) = tokio_tungstenite::client_async(req, tunneled_stream). await .map_err(| e
|
||
| anyhow::Error::from(e).context("WebSocket handshake via proxy failed")) ?;
|
||
Ok((ws, resp)) } else { connect_async(req). await .map_err(| e |
|
||
anyhow::Error::from(e).context("WebSocket connection failed")) } }) => { match
|
||
result { Ok(Ok((ws, resp))) => { if let Some(proto) = resp.headers()
|
||
.get("Sec-WebSocket-Protocol") { info!(subprotocol = ? proto,
|
||
"WS subprotocol negotiated"); } Ok(ws) } Ok(Err(e)) => Err(e), Err(_) =>
|
||
anyhow::bail!("WebSocket connection timed out after {} seconds",
|
||
CONNECT_TIMEOUT_SECS), } }
|
||
}
|
||
}
|
||
/// Run a single WebSocket session, handling messages until disconnection.
|
||
pub(crate) async fn run_websocket_session<S>(
|
||
ws: tokio_tungstenite::WebSocketStream<S>,
|
||
to_agent_tx: &mpsc::UnboundedSender<String>,
|
||
from_agent_rx: &mut mpsc::UnboundedReceiver<String>,
|
||
cancel: &CancellationToken,
|
||
) -> anyhow::Result<SessionEndReason>
|
||
where
|
||
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + 'static,
|
||
{
|
||
run_websocket_session_with_liveness(
|
||
ws,
|
||
to_agent_tx,
|
||
from_agent_rx,
|
||
cancel,
|
||
Duration::from_secs(READ_LIVENESS_TIMEOUT_SECS),
|
||
)
|
||
.await
|
||
}
|
||
/// [`run_websocket_session`] with an explicit read-liveness window
|
||
/// (separate entry point so tests can use a short deadline).
|
||
pub(crate) async fn run_websocket_session_with_liveness<S>(
|
||
ws: tokio_tungstenite::WebSocketStream<S>,
|
||
to_agent_tx: &mpsc::UnboundedSender<String>,
|
||
from_agent_rx: &mut mpsc::UnboundedReceiver<String>,
|
||
cancel: &CancellationToken,
|
||
liveness: Duration,
|
||
) -> anyhow::Result<SessionEndReason>
|
||
where
|
||
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + 'static,
|
||
{
|
||
let (mut ws_outbound, mut ws_inbound) = ws.split();
|
||
let (auth_error_tx, mut auth_error_rx) = mpsc::channel::<()>(1);
|
||
let cancel_read = cancel.clone();
|
||
let read_from_ws = async move {
|
||
loop {
|
||
tokio::select! {
|
||
_ = cancel_read.cancelled() => break, msg_res =
|
||
tokio::time::timeout(liveness, ws_inbound.next()) => { let Ok(msg_opt) =
|
||
msg_res else { tprintln!("ws_inbound::liveness_timeout");
|
||
warn!(timeout_secs = liveness.as_secs(),
|
||
"no WS traffic within liveness window, treating connection as dead");
|
||
xai_grok_telemetry::unified_log::warn("relay: read liveness timeout, reconnecting",
|
||
None, Some(serde_json::json!({ "timeout_secs" : liveness.as_secs(),
|
||
})),); break; }; let Some(msg) = msg_opt else { break }; match msg {
|
||
Ok(Message::Text(text)) => { let trimmed_end = text
|
||
.trim_end_matches(['\r', '\n']); if trimmed_end.is_empty() {
|
||
debug!("received empty/whitespace WS text frame - skipping"); continue; }
|
||
let json : serde_json::Value = match serde_json::from_str(trimmed_end) {
|
||
Ok(v) => v, Err(_) => { debug!("failed to parse WS message as JSON");
|
||
continue; } }; if let Some(err) = json.get("error") { let code = err
|
||
.get("code").and_then(| c | c.as_i64()).unwrap_or(0); if code ==
|
||
AUTH_ERROR_CODE { let _ = auth_error_tx.send(()). await; return (false,
|
||
true); } tracing::warn!(error_code = code, "Server error (skipping)");
|
||
continue; } match json.get("method").and_then(| m | m.as_str()) {
|
||
Some(method) => tprintln!("acp_inbound::{}", method), None =>
|
||
tprintln!("ws_inbound::text"), } debug!(bytes = trimmed_end.len(),
|
||
"received WS text -> agent"); if to_agent_tx.send(trimmed_end
|
||
.to_string()).is_err() {
|
||
warn!("Failed to forward message to agent - channel closed"); break; } }
|
||
Ok(Message::Binary(bin)) => { tprintln!("ws_inbound::binary"); if let
|
||
Ok(s) = std::str::from_utf8(& bin) { let s = s.trim_end_matches(['\r',
|
||
'\n']); if s.is_empty() {
|
||
debug!("received empty WS binary frame - skipping"); continue; }
|
||
debug!(bytes = s.len(), "received WS binary(utf8) -> agent"); if
|
||
to_agent_tx.send(s.to_string()).is_err() { break; } } else {
|
||
debug!("received non-utf8 WS binary frame - skipping"); } }
|
||
Ok(Message::Close(frame_opt)) => { tprintln!("ws_inbound::close"); if let
|
||
Some(frame) = frame_opt { info!(code = ? frame.code, reason = % frame
|
||
.reason, "WS close received"); } else {
|
||
info!("WS close received (no frame)"); } break; } Ok(Message::Ping(p)) =>
|
||
{ tprintln!("ws_inbound::ping"); debug!(len = p.len(),
|
||
"received WS Ping"); } Ok(Message::Pong(p)) => {
|
||
tprintln!("ws_inbound::pong"); debug!(len = p.len(), "received WS Pong");
|
||
} Ok(Message::Frame(_)) => { tprintln!("ws_inbound::frame"); } Err(e) =>
|
||
{ tprintln!("ws_inbound::error::{:?}", & e); warn!(error = ? e,
|
||
"WS read error"); break; } } }
|
||
}
|
||
}
|
||
(true, false)
|
||
};
|
||
let cancel_write = cancel.clone();
|
||
let write_to_ws = async move {
|
||
let mut keepalive = tokio::time::interval(Duration::from_secs(KEEPALIVE_INTERVAL_SECS));
|
||
loop {
|
||
tokio::select! {
|
||
_ = cancel_write.cancelled() => break, msg_opt = from_agent_rx.recv() =>
|
||
{ match msg_opt { Some(msg) => { if
|
||
tracing::enabled!(tracing::Level::DEBUG) { if let Ok(json_val) =
|
||
serde_json::from_str::< serde_json::Value > (& msg) { let method =
|
||
json_val.get("method").and_then(| m | m.as_str()); let line_to_print =
|
||
match method { Some("session/update") => { let params = json_val
|
||
.get("params").unwrap_or(& serde_json::Value::Null);
|
||
format!("acp_outbound::session/update::{params}") } Some(m) =>
|
||
format!("acp_outbound::{m}"), None => "acp_outbound::response"
|
||
.to_string(), }; debug!("{line_to_print}"); } else {
|
||
debug!("acp_outbound::response"); } } if ! msg.is_empty() && let Err(e) =
|
||
ws_outbound.send(Message::Text(Utf8Bytes::from(msg))). await {
|
||
warn!(error = ? e, "failed to send to WS"); break; } } None => {
|
||
info!("Agent outbound channel closed"); break; } } } _ = keepalive.tick()
|
||
=> { tprintln!("ws::keep_alive_tick"); if let Err(e) = ws_outbound
|
||
.send(Message::Ping(Vec::new().into())). await {
|
||
tprintln!("ws::keep_alive::error::{:?}", & e); break; } }
|
||
}
|
||
}
|
||
anyhow::Ok(())
|
||
};
|
||
tokio::select! {
|
||
(_, auth_error) = read_from_ws => {
|
||
info!("WebSocket read task completed (connection closed)"); if auth_error {
|
||
return Ok(SessionEndReason::AuthError); } } res = write_to_ws => {
|
||
info!("WebSocket write task completed"); res ?; }
|
||
}
|
||
if auth_error_rx.try_recv().is_ok() {
|
||
return Ok(SessionEndReason::AuthError);
|
||
}
|
||
Ok(SessionEndReason::Normal)
|
||
}
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::auth::AuthMode;
|
||
use serde_json::json;
|
||
use std::sync::atomic::{AtomicU32, Ordering};
|
||
use tokio_tungstenite::tungstenite::{Utf8Bytes, protocol::Role};
|
||
/// Create an in-memory WebSocket pair (no network, no handshake needed).
|
||
async fn ws_pair() -> (
|
||
tokio_tungstenite::WebSocketStream<tokio::io::DuplexStream>,
|
||
tokio_tungstenite::WebSocketStream<tokio::io::DuplexStream>,
|
||
) {
|
||
let (client, server) = tokio::io::duplex(64 * 1024);
|
||
let client_ws =
|
||
tokio_tungstenite::WebSocketStream::from_raw_socket(client, Role::Client, None).await;
|
||
let server_ws =
|
||
tokio_tungstenite::WebSocketStream::from_raw_socket(server, Role::Server, None).await;
|
||
(client_ws, server_ws)
|
||
}
|
||
#[test]
|
||
fn test_handshake_401_detected_through_anyhow_context() {
|
||
use tokio_tungstenite::tungstenite::Error as WsError;
|
||
let resp = axum::http::Response::builder()
|
||
.status(401)
|
||
.body(None::<Vec<u8>>)
|
||
.unwrap();
|
||
let err = anyhow::Error::from(WsError::Http(resp)).context("WebSocket connection failed");
|
||
assert!(is_handshake_unauthorized(&err));
|
||
}
|
||
#[test]
|
||
fn test_handshake_non_401_and_non_ws_errors_rejected() {
|
||
use tokio_tungstenite::tungstenite::Error as WsError;
|
||
let resp = axum::http::Response::builder()
|
||
.status(403)
|
||
.body(None::<Vec<u8>>)
|
||
.unwrap();
|
||
let err = anyhow::Error::from(WsError::Http(resp)).context("WebSocket connection failed");
|
||
assert!(!is_handshake_unauthorized(&err));
|
||
let err = anyhow::anyhow!("some random error");
|
||
assert!(!is_handshake_unauthorized(&err));
|
||
}
|
||
#[tokio::test]
|
||
async fn test_ws_session_auth_error_returns_auth_error() {
|
||
let (client_ws, server_ws) = ws_pair().await;
|
||
let (mut server_tx, _server_rx) = server_ws.split();
|
||
let (to_agent_tx, _to_agent_rx) = mpsc::unbounded_channel::<String>();
|
||
let (_agent_out_tx, mut agent_out_rx) = mpsc::unbounded_channel::<String>();
|
||
let cancel = CancellationToken::new();
|
||
tokio::spawn(async move {
|
||
let auth_error = json!(
|
||
{ "jsonrpc" : "2.0", "id" : 1, "error" : { "code" : - 32000, "message" :
|
||
"Authentication required" } }
|
||
);
|
||
let _ = server_tx
|
||
.send(Message::Text(Utf8Bytes::from(auth_error.to_string())))
|
||
.await;
|
||
let _ = server_tx.close().await;
|
||
});
|
||
let result = tokio::time::timeout(
|
||
Duration::from_secs(5),
|
||
run_websocket_session(client_ws, &to_agent_tx, &mut agent_out_rx, &cancel),
|
||
)
|
||
.await
|
||
.expect("test timed out")
|
||
.expect("session should not error");
|
||
assert_eq!(result, SessionEndReason::AuthError);
|
||
}
|
||
#[tokio::test]
|
||
async fn test_ws_session_non_auth_error_skipped() {
|
||
let (client_ws, server_ws) = ws_pair().await;
|
||
let (mut server_tx, _server_rx) = server_ws.split();
|
||
let (to_agent_tx, _to_agent_rx) = mpsc::unbounded_channel::<String>();
|
||
let (_agent_out_tx, mut agent_out_rx) = mpsc::unbounded_channel::<String>();
|
||
let cancel = CancellationToken::new();
|
||
tokio::spawn(async move {
|
||
let other_error = json!(
|
||
{ "jsonrpc" : "2.0", "id" : 1, "error" : { "code" : - 32600, "message" :
|
||
"Invalid Request" } }
|
||
);
|
||
let _ = server_tx
|
||
.send(Message::Text(Utf8Bytes::from(other_error.to_string())))
|
||
.await;
|
||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||
let _ = server_tx.close().await;
|
||
});
|
||
let result = tokio::time::timeout(
|
||
Duration::from_secs(5),
|
||
run_websocket_session(client_ws, &to_agent_tx, &mut agent_out_rx, &cancel),
|
||
)
|
||
.await
|
||
.expect("test timed out")
|
||
.expect("session should not error");
|
||
assert_eq!(result, SessionEndReason::Normal);
|
||
}
|
||
#[tokio::test]
|
||
async fn test_ws_session_normal_close_returns_normal() {
|
||
let (client_ws, server_ws) = ws_pair().await;
|
||
let (mut server_tx, _server_rx) = server_ws.split();
|
||
let (to_agent_tx, _to_agent_rx) = mpsc::unbounded_channel::<String>();
|
||
let (_agent_out_tx, mut agent_out_rx) = mpsc::unbounded_channel::<String>();
|
||
let cancel = CancellationToken::new();
|
||
tokio::spawn(async move {
|
||
let _ = server_tx.send(Message::Close(None)).await;
|
||
});
|
||
let result = tokio::time::timeout(
|
||
Duration::from_secs(5),
|
||
run_websocket_session(client_ws, &to_agent_tx, &mut agent_out_rx, &cancel),
|
||
)
|
||
.await
|
||
.expect("test timed out")
|
||
.expect("session should not error");
|
||
assert_eq!(result, SessionEndReason::Normal);
|
||
}
|
||
#[tokio::test]
|
||
async fn test_ws_session_read_liveness_timeout_ends_session() {
|
||
let (client_ws, server_ws) = ws_pair().await;
|
||
let _silent_server = server_ws;
|
||
let (to_agent_tx, _to_agent_rx) = mpsc::unbounded_channel::<String>();
|
||
let (_agent_out_tx, mut agent_out_rx) = mpsc::unbounded_channel::<String>();
|
||
let cancel = CancellationToken::new();
|
||
let result = tokio::time::timeout(
|
||
Duration::from_secs(5),
|
||
run_websocket_session_with_liveness(
|
||
client_ws,
|
||
&to_agent_tx,
|
||
&mut agent_out_rx,
|
||
&cancel,
|
||
Duration::from_millis(100),
|
||
),
|
||
)
|
||
.await
|
||
.expect("session must end via read-liveness timeout instead of hanging")
|
||
.expect("session should not error");
|
||
assert_eq!(result, SessionEndReason::Normal);
|
||
}
|
||
#[tokio::test]
|
||
async fn test_ws_session_inbound_traffic_resets_liveness_window() {
|
||
let (client_ws, server_ws) = ws_pair().await;
|
||
let (mut server_tx, _server_rx) = server_ws.split();
|
||
let (to_agent_tx, mut to_agent_rx) = mpsc::unbounded_channel::<String>();
|
||
let (_agent_out_tx, mut agent_out_rx) = mpsc::unbounded_channel::<String>();
|
||
let cancel = CancellationToken::new();
|
||
tokio::spawn(async move {
|
||
for i in 0..12 {
|
||
let msg = json!({ "jsonrpc" : "2.0", "method" : "ping", "id" : i });
|
||
if server_tx
|
||
.send(Message::Text(Utf8Bytes::from(msg.to_string())))
|
||
.await
|
||
.is_err()
|
||
{
|
||
return;
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||
}
|
||
let _ = server_tx.close().await;
|
||
});
|
||
let result = tokio::time::timeout(
|
||
Duration::from_secs(5),
|
||
run_websocket_session_with_liveness(
|
||
client_ws,
|
||
&to_agent_tx,
|
||
&mut agent_out_rx,
|
||
&cancel,
|
||
Duration::from_millis(200),
|
||
),
|
||
)
|
||
.await
|
||
.expect("test timed out")
|
||
.expect("session should not error");
|
||
assert_eq!(result, SessionEndReason::Normal);
|
||
let mut forwarded = 0;
|
||
while to_agent_rx.try_recv().is_ok() {
|
||
forwarded += 1;
|
||
}
|
||
assert_eq!(forwarded, 12);
|
||
}
|
||
#[tokio::test]
|
||
async fn test_ws_session_forwards_text_to_agent() {
|
||
let (client_ws, server_ws) = ws_pair().await;
|
||
let (mut server_tx, _server_rx) = server_ws.split();
|
||
let (to_agent_tx, mut to_agent_rx) = mpsc::unbounded_channel::<String>();
|
||
let (_agent_out_tx, mut agent_out_rx) = mpsc::unbounded_channel::<String>();
|
||
let cancel = CancellationToken::new();
|
||
let test_msg = json!(
|
||
{ "jsonrpc" : "2.0", "id" : 1, "method" : "initialize", "params" : {} }
|
||
);
|
||
let msg_str = test_msg.to_string();
|
||
tokio::spawn(async move {
|
||
let _ = server_tx
|
||
.send(Message::Text(Utf8Bytes::from(msg_str)))
|
||
.await;
|
||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||
let _ = server_tx.close().await;
|
||
});
|
||
let _result = tokio::time::timeout(
|
||
Duration::from_secs(5),
|
||
run_websocket_session(client_ws, &to_agent_tx, &mut agent_out_rx, &cancel),
|
||
)
|
||
.await
|
||
.expect("test timed out");
|
||
let received = to_agent_rx
|
||
.try_recv()
|
||
.expect("should have forwarded message to agent");
|
||
let received_json: serde_json::Value = serde_json::from_str(&received).unwrap();
|
||
assert_eq!(received_json["method"], "initialize");
|
||
}
|
||
#[tokio::test]
|
||
async fn test_ws_session_cancel_stops_session() {
|
||
let (client_ws, server_ws) = ws_pair().await;
|
||
let _server_ws = server_ws;
|
||
let (to_agent_tx, _to_agent_rx) = mpsc::unbounded_channel::<String>();
|
||
let (_agent_out_tx, mut agent_out_rx) = mpsc::unbounded_channel::<String>();
|
||
let cancel = CancellationToken::new();
|
||
let cancel_clone = cancel.clone();
|
||
tokio::spawn(async move {
|
||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||
cancel_clone.cancel();
|
||
});
|
||
let result = tokio::time::timeout(
|
||
Duration::from_secs(5),
|
||
run_websocket_session(client_ws, &to_agent_tx, &mut agent_out_rx, &cancel),
|
||
)
|
||
.await
|
||
.expect("test timed out")
|
||
.expect("session should not error");
|
||
assert_eq!(result, SessionEndReason::Normal);
|
||
}
|
||
/// Helper to create a test GrokAuth with the given key.
|
||
fn test_auth(key: &str) -> GrokAuth {
|
||
GrokAuth {
|
||
key: key.to_string(),
|
||
refresh_token: Some("rt".to_string()),
|
||
..GrokAuth::test_default()
|
||
}
|
||
}
|
||
#[test]
|
||
fn for_session_builds_only_for_xai_issuer() {
|
||
use crate::auth::XAI_OAUTH2_ISSUER;
|
||
let cfg = GrokComConfig::default();
|
||
let builds = |a: &GrokAuth| RelayConfig::for_session(a, &cfg, None, None).is_some();
|
||
let xai = GrokAuth {
|
||
auth_mode: AuthMode::Oidc,
|
||
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
|
||
..test_auth("xai-bearer")
|
||
};
|
||
assert!(xai.is_xai_auth(), "precondition: is_xai_auth");
|
||
assert!(builds(&xai));
|
||
let external_xai = GrokAuth {
|
||
auth_mode: AuthMode::External,
|
||
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
|
||
..test_auth("ext-bearer")
|
||
};
|
||
assert!(external_xai.is_xai_auth(), "precondition: is_xai_auth");
|
||
assert!(builds(&external_xai));
|
||
assert!(!builds(&GrokAuth {
|
||
key: String::new(),
|
||
..xai.clone()
|
||
}));
|
||
assert!(!builds(&GrokAuth {
|
||
auth_mode: AuthMode::ApiKey,
|
||
..test_auth("k")
|
||
}));
|
||
assert!(!builds(&GrokAuth {
|
||
auth_mode: AuthMode::External,
|
||
..test_auth("k")
|
||
}));
|
||
assert!(!builds(&GrokAuth {
|
||
auth_mode: AuthMode::WebLogin,
|
||
..test_auth("k")
|
||
}));
|
||
assert!(!builds(&GrokAuth {
|
||
auth_mode: AuthMode::Oidc,
|
||
oidc_issuer: Some("https://login.acme-corp.example/oauth2".to_string()),
|
||
..test_auth("k")
|
||
}));
|
||
assert!(!builds(&GrokAuth {
|
||
auth_mode: AuthMode::External,
|
||
oidc_issuer: Some("https://login.acme-corp.example/oauth2".to_string()),
|
||
..test_auth("k")
|
||
}));
|
||
}
|
||
/// Helper: write a GrokAuth to disk under the given scope.
|
||
fn write_test_auth_to_disk(dir: &std::path::Path, scope: &str, auth: &GrokAuth) {
|
||
let path = dir.join("auth.json");
|
||
let mut map = crate::auth::read_auth_json(&path).unwrap_or_default();
|
||
map.insert(scope.to_owned(), auth.clone());
|
||
let json = serde_json::to_string_pretty(&map).unwrap();
|
||
std::fs::write(&path, json).unwrap();
|
||
}
|
||
/// Regression: `auth.json` vanishes (deleted/corrupt/externally
|
||
/// removed) while the process still holds an expired access token +
|
||
/// valid refresh token in `AuthManager` memory. Relay 401 recovery
|
||
/// must drive the full refresh chain — mint a fresh token via the
|
||
/// refresher and REWRITE `auth.json` — instead of dead-ending. A
|
||
/// relay holding a private, refresher-less `AuthManager` fails this:
|
||
/// it can only adopt sibling disk tokens, and there are none.
|
||
#[tokio::test]
|
||
async fn auth_recovery_refreshes_and_heals_missing_auth_json() {
|
||
use crate::auth::XAI_OAUTH2_ISSUER;
|
||
use crate::auth::refresh::{RefreshOutcome, TokenRefresher};
|
||
use std::sync::atomic::AtomicU32;
|
||
struct CountingRefresher {
|
||
calls: Arc<AtomicU32>,
|
||
}
|
||
#[async_trait::async_trait]
|
||
impl TokenRefresher for CountingRefresher {
|
||
async fn refresh(
|
||
&self,
|
||
_reason: crate::auth::manager::RefreshReason,
|
||
) -> RefreshOutcome {
|
||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||
RefreshOutcome::Success(Box::new(GrokAuth {
|
||
key: "fresh-from-authority".into(),
|
||
auth_mode: AuthMode::Oidc,
|
||
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
|
||
refresh_token: Some("rt-rotated".into()),
|
||
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
|
||
..GrokAuth::test_default()
|
||
}))
|
||
}
|
||
}
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let cfg = crate::auth::GrokComConfig::default();
|
||
let scope = cfg.auth_scope();
|
||
let am = Arc::new(
|
||
AuthManager::new(dir.path(), cfg.clone()).with_proxy_base_url("http://127.0.0.1:1"),
|
||
);
|
||
let expired_session = GrokAuth {
|
||
auth_mode: AuthMode::Oidc,
|
||
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
|
||
refresh_token: Some("rt-valid-unconsumed".into()),
|
||
expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(14)),
|
||
..test_auth("expired-overnight")
|
||
};
|
||
am.hot_swap(expired_session.clone());
|
||
assert!(
|
||
!dir.path().join("auth.json").exists(),
|
||
"precondition: no auth.json on disk"
|
||
);
|
||
let calls = Arc::new(AtomicU32::new(0));
|
||
am.set_refresher(Arc::new(CountingRefresher {
|
||
calls: calls.clone(),
|
||
}));
|
||
let mut config = RelayConfig::for_session(&expired_session, &cfg, None, Some(am.clone()))
|
||
.expect("x.ai OIDC session is relay-eligible");
|
||
let cancel = CancellationToken::new();
|
||
let recovered = attempt_auth_recovery(&mut config, &cancel, "test 401").await;
|
||
assert!(recovered, "recovery must succeed via the shared refresher");
|
||
assert!(!cancel.is_cancelled(), "relay must keep running");
|
||
assert_eq!(calls.load(Ordering::SeqCst), 1, "exactly one IdP refresh");
|
||
assert_eq!(config.auth.key, "fresh-from-authority");
|
||
let store = crate::auth::read_auth_json(&dir.path().join("auth.json"))
|
||
.expect("auth.json must be recreated");
|
||
let healed = store.get(&scope).expect("scope entry restored");
|
||
assert_eq!(healed.key, "fresh-from-authority");
|
||
assert_eq!(healed.refresh_token.as_deref(), Some("rt-rotated"));
|
||
}
|
||
/// Recovery returning the *unchanged* token (fresh-mint guard) must report
|
||
/// no recovery — the caller then backs off before reconnecting instead of
|
||
/// tight-looping — without cancelling the relay or touching the IdP.
|
||
#[tokio::test]
|
||
async fn attempt_auth_recovery_same_key_backs_off_without_cancel() {
|
||
use crate::auth::XAI_OAUTH2_ISSUER;
|
||
use crate::auth::refresh::{RefreshOutcome, TokenRefresher};
|
||
struct PanicRefresher;
|
||
#[async_trait::async_trait]
|
||
impl TokenRefresher for PanicRefresher {
|
||
async fn refresh(
|
||
&self,
|
||
_reason: crate::auth::manager::RefreshReason,
|
||
) -> RefreshOutcome {
|
||
panic!("fresh-mint guard must keep recovery away from the IdP");
|
||
}
|
||
}
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let cfg = crate::auth::GrokComConfig::default();
|
||
let am = Arc::new(AuthManager::new(dir.path(), cfg.clone()));
|
||
let fresh_session = GrokAuth {
|
||
auth_mode: AuthMode::Oidc,
|
||
oidc_issuer: Some(XAI_OAUTH2_ISSUER.to_string()),
|
||
refresh_token: Some("rt-valid".into()),
|
||
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
|
||
..test_auth("fresh-key")
|
||
};
|
||
am.hot_swap(fresh_session.clone());
|
||
am.set_refresher(Arc::new(PanicRefresher));
|
||
let mut config = RelayConfig::for_session(&fresh_session, &cfg, None, Some(am.clone()))
|
||
.expect("x.ai OIDC session is relay-eligible");
|
||
let cancel = CancellationToken::new();
|
||
let recovered = attempt_auth_recovery(&mut config, &cancel, "test 401").await;
|
||
assert!(!recovered, "same-key recovery must take the backoff path");
|
||
assert!(!cancel.is_cancelled(), "relay must keep reconnecting");
|
||
assert_eq!(config.auth.key, "fresh-key", "config auth stays unchanged");
|
||
}
|
||
#[tokio::test]
|
||
async fn test_auth_refresh_via_auth_manager_on_auth_error() {
|
||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let addr = listener.local_addr().unwrap();
|
||
let connection_count = Arc::new(AtomicU32::new(0));
|
||
let count_clone = connection_count.clone();
|
||
tokio::spawn(async move {
|
||
loop {
|
||
let Ok((stream, _)) = listener.accept().await else {
|
||
break;
|
||
};
|
||
let count = count_clone.clone();
|
||
tokio::spawn(async move {
|
||
let Ok(ws) = tokio_tungstenite::accept_async(stream).await else {
|
||
return;
|
||
};
|
||
let (mut tx, _rx) = ws.split();
|
||
let n = count.fetch_add(1, Ordering::SeqCst);
|
||
if n == 0 {
|
||
let auth_err = json!(
|
||
{ "jsonrpc" : "2.0", "id" : 1, "error" : { "code" : - 32000,
|
||
"message" : "Token expired" } }
|
||
);
|
||
let _ = tx
|
||
.send(Message::Text(Utf8Bytes::from(auth_err.to_string())))
|
||
.await;
|
||
} else {
|
||
tokio::time::sleep(Duration::from_secs(30)).await;
|
||
}
|
||
});
|
||
}
|
||
});
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let cfg = crate::auth::GrokComConfig::default();
|
||
let scope = cfg.auth_scope();
|
||
let am = Arc::new(AuthManager::new(dir.path(), cfg));
|
||
am.hot_swap(test_auth("old-key"));
|
||
write_test_auth_to_disk(dir.path(), &scope, &test_auth("new-key"));
|
||
let config = RelayConfig {
|
||
ws_url: format!("ws://{}", addr),
|
||
ws_origin: format!("http://{}", addr),
|
||
token_header: "test-token".to_string(),
|
||
auth: test_auth("old-key"),
|
||
auth_manager: Some(am),
|
||
};
|
||
let cancel = CancellationToken::new();
|
||
let (from_relay_tx, _from_relay_rx) = mpsc::unbounded_channel();
|
||
let (_to_relay_tx, _handle) = spawn_relay_connection(config, from_relay_tx, cancel.clone());
|
||
tokio::time::sleep(Duration::from_secs(3)).await;
|
||
cancel.cancel();
|
||
assert!(
|
||
connection_count.load(Ordering::SeqCst) >= 2,
|
||
"should have connected at least twice (original + after refresh), got {}",
|
||
connection_count.load(Ordering::SeqCst)
|
||
);
|
||
}
|
||
#[tokio::test]
|
||
async fn test_auth_refresh_failure_continues_with_backoff() {
|
||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let addr = listener.local_addr().unwrap();
|
||
let connection_count = Arc::new(AtomicU32::new(0));
|
||
let count_clone = connection_count.clone();
|
||
tokio::spawn(async move {
|
||
loop {
|
||
let Ok((stream, _)) = listener.accept().await else {
|
||
break;
|
||
};
|
||
let count = count_clone.clone();
|
||
tokio::spawn(async move {
|
||
let Ok(ws) = tokio_tungstenite::accept_async(stream).await else {
|
||
return;
|
||
};
|
||
let (mut tx, _rx) = ws.split();
|
||
count.fetch_add(1, Ordering::SeqCst);
|
||
let auth_err = json!(
|
||
{ "jsonrpc" : "2.0", "id" : 1, "error" : { "code" : - 32000,
|
||
"message" : "Token expired" } }
|
||
);
|
||
let _ = tx
|
||
.send(Message::Text(Utf8Bytes::from(auth_err.to_string())))
|
||
.await;
|
||
});
|
||
}
|
||
});
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let cfg = crate::auth::GrokComConfig::default();
|
||
let scope = cfg.auth_scope();
|
||
let am = Arc::new(AuthManager::new(dir.path(), cfg));
|
||
am.hot_swap(test_auth("old-key"));
|
||
write_test_auth_to_disk(dir.path(), &scope, &test_auth("old-key"));
|
||
let config = RelayConfig {
|
||
ws_url: format!("ws://{}", addr),
|
||
ws_origin: format!("http://{}", addr),
|
||
token_header: "test-token".to_string(),
|
||
auth: test_auth("old-key"),
|
||
auth_manager: Some(am),
|
||
};
|
||
let cancel = CancellationToken::new();
|
||
let (from_relay_tx, _from_relay_rx) = mpsc::unbounded_channel();
|
||
let (_to_relay_tx, _handle) = spawn_relay_connection(config, from_relay_tx, cancel.clone());
|
||
tokio::time::sleep(Duration::from_secs(4)).await;
|
||
cancel.cancel();
|
||
assert!(
|
||
connection_count.load(Ordering::SeqCst) >= 2,
|
||
"should have retried after failed refresh, got {}",
|
||
connection_count.load(Ordering::SeqCst)
|
||
);
|
||
}
|
||
}
|