1use std::any::Any;
15use std::collections::HashMap;
16use std::sync::atomic::{AtomicU64, Ordering};
17use std::sync::Mutex;
18
19use bao_cdp::servo_bridge::{BridgeCommand, BridgeSender};
20use cdp_server::{CdpError, CdpMessage, EventSender, RegistryDispatch};
21use serde_json::{json, Value};
22
23const ERR_SESSION_NOT_FOUND: i64 = -32001;
26
27const BROWSER_TARGET: &str = "__browser__";
30
31const SERVED_DOMAINS: [&str; 21] = [
33 "Target",
34 "Page",
35 "Runtime",
36 "DOM",
37 "Network",
38 "CSS",
39 "Emulation",
40 "Input",
41 "Overlay",
42 "Debugger",
43 "Log",
44 "Fetch",
45 "Storage",
46 "Security",
47 "Profiler",
48 "HeapProfiler",
49 "Memory",
50 "Performance",
51 "SystemInfo",
52 "ServiceWorker",
53 "Browser",
54];
55
56static SESSION_COUNTER: AtomicU64 = AtomicU64::new(1);
57static CONTEXT_COUNTER: AtomicU64 = AtomicU64::new(1);
58
59fn next_session_id() -> String {
60 let n = SESSION_COUNTER.fetch_add(1, Ordering::Relaxed);
61 format!("bao-session-{n:016x}")
62}
63
64pub struct BaoWsRegistry {
73 bridge: BridgeSender,
74 attached_sessions: Mutex<HashMap<String, String>>,
76 auto_attach: Mutex<bool>,
79}
80
81impl BaoWsRegistry {
82 pub fn new(bridge: BridgeSender) -> Self {
83 BaoWsRegistry {
84 bridge,
85 attached_sessions: Mutex::new(HashMap::new()),
86 auto_attach: Mutex::new(false),
87 }
88 }
89
90 fn dispatch_session_command(
94 &self,
95 method: &str,
96 params: &Option<Value>,
97 msg: &CdpMessage,
98 event_sender: &dyn EventSender,
99 ) -> Option<Result<Value, CdpError>> {
100 match method {
101 "Target.attachToTarget" => Some(self.attach_to_target(params)),
102 "Target.detachFromTarget" => Some(self.detach_from_target(params)),
103 "Page.createIsolatedWorld" => {
107 Some(self.create_isolated_world(params, msg, event_sender))
108 }
109 "Target.setAutoAttach" => {
110 Some(self.set_auto_attach(params, msg.session_id.is_none(), event_sender))
115 }
116 _ => None,
117 }
118 }
119
120 fn attach_to_target(&self, params: &Option<Value>) -> Result<Value, CdpError> {
121 let target_id = require_param(params, "targetId")?;
122 let flatten = params
126 .as_ref()
127 .and_then(|p| p.get("flatten"))
128 .and_then(|v| v.as_bool())
129 .unwrap_or(false);
130 if !flatten {
131 return Err(CdpError {
132 code: -32000,
133 message: "'Target.attachToTarget' not supported: only flatten=true (Playwright/Puppeteer mode) is implemented".into(),
134 });
135 }
136 Ok(json!({ "sessionId": self.mint_session(&target_id) }))
137 }
138
139 fn detach_from_target(&self, params: &Option<Value>) -> Result<Value, CdpError> {
140 let session_id = require_param(params, "sessionId")?;
141 if let Ok(mut table) = self.attached_sessions.lock() {
142 table.remove(session_id.as_str());
143 }
144 Ok(json!({}))
145 }
146
147 fn set_auto_attach(
152 &self,
153 params: &Option<Value>,
154 from_browser_session: bool,
155 event_sender: &dyn EventSender,
156 ) -> Result<Value, CdpError> {
157 let auto_attach = params
158 .as_ref()
159 .and_then(|p| p.get("autoAttach"))
160 .and_then(|v| v.as_bool())
161 .unwrap_or(false);
162 if let Ok(mut flag) = self.auto_attach.lock() {
163 *flag = auto_attach;
164 }
165 if !auto_attach || !from_browser_session {
166 return Ok(json!({}));
167 }
168 let listed = self
171 .bridge
172 .send(BridgeCommand::ListTargets)
173 .result
174 .ok()
175 .and_then(|v| v.as_array().cloned());
176 if let Some(entries) = listed {
177 for entry in entries {
178 let Some(id) = entry.get("id").and_then(|v| v.as_str()) else {
179 continue;
180 };
181 let session_id = self.mint_session(id);
182 event_sender.send_event(
183 "Target.attachedToTarget",
184 json!({
185 "sessionId": session_id,
186 "targetInfo": {
187 "targetId": id,
188 "type": "page",
189 "title": entry.get("title").cloned().unwrap_or(json!("")),
190 "url": entry.get("url").cloned().unwrap_or(json!("about:blank")),
191 "attached": true,
192 "browserContextId": "bao-default-context",
193 },
194 }),
195 );
196 }
197 }
198 Ok(json!({}))
199 }
200
201 fn create_isolated_world(
210 &self,
211 params: &Option<Value>,
212 msg: &CdpMessage,
213 event_sender: &dyn EventSender,
214 ) -> Result<Value, CdpError> {
215 let world_name = params
216 .as_ref()
217 .and_then(|p| p.get("worldName"))
218 .and_then(|v| v.as_str())
219 .unwrap_or("")
220 .to_string();
221 let context_id = CONTEXT_COUNTER.fetch_add(1, Ordering::Relaxed);
222 if let Some(sid) = msg.session_id.as_deref() {
223 event_sender.send_session_event(
224 sid,
225 "Runtime.executionContextCreated",
226 json!({
227 "context": {
228 "id": context_id,
229 "origin": "-",
230 "name": world_name,
231 "auxData": { "isDefault": false },
232 }
233 }),
234 );
235 }
236 Ok(json!({ "executionContextId": context_id }))
237 }
238
239 fn mint_session(&self, target_id: &str) -> String {
241 let session_id = next_session_id();
242 if let Ok(mut table) = self.attached_sessions.lock() {
243 table.insert(session_id.clone(), target_id.to_string());
244 }
245 session_id
246 }
247
248 fn emit(
252 &self,
253 event_sender: &dyn EventSender,
254 session_id: Option<&str>,
255 method: &str,
256 params: Value,
257 ) {
258 match session_id {
259 Some(sid) => event_sender.send_session_event(sid, method, params),
260 None => event_sender.send_event(method, params),
261 }
262 }
263}
264
265fn require_param(params: &Option<Value>, key: &str) -> Result<String, CdpError> {
266 params
267 .as_ref()
268 .and_then(|p| p.get(key))
269 .and_then(|v| v.as_str())
270 .filter(|s| !s.is_empty())
271 .map(|s| s.to_string())
272 .ok_or_else(|| CdpError {
273 code: -32602,
274 message: format!("requires a non-empty {key} param"),
275 })
276}
277
278impl RegistryDispatch for BaoWsRegistry {
279 fn dispatch_command(
280 &self,
281 method: &str,
282 params: Value,
283 event_sender: &dyn EventSender,
284 ) -> Option<Result<Value, CdpError>> {
285 let msg = CdpMessage {
288 id: None,
289 method: method.to_string(),
290 params: Some(params),
291 session_id: None,
292 };
293 self.dispatch_message(&msg, BROWSER_TARGET, event_sender)
294 }
295
296 fn dispatch_message(
297 &self,
298 msg: &CdpMessage,
299 ws_target_id: &str,
300 event_sender: &dyn EventSender,
301 ) -> Option<Result<Value, CdpError>> {
302 if let Some(result) =
304 self.dispatch_session_command(&msg.method, &msg.params, msg, event_sender)
305 {
306 return Some(result);
307 }
308
309 let target_id = match &msg.session_id {
313 Some(sid) => match self
314 .attached_sessions
315 .lock()
316 .ok()
317 .and_then(|t| t.get(sid).cloned())
318 {
319 Some(t) => t,
320 None => {
321 return Some(Err(CdpError {
322 code: ERR_SESSION_NOT_FOUND,
323 message: format!("Session with given id not found: {sid}"),
324 }))
325 }
326 },
327 None => ws_target_id.to_string(),
328 };
329
330 let response = bao_cdp::handle_command(
332 msg.clone(),
333 &target_id,
334 &msg.params,
335 Some(&self.bridge),
336 );
337 let result = match (response.result, response.error) {
338 (Some(result), _) => Ok(result),
339 (None, Some(err)) => Err(err),
340 (None, None) => Ok(json!({})),
341 };
342
343 if result.is_ok() {
346 let sid = msg.session_id.as_deref();
347 match msg.method.as_str() {
348 "Runtime.enable" => {
352 let context_id = CONTEXT_COUNTER.fetch_add(1, Ordering::Relaxed);
356 self.emit(
357 event_sender,
358 sid,
359 "Runtime.executionContextCreated",
360 json!({
361 "context": {
362 "id": context_id,
363 "origin": "-",
364 "name": "",
365 "auxData": {
366 "isDefault": true,
367 "type": "default",
368 "frameId": target_id,
369 },
370 }
371 }),
372 );
373 }
374 "Page.navigate" => {
377 if let Ok(ref r) = result {
378 let fid = r
379 .get("frameId")
380 .and_then(|v| v.as_str())
381 .unwrap_or(&target_id)
382 .to_string();
383 let loader = r
384 .get("loaderId")
385 .and_then(|v| v.as_str())
386 .unwrap_or("")
387 .to_string();
388 let url = msg
389 .params
390 .as_ref()
391 .and_then(|p| p.get("url"))
392 .and_then(|v| v.as_str())
393 .unwrap_or("about:blank")
394 .to_string();
395 self.emit(
396 event_sender,
397 sid,
398 "Page.frameStartedLoading",
399 json!({ "frameId": fid }),
400 );
401 self.emit(
407 event_sender,
408 sid,
409 "Runtime.executionContextsCleared",
410 json!({}),
411 );
412 self.emit(
413 event_sender,
414 sid,
415 "Page.frameNavigated",
416 json!({
417 "frame": {
418 "id": fid,
419 "loaderId": loader,
420 "url": url,
421 "mimeType": "text/html",
422 "securityOrigin": "",
423 },
424 }),
425 );
426 let context_id = CONTEXT_COUNTER.fetch_add(1, Ordering::Relaxed);
427 self.emit(
428 event_sender,
429 sid,
430 "Runtime.executionContextCreated",
431 json!({
432 "context": {
433 "id": context_id,
434 "origin": "-",
435 "name": "",
436 "auxData": {
437 "isDefault": true,
438 "type": "default",
439 "frameId": fid,
440 },
441 }
442 }),
443 );
444 }
445 }
446 "Target.createTarget" => {
450 let auto = self.auto_attach.lock().map(|f| *f).unwrap_or(false);
451 if auto {
452 if let Ok(ref r) = result {
453 if let Some(new_id) = r.get("targetId").and_then(|v| v.as_str()) {
454 let session_id = self.mint_session(new_id);
455 event_sender.send_event(
456 "Target.attachedToTarget",
457 json!({
458 "sessionId": session_id,
459 "targetInfo": {
460 "targetId": new_id,
461 "type": "page",
462 "title": "",
463 "url": "about:blank",
464 "attached": true,
465 "browserContextId": "bao-default-context",
466 },
467 }),
468 );
469 }
470 }
471 }
472 }
473 _ => {}
474 }
475 }
476
477 Some(result)
478 }
479
480 fn notify_session_created(&self, _domain: &str, _session_id: &str) {
481 }
483
484 fn notify_session_destroyed(&self, _domains: &[String], _session_id: &str) {
485 }
489
490 fn has_domain(&self, domain: &str) -> bool {
491 SERVED_DOMAINS.contains(&domain)
492 }
493
494 fn as_any(&self) -> &dyn Any {
495 self
496 }
497}
498
499#[cfg(test)]
500mod tests {
501 use super::*;
502 use bao_cdp::servo_bridge::{bridge_channel, BridgeResponse};
503 use std::sync::Arc;
504 use std::time::Duration;
505
506 struct NopSender;
507 impl EventSender for NopSender {
508 fn send_event(&self, _: &str, _: Value) {}
509 }
510
511 struct CapturingSender {
512 events: Mutex<Vec<(String, Value)>>,
513 session_events: Mutex<Vec<(String, String, Value)>>,
514 }
515 impl CapturingSender {
516 fn new() -> Arc<Self> {
517 Arc::new(CapturingSender {
518 events: Mutex::new(Vec::new()),
519 session_events: Mutex::new(Vec::new()),
520 })
521 }
522 }
523 impl EventSender for CapturingSender {
524 fn send_event(&self, method: &str, params: Value) {
525 self.events
526 .lock()
527 .unwrap()
528 .push((method.to_string(), params));
529 }
530 fn send_session_event(&self, session_id: &str, method: &str, params: Value) {
531 self.session_events
532 .lock()
533 .unwrap()
534 .push((session_id.to_string(), method.to_string(), params));
535 }
536 }
537
538 fn page_responder(rx: bao_cdp::servo_bridge::BridgeReceiver) -> std::thread::JoinHandle<()> {
539 std::thread::spawn(move || {
540 loop {
541 let handled = rx.try_process(|cmd| match cmd {
542 BridgeCommand::ListTargets => BridgeResponse {
543 result: Ok(json!([
544 { "id": "1", "title": "Page 1", "url": "about:blank" }
545 ])),
546 },
547 BridgeCommand::Navigate { .. } => BridgeResponse {
548 result: Ok(json!({ "frameId": "1", "loaderId": "loader-1" })),
549 },
550 BridgeCommand::EvaluateJs { expression, .. } => BridgeResponse {
551 result: Ok(json!({ "result": { "type": "string", "value": expression } })),
552 },
553 _ => BridgeResponse {
554 result: Ok(json!({})),
555 },
556 });
557 if !handled {
558 std::thread::sleep(Duration::from_millis(1));
559 }
560 }
561 })
562 }
563
564 fn msg(method: &str, params: Value, session_id: Option<String>) -> CdpMessage {
565 CdpMessage {
566 id: Some(1),
567 method: method.to_string(),
568 params: Some(params),
569 session_id,
570 }
571 }
572
573 #[test]
575 fn attach_to_target_mints_real_unique_sessions() {
576 let (tx, rx) = bridge_channel(Duration::from_secs(2));
577 let _keeper = page_responder(rx);
578 let reg = BaoWsRegistry::new(tx);
579 let sender = NopSender;
580
581 let r1 = reg
582 .dispatch_message(
583 &msg(
584 "Target.attachToTarget",
585 json!({"targetId": "1", "flatten": true}),
586 None,
587 ),
588 BROWSER_TARGET,
589 &sender,
590 )
591 .unwrap()
592 .unwrap();
593 let r2 = reg
594 .dispatch_message(
595 &msg(
596 "Target.attachToTarget",
597 json!({"targetId": "1", "flatten": true}),
598 None,
599 ),
600 BROWSER_TARGET,
601 &sender,
602 )
603 .unwrap()
604 .unwrap();
605 let s1 = r1["sessionId"].as_str().unwrap().to_string();
606 let s2 = r2["sessionId"].as_str().unwrap().to_string();
607 assert_ne!(s1, s2, "each attach mints a fresh session id");
608 }
609
610 #[test]
611 fn attach_to_target_requires_flatten() {
612 let (tx, _rx) = bridge_channel(Duration::from_millis(100));
613 let reg = BaoWsRegistry::new(tx);
614 let sender = NopSender;
615 let err = reg
616 .dispatch_message(
617 &msg("Target.attachToTarget", json!({"targetId": "1"}), None),
618 BROWSER_TARGET,
619 &sender,
620 )
621 .unwrap()
622 .unwrap_err();
623 assert_eq!(err.code, -32000);
624 assert!(err.message.contains("flatten"));
625 }
626
627 #[test]
628 fn flattened_session_routes_to_attached_target() {
629 let (tx, rx) = bridge_channel(Duration::from_secs(2));
630 let _keeper = page_responder(rx);
631 let reg = BaoWsRegistry::new(tx);
632 let sender = NopSender;
633
634 let r = reg
635 .dispatch_message(
636 &msg(
637 "Target.attachToTarget",
638 json!({"targetId": "1", "flatten": true}),
639 None,
640 ),
641 BROWSER_TARGET,
642 &sender,
643 )
644 .unwrap()
645 .unwrap();
646 let sid = r["sessionId"].as_str().unwrap().to_string();
647
648 let nav = reg
651 .dispatch_message(
652 &msg("Page.navigate", json!({"url": "about:blank"}), Some(sid.clone())),
653 BROWSER_TARGET,
654 &sender,
655 )
656 .unwrap()
657 .unwrap();
658 assert_eq!(nav["frameId"], "1");
659 }
660
661 #[test]
662 fn unknown_session_id_is_explicit_error() {
663 let (tx, _rx) = bridge_channel(Duration::from_millis(100));
664 let reg = BaoWsRegistry::new(tx);
665 let sender = NopSender;
666 let err = reg
667 .dispatch_message(
668 &msg("Page.navigate", json!({"url": "about:blank"}), Some("nope".into())),
669 BROWSER_TARGET,
670 &sender,
671 )
672 .unwrap()
673 .unwrap_err();
674 assert_eq!(err.code, -32001);
675 assert!(err.message.contains("not found"));
676 }
677
678 #[test]
679 fn detach_removes_routing_entry() {
680 let (tx, _rx) = bridge_channel(Duration::from_millis(100));
681 let reg = BaoWsRegistry::new(tx);
682 let sender = NopSender;
683 let r = reg
685 .dispatch_message(
686 &msg(
687 "Target.attachToTarget",
688 json!({"targetId": "7", "flatten": true}),
689 None,
690 ),
691 BROWSER_TARGET,
692 &sender,
693 )
694 .unwrap()
695 .unwrap();
696 let sid = r["sessionId"].as_str().unwrap().to_string();
697 reg.dispatch_message(
698 &msg(
699 "Target.detachFromTarget",
700 json!({"sessionId": sid.clone()}),
701 None,
702 ),
703 BROWSER_TARGET,
704 &sender,
705 )
706 .unwrap()
707 .unwrap();
708 let err = reg
709 .dispatch_message(&msg("Page.enable", json!({}), Some(sid)), BROWSER_TARGET, &sender)
710 .unwrap()
711 .unwrap_err();
712 assert_eq!(err.code, -32001);
713 }
714
715 #[test]
716 fn page_session_target_used_without_session_id() {
717 let (tx, rx) = bridge_channel(Duration::from_secs(2));
718 let _keeper = page_responder(rx);
719 let reg = BaoWsRegistry::new(tx);
720 let sender = NopSender;
721 let r = reg
723 .dispatch_message(
724 &msg("Runtime.evaluate", json!({"expression": "1+1"}), None),
725 "1",
726 &sender,
727 )
728 .unwrap()
729 .unwrap();
730 assert_eq!(r["result"]["value"], "1+1");
731 }
732
733 #[test]
734 fn fetch_domain_is_explicit_error() {
735 let (tx, _rx) = bridge_channel(Duration::from_millis(100));
736 let reg = BaoWsRegistry::new(tx);
737 let sender = NopSender;
738 let err = reg
739 .dispatch_message(
740 &msg(
741 "Fetch.enable",
742 json!({"patterns": [{"urlPattern": "*"}]}),
743 None,
744 ),
745 "1",
746 &sender,
747 )
748 .unwrap()
749 .unwrap_err();
750 assert!(err.message.contains("no request interception facility"));
751 }
752
753 #[test]
754 fn has_domain_served_domains() {
755 let (tx, _rx) = bridge_channel(Duration::from_millis(100));
756 let reg = BaoWsRegistry::new(tx);
757 assert!(reg.has_domain("Page"));
758 assert!(reg.has_domain("Runtime"));
759 assert!(reg.has_domain("Target"));
760 assert!(reg.has_domain("Browser"));
761 assert!(!reg.has_domain("NotADomain"));
762 }
763
764 #[test]
765 fn set_auto_attach_emits_attached_to_target_for_existing_pages() {
766 let (tx, rx) = bridge_channel(Duration::from_secs(2));
767 let _keeper = page_responder(rx);
768 let reg = BaoWsRegistry::new(tx);
769 let sender = CapturingSender::new();
770
771 reg.dispatch_message(
772 &msg("Target.setAutoAttach", json!({"autoAttach": true, "flatten": true}), None),
773 BROWSER_TARGET,
774 &*sender,
775 )
776 .unwrap()
777 .unwrap();
778
779 let events = sender.events.lock().unwrap();
780 let attach_events: Vec<_> = events
781 .iter()
782 .filter(|(m, _)| m == "Target.attachedToTarget")
783 .collect();
784 assert_eq!(attach_events.len(), 1, "one event per listed page");
785 let (_, params) = &attach_events[0];
786 assert_eq!(params["targetInfo"]["targetId"], "1");
787 assert!(params["sessionId"].as_str().is_some());
788
789 let sid = params["sessionId"].as_str().unwrap().to_string();
791 drop(events);
792 let r = reg
793 .dispatch_message(
794 &msg("Runtime.evaluate", json!({"expression": "x"}), Some(sid)),
795 BROWSER_TARGET,
796 &*sender,
797 )
798 .unwrap()
799 .unwrap();
800 assert!(r["result"].is_object());
801 }
802
803 #[test]
804 fn runtime_enable_emits_execution_context_created_on_session() {
805 let (tx, rx) = bridge_channel(Duration::from_secs(2));
806 let _keeper = page_responder(rx);
807 let reg = BaoWsRegistry::new(tx);
808 let sender = CapturingSender::new();
809
810 let attach = reg
811 .dispatch_message(
812 &msg(
813 "Target.attachToTarget",
814 json!({"targetId": "1", "flatten": true}),
815 None,
816 ),
817 BROWSER_TARGET,
818 &*sender,
819 )
820 .unwrap()
821 .unwrap();
822 let sid = attach["sessionId"].as_str().unwrap().to_string();
823
824 reg.dispatch_message(
825 &msg("Runtime.enable", json!({}), Some(sid.clone())),
826 BROWSER_TARGET,
827 &*sender,
828 )
829 .unwrap()
830 .unwrap();
831
832 let session_events = sender.session_events.lock().unwrap();
833 let ctx_events: Vec<_> = session_events
834 .iter()
835 .filter(|(s, m, _)| s == &sid && m == "Runtime.executionContextCreated")
836 .collect();
837 assert_eq!(ctx_events.len(), 1);
838 assert!(ctx_events[0].2["context"]["id"].as_u64().is_some());
839 }
840
841 #[test]
842 fn navigate_emits_frame_lifecycle_events_on_session() {
843 let (tx, rx) = bridge_channel(Duration::from_secs(2));
844 let _keeper = page_responder(rx);
845 let reg = BaoWsRegistry::new(tx);
846 let sender = CapturingSender::new();
847
848 let attach = reg
849 .dispatch_message(
850 &msg(
851 "Target.attachToTarget",
852 json!({"targetId": "1", "flatten": true}),
853 None,
854 ),
855 BROWSER_TARGET,
856 &*sender,
857 )
858 .unwrap()
859 .unwrap();
860 let sid = attach["sessionId"].as_str().unwrap().to_string();
861
862 reg.dispatch_message(
863 &msg("Page.navigate", json!({"url": "https://example.com"}), Some(sid.clone())),
864 BROWSER_TARGET,
865 &*sender,
866 )
867 .unwrap()
868 .unwrap();
869
870 let session_events = sender.session_events.lock().unwrap();
871 let methods: Vec<&str> = session_events
872 .iter()
873 .filter(|(s, _, _)| s == &sid)
874 .map(|(_, m, _)| m.as_str())
875 .collect();
876 assert!(methods.contains(&"Page.frameStartedLoading"));
877 assert!(methods.contains(&"Page.frameNavigated"));
878 let nav = session_events
879 .iter()
880 .find(|(_, m, _)| m == "Page.frameNavigated")
881 .unwrap();
882 assert_eq!(nav.2["frame"]["url"], "https://example.com");
883 assert_eq!(nav.2["frame"]["id"], "1");
884 }
885}