diff --git a/agent/src/main.rs b/agent/src/main.rs index df4b956..47d0ca0 100644 --- a/agent/src/main.rs +++ b/agent/src/main.rs @@ -76,6 +76,8 @@ const PUBLIC_IP_V6_URL: &str = "https://ipv6.icanhazip.com/"; const RUNTIME_STATS_INTERVAL: Duration = Duration::from_secs(5 * 60); const CLOCK_CALIBRATION_MAX_AGE: Duration = Duration::from_secs(24 * 60 * 60); const CLOCK_CALIBRATION_MIN_CHANGE_MS: i64 = 20_000; +const AGENT_PROTOCOL_VERSION: &str = "1"; +const AGENT_CAPABILITIES: &str = "metrics-v1,config-v1,remote-exec-v1,task-ack-v1"; type Error = Box; type Result = std::result::Result; @@ -1672,6 +1674,14 @@ fn connect_live(endpoint: &str, token: &str) -> Result<(LiveSocket, Option) request .headers_mut() .insert("User-Agent", format!("nodeflare-agent/{VERSION}").parse()?); + request.headers_mut().insert( + "X-NodeFlare-Agent-Protocol", + AGENT_PROTOCOL_VERSION.parse()?, + ); + request.headers_mut().insert( + "X-NodeFlare-Agent-Capabilities", + AGENT_CAPABILITIES.parse()?, + ); let stream = connect_live_stream(&url, deadline)?; match client_tls(request, stream) { Ok(connected) => break connected, @@ -1700,6 +1710,16 @@ fn connect_live(endpoint: &str, token: &str) -> Result<(LiveSocket, Option) .get("date") .and_then(|value| value.to_str().ok()) .and_then(|value| clock_offset_from_http_date(value, started_ms, ended_ms)); + if !response + .headers() + .get("x-nodeflare-agent-protocol") + .is_some_and(|value| value == AGENT_PROTOCOL_VERSION) + { + return Err(format!( + "backend protocol mismatch; update the backend to protocol {AGENT_PROTOCOL_VERSION}" + ) + .into()); + } Ok((socket, clock_offset_ms)) } diff --git a/backend/src/websocket/agent.rs b/backend/src/websocket/agent.rs index 7cf5deb..1a2781c 100644 --- a/backend/src/websocket/agent.rs +++ b/backend/src/websocket/agent.rs @@ -16,6 +16,11 @@ use tokio::sync::mpsc; const MAX_AGENT_MESSAGE_BYTES: usize = 2 * 1024 * 1024; const MAX_TASK_RESULT_BYTES: usize = 1024 * 1024; +const AGENT_PROTOCOL_VERSION: &str = "1"; +const AGENT_PROTOCOL_HEADER: &str = "x-nodeflare-agent-protocol"; +const AGENT_CAPABILITIES_HEADER: &str = "x-nodeflare-agent-capabilities"; +const REQUIRED_AGENT_CAPABILITIES: [&str; 4] = + ["metrics-v1", "config-v1", "remote-exec-v1", "task-ack-v1"]; #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -46,6 +51,13 @@ pub async fn handle( Ok(access) => access, Err(error) => return error.into_response(), }; + if !agent_protocol_supported(&headers) { + return ApiResponse::error( + axum::http::StatusCode::UPGRADE_REQUIRED, + format!("Agent 协议不兼容，需要协议版本 {AGENT_PROTOCOL_VERSION}"), + ) + .into_response(); + } let token = headers .get(axum::http::header::AUTHORIZATION) .and_then(|value| value.to_str().ok()) @@ -62,9 +74,15 @@ pub async fn handle( let remote_ip = client_ip(&headers, peer, &state.config.trusted_proxies); let token = token.to_string(); let connection_state = Arc::clone(&state); - ws.max_message_size(MAX_AGENT_MESSAGE_BYTES) + let mut response = ws + .max_message_size(MAX_AGENT_MESSAGE_BYTES) .max_frame_size(MAX_AGENT_MESSAGE_BYTES) - .on_upgrade(move |socket| run(socket, connection_state, identity, remote_ip, token)) + .on_upgrade(move |socket| run(socket, connection_state, identity, remote_ip, token)); + response.headers_mut().insert( + AGENT_PROTOCOL_HEADER, + axum::http::HeaderValue::from_static(AGENT_PROTOCOL_VERSION), + ); + response } async fn run( @@ -418,6 +436,26 @@ fn valid_task_id(value: &str) -> bool { !value.is_empty() && value.len() <= 80 && uuid::Uuid::parse_str(value).is_ok() } +fn agent_protocol_supported(headers: &HeaderMap) -> bool { + let protocol = headers + .get(AGENT_PROTOCOL_HEADER) + .and_then(|value| value.to_str().ok()) + .map(str::trim); + if protocol != Some(AGENT_PROTOCOL_VERSION) { + return false; + } + let capabilities = headers + .get(AGENT_CAPABILITIES_HEADER) + .and_then(|value| value.to_str().ok()) + .unwrap_or("") + .split(',') + .map(str::trim) + .collect::>(); + REQUIRED_AGENT_CAPABILITIES + .iter() + .all(|capability| capabilities.contains(capability)) +} + pub fn queue_remote_task( connection: &AgentConnection, task: &RemoteTaskInfo, @@ -441,6 +479,7 @@ pub fn queue_remote_task( #[cfg(test)] mod tests { use super::*; + use axum::http::HeaderValue; #[test] fn remote_dispatch_reports_backpressure_and_disconnects() { @@ -484,6 +523,31 @@ mod tests { ); } + #[test] + fn requires_the_current_agent_protocol_and_capabilities() { + let mut headers = HeaderMap::new(); + assert!(!agent_protocol_supported(&headers)); + headers.insert( + AGENT_PROTOCOL_HEADER, + HeaderValue::from_static(AGENT_PROTOCOL_VERSION), + ); + headers.insert( + AGENT_CAPABILITIES_HEADER, + HeaderValue::from_static("metrics-v1,config-v1,remote-exec-v1,task-ack-v1"), + ); + assert!(agent_protocol_supported(&headers)); + headers.remove(AGENT_CAPABILITIES_HEADER); + assert!(!agent_protocol_supported(&headers)); + headers.insert( + AGENT_CAPABILITIES_HEADER, + HeaderValue::from_static("metrics-v1,config-v1,remote-exec-v1,task-ack-v1"), + ); + headers.insert(AGENT_PROTOCOL_HEADER, HeaderValue::from_static("0")); + assert!(!agent_protocol_supported(&headers)); + headers.insert(AGENT_PROTOCOL_HEADER, HeaderValue::from_static("999")); + assert!(!agent_protocol_supported(&headers)); + } + #[test] fn updates_require_an_explicit_persistence_flag() { let mut message = serde_json::json!({"samples": []}); diff --git a/frontend/src/AppProvider.tsx b/frontend/src/AppProvider.tsx index 99a997d..e2feec1 100644 --- a/frontend/src/AppProvider.tsx +++ b/frontend/src/AppProvider.tsx @@ -14,6 +14,7 @@ import { resolveBackground, themeToggle } from "./theme"; import { BOOTSTRAP_POLL_INTERVAL_MS, createRefreshQueue, + LIVE_CARD_REFRESH_INTERVAL_MS, shouldSyncBootstrap, } from "./refresh"; import { connectLive } from "./transport"; @@ -231,7 +232,7 @@ export function AppProvider({ children }: { children: ReactNode }) { timer = undefined; if (document.hidden || navigator.onLine === false) return; tick(); - timer = window.setInterval(tick, 1_000); + timer = window.setInterval(tick, LIVE_CARD_REFRESH_INTERVAL_MS); }; sync(); document.addEventListener("visibilitychange", sync); diff --git a/frontend/src/refresh.test.ts b/frontend/src/refresh.test.ts index 28e41e0..6358134 100644 --- a/frontend/src/refresh.test.ts +++ b/frontend/src/refresh.test.ts @@ -4,9 +4,14 @@ import { createRefreshQueue, hasActiveRemoteTasks, isRemoteTaskActive, + LIVE_CARD_REFRESH_INTERVAL_MS, shouldSyncBootstrap, } from "./refresh"; +test("keeps server card refreshes at one second", () => { + expect(LIVE_CARD_REFRESH_INTERVAL_MS).toBe(1_000); +}); + describe("createRefreshQueue", () => { test("serializes requests and collapses overlap into one trailing refresh", async () => { const releases: Array<() => void> = []; diff --git a/frontend/src/refresh.ts b/frontend/src/refresh.ts index c95bd89..016df02 100644 --- a/frontend/src/refresh.ts +++ b/frontend/src/refresh.ts @@ -2,6 +2,7 @@ import type { RemoteTask } from "./types"; export const BOOTSTRAP_POLL_INTERVAL_MS = 15_000; export const BOOTSTRAP_LIVE_SYNC_INTERVAL_MS = 5 * 60_000; +export const LIVE_CARD_REFRESH_INTERVAL_MS = 1_000; export const REMOTE_TASK_POLL_INTERVAL_MS = 2_000; export const REMOTE_TASK_POLL_TIMEOUT_MS = 60_000; diff --git a/scripts/agent-runtime-test.mjs b/scripts/agent-runtime-test.mjs index e624d4a..b7fe1c6 100644 --- a/scripts/agent-runtime-test.mjs +++ b/scripts/agent-runtime-test.mjs @@ -16,6 +16,11 @@ const directory = mkdtempSync(join(tmpdir(), "nodeflare-agent-runtime-")); const agents = []; const sockets = new Set(); const releaseRequests = []; +const agentProtocolHeaders = { + "x-nodeflare-agent-protocol": "1", + "x-nodeflare-agent-capabilities": + "metrics-v1,config-v1,remote-exec-v1,task-ack-v1", +}; const proxy = createHttpServer((_request, response) => { response.writeHead(502); response.end(); @@ -70,9 +75,15 @@ try { console.log("agent runtime: automatic update during three-second sampling"); websocketServer = new WebSocketServer({ noServer: true }); + websocketServer.on("headers", (headers) => { + headers.push("X-NodeFlare-Agent-Protocol: 1"); + }); upgradeServer = createHttpServer(); let redirected = false; upgradeServer.on("upgrade", (request, socket, head) => { + for (const [name, expected] of Object.entries(agentProtocolHeaders)) { + assert.equal(request.headers[name], expected, `Agent must send ${name}`); + } if (request.url === "/api/agent/ws") { redirected = true; socket.end("HTTP/1.1 307 Temporary Redirect\r\nLocation: /live\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"); diff --git a/scripts/smoke-test.sh b/scripts/smoke-test.sh index 4c102c9..b574c8a 100644 --- a/scripts/smoke-test.sh +++ b/scripts/smoke-test.sh @@ -356,6 +356,8 @@ if [ -n "${MONITOR_MIGRATION_URL:-}" ]; then [ "$frozen_delete_status" = "503" ] frozen_agent_status=$(monitor_curl --silent --show-error --output /dev/null --write-out '%{http_code}' \ -H "Authorization: Bearer $agent_token" \ + -H 'X-NodeFlare-Agent-Protocol: 1' \ + -H 'X-NodeFlare-Agent-Capabilities: metrics-v1,config-v1,remote-exec-v1,task-ack-v1' \ -H 'Connection: Upgrade' -H 'Upgrade: websocket' -H 'Sec-WebSocket-Version: 13' \ -H 'Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==' "$MONITOR_BASE_URL/api/agent/ws") [ "$frozen_agent_status" = "503" ] diff --git a/scripts/websocket-smoke.mjs b/scripts/websocket-smoke.mjs index c179a91..7ce629b 100644 --- a/scripts/websocket-smoke.mjs +++ b/scripts/websocket-smoke.mjs @@ -27,6 +27,11 @@ const serverId = process.env.MONITOR_SERVER_ID; const latencyTaskId = process.env.MONITOR_LATENCY_TASK_ID; const expectTaskAssigned = process.env.MONITOR_EXPECT_TASK_ASSIGNED !== "0"; const configOnly = process.env.MONITOR_CONFIG_ONLY === "1"; +const agentProtocolHeaders = { + "X-NodeFlare-Agent-Protocol": "1", + "X-NodeFlare-Agent-Capabilities": + "metrics-v1,config-v1,remote-exec-v1,task-ack-v1", +}; if (!baseUrl || !adminToken || !agentToken || !serverId || !latencyTaskId) { throw new Error("MONITOR_BASE_URL, MONITOR_ADMIN_TOKEN, MONITOR_AGENT_TOKEN, MONITOR_SERVER_ID and MONITOR_LATENCY_TASK_ID are required"); @@ -115,25 +120,35 @@ async function waitForRemoteTask(taskId, expectedStatus) { throw new Error(`Remote task ${taskId} did not reach ${expectedStatus}`); } -function socketHeaders(token, headers = {}) { +function socketHeaders(path, token, headers = {}, includeAgentProtocol = true) { return { Authorization: `Bearer ${token}`, + ...(includeAgentProtocol && path === "/api/agent/ws" ? agentProtocolHeaders : {}), ...headers, }; } function openSocket(path, token, headers = {}, observeMessage) { return new Promise((resolve, reject) => { + let serverProtocol; const socket = new WebSocket(websocketUrl(path), { - headers: socketHeaders(token, headers), + headers: socketHeaders(path, token, headers), }); if (observeMessage) socket.on("message", observeMessage); const timer = setTimeout(() => { socket.terminate(); reject(new Error(`${path} WebSocket handshake timed out`)); }, 5_000); + socket.once("upgrade", (response) => { + serverProtocol = response.headers["x-nodeflare-agent-protocol"]; + }); socket.once("open", () => { clearTimeout(timer); + if (path === "/api/agent/ws" && serverProtocol !== "1") { + socket.terminate(); + reject(new Error(`Invalid backend Agent protocol: ${serverProtocol}`)); + return; + } resolve(socket); }); socket.once("unexpected-response", (_request, response) => { @@ -147,10 +162,10 @@ function openSocket(path, token, headers = {}, observeMessage) { }); } -function expectSocketStatus(path, token, statusCode, headers = {}) { +function expectSocketStatus(path, token, statusCode, headers = {}, includeAgentProtocol = true) { return new Promise((resolve, reject) => { const socket = new WebSocket(websocketUrl(path), { - headers: socketHeaders(token, headers), + headers: socketHeaders(path, token, headers, includeAgentProtocol), }); const timer = setTimeout(() => { socket.terminate(); @@ -286,6 +301,22 @@ function waitForSocketClose(socket, expected, timeoutMs = 5_000) { }); } +await expectSocketStatus("/api/agent/ws", "invalid-agent-token", 426, {}, false); +await expectSocketStatus( + "/api/agent/ws", + "invalid-agent-token", + 426, + { "X-NodeFlare-Agent-Protocol": "0" }, +); +await expectSocketStatus( + "/api/agent/ws", + "invalid-agent-token", + 426, + { + "X-NodeFlare-Agent-Capabilities": + "metrics-v1,config-v1,remote-exec-v1", + }, +); await expectSocketStatus("/api/agent/ws", "invalid-agent-token", 401); await expectSocketStatus("/api/agent/ws", "", 401); await expectSocketStatus("/api/agent/ws", "x".repeat(513), 401);