1use super::{BootstrapExtensionContext, HandlerPlugin};
14use crate::{ResponseWriter, Result};
15use std::collections::HashSet;
16use webui_protocol::WebUIProtocol;
17
18const REPEAT_START: &str = "<!--wr-->";
19const REPEAT_END: &str = "<!--/wr-->";
20const REPEAT_ITEM: &str = "<!--wi-->";
21const COND_START: &str = "<!--wc-->";
22const COND_END: &str = "<!--/wc-->";
23const RAW_START_PREFIX: &str = "<!--w";
24const RAW_END_PREFIX: &str = "<!--/w";
25const MARKER_SUFFIX: &str = "-->";
26
27pub struct WebUIHydrationPlugin {
33 raw_index: usize,
34 raw_marker: String,
35}
36
37impl WebUIHydrationPlugin {
38 #[must_use]
40 pub fn new() -> Self {
41 Self {
42 raw_index: 0,
43 raw_marker: String::new(),
44 }
45 }
46}
47
48fn write_raw_marker(
49 marker: &mut String,
50 writer: &mut dyn ResponseWriter,
51 index: usize,
52 closing: bool,
53) -> Result<()> {
54 let mut reversed_digits = [0_u8; 20];
55 let mut digit_count = 0;
56 let mut remaining = index;
57 loop {
58 reversed_digits[digit_count] = b"0123456789"[remaining % 10];
59 digit_count += 1;
60 remaining /= 10;
61 if remaining == 0 {
62 break;
63 }
64 }
65
66 marker.clear();
67 if marker.capacity() == 0 {
68 marker.reserve(32);
69 }
70 marker.push_str(if closing {
71 RAW_END_PREFIX
72 } else {
73 RAW_START_PREFIX
74 });
75 for index in (0..digit_count).rev() {
76 marker.push(char::from(reversed_digits[index]));
77 }
78 marker.push_str(MARKER_SUFFIX);
79 writer.write(marker)
80}
81
82impl Default for WebUIHydrationPlugin {
83 fn default() -> Self {
84 Self::new()
85 }
86}
87
88impl HandlerPlugin for WebUIHydrationPlugin {
89 fn push_scope(&mut self) {}
90
91 fn pop_scope(&mut self) {}
92
93 fn on_binding_start(
94 &mut self,
95 _name: &str,
96 raw: bool,
97 writer: &mut dyn ResponseWriter,
98 ) -> Result<()> {
99 if raw {
100 write_raw_marker(&mut self.raw_marker, writer, self.raw_index, false)
101 } else {
102 Ok(())
103 }
104 }
105
106 fn on_binding_end(
107 &mut self,
108 _name: &str,
109 raw: bool,
110 writer: &mut dyn ResponseWriter,
111 ) -> Result<()> {
112 if raw {
113 write_raw_marker(&mut self.raw_marker, writer, self.raw_index, true)?;
114 self.raw_index += 1;
115 }
116 Ok(())
117 }
118
119 fn on_for_start(&mut self, _name: &str, writer: &mut dyn ResponseWriter) -> Result<()> {
120 writer.write(REPEAT_START)
121 }
122
123 fn on_for_end(&mut self, _name: &str, writer: &mut dyn ResponseWriter) -> Result<()> {
124 writer.write(REPEAT_END)
125 }
126
127 fn on_if_start(&mut self, _name: &str, writer: &mut dyn ResponseWriter) -> Result<()> {
128 writer.write(COND_START)
129 }
130
131 fn on_if_end(&mut self, _name: &str, writer: &mut dyn ResponseWriter) -> Result<()> {
132 writer.write(COND_END)
133 }
134
135 fn on_repeat_item_start(
136 &mut self,
137 _index: usize,
138 writer: &mut dyn ResponseWriter,
139 ) -> Result<()> {
140 writer.write(REPEAT_ITEM)
141 }
142
143 fn on_repeat_item_end(
145 &mut self,
146 _index: usize,
147 _writer: &mut dyn ResponseWriter,
148 ) -> Result<()> {
149 Ok(())
150 }
151
152 fn on_element_data(&mut self, _data: &[u8], _writer: &mut dyn ResponseWriter) -> Result<()> {
153 Ok(())
154 }
155
156 fn emit_templates(
163 &self,
164 protocol: &WebUIProtocol,
165 components: &HashSet<String>,
166 nonce: Option<&str>,
167 writer: &mut dyn ResponseWriter,
168 ) -> Result<()> {
169 webui_emit_templates(
170 protocol,
171 components.iter().map(String::as_str),
172 nonce,
173 writer,
174 )
175 }
176
177 fn emit_templates_slice(
178 &self,
179 protocol: &WebUIProtocol,
180 tags: &[&str],
181 nonce: Option<&str>,
182 writer: &mut dyn ResponseWriter,
183 ) -> Result<()> {
184 webui_emit_templates(protocol, tags.iter().copied(), nonce, writer)
185 }
186
187 fn collect_template_payloads<'a>(
189 &self,
190 protocol: &'a WebUIProtocol,
191 components: &HashSet<String>,
192 ) -> Option<Vec<super::WebUiTemplatePayload<'a>>> {
193 webui_collect_payloads(protocol, components.iter().map(String::as_str))
194 }
195
196 fn collect_template_payloads_slice<'a>(
197 &self,
198 protocol: &'a WebUIProtocol,
199 tags: &[&str],
200 ) -> Option<Vec<super::WebUiTemplatePayload<'a>>> {
201 webui_collect_payloads(protocol, tags.iter().copied())
202 }
203
204 fn emit_bootstrap_extension(
205 &self,
206 context: BootstrapExtensionContext<'_>,
207 writer: &mut dyn ResponseWriter,
208 ) -> Result<()> {
209 webui_emit_bootstrap_fns(context.payloads, context.nonce, writer)
210 }
211
212 fn emit_bootstrap_extension_payloads(
213 &self,
214 payloads: &[super::WebUiTemplatePayload<'_>],
215 nonce: Option<&str>,
216 writer: &mut dyn ResponseWriter,
217 ) -> Result<()> {
218 webui_emit_bootstrap_fns(payloads, nonce, writer)
219 }
220}
221
222fn webui_emit_templates<'b>(
228 protocol: &WebUIProtocol,
229 tags: impl Iterator<Item = &'b str>,
230 nonce: Option<&str>,
231 writer: &mut dyn ResponseWriter,
232) -> Result<()> {
233 let mut templates: Vec<&str> = Vec::new();
234
235 for name in tags {
236 if let Some(template) = protocol
237 .components
238 .get(name)
239 .map(|component| component.template.as_str())
240 .filter(|t| !t.is_empty())
241 {
242 templates.push(template);
243 }
244 }
245
246 if templates.is_empty() {
247 return Ok(());
248 }
249
250 if let Some(nonce) = nonce {
251 writer.write("<script nonce=\"")?;
252 writer.write(nonce)?;
253 writer.write("\">\n")?;
254 } else {
255 writer.write("<script>\n")?;
256 }
257 for tmpl in &templates {
258 writer.write(tmpl)?;
259 }
260 writer.write("</script>\n")?;
261
262 Ok(())
263}
264
265fn webui_collect_payloads<'a, 'b>(
270 protocol: &'a WebUIProtocol,
271 tags: impl Iterator<Item = &'b str>,
272) -> Option<Vec<super::WebUiTemplatePayload<'a>>> {
273 let mut templates: Vec<super::WebUiTemplatePayload<'a>> = Vec::new();
274 for name in tags {
275 if let Some((tag_name, component)) = protocol.components.get_key_value(name) {
276 if !component.template_json.is_empty() {
277 templates.push(super::WebUiTemplatePayload {
278 tag_name: tag_name.as_str(),
279 template_json: component.template_json.as_str(),
280 template_functions: component.template_functions.as_str(),
281 });
282 }
283 }
284 }
285 if templates.is_empty() {
286 None
287 } else {
288 Some(templates)
289 }
290}
291
292fn webui_emit_bootstrap_fns(
296 payloads: &[super::WebUiTemplatePayload<'_>],
297 nonce: Option<&str>,
298 writer: &mut dyn ResponseWriter,
299) -> Result<()> {
300 let has_functions = payloads
301 .iter()
302 .any(|payload| !payload.template_functions.is_empty());
303 if !has_functions {
304 return Ok(());
305 }
306
307 if let Some(nonce) = nonce {
308 writer.write("<script nonce=\"")?;
309 writer.write(nonce)?;
310 writer.write("\">")?;
311 } else {
312 writer.write("<script>")?;
313 }
314 writer.write("(function(){var w=window.__webui||(window.__webui={});")?;
315 writer.write("var f=w.templateFns||(w.templateFns={});")?;
316 let mut json_scratch = Vec::new();
319 for payload in payloads {
320 if payload.template_functions.is_empty() {
321 continue;
322 }
323 writer.write("f[")?;
324 crate::write_script_safe_json(writer, &mut json_scratch, payload.tag_name)?;
325 writer.write("]=")?;
326 writer.write(payload.template_functions)?;
327 writer.write(";")?;
328 }
329 writer.write("})();</script>\n")
330}
331
332#[cfg(test)]
333mod tests {
334 use super::*;
335
336 struct TestWriter {
337 output: String,
338 }
339
340 impl TestWriter {
341 fn new() -> Self {
342 Self {
343 output: String::new(),
344 }
345 }
346 }
347
348 impl ResponseWriter for TestWriter {
349 fn write(&mut self, content: &str) -> Result<()> {
350 self.output.push_str(content);
351 Ok(())
352 }
353 fn end(&mut self) -> Result<()> {
354 Ok(())
355 }
356 }
357
358 #[test]
359 fn test_default_creates_instance() {
360 let _plugin = WebUIHydrationPlugin::default();
361 }
362
363 #[test]
364 fn test_signal_binding_emits_no_output() {
365 let mut plugin = WebUIHydrationPlugin::new();
366 let mut writer = TestWriter::new();
367 plugin.on_binding_start("x", false, &mut writer).unwrap();
368 plugin.on_binding_end("x", false, &mut writer).unwrap();
369 assert_eq!(
370 writer.output, "",
371 "signal binding hooks must not emit output"
372 );
373 }
374
375 #[test]
376 fn test_raw_signal_binding_emits_range_markers() {
377 let mut plugin = WebUIHydrationPlugin::new();
378 let mut writer = TestWriter::new();
379 plugin.on_binding_start("html", true, &mut writer).unwrap();
380 writer.write("<b>trusted</b>").unwrap();
381 plugin.on_binding_end("html", true, &mut writer).unwrap();
382 assert_eq!(writer.output, "<!--w0--><b>trusted</b><!--/w0-->");
383 }
384
385 #[test]
386 fn test_for_loop_emits_repeat_markers() {
387 let mut plugin = WebUIHydrationPlugin::new();
388 let mut writer = TestWriter::new();
389 plugin.on_for_start("items", &mut writer).unwrap();
390 plugin.on_repeat_item_start(0, &mut writer).unwrap();
391 plugin.push_scope();
392 writer.write("<div>A</div>").unwrap();
393 plugin.pop_scope();
394 plugin.on_repeat_item_end(0, &mut writer).unwrap();
395 plugin.on_repeat_item_start(1, &mut writer).unwrap();
396 plugin.push_scope();
397 writer.write("<div>B</div>").unwrap();
398 plugin.pop_scope();
399 plugin.on_repeat_item_end(1, &mut writer).unwrap();
400 plugin.on_for_end("items", &mut writer).unwrap();
401 assert_eq!(
402 writer.output,
403 "<!--wr--><!--wi--><div>A</div><!--wi--><div>B</div><!--/wr-->"
404 );
405 }
406
407 #[test]
408 fn test_empty_for_loop_emits_boundary_markers() {
409 let mut plugin = WebUIHydrationPlugin::new();
410 let mut writer = TestWriter::new();
411 plugin.on_for_start("items", &mut writer).unwrap();
412 plugin.on_for_end("items", &mut writer).unwrap();
413 assert_eq!(writer.output, "<!--wr--><!--/wr-->");
414 }
415
416 #[test]
417 fn test_if_true_emits_cond_markers() {
418 let mut plugin = WebUIHydrationPlugin::new();
419 let mut writer = TestWriter::new();
420 plugin.on_if_start("show", &mut writer).unwrap();
421 plugin.push_scope();
422 writer.write("<p>visible</p>").unwrap();
423 plugin.pop_scope();
424 plugin.on_if_end("show", &mut writer).unwrap();
425 assert_eq!(writer.output, "<!--wc--><p>visible</p><!--/wc-->");
426 }
427
428 #[test]
429 fn test_if_false_emits_empty_cond_markers() {
430 let mut plugin = WebUIHydrationPlugin::new();
431 let mut writer = TestWriter::new();
432 plugin.on_if_start("show", &mut writer).unwrap();
433 plugin.on_if_end("show", &mut writer).unwrap();
435 assert_eq!(writer.output, "<!--wc--><!--/wc-->");
436 }
437
438 #[test]
439 fn test_nested_repeat_inside_conditional() {
440 let mut plugin = WebUIHydrationPlugin::new();
441 let mut writer = TestWriter::new();
442 plugin.on_if_start("show", &mut writer).unwrap();
443 plugin.push_scope();
444 plugin.on_for_start("items", &mut writer).unwrap();
445 plugin.on_repeat_item_start(0, &mut writer).unwrap();
446 plugin.push_scope();
447 writer.write("<li>X</li>").unwrap();
448 plugin.pop_scope();
449 plugin.on_repeat_item_end(0, &mut writer).unwrap();
450 plugin.on_for_end("items", &mut writer).unwrap();
451 plugin.pop_scope();
452 plugin.on_if_end("show", &mut writer).unwrap();
453 assert_eq!(
454 writer.output,
455 "<!--wc--><!--wr--><!--wi--><li>X</li><!--/wr--><!--/wc-->"
456 );
457 }
458
459 #[test]
460 fn test_element_data_emits_no_output() {
461 let mut plugin = WebUIHydrationPlugin::new();
462 let mut writer = TestWriter::new();
463 plugin.push_scope();
464 let data = 3u32.to_le_bytes();
465 plugin.on_element_data(&data, &mut writer).unwrap();
466 assert_eq!(writer.output, "", "element data must not emit output");
467 }
468
469 #[test]
470 fn test_scope_push_pop_is_noop() {
471 let mut plugin = WebUIHydrationPlugin::new();
472 plugin.push_scope();
473 plugin.push_scope();
474 plugin.pop_scope();
475 plugin.pop_scope();
476 }
477
478 #[test]
479 fn test_route_component_state_emits_no_output() {
480 let plugin = WebUIHydrationPlugin::new();
481 let mut writer = TestWriter::new();
482 let state = serde_json::json!({
483 "title": "Hello",
484 "cartOpen": true,
485 "items": [{"name": "A&B"}],
486 });
487 plugin
488 .write_route_component_state(&state, &mut writer)
489 .unwrap();
490 assert_eq!(
491 writer.output, "",
492 "route component state must not emit output"
493 );
494 }
495
496 #[test]
497 fn test_full_lifecycle_with_markers() {
498 let mut plugin = WebUIHydrationPlugin::new();
499 let mut writer = TestWriter::new();
500
501 plugin.push_scope();
503 plugin.on_binding_start("a", false, &mut writer).unwrap();
505 writer.write("hello").unwrap();
506 plugin.on_binding_end("a", false, &mut writer).unwrap();
507 plugin.on_for_start("list", &mut writer).unwrap();
509 plugin.on_repeat_item_start(0, &mut writer).unwrap();
510 plugin.push_scope();
511 writer.write("<x-item>1</x-item>").unwrap();
512 plugin.pop_scope();
513 plugin.on_repeat_item_end(0, &mut writer).unwrap();
514 plugin.on_for_end("list", &mut writer).unwrap();
515 plugin.on_if_start("flag", &mut writer).unwrap();
517 plugin.push_scope();
518 writer.write("<p>yes</p>").unwrap();
519 plugin.pop_scope();
520 plugin.on_if_end("flag", &mut writer).unwrap();
521 plugin.pop_scope();
522
523 assert_eq!(
524 writer.output,
525 "hello<!--wr--><!--wi--><x-item>1</x-item><!--/wr--><!--wc--><p>yes</p><!--/wc-->"
526 );
527 }
528
529 fn iife_template(tag: &str, body: &str) -> String {
530 format!("(function(){{var w=(window.__webui||(window.__webui={{}})).templates||(window.__webui.templates={{}});w['{tag}']={{{body}}}}})();\n")
531 }
532
533 #[test]
534 fn test_emit_templates_only_emits_provided_components() {
535 let mut writer = TestWriter::new();
536
537 let mut protocol = webui_protocol::WebUIProtocol::new(std::collections::HashMap::new());
538 protocol
539 .components
540 .entry("comp-a".to_string())
541 .or_default()
542 .template = iife_template("comp-a", "h:\"a\"");
543 protocol
544 .components
545 .entry("comp-b".to_string())
546 .or_default()
547 .template = iife_template("comp-b", "h:\"b\"");
548 protocol
549 .components
550 .entry("comp-c".to_string())
551 .or_default()
552 .template = iife_template("comp-c", "h:\"c\"");
553
554 let mut components = std::collections::HashSet::new();
555 components.insert("comp-a".to_string());
556
557 let plugin = WebUIHydrationPlugin::new();
558 plugin
559 .emit_templates(&protocol, &components, None, &mut writer)
560 .unwrap();
561
562 assert!(
563 writer.output.contains("comp-a"),
564 "rendered component should be emitted: {}",
565 writer.output
566 );
567 assert!(
568 !writer.output.contains("comp-b"),
569 "non-rendered component should NOT be emitted: {}",
570 writer.output
571 );
572 assert!(
573 !writer.output.contains("comp-c"),
574 "non-rendered component should NOT be emitted: {}",
575 writer.output
576 );
577 }
578
579 #[test]
580 fn test_emit_templates_empty_set() {
581 let mut writer = TestWriter::new();
582 let mut protocol = webui_protocol::WebUIProtocol::new(std::collections::HashMap::new());
583 protocol
584 .components
585 .entry("comp-a".to_string())
586 .or_default()
587 .template = iife_template("comp-a", "h:\"a\"");
588 let components = std::collections::HashSet::new();
589 let plugin = WebUIHydrationPlugin::new();
590 plugin
591 .emit_templates(&protocol, &components, None, &mut writer)
592 .unwrap();
593 assert_eq!(writer.output, "", "empty set should emit nothing");
594 }
595
596 #[test]
597 fn test_on_render_complete_consolidates_into_single_script() {
598 let mut writer = TestWriter::new();
599
600 let mut protocol = webui_protocol::WebUIProtocol::new(std::collections::HashMap::new());
601 protocol
602 .components
603 .entry("comp-a".to_string())
604 .or_default()
605 .template = iife_template("comp-a", "h:\"a\"");
606 protocol
607 .components
608 .entry("comp-b".to_string())
609 .or_default()
610 .template = iife_template("comp-b", "h:\"b\"");
611
612 let mut rendered = std::collections::HashSet::new();
613 rendered.insert("comp-a".to_string());
614 rendered.insert("comp-b".to_string());
615
616 let plugin = WebUIHydrationPlugin::new();
617 plugin
618 .emit_templates(&protocol, &rendered, None, &mut writer)
619 .unwrap();
620
621 let script_count = writer.output.matches("<script").count();
622 assert_eq!(
623 script_count, 1,
624 "should emit exactly one <script> tag, got {script_count}: {}",
625 writer.output
626 );
627 assert!(writer.output.contains("comp-a"));
628 assert!(writer.output.contains("comp-b"));
629 }
630
631 #[test]
632 fn test_on_render_complete_skips_empty_template() {
633 let mut writer = TestWriter::new();
634 let mut protocol = webui_protocol::WebUIProtocol::new(std::collections::HashMap::new());
635 protocol
636 .components
637 .entry("comp-a".to_string())
638 .or_default()
639 .template = String::new();
640 let mut rendered = std::collections::HashSet::new();
641 rendered.insert("comp-a".to_string());
642 let plugin = WebUIHydrationPlugin::new();
643 plugin
644 .emit_templates(&protocol, &rendered, None, &mut writer)
645 .unwrap();
646 assert_eq!(writer.output, "", "empty template should not be emitted");
647 }
648
649 #[test]
650 fn test_on_render_complete_unknown_component() {
651 let mut writer = TestWriter::new();
652 let protocol = webui_protocol::WebUIProtocol::new(std::collections::HashMap::new());
653 let mut rendered = std::collections::HashSet::new();
654 rendered.insert("nonexistent-comp".to_string());
655 let plugin = WebUIHydrationPlugin::new();
656 plugin
657 .emit_templates(&protocol, &rendered, None, &mut writer)
658 .unwrap();
659 assert_eq!(
660 writer.output, "",
661 "unknown component should not cause error"
662 );
663 }
664
665 use std::collections::HashMap;
668 use webui_protocol::{ConditionExpr, FragmentList, WebUIFragment, WebUIProtocol};
669 use webui_test_utils::test_json;
670
671 use crate::{RenderOptions, WebUIHandler};
672
673 fn render_with_webui_plugin(protocol: &WebUIProtocol, state: &serde_json::Value) -> String {
674 let mut writer = TestWriter::new();
675 let handler = WebUIHandler::with_plugin(|| Box::new(WebUIHydrationPlugin::new()));
676 handler
677 .handle(
678 protocol,
679 state,
680 &RenderOptions::new("index.html", "/"),
681 &mut writer,
682 )
683 .unwrap();
684 writer.output
685 }
686
687 #[test]
688 fn test_handler_emits_hydration_markers_for_loop_and_if() {
689 let mut fragments = HashMap::new();
691 fragments.insert(
692 "index.html".to_string(),
693 FragmentList {
694 fragments: vec![
695 WebUIFragment::for_loop("item", "items", "for-body"),
696 WebUIFragment::if_cond(ConditionExpr::identifier("show"), "if-body"),
697 WebUIFragment::signal("html", true),
698 ],
699 contains_boundary: false,
700 },
701 );
702 fragments.insert(
703 "for-body".to_string(),
704 FragmentList {
705 fragments: vec![WebUIFragment::signal("item", false)],
706 contains_boundary: false,
707 },
708 );
709 fragments.insert(
710 "if-body".to_string(),
711 FragmentList {
712 fragments: vec![WebUIFragment::raw("<span>yes</span>")],
713 contains_boundary: false,
714 },
715 );
716 let protocol = WebUIProtocol::new(fragments);
717 let state = test_json!({"items": ["a", "b"], "show": true, "html": "<i>raw</i>"});
718 let output = render_with_webui_plugin(&protocol, &state);
719
720 assert!(
722 output.contains("<!--wr-->"),
723 "Expected repeat-start marker, got: {output}"
724 );
725 assert!(
726 output.contains("<!--/wr-->"),
727 "Expected repeat-end marker, got: {output}"
728 );
729 assert!(
730 output.contains("<!--wi-->"),
731 "Expected repeat-item marker, got: {output}"
732 );
733 assert_eq!(
735 output.matches("<!--wi-->").count(),
736 2,
737 "Expected 2 repeat-item markers, got: {output}"
738 );
739
740 assert!(
742 output.contains("<!--wc-->"),
743 "Expected cond-start marker, got: {output}"
744 );
745 assert!(
746 output.contains("<!--/wc-->"),
747 "Expected cond-end marker, got: {output}"
748 );
749
750 assert!(
752 output.contains("a"),
753 "Expected for-loop items in output, got: {output}"
754 );
755 assert!(
756 output.contains("b"),
757 "Expected for-loop items in output, got: {output}"
758 );
759 assert!(
760 output.contains("<span>yes</span>"),
761 "Expected if-condition body, got: {output}"
762 );
763 assert!(
764 output.contains("<!--w0--><i>raw</i><!--/w0-->"),
765 "Expected bounded raw HTML markers, got: {output}"
766 );
767 }
768
769 #[test]
770 fn test_handler_keeps_raw_text_signal_marker_free() {
771 let mut fragments = HashMap::new();
772 fragments.insert(
773 "index.html".to_string(),
774 FragmentList {
775 fragments: vec![
776 WebUIFragment::raw("<style>:root{"),
777 WebUIFragment::raw_text_signal("tokens", true),
778 WebUIFragment::raw("}</style>"),
779 ],
780 contains_boundary: false,
781 },
782 );
783 let protocol = WebUIProtocol::new(fragments);
784 let state = test_json!({"tokens": "--accent:red"});
785 let output = render_with_webui_plugin(&protocol, &state);
786
787 assert!(output.contains("<style>:root{--accent:red}</style>"));
788 assert!(!output.contains("<!--w0-->"));
789 assert!(!output.contains("<!--/w0-->"));
790 }
791}