| 1 | //! Selected Native proposals, the real existing pool and local HTTP/SSE peers. |
| 2 | use super::*; |
| 3 | use crate::extension_host::ExtensionHostManager; |
| 4 | use crate::extension_host::supervisor::HostRequestContext; |
| 5 | use serde_json::Value; |
| 6 | use std::sync::Mutex; |
| 7 | use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; |
| 8 | use tokio_util::sync::CancellationToken; |
| 9 | |
| 10 | struct RemotePeer { |
| 11 | base: String, |
| 12 | requests: Arc<Mutex<Vec<Value>>>, |
| 13 | cancel: CancellationToken, |
| 14 | task: tokio::task::JoinHandle<()>, |
| 15 | } |
| 16 | impl Drop for RemotePeer { |
| 17 | fn drop(&mut self) { |
| 18 | self.cancel.cancel(); |
| 19 | self.task.abort(); |
| 20 | } |
| 21 | } |
| 22 | fn reply(request: &Value) -> Option<Value> { |
| 23 | let id = request.get("id")?.clone(); |
| 24 | assert!( |
| 25 | id.is_string(), |
| 26 | "both backends must preserve Core's actual string ID" |
| 27 | ); |
| 28 | let result = match request["method"].as_str().unwrap() { |
| 29 | "initialize" => { |
| 30 | json!({"protocolVersion":request["params"]["protocolVersion"],"capabilities":{"tools":{}},"serverInfo":{"name":"reviewed-local-peer","version":"1"}}) |
| 31 | } |
| 32 | "tools/list" => { |
| 33 | json!({"tools":[{"name":"remote_echo","description":"local read-only fixture","inputSchema":{"type":"object"},"annotations":{"readOnlyHint":true}}]}) |
| 34 | } |
| 35 | "tools/call" => { |
| 36 | json!({"content":[{"type":"text","text":format!("remote:{}",request["params"]["arguments"])}]}) |
| 37 | } |
| 38 | "ping" => json!({}), |
| 39 | "resources/list" => json!({"resources":[]}), |
| 40 | "resources/templates/list" => json!({"resourceTemplates":[]}), |
| 41 | "prompts/list" => json!({"prompts":[]}), |
| 42 | other => panic!("unexpected local fixture method {other}"), |
| 43 | }; |
| 44 | Some(json!({"jsonrpc":"2.0","id":id,"result":result})) |
| 45 | } |
| 46 | async fn peer() -> RemotePeer { |
| 47 | let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); |
| 48 | let base = format!("http://{}", listener.local_addr().unwrap()); |
| 49 | let cancel = CancellationToken::new(); |
| 50 | let server_cancel = cancel.clone(); |
| 51 | let requests = Arc::new(Mutex::new(Vec::new())); |
| 52 | let observed = Arc::clone(&requests); |
| 53 | let (events, _) = tokio::sync::broadcast::channel::<Value>(32); |
| 54 | let task = tokio::spawn(async move { |
| 55 | let mut connections = tokio::task::JoinSet::new(); |
| 56 | loop { |
| 57 | tokio::select! { |
| 58 | biased; |
| 59 | _ = server_cancel.cancelled() => break, |
| 60 | connection = listener.accept() => { |
| 61 | let (socket, _) = connection.unwrap(); |
| 62 | let cancel = server_cancel.clone(); |
| 63 | let events = events.clone(); |
| 64 | let observed = Arc::clone(&observed); |
| 65 | connections.spawn(async move { |
| 66 | let mut reader = BufReader::new(socket); |
| 67 | let mut first = String::new(); |
| 68 | if reader.read_line(&mut first).await.unwrap() == 0 { return; } |
| 69 | let mut length = 0; |
| 70 | let mut header_bytes = first.len(); |
| 71 | loop { |
| 72 | let mut line = String::new(); |
| 73 | if reader.read_line(&mut line).await.unwrap() == 0 { return; } |
| 74 | header_bytes += line.len(); |
| 75 | assert!(header_bytes <= 64 * 1024); |
| 76 | if line == "\r\n" { break; } |
| 77 | if let Some(value) = line.to_ascii_lowercase().strip_prefix("content-length:") { length = value.trim().parse::<usize>().unwrap(); } |
| 78 | } |
| 79 | assert!(length <= 64 * 1024); |
| 80 | let mut body = vec![0; length]; |
| 81 | reader.read_exact(&mut body).await.unwrap(); |
| 82 | let mut socket = reader.into_inner(); |
| 83 | if first.starts_with("GET /sse ") { |
| 84 | let mut incoming = events.subscribe(); |
| 85 | socket.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\nevent: endpoint\ndata: /messages\n\n").await.unwrap(); |
| 86 | loop { |
| 87 | tokio::select! { |
| 88 | biased; |
| 89 | _ = cancel.cancelled() => break, |
| 90 | event = incoming.recv() => { |
| 91 | let Ok(event) = event else { break }; |
| 92 | let bytes = format!("event: message\ndata: {event}\n\n"); |
| 93 | if socket.write_all(bytes.as_bytes()).await.is_err() { break; } |
| 94 | } |
| 95 | } |
| 96 | } |
| 97 | } else if first.starts_with("POST /mcp ") || first.starts_with("POST /messages ") { |
| 98 | let request: Value = serde_json::from_slice(&body).unwrap(); |
| 99 | { |
| 100 | let mut requests = observed.lock().unwrap(); |
| 101 | assert!(requests.len() < 256); |
| 102 | requests.push(request.clone()); |
| 103 | } |
| 104 | let response = reply(&request); |
| 105 | if first.starts_with("POST /messages ") { |
| 106 | if let Some(response) = response { let _ = events.send(response); } |
| 107 | socket.write_all(b"HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection: close\r\n\r\n").await.unwrap(); |
| 108 | } else if let Some(response) = response { |
| 109 | let body = response.to_string(); |
| 110 | socket.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",body.len()).as_bytes()).await.unwrap(); |
| 111 | } else { |
| 112 | socket.write_all(b"HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection: close\r\n\r\n").await.unwrap(); |
| 113 | } |
| 114 | } else { |
| 115 | socket.write_all(b"HTTP/1.1 405 Method Not Allowed\r\nContent-Length: 0\r\nConnection: close\r\n\r\n").await.unwrap(); |
| 116 | } |
| 117 | }); |
| 118 | } |
| 119 | } |
| 120 | } |
| 121 | connections.abort_all(); |
| 122 | while connections.join_next().await.is_some() {} |
| 123 | }); |
| 124 | RemotePeer { |
| 125 | base, |
| 126 | requests, |
| 127 | cancel, |
| 128 | task, |
| 129 | } |
| 130 | } |
| 131 | |
| 132 | async fn admit_remote( |
| 133 | manager: &Arc<ExtensionHostManager>, |
| 134 | caller: &HostAttachment, |
| 135 | transport: &str, |
| 136 | url: &str, |
| 137 | name: &str, |
| 138 | ) -> u64 { |
| 139 | let current = for_plugins(&caller.plugin_view()).unwrap(); |
| 140 | let owner = current[0].3.registration.owner.clone(); |
| 141 | let scope = current[0].3.registration.scope.clone(); |
| 142 | let mut proposal = params(&owner, scope, name.into()); |
| 143 | proposal.spec.description = json!({"type":transport,"url":url}).to_string(); |
| 144 | let generation = manager.shared.plugin.host_generation.load(Ordering::SeqCst); |
| 145 | let (cx, _violations, _cancel) = HostRequestContext::for_test(1); |
| 146 | match admit(&manager.shared, HostTier::Plugin, generation, proposal, &cx).await { |
| 147 | RegisterResult::Admitted { handle } => handle, |
| 148 | RegisterResult::Refused { refused } => panic!("remote proposal refused: {refused}"), |
| 149 | } |
| 150 | } |
| 151 | |
| 152 | #[tokio::test(flavor = "current_thread")] |
| 153 | async fn selected_native_literal_remote_http_and_sse_admit_and_execute_on_both_backends() { |
| 154 | let _home = crate::test_support::SealedHome::new(); |
| 155 | let _policy = TestPolicyGuard::extension_host(true); |
| 156 | let Some(node) = node_for_tests("selected_native_remote_http_and_sse") else { |
| 157 | return; |
| 158 | }; |
| 159 | let fixture = FixturePlugins::new(&["raw-dsh-mcp"]).await; |
| 160 | let manager = fixture.manager(node); |
| 161 | let _manager = TestManagerGuard::install(Arc::clone(&manager)); |
| 162 | let caller = manager.attach(selected(&fixture, "a")); |
| 163 | caller.reconcile().await.unwrap(); |
| 164 | assert_reviewed_owner_active(&manager, &caller); |
| 165 | // The wrapper has only a broad reviewed Native receipt; no duplicated |
| 166 | // declarative endpoint capability inventory is invented for dynamic rows. |
| 167 | assert!( |
| 168 | fixture |
| 169 | .registry() |
| 170 | .get("raw-dsh-mcp") |
| 171 | .unwrap() |
| 172 | .inventory |
| 173 | .network_hosts |
| 174 | .is_empty() |
| 175 | ); |
| 176 | let peer = peer().await; |
| 177 | for (kind, suffix) in [("streamable-http", "mcp"), ("sse", "sse")] { |
| 178 | let public_name = format!("remote_{suffix}"); |
| 179 | let handle = admit_remote( |
| 180 | &manager, |
| 181 | &caller, |
| 182 | kind, |
| 183 | &format!("{}/{suffix}", peer.base), |
| 184 | &public_name, |
| 185 | ) |
| 186 | .await; |
| 187 | for backend in [McpBackend::Rust, McpBackend::Host] { |
| 188 | let mut active = pool(&fixture, &caller, backend); |
| 189 | let errors = active.connect_all().await; |
| 190 | assert!(errors.is_empty(), "{kind} {backend:?}: {errors:?}"); |
| 191 | let name = active |
| 192 | .all_tools() |
| 193 | .into_iter() |
| 194 | .find(|(_, tool)| tool.name == "remote_echo") |
| 195 | .unwrap() |
| 196 | .0; |
| 197 | let result = active |
| 198 | .call_tool(&name, json!({"backend":format!("{backend:?}")})) |
| 199 | .await |
| 200 | .unwrap(); |
| 201 | assert!( |
| 202 | result["content"][0]["text"] |
| 203 | .as_str() |
| 204 | .unwrap() |
| 205 | .starts_with("remote:") |
| 206 | ); |
| 207 | assert!( |
| 208 | peer.requests |
| 209 | .lock() |
| 210 | .unwrap() |
| 211 | .iter() |
| 212 | .any(|request| request["method"] == "tools/call" && request["id"].is_string()) |
| 213 | ); |
| 214 | active.shutdown_all().await; |
| 215 | } |
| 216 | let definition = for_plugins(&caller.plugin_view()) |
| 217 | .unwrap() |
| 218 | .into_iter() |
| 219 | .find(|row| row.0 == public_name) |
| 220 | .unwrap() |
| 221 | .3; |
| 222 | manager |
| 223 | .shared |
| 224 | .registry |
| 225 | .lock() |
| 226 | .unwrap() |
| 227 | .unregister(&definition.registration.owner, handle); |
| 228 | assert!(definition.validate().is_err()); |
| 229 | assert!(definition.registration.cancel.is_cancelled()); |
| 230 | } |
| 231 | manager.shutdown().await; |
| 232 | } |
| 233 | |
| 234 | #[tokio::test(flavor = "current_thread")] |
| 235 | async fn native_remote_network_deny_and_credential_refusal_prevent_requests() { |
| 236 | let _home = crate::test_support::SealedHome::new(); |
| 237 | let _policy = TestPolicyGuard::extension_host(true); |
| 238 | let Some(node) = node_for_tests("native_remote_deny_and_credentials") else { |
| 239 | return; |
| 240 | }; |
| 241 | let fixture = FixturePlugins::new(&["raw-dsh-mcp"]).await; |
| 242 | let manager = fixture.manager(node); |
| 243 | let _manager = TestManagerGuard::install(Arc::clone(&manager)); |
| 244 | let caller = manager.attach(selected(&fixture, "a")); |
| 245 | caller.reconcile().await.unwrap(); |
| 246 | assert_reviewed_owner_active(&manager, &caller); |
| 247 | let peer = peer().await; |
| 248 | admit_remote( |
| 249 | &manager, |
| 250 | &caller, |
| 251 | "streamable-http", |
| 252 | &format!("{}/mcp", peer.base), |
| 253 | "remote", |
| 254 | ) |
| 255 | .await; |
| 256 | for backend in [McpBackend::Rust, McpBackend::Host] { |
| 257 | let policy: crate::network_policy::NetworkPolicy = |
| 258 | serde_json::from_value(json!({"default":"deny","audit":false})).unwrap(); |
| 259 | let mut denied = pool(&fixture, &caller, backend).with_network_policy( |
| 260 | crate::network_policy::NetworkPolicyDecider::new(policy, None), |
| 261 | ); |
| 262 | assert!(!denied.connect_all().await.is_empty()); |
| 263 | denied.shutdown_all().await; |
| 264 | } |
| 265 | assert!(peer.requests.lock().unwrap().is_empty()); |
| 266 | let definition = for_plugins(&caller.plugin_view()) |
| 267 | .unwrap() |
| 268 | .into_iter() |
| 269 | .find(|row| row.0 == "remote") |
| 270 | .unwrap() |
| 271 | .3; |
| 272 | for value in [ |
| 273 | json!({"type":"streamable-http","url":format!("{}/mcp",peer.base),"headers":{"Authorization":"literal-fixture-only"}}), |
| 274 | json!({"type":"sse","url":format!("{}/sse?token=fixture",peer.base)}), |
| 275 | ] { |
| 276 | let mut proposal = params( |
| 277 | &definition.registration.owner, |
| 278 | definition.registration.scope.clone(), |
| 279 | "refused".into(), |
| 280 | ); |
| 281 | proposal.spec.description = value.to_string(); |
| 282 | let (cx, _violations, _cancel) = HostRequestContext::for_test(2); |
| 283 | assert!(matches!( |
| 284 | admit( |
| 285 | &manager.shared, |
| 286 | HostTier::Plugin, |
| 287 | definition.registration.host_generation, |
| 288 | proposal, |
| 289 | &cx |
| 290 | ) |
| 291 | .await, |
| 292 | RegisterResult::Refused { .. } |
| 293 | )); |
| 294 | } |
| 295 | assert!(peer.requests.lock().unwrap().is_empty()); |
| 296 | manager.shutdown().await; |
| 297 | } |
| 298 |