1#[cfg(feature = "anomaly-detection")]
7use crate::anomaly::AnomalyDetector;
8use crate::config::{NodeConfig, NodeRole, ProxyConfig};
9#[cfg(feature = "edge-proxy")]
10use crate::edge::{EdgeCache, EdgeRegistry, InvalidationEvent};
11#[cfg(feature = "wasm-plugins")]
12use crate::plugins::PluginManager;
13#[cfg(feature = "ha-tr")]
14use crate::replay::{ReplayEngine, TimeTravelRequest};
15use crate::server::{NodeHealth, ServerMetricsSnapshot};
16use crate::{ProxyError, Result};
17#[cfg(feature = "ha-tr")]
18use chrono::{DateTime, Utc};
19use serde::{Deserialize, Serialize};
20use std::collections::HashMap;
21use std::net::SocketAddr;
22use std::sync::atomic::{AtomicUsize, Ordering};
23use std::sync::Arc;
24use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
25use tokio::net::TcpStream;
26use tokio::sync::{broadcast, RwLock};
27
28const ADMIN_UI_HTML: &str = include_str!("admin_ui.html");
32
33fn constant_time_eq_str(a: &str, b: &str) -> bool {
36 let (a, b) = (a.as_bytes(), b.as_bytes());
37 if a.len() != b.len() {
38 return false;
39 }
40 let mut diff = 0u8;
41 for i in 0..a.len() {
42 diff |= a[i] ^ b[i];
43 }
44 diff == 0
45}
46
47pub struct AdminServer {
48 listen_address: String,
50 state: Arc<AdminState>,
52 shutdown_tx: broadcast::Sender<()>,
54}
55
56pub struct AdminState {
58 pub node_health: RwLock<HashMap<String, NodeHealth>>,
60 pub metrics: RwLock<ServerMetricsSnapshot>,
62 pub active_sessions: RwLock<u64>,
64 pub config_snapshot: RwLock<ConfigSnapshot>,
66 pub proxy_config: RwLock<Option<ProxyConfig>>,
68 read_lb_counter: AtomicUsize,
70 commands: RwLock<HashMap<String, CommandHandler>>,
72 #[cfg(feature = "pool-modes")]
76 pub pool_manager: RwLock<Option<Arc<crate::pool::ConnectionPoolManager>>>,
77 #[cfg(feature = "circuit-breaker")]
80 pub circuit_breaker: RwLock<Option<Arc<crate::circuit_breaker::CircuitBreakerManager>>>,
81 #[cfg(feature = "ha-tr")]
85 pub replay_engine: RwLock<Option<Arc<ReplayEngine>>>,
86 #[cfg(feature = "wasm-plugins")]
91 pub plugin_manager: RwLock<Option<Arc<PluginManager>>>,
92 pub chaos_overrides: RwLock<HashMap<String, ChaosOverride>>,
97 #[cfg(feature = "anomaly-detection")]
101 pub anomaly_detector: RwLock<Option<Arc<AnomalyDetector>>>,
102 #[cfg(feature = "query-analytics")]
105 pub analytics: RwLock<Option<Arc<crate::analytics::QueryAnalytics>>>,
106 #[cfg(feature = "edge-proxy")]
109 pub edge_cache: RwLock<Option<Arc<EdgeCache>>>,
110 #[cfg(feature = "edge-proxy")]
111 pub edge_registry: RwLock<Option<Arc<EdgeRegistry>>>,
112 pub auth_token: RwLock<Option<String>>,
115 pub migration: RwLock<Option<MigrationInfo>>,
118 pub branch: RwLock<Option<crate::config::BranchConfig>>,
121}
122
123#[derive(Clone)]
126pub struct MigrationInfo {
127 pub target: String,
128 pub writes_only: bool,
129 pub metrics: Arc<crate::mirror::MirrorMetrics>,
130 pub config: crate::config::MirrorConfig,
132 pub cutover: Arc<arc_swap::ArcSwap<Option<Arc<crate::mirror::CutoverTarget>>>>,
134 pub cutover_target: crate::mirror::CutoverTarget,
135}
136
137#[derive(Debug, Clone, Serialize)]
142pub struct ChaosOverride {
143 pub since: String,
145 pub kind: String,
147 pub note: String,
149}
150
151type CommandHandler = Arc<dyn Fn(&[&str]) -> Result<String> + Send + Sync>;
153
154#[derive(Debug, Clone, Serialize, Deserialize)]
156pub struct ConfigSnapshot {
157 pub listen_address: String,
158 pub admin_address: String,
159 pub tr_enabled: bool,
160 pub tr_mode: String,
161 pub pool_min_connections: usize,
162 pub pool_max_connections: usize,
163 pub nodes: Vec<NodeSnapshot>,
164}
165
166#[derive(Debug, Clone, Serialize, Deserialize)]
168pub struct NodeSnapshot {
169 pub address: String,
170 pub role: String,
171 pub weight: u32,
172 pub enabled: bool,
173}
174
175impl AdminServer {
176 pub fn new(listen_address: String, state: Arc<AdminState>) -> Self {
178 let (shutdown_tx, _) = broadcast::channel(1);
179
180 Self {
181 listen_address,
182 state,
183 shutdown_tx,
184 }
185 }
186
187 pub async fn run(&self) -> Result<()> {
189 let listener = crate::server::bind_reuseport(&self.listen_address)?;
192
193 tracing::info!(
194 "Admin API listening on {} (SO_REUSEPORT)",
195 self.listen_address
196 );
197
198 let mut shutdown_rx = self.shutdown_tx.subscribe();
199 let conn_limit = std::sync::Arc::new(tokio::sync::Semaphore::new(Self::MAX_ADMIN_CONNS));
203
204 loop {
205 tokio::select! {
206 accept_result = listener.accept() => {
207 match accept_result {
208 Ok((stream, addr)) => {
209 let permit = match conn_limit.clone().try_acquire_owned() {
210 Ok(p) => p,
211 Err(_) => {
212 tracing::warn!(%addr, "admin connection limit reached; dropping");
213 drop(stream);
214 continue;
215 }
216 };
217 let state = self.state.clone();
218 tokio::spawn(async move {
219 let _permit = permit; if let Err(e) = Self::handle_connection(stream, addr, state).await {
221 tracing::error!("Admin connection error: {}", e);
222 }
223 });
224 }
225 Err(e) => {
226 tracing::error!("Admin accept error: {}", e);
227 }
228 }
229 }
230 _ = shutdown_rx.recv() => {
231 tracing::info!("Admin server shutting down");
232 break;
233 }
234 }
235 }
236
237 Ok(())
238 }
239
240 const MAX_ADMIN_CONNS: usize = 256;
244 const ADMIN_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15);
245 #[cfg(feature = "edge-proxy")]
251 const ADMIN_SSE_WRITE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
252 const MAX_ADMIN_HEADERS: usize = 100;
254 const MAX_ADMIN_HEADER_BYTES: usize = 64 * 1024;
256 const MAX_ADMIN_BODY_BYTES: usize = 8 * 1024 * 1024;
259
260 async fn handle_connection(
262 mut stream: TcpStream,
263 addr: SocketAddr,
264 state: Arc<AdminState>,
265 ) -> Result<()> {
266 tracing::debug!("Admin connection from {}", addr);
267
268 let (reader, mut writer) = stream.split();
269 let mut reader = BufReader::new(reader);
270 let mut line = String::new();
271
272 let mut headers = Vec::new();
274 let mut content_length: usize = 0;
275 let mut header_bytes: usize = 0;
276
277 loop {
278 line.clear();
279 let bytes_read =
283 match tokio::time::timeout(Self::ADMIN_READ_TIMEOUT, reader.read_line(&mut line))
284 .await
285 {
286 Ok(r) => r.map_err(|e| ProxyError::Network(format!("Read error: {}", e)))?,
287 Err(_) => return Ok(()), };
289
290 if bytes_read == 0 || line == "\r\n" {
291 break;
292 }
293
294 header_bytes += bytes_read;
297 if headers.len() >= Self::MAX_ADMIN_HEADERS
298 || header_bytes > Self::MAX_ADMIN_HEADER_BYTES
299 {
300 Self::send_response(
301 &mut writer,
302 431,
303 "Request Header Fields Too Large",
304 "header section too large",
305 )
306 .await?;
307 return Ok(());
308 }
309
310 let trimmed = line.trim();
312 if trimmed.to_lowercase().starts_with("content-length:") {
313 if let Some(len_str) = trimmed.split(':').nth(1) {
314 content_length = len_str.trim().parse().unwrap_or(0);
315 }
316 }
317 headers.push(trimmed.to_string());
318 }
319
320 if headers.is_empty() {
321 return Ok(());
322 }
323
324 if content_length > Self::MAX_ADMIN_BODY_BYTES {
329 Self::send_response(
330 &mut writer,
331 413,
332 "Payload Too Large",
333 "request body exceeds admin size limit",
334 )
335 .await?;
336 return Ok(());
337 }
338
339 let request_line = &headers[0];
341 let parts: Vec<&str> = request_line.split_whitespace().collect();
342
343 if parts.len() < 2 {
344 Self::send_response(&mut writer, 400, "Bad Request", "Invalid request line").await?;
345 return Ok(());
346 }
347
348 let method = parts[0];
349 let path = parts[1];
350
351 {
359 let required = state.auth_token.read().await.clone();
360 if let Some(token) = required {
361 let path_only = path.split('?').next().unwrap_or(path);
362 let is_open = method == "GET"
363 && matches!(
364 path_only,
365 "/health" | "/healthz" | "/livez" | "/readyz" | "/" | "/ui" | "/ui/"
366 );
367 if !is_open && !Self::admin_authorized(&headers, &token) {
368 Self::send_response(
369 &mut writer,
370 401,
371 "Unauthorized",
372 "{\"error\":\"missing or invalid admin bearer token\"}",
373 )
374 .await?;
375 return Ok(());
376 }
377 }
378 }
379
380 #[cfg(feature = "edge-proxy")]
393 if method == "GET" && path.split('?').next().unwrap_or(path) == "/api/edge/subscribe" {
394 let params = parse_query_params(path);
395 let edge_id = params.get("edge_id").map(String::as_str).unwrap_or("");
396 if edge_id.is_empty() {
397 Self::send_json_response(
398 &mut writer,
399 400,
400 &serde_json::json!({ "error": "edge_id query parameter is required" }),
401 )
402 .await?;
403 return Ok(());
404 }
405 let region = params.get("region").map(String::as_str).unwrap_or("");
406 let base_url = params.get("base_url").map(String::as_str).unwrap_or("");
407 return Self::handle_edge_subscribe(&mut writer, &state, edge_id, region, base_url)
408 .await;
409 }
410
411 let body = if content_length > 0 && (method == "POST" || method == "PUT") {
413 let mut body_buf = vec![0u8; content_length];
414 match tokio::time::timeout(Self::ADMIN_READ_TIMEOUT, reader.read_exact(&mut body_buf))
415 .await
416 {
417 Ok(r) => {
418 r.map_err(|e| ProxyError::Network(format!("Body read error: {}", e)))?;
419 }
420 Err(_) => return Ok(()), }
422 Some(String::from_utf8_lossy(&body_buf).to_string())
423 } else {
424 None
425 };
426
427 if method == "GET" && (path == "/" || path == "/ui" || path == "/ui/") {
430 Self::send_html_response(&mut writer, 200, ADMIN_UI_HTML).await?;
431 return Ok(());
432 }
433
434 let response = Self::route_request(method, path, body.as_deref(), &state).await;
436
437 match response {
438 Ok((status, body)) => {
439 Self::send_json_response(&mut writer, status, &body).await?;
440 }
441 Err(e) => {
442 let error = ErrorResponse {
443 error: e.to_string(),
444 };
445 Self::send_json_response(&mut writer, 500, &error).await?;
446 }
447 }
448
449 Ok(())
450 }
451
452 fn admin_authorized(headers: &[String], token: &str) -> bool {
455 let expected = format!("Bearer {}", token);
456 for h in headers {
457 let mut sp = h.splitn(2, ':');
458 let name = sp.next().unwrap_or("").trim();
459 if name.eq_ignore_ascii_case("authorization") {
460 let value = sp.next().unwrap_or("").trim();
461 return constant_time_eq_str(value, &expected);
462 }
463 }
464 false
465 }
466
467 async fn send_html_response(
469 writer: &mut tokio::net::tcp::WriteHalf<'_>,
470 status: u16,
471 html: &str,
472 ) -> Result<()> {
473 let status_text = match status {
474 200 => "OK",
475 404 => "Not Found",
476 _ => "Unknown",
477 };
478 let response = format!(
479 "HTTP/1.1 {} {}\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
480 status,
481 status_text,
482 html.len(),
483 html
484 );
485 writer
486 .write_all(response.as_bytes())
487 .await
488 .map_err(|e| ProxyError::Network(format!("Write error: {}", e)))?;
489 Ok(())
490 }
491
492 async fn route_request(
494 method: &str,
495 path: &str,
496 body: Option<&str>,
497 state: &Arc<AdminState>,
498 ) -> Result<(u16, serde_json::Value)> {
499 match (method, path) {
500 ("POST", "/api/sql") => Self::handle_sql_request(body, state).await,
502
503 ("GET", "/health") | ("GET", "/healthz") => {
505 let health = HealthResponse { status: "ok" };
506 Ok((200, serde_json::to_value(health)?))
507 }
508 ("GET", "/health/ready") | ("GET", "/readyz") => {
509 let ready = Self::check_readiness(state).await;
510 let response = ReadinessResponse {
511 ready,
512 message: if ready {
513 "Proxy is ready"
514 } else {
515 "Proxy is not ready"
516 },
517 };
518 let status = if ready { 200 } else { 503 };
519 Ok((status, serde_json::to_value(response)?))
520 }
521 ("GET", "/health/live") | ("GET", "/livez") => {
522 let response = LivenessResponse { alive: true };
523 Ok((200, serde_json::to_value(response)?))
524 }
525
526 ("GET", "/metrics") => {
528 let metrics = state.metrics.read().await.clone();
529 Ok((200, serde_json::to_value(MetricsResponse::from(metrics))?))
530 }
531 ("GET", "/metrics/prometheus") => {
532 let metrics = state.metrics.read().await.clone();
533 let prometheus = Self::format_prometheus_metrics(&metrics);
534 Ok((200, serde_json::json!({ "text": prometheus })))
535 }
536
537 ("GET", "/nodes") => {
539 let health = state.node_health.read().await;
540 let nodes: Vec<NodeHealthResponse> = health
541 .values()
542 .map(|h| NodeHealthResponse::from(h.clone()))
543 .collect();
544 Ok((200, serde_json::to_value(nodes)?))
545 }
546 ("GET", path) if path.starts_with("/nodes/") => {
547 let node_addr = path.trim_start_matches("/nodes/");
548 let health = state.node_health.read().await;
549 match health.get(node_addr) {
550 Some(h) => Ok((
551 200,
552 serde_json::to_value(NodeHealthResponse::from(h.clone()))?,
553 )),
554 None => Ok((404, serde_json::json!({ "error": "Node not found" }))),
555 }
556 }
557 ("POST", path) if path.starts_with("/nodes/") && path.ends_with("/enable") => {
558 let node_addr = path
559 .trim_start_matches("/nodes/")
560 .trim_end_matches("/enable");
561 Self::set_node_enabled(state, node_addr, true).await?;
562 Ok((200, serde_json::json!({ "status": "enabled" })))
563 }
564 ("POST", path) if path.starts_with("/nodes/") && path.ends_with("/disable") => {
565 let node_addr = path
566 .trim_start_matches("/nodes/")
567 .trim_end_matches("/disable");
568 Self::set_node_enabled(state, node_addr, false).await?;
569 Ok((200, serde_json::json!({ "status": "disabled" })))
570 }
571
572 ("GET", "/topology") => {
578 let topo = Self::compute_topology(state).await;
579 Ok((200, serde_json::to_value(topo)?))
580 }
581
582 #[cfg(feature = "ha-tr")]
586 ("POST", "/api/replay") => Self::handle_replay_request(body, state).await,
587 #[cfg(not(feature = "ha-tr"))]
588 ("POST", "/api/replay") => Ok((
589 503,
590 serde_json::json!({ "error": "ha-tr feature not compiled in" }),
591 )),
592
593 #[cfg(feature = "ha-tr")]
599 ("POST", "/api/shadow") => Self::handle_shadow_request(body).await,
600 #[cfg(not(feature = "ha-tr"))]
601 ("POST", "/api/shadow") => Ok((
602 503,
603 serde_json::json!({ "error": "ha-tr feature not compiled in" }),
604 )),
605
606 ("GET", "/plugins") => Self::handle_plugins_list(state).await,
611
612 #[cfg(feature = "anomaly-detection")]
616 ("GET", p) if p == "/anomalies" || p.starts_with("/anomalies?") => {
617 Self::handle_anomalies_list(p, state).await
618 }
619 #[cfg(not(feature = "anomaly-detection"))]
620 ("GET", p) if p == "/anomalies" || p.starts_with("/anomalies?") => Ok((
621 503,
622 serde_json::json!({ "error": "anomaly-detection feature not compiled in" }),
623 )),
624
625 #[cfg(feature = "query-analytics")]
627 ("GET", p)
628 if p == "/api/analytics"
629 || p == "/analytics"
630 || p.starts_with("/api/analytics?")
631 || p.starts_with("/analytics?") =>
632 {
633 Self::handle_analytics(p, state).await
634 }
635 #[cfg(not(feature = "query-analytics"))]
636 ("GET", p)
637 if p == "/api/analytics"
638 || p == "/analytics"
639 || p.starts_with("/api/analytics?")
640 || p.starts_with("/analytics?") =>
641 {
642 Ok((
643 503,
644 serde_json::json!({ "error": "query-analytics feature not compiled in" }),
645 ))
646 }
647
648 #[cfg(feature = "edge-proxy")]
656 ("GET", "/api/edge") => Self::handle_edge_status(state).await,
657 #[cfg(feature = "edge-proxy")]
658 ("POST", "/api/edge/register") => Self::handle_edge_register(body, state).await,
659 #[cfg(feature = "edge-proxy")]
660 ("POST", "/api/edge/invalidate") => Self::handle_edge_invalidate(body, state).await,
661 #[cfg(not(feature = "edge-proxy"))]
662 ("GET", "/api/edge")
663 | ("POST", "/api/edge/register")
664 | ("POST", "/api/edge/invalidate") => Ok((
665 503,
666 serde_json::json!({ "error": "edge-proxy feature not compiled in" }),
667 )),
668 #[cfg(not(feature = "edge-proxy"))]
672 ("GET", p) if p == "/api/edge/subscribe" || p.starts_with("/api/edge/subscribe?") => {
673 Ok((
674 503,
675 serde_json::json!({ "error": "edge-proxy feature not compiled in" }),
676 ))
677 }
678
679 ("POST", "/api/chaos") => Self::handle_chaos_request(body, state).await,
683 ("GET", "/api/chaos") => {
686 let overrides = state.chaos_overrides.read().await.clone();
687 Ok((200, serde_json::to_value(overrides)?))
688 }
689
690 #[cfg(feature = "circuit-breaker")]
693 ("GET", "/api/circuit") => Self::handle_circuit_status(state).await,
694 #[cfg(not(feature = "circuit-breaker"))]
695 ("GET", "/api/circuit") => Ok((
696 503,
697 serde_json::json!({ "error": "circuit-breaker feature not enabled" }),
698 )),
699
700 ("GET", "/api/migration/status") | ("GET", "/migration/status") => {
702 match state.migration.read().await.as_ref() {
703 Some(info) => {
704 let st =
705 crate::mirror::status(&info.target, info.writes_only, &info.metrics);
706 let mut v = serde_json::to_value(st)?;
707 let cut = info.cutover.load_full().is_some();
708 v["cutover_active"] = serde_json::json!(cut);
709 Ok((200, v))
710 }
711 None => Ok((
712 503,
713 serde_json::json!({ "error": "traffic mirroring not enabled" }),
714 )),
715 }
716 }
717
718 ("POST", "/api/migration/cutover") | ("POST", "/migration/cutover") => {
720 let info = state.migration.read().await.clone();
721 let Some(info) = info else {
722 return Ok((
723 503,
724 serde_json::json!({ "error": "traffic mirroring not enabled" }),
725 ));
726 };
727 let force = path.contains("force=true")
728 || body.map(|b| b.contains("\"force\":true")).unwrap_or(false);
729 let st = crate::mirror::status(&info.target, info.writes_only, &info.metrics);
730 if !st.migration_ready && !force {
731 return Ok((
732 409,
733 serde_json::json!({
734 "ok": false,
735 "error": "not migration_ready (backlog/drops present); pass force=true to override",
736 "status": st,
737 }),
738 ));
739 }
740 info.cutover
741 .store(Arc::new(Some(Arc::new(info.cutover_target.clone()))));
742 tracing::warn!(target = %info.cutover_target.addr, "migration cutover: new connections now route to the promoted target");
743 Ok((
744 200,
745 serde_json::json!({ "ok": true, "promoted_to": info.cutover_target.addr }),
746 ))
747 }
748
749 ("POST", "/api/migration/cutover/rollback")
751 | ("POST", "/migration/cutover/rollback") => {
752 let info = state.migration.read().await.clone();
753 let Some(info) = info else {
754 return Ok((
755 503,
756 serde_json::json!({ "error": "traffic mirroring not enabled" }),
757 ));
758 };
759 info.cutover.store(Arc::new(None));
760 Ok((200, serde_json::json!({ "ok": true, "rolled_back": true })))
761 }
762
763 ("POST", "/api/migration/snapshot") | ("POST", "/migration/snapshot") => {
765 let info = state.migration.read().await.clone();
766 let Some(info) = info else {
767 return Ok((
768 503,
769 serde_json::json!({ "error": "traffic mirroring not enabled" }),
770 ));
771 };
772 let body = body.unwrap_or("{}");
773 let req: serde_json::Value = serde_json::from_str(body)
774 .map_err(|e| ProxyError::Internal(format!("invalid JSON: {}", e)))?;
775 let tables: Vec<String> = req
776 .get("tables")
777 .and_then(|t| t.as_array())
778 .map(|a| {
779 a.iter()
780 .filter_map(|v| v.as_str().map(String::from))
781 .collect()
782 })
783 .unwrap_or_default();
784 if tables.is_empty() {
785 return Ok((
786 400,
787 serde_json::json!({ "error": "provide a non-empty 'tables' array" }),
788 ));
789 }
790 match crate::mirror::snapshot_tables(&info.config, &tables).await {
791 Ok(rep) => {
792 let total: u64 = rep.iter().map(|t| t.copied).sum();
793 Ok((
794 200,
795 serde_json::json!({ "ok": true, "tables": rep, "rows_copied": total }),
796 ))
797 }
798 Err(e) => Ok((500, serde_json::json!({ "ok": false, "error": e }))),
799 }
800 }
801
802 ("GET", p) if p == "/api/branch" || p == "/branch" || p.starts_with("/api/branch?") => {
804 let cfg = state.branch.read().await.clone();
805 let Some(cfg) = cfg else {
806 return Ok((
807 503,
808 serde_json::json!({ "error": "branch databases not enabled" }),
809 ));
810 };
811 match crate::branch::list(&cfg).await {
812 Ok(branches) => Ok((200, serde_json::json!({ "branches": branches }))),
813 Err(e) => Ok((500, serde_json::json!({ "error": e }))),
814 }
815 }
816 ("POST", p) if p == "/api/branch" || p == "/branch" => {
817 let cfg = state.branch.read().await.clone();
818 let Some(cfg) = cfg else {
819 return Ok((
820 503,
821 serde_json::json!({ "error": "branch databases not enabled" }),
822 ));
823 };
824 let req: serde_json::Value = serde_json::from_str(body.unwrap_or("{}"))
825 .map_err(|e| ProxyError::Internal(format!("invalid JSON: {}", e)))?;
826 let name = req.get("name").and_then(|v| v.as_str()).unwrap_or("");
827 if name.is_empty() {
828 return Ok((400, serde_json::json!({ "error": "provide 'name'" })));
829 }
830 let base = req.get("base").and_then(|v| v.as_str());
831 match crate::branch::create(&cfg, name, base).await {
832 Ok(()) => Ok((
833 200,
834 serde_json::json!({ "ok": true, "branch": name,
835 "base": base.unwrap_or(&cfg.base_database) }),
836 )),
837 Err(e) => Ok((500, serde_json::json!({ "ok": false, "error": e }))),
838 }
839 }
840 ("DELETE", p) if p.starts_with("/api/branch") || p.starts_with("/branch") => {
841 let cfg = state.branch.read().await.clone();
842 let Some(cfg) = cfg else {
843 return Ok((
844 503,
845 serde_json::json!({ "error": "branch databases not enabled" }),
846 ));
847 };
848 let name = p.find("name=").map(|i| &p[i + 5..]).unwrap_or("");
849 if name.is_empty() {
850 return Ok((
851 400,
852 serde_json::json!({ "error": "provide ?name=<branch>" }),
853 ));
854 }
855 match crate::branch::drop(&cfg, name).await {
856 Ok(()) => Ok((200, serde_json::json!({ "ok": true, "dropped": name }))),
857 Err(e) => Ok((500, serde_json::json!({ "ok": false, "error": e }))),
858 }
859 }
860
861 ("GET", "/config") => {
863 let config = state.config_snapshot.read().await.clone();
864 Ok((200, serde_json::to_value(config)?))
865 }
866
867 ("GET", "/sessions") => {
869 let count = *state.active_sessions.read().await;
870 let response = SessionsResponse {
871 active_sessions: count,
872 };
873 Ok((200, serde_json::to_value(response)?))
874 }
875
876 ("GET", "/pools") => {
878 let pools = Self::get_pool_stats(state).await;
879 Ok((200, serde_json::to_value(pools)?))
880 }
881
882 ("GET", "/version") => {
884 let response = VersionResponse {
885 version: crate::VERSION.to_string(),
886 build_time: env!("CARGO_PKG_VERSION").to_string(),
887 };
888 Ok((200, serde_json::to_value(response)?))
889 }
890
891 #[cfg(feature = "wasm-plugins")]
897 (m, p) if p.starts_with("/admin/kv/") => Self::handle_kv(m, p, body, state).await,
898 #[cfg(not(feature = "wasm-plugins"))]
899 (_m, p) if p.starts_with("/admin/kv/") => Ok((
900 501,
901 serde_json::json!({ "error": "proxy built without the wasm-plugins feature" }),
902 )),
903
904 _ => Ok((404, serde_json::json!({ "error": "Not found" }))),
906 }
907 }
908
909 async fn handle_sql_request(
911 body: Option<&str>,
912 state: &Arc<AdminState>,
913 ) -> Result<(u16, serde_json::Value)> {
914 let body = body.ok_or_else(|| ProxyError::Internal("Missing request body".to_string()))?;
916 let request: SqlRequest = serde_json::from_str(body)
917 .map_err(|e| ProxyError::Internal(format!("Invalid JSON: {}", e)))?;
918
919 let sql = request.query.trim();
920 if sql.is_empty() {
921 return Ok((400, serde_json::json!({ "error": "Empty query" })));
922 }
923
924 let is_write = Self::is_write_query(sql);
926 let query_type = if is_write { "write" } else { "read" };
927
928 let proxy_config = state.proxy_config.read().await;
930 let config = proxy_config
931 .as_ref()
932 .ok_or_else(|| ProxyError::Internal("Proxy config not initialized".to_string()))?;
933
934 let health = state.node_health.read().await;
936
937 let target_node = if is_write {
939 Self::select_primary_node(config, &health)?
941 } else {
942 Self::select_read_node(config, &health, state)?
944 };
945
946 let target_address = format!("{}:{}", target_node.host, target_node.port);
947 let http_port = target_node.http_port;
949 let http_url = format!("http://{}:{}/api/sql", target_node.host, http_port);
950
951 tracing::debug!(
952 "Routing {} query to {} ({})",
953 query_type,
954 target_address,
955 match target_node.role {
956 NodeRole::Primary => "primary",
957 NodeRole::Standby => "standby",
958 NodeRole::ReadReplica => "replica",
959 }
960 );
961
962 let result = Self::forward_sql_request(&http_url, sql).await?;
964
965 let response = SqlResponse {
967 query_type: query_type.to_string(),
968 routed_to: target_address,
969 node_role: format!("{:?}", target_node.role).to_lowercase(),
970 result,
971 };
972
973 Ok((200, serde_json::to_value(response)?))
974 }
975
976 fn is_write_query(sql: &str) -> bool {
978 let upper = sql.trim().to_uppercase();
979
980 if upper.starts_with("INSERT")
982 || upper.starts_with("UPDATE")
983 || upper.starts_with("DELETE")
984 || upper.starts_with("CREATE")
985 || upper.starts_with("ALTER")
986 || upper.starts_with("DROP")
987 || upper.starts_with("TRUNCATE")
988 || upper.starts_with("GRANT")
989 || upper.starts_with("REVOKE")
990 || upper.starts_with("VACUUM")
991 || upper.starts_with("REINDEX")
992 || upper.starts_with("MERGE")
993 || upper.starts_with("UPSERT")
994 {
995 return true;
996 }
997
998 if upper.starts_with("BEGIN")
1000 || upper.starts_with("COMMIT")
1001 || upper.starts_with("ROLLBACK")
1002 || upper.starts_with("SAVEPOINT")
1003 {
1004 return true;
1006 }
1007
1008 false
1010 }
1011
1012 fn select_primary_node<'a>(
1014 config: &'a ProxyConfig,
1015 health: &HashMap<String, NodeHealth>,
1016 ) -> Result<&'a NodeConfig> {
1017 config
1018 .nodes
1019 .iter()
1020 .find(|n| {
1021 n.role == NodeRole::Primary
1022 && n.enabled
1023 && health.get(&n.address()).map(|h| h.healthy).unwrap_or(false)
1024 })
1025 .ok_or_else(|| ProxyError::Internal("No healthy primary node available".to_string()))
1026 }
1027
1028 fn select_read_node<'a>(
1030 config: &'a ProxyConfig,
1031 health: &HashMap<String, NodeHealth>,
1032 state: &AdminState,
1033 ) -> Result<&'a NodeConfig> {
1034 let healthy_nodes: Vec<&NodeConfig> = config
1036 .nodes
1037 .iter()
1038 .filter(|n| n.enabled && health.get(&n.address()).map(|h| h.healthy).unwrap_or(false))
1039 .collect();
1040
1041 if healthy_nodes.is_empty() {
1042 return Err(ProxyError::Internal(
1043 "No healthy nodes available".to_string(),
1044 ));
1045 }
1046
1047 if config.load_balancer.read_write_split {
1049 let read_nodes: Vec<&NodeConfig> = healthy_nodes
1050 .iter()
1051 .filter(|n| n.role == NodeRole::Standby || n.role == NodeRole::ReadReplica)
1052 .copied()
1053 .collect();
1054
1055 if !read_nodes.is_empty() {
1056 let counter = state.read_lb_counter.fetch_add(1, Ordering::Relaxed);
1058 let index = counter % read_nodes.len();
1059 return Ok(read_nodes[index]);
1060 }
1061 }
1062
1063 let counter = state.read_lb_counter.fetch_add(1, Ordering::Relaxed);
1065 let index = counter % healthy_nodes.len();
1066 Ok(healthy_nodes[index])
1067 }
1068
1069 async fn forward_sql_request(url: &str, sql: &str) -> Result<serde_json::Value> {
1071 let request_body = serde_json::json!({ "query": sql });
1073 let body_bytes = serde_json::to_vec(&request_body)
1074 .map_err(|e| ProxyError::Internal(format!("JSON serialization error: {}", e)))?;
1075
1076 let url_parts: Vec<&str> = url.trim_start_matches("http://").splitn(2, '/').collect();
1078 if url_parts.is_empty() {
1079 return Err(ProxyError::Internal("Invalid URL".to_string()));
1080 }
1081
1082 let host_port = url_parts[0];
1083 let path = if url_parts.len() > 1 {
1084 format!("/{}", url_parts[1])
1085 } else {
1086 "/".to_string()
1087 };
1088
1089 let stream = TcpStream::connect(host_port).await.map_err(|e| {
1091 ProxyError::Network(format!("Failed to connect to {}: {}", host_port, e))
1092 })?;
1093
1094 let (reader, mut writer) = stream.into_split();
1095 let mut reader = BufReader::new(reader);
1096
1097 let request = format!(
1099 "POST {} HTTP/1.1\r\nHost: {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
1100 path,
1101 host_port,
1102 body_bytes.len()
1103 );
1104
1105 writer
1106 .write_all(request.as_bytes())
1107 .await
1108 .map_err(|e| ProxyError::Network(format!("Write error: {}", e)))?;
1109 writer
1110 .write_all(&body_bytes)
1111 .await
1112 .map_err(|e| ProxyError::Network(format!("Write body error: {}", e)))?;
1113
1114 let mut response_headers = Vec::new();
1116 let mut line = String::new();
1117 let mut content_length: usize = 0;
1118
1119 loop {
1120 line.clear();
1121 let bytes_read = reader
1122 .read_line(&mut line)
1123 .await
1124 .map_err(|e| ProxyError::Network(format!("Response read error: {}", e)))?;
1125
1126 if bytes_read == 0 || line == "\r\n" {
1127 break;
1128 }
1129
1130 let trimmed = line.trim();
1131 if trimmed.to_lowercase().starts_with("content-length:") {
1132 if let Some(len_str) = trimmed.split(':').nth(1) {
1133 content_length = len_str.trim().parse().unwrap_or(0);
1134 }
1135 }
1136 response_headers.push(trimmed.to_string());
1137 }
1138
1139 let mut body_buf = vec![0u8; content_length];
1141 if content_length > 0 {
1142 reader
1143 .read_exact(&mut body_buf)
1144 .await
1145 .map_err(|e| ProxyError::Network(format!("Response body read error: {}", e)))?;
1146 }
1147
1148 let response_body = String::from_utf8_lossy(&body_buf);
1149
1150 serde_json::from_str(&response_body).map_err(|e| {
1152 ProxyError::Internal(format!(
1153 "Invalid JSON response: {} - body: {}",
1154 e, response_body
1155 ))
1156 })
1157 }
1158
1159 async fn check_readiness(state: &Arc<AdminState>) -> bool {
1161 let health = state.node_health.read().await;
1162
1163 health.values().any(|h| h.healthy)
1165 }
1166
1167 async fn set_node_enabled(
1169 state: &Arc<AdminState>,
1170 node_addr: &str,
1171 enabled: bool,
1172 ) -> Result<()> {
1173 let mut health = state.node_health.write().await;
1174
1175 if let Some(node_health) = health.get_mut(node_addr) {
1176 node_health.healthy = enabled;
1177 Ok(())
1178 } else {
1179 Err(ProxyError::Config(format!("Node not found: {}", node_addr)))
1180 }
1181 }
1182
1183 async fn get_pool_stats(_state: &Arc<AdminState>) -> Vec<PoolStatsResponse> {
1185 #[cfg(feature = "pool-modes")]
1188 if let Some(mgr) = _state.pool_manager.read().await.clone() {
1189 let stats = mgr.get_stats().await;
1190 return stats
1191 .node_stats
1192 .into_iter()
1193 .map(|ns| PoolStatsResponse {
1194 node: ns.node_id.0.to_string(),
1195 active_connections: ns.active as u64,
1196 idle_connections: ns.idle as u64,
1197 pending_requests: 0,
1200 total_connections_created: ns.total as u64,
1201 total_connections_closed: 0,
1202 })
1203 .collect();
1204 }
1205 Vec::new()
1206 }
1207
1208 #[cfg(feature = "ha-tr")]
1212 async fn handle_replay_request(
1213 body: Option<&str>,
1214 state: &Arc<AdminState>,
1215 ) -> Result<(u16, serde_json::Value)> {
1216 let raw =
1217 body.ok_or_else(|| ProxyError::Internal("replay: empty request body".to_string()))?;
1218 let req: ReplayRequestBody = match serde_json::from_str(raw) {
1219 Ok(r) => r,
1220 Err(e) => {
1221 return Ok((
1222 400,
1223 serde_json::json!({ "error": format!("invalid body: {}", e) }),
1224 ));
1225 }
1226 };
1227 let engine = match state.replay_engine.read().await.clone() {
1228 Some(e) => e,
1229 None => {
1230 return Ok((
1231 503,
1232 serde_json::json!({ "error": "replay engine not attached" }),
1233 ));
1234 }
1235 };
1236 let tt = TimeTravelRequest {
1237 from: req.from,
1238 to: req.to,
1239 target_host: req.target_host,
1240 target_port: req.target_port,
1241 target_user: req.target_user,
1242 target_password: req.target_password,
1243 target_database: req.target_database,
1244 };
1245 match engine.replay_window(&tt).await {
1246 Ok(summary) => Ok((200, serde_json::to_value(summary)?)),
1247 Err(e) => Ok((
1248 500,
1249 serde_json::json!({ "error": format!("replay failed: {}", e) }),
1250 )),
1251 }
1252 }
1253
1254 #[cfg(feature = "edge-proxy")]
1257 async fn handle_edge_status(state: &Arc<AdminState>) -> Result<(u16, serde_json::Value)> {
1258 let cache_stats = state.edge_cache.read().await.clone().map(|c| c.stats());
1259 let edges = match state.edge_registry.read().await.clone() {
1260 Some(r) => r.list(),
1261 None => Vec::new(),
1262 };
1263 Ok((
1264 200,
1265 serde_json::json!({
1266 "cache": cache_stats,
1267 "registered": edges,
1268 "edge_count": edges.len(),
1269 }),
1270 ))
1271 }
1272
1273 #[cfg(feature = "edge-proxy")]
1280 async fn handle_edge_register(
1281 body: Option<&str>,
1282 state: &Arc<AdminState>,
1283 ) -> Result<(u16, serde_json::Value)> {
1284 let raw =
1285 body.ok_or_else(|| ProxyError::Internal("edge register: empty body".to_string()))?;
1286 let req: EdgeRegisterBody = match serde_json::from_str(raw) {
1287 Ok(r) => r,
1288 Err(e) => {
1289 return Ok((
1290 400,
1291 serde_json::json!({ "error": format!("invalid body: {}", e) }),
1292 ));
1293 }
1294 };
1295 let registry = match state.edge_registry.read().await.clone() {
1296 Some(r) => r,
1297 None => {
1298 return Ok((
1299 503,
1300 serde_json::json!({ "error": "edge registry not attached" }),
1301 ));
1302 }
1303 };
1304 let now = chrono::Utc::now().to_rfc3339();
1305 match registry.register(&req.edge_id, &req.region, &req.base_url, &now) {
1306 Ok(_rx) => {
1307 Ok((
1314 201,
1315 serde_json::json!({
1316 "edge_id": req.edge_id,
1317 "region": req.region,
1318 "base_url": req.base_url,
1319 "registered_at": now,
1320 }),
1321 ))
1322 }
1323 Err(e) => Ok((503, serde_json::json!({ "error": e.to_string() }))),
1324 }
1325 }
1326
1327 #[cfg(feature = "edge-proxy")]
1341 async fn handle_edge_subscribe(
1342 writer: &mut tokio::net::tcp::WriteHalf<'_>,
1343 state: &Arc<AdminState>,
1344 edge_id: &str,
1345 region: &str,
1346 base_url: &str,
1347 ) -> Result<()> {
1348 let registry = match state.edge_registry.read().await.clone() {
1349 Some(r) => r,
1350 None => {
1351 Self::send_json_response(
1352 writer,
1353 503,
1354 &serde_json::json!({ "error": "edge registry not attached" }),
1355 )
1356 .await?;
1357 return Ok(());
1358 }
1359 };
1360 let now = chrono::Utc::now().to_rfc3339();
1361 let mut rx = match registry.register(edge_id, region, base_url, &now) {
1362 Ok(rx) => rx,
1363 Err(e) => {
1365 Self::send_json_response(
1366 writer,
1367 503,
1368 &serde_json::json!({ "error": e.to_string() }),
1369 )
1370 .await?;
1371 return Ok(());
1372 }
1373 };
1374
1375 if !Self::write_sse(
1383 writer,
1384 b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: keep-alive\r\n\r\n",
1385 )
1386 .await
1387 {
1388 return Ok(());
1389 }
1390
1391 if let Some(cache) = state.edge_cache.read().await.clone() {
1398 let hello = InvalidationEvent {
1399 up_to_version: cache.current_version(),
1400 tables: Vec::new(),
1401 committed_at: now.clone(),
1402 epoch: cache.epoch(),
1403 };
1404 let json = serde_json::to_string(&hello)
1405 .map_err(|e| ProxyError::Internal(format!("JSON error: {}", e)))?;
1406 let frame = format!("event: invalidate\ndata: {}\n\n", json);
1407 if !Self::write_sse(writer, frame.as_bytes()).await {
1408 return Ok(());
1409 }
1410 }
1411
1412 tracing::info!(edge_id, region, "edge subscribed to invalidation stream");
1413
1414 let mut heartbeat = tokio::time::interval(std::time::Duration::from_secs(15));
1415 heartbeat.tick().await;
1418
1419 loop {
1420 tokio::select! {
1421 ev = rx.recv() => {
1422 match ev {
1423 Some(ev) => {
1424 let json = serde_json::to_string(&ev)
1427 .map_err(|e| ProxyError::Internal(format!("JSON error: {}", e)))?;
1428 let frame = format!("event: invalidate\ndata: {}\n\n", json);
1429 if !Self::write_sse(writer, frame.as_bytes()).await {
1430 tracing::debug!(edge_id, "edge SSE write failed; closing stream");
1434 return Ok(());
1435 }
1436 }
1437 None => {
1441 tracing::debug!(edge_id, "edge subscription evicted by registry");
1442 return Ok(());
1443 }
1444 }
1445 }
1446 _ = heartbeat.tick() => {
1447 if !Self::write_sse(writer, b": keepalive\n\n").await {
1451 tracing::debug!(edge_id, "edge SSE heartbeat failed; closing stream");
1452 return Ok(());
1453 }
1454 registry.touch(edge_id, &chrono::Utc::now().to_rfc3339());
1459 }
1460 }
1461 }
1462 }
1463
1464 #[cfg(feature = "edge-proxy")]
1467 async fn write_sse(writer: &mut tokio::net::tcp::WriteHalf<'_>, bytes: &[u8]) -> bool {
1468 let write = async {
1469 writer.write_all(bytes).await?;
1470 writer.flush().await
1471 };
1472 matches!(
1473 tokio::time::timeout(Self::ADMIN_SSE_WRITE_TIMEOUT, write).await,
1474 Ok(Ok(()))
1475 )
1476 }
1477
1478 #[cfg(feature = "edge-proxy")]
1485 async fn handle_edge_invalidate(
1486 body: Option<&str>,
1487 state: &Arc<AdminState>,
1488 ) -> Result<(u16, serde_json::Value)> {
1489 let raw =
1490 body.ok_or_else(|| ProxyError::Internal("edge invalidate: empty body".to_string()))?;
1491 let req: EdgeInvalidateBody = match serde_json::from_str(raw) {
1492 Ok(r) => r,
1493 Err(e) => {
1494 return Ok((
1495 400,
1496 serde_json::json!({ "error": format!("invalid body: {}", e) }),
1497 ));
1498 }
1499 };
1500 let cache = match state.edge_cache.read().await.clone() {
1501 Some(c) => c,
1502 None => {
1503 return Ok((
1504 503,
1505 serde_json::json!({ "error": "edge cache not attached" }),
1506 ));
1507 }
1508 };
1509 let registry = match state.edge_registry.read().await.clone() {
1510 Some(r) => r,
1511 None => {
1512 return Ok((
1513 503,
1514 serde_json::json!({ "error": "edge registry not attached" }),
1515 ));
1516 }
1517 };
1518 let version = req
1528 .up_to_version
1529 .map(|v| v.min(cache.current_version()))
1530 .unwrap_or_else(|| cache.current_version());
1531 let dropped_local = cache.invalidate(version, &req.tables);
1533 let ev = InvalidationEvent {
1535 up_to_version: version,
1536 tables: req.tables.clone(),
1537 committed_at: chrono::Utc::now().to_rfc3339(),
1538 epoch: cache.epoch(),
1539 };
1540 let (sent, pruned) = registry.broadcast(ev).await;
1541 Ok((
1542 200,
1543 serde_json::json!({
1544 "version": version,
1545 "tables": req.tables,
1546 "dropped_local": dropped_local,
1547 "edges_notified": sent,
1548 "edges_pruned": pruned,
1549 }),
1550 ))
1551 }
1552
1553 #[cfg(feature = "anomaly-detection")]
1558 async fn handle_anomalies_list(
1559 path: &str,
1560 state: &Arc<AdminState>,
1561 ) -> Result<(u16, serde_json::Value)> {
1562 let limit = parse_limit_query(path, 100, 1024);
1563 let det = match state.anomaly_detector.read().await.clone() {
1564 Some(d) => d,
1565 None => {
1566 return Ok((
1567 503,
1568 serde_json::json!({ "error": "anomaly detector not attached" }),
1569 ));
1570 }
1571 };
1572 let events = det.recent_events(limit);
1573 Ok((
1574 200,
1575 serde_json::json!({
1576 "count": events.len(),
1577 "limit": limit,
1578 "events": events,
1579 "buffer_total": det.event_count(),
1580 }),
1581 ))
1582 }
1583
1584 #[cfg(feature = "query-analytics")]
1587 async fn handle_analytics(
1588 path: &str,
1589 state: &Arc<AdminState>,
1590 ) -> Result<(u16, serde_json::Value)> {
1591 use crate::analytics::OrderBy;
1592 let limit = parse_limit_query(path, 50, 1024);
1593 let a = match state.analytics.read().await.clone() {
1594 Some(a) => a,
1595 None => {
1596 return Ok((503, serde_json::json!({ "error": "analytics not enabled" })));
1597 }
1598 };
1599 let top: Vec<serde_json::Value> = a
1600 .top_queries(OrderBy::Calls, limit)
1601 .into_iter()
1602 .map(|s| {
1603 serde_json::json!({
1604 "fingerprint": s.fingerprint_hash,
1605 "normalized": s.normalized,
1606 "calls": s.calls,
1607 "avg_ms": s.avg_time.as_secs_f64() * 1000.0,
1608 "p99_ms": s.p99.as_secs_f64() * 1000.0,
1609 "rows": s.rows,
1610 "errors": s.errors,
1611 })
1612 })
1613 .collect();
1614 let slow_count = a.slow_queries(limit).len();
1615 Ok((
1616 200,
1617 serde_json::json!({
1618 "limit": limit,
1619 "top_queries": top,
1620 "slow_query_count": slow_count,
1621 }),
1622 ))
1623 }
1624
1625 #[cfg(feature = "ha-tr")]
1635 async fn handle_shadow_request(body: Option<&str>) -> Result<(u16, serde_json::Value)> {
1636 use crate::backend::{
1637 tls::default_client_config, BackendClient, BackendConfig, ParamValue, TlsMode,
1638 };
1639 use crate::shadow_execute::shadow_execute;
1640
1641 let raw =
1642 body.ok_or_else(|| ProxyError::Internal("shadow: empty request body".to_string()))?;
1643 let req: ShadowRequestBody = match serde_json::from_str(raw) {
1644 Ok(r) => r,
1645 Err(e) => {
1646 return Ok((
1647 400,
1648 serde_json::json!({ "error": format!("invalid body: {}", e) }),
1649 ));
1650 }
1651 };
1652
1653 let mk_cfg = |host: String,
1656 port: u16,
1657 user: Option<String>,
1658 password: Option<String>,
1659 database: Option<String>| BackendConfig {
1660 host,
1661 port,
1662 user: user.unwrap_or_else(|| "postgres".into()),
1663 password,
1664 database,
1665 application_name: Some("heliosdb-proxy-shadow".into()),
1666 tls_mode: TlsMode::Disable,
1667 connect_timeout: std::time::Duration::from_secs(5),
1668 query_timeout: std::time::Duration::from_secs(30),
1669 tls_config: default_client_config(),
1670 };
1671 let source_cfg = mk_cfg(
1672 req.source_host,
1673 req.source_port,
1674 req.source_user,
1675 req.source_password,
1676 req.source_database,
1677 );
1678 let shadow_cfg = mk_cfg(
1679 req.shadow_host,
1680 req.shadow_port,
1681 req.shadow_user,
1682 req.shadow_password,
1683 req.shadow_database,
1684 );
1685
1686 let mut source = match BackendClient::connect(&source_cfg).await {
1690 Ok(c) => c,
1691 Err(e) => {
1692 return Ok((
1693 500,
1694 serde_json::json!({ "error": format!("source connect: {}", e) }),
1695 ));
1696 }
1697 };
1698
1699 let params: Vec<ParamValue> = req
1700 .params
1701 .unwrap_or_default()
1702 .into_iter()
1703 .map(ParamValue::Text)
1704 .collect();
1705
1706 let outcome = shadow_execute(&mut source, &shadow_cfg, &req.sql, ¶ms).await;
1707 source.close().await;
1708
1709 match outcome {
1710 Ok((_qr, report)) => Ok((
1711 200,
1712 serde_json::json!({
1713 "sql": report.sql,
1714 "both_succeeded": report.both_succeeded,
1715 "row_count_match": report.row_count_match,
1716 "row_hash_match": report.row_hash_match,
1717 "primary_elapsed_us": report.primary_elapsed_us,
1718 "shadow_elapsed_us": report.shadow_elapsed_us,
1719 "primary_error": report.primary_error,
1720 "shadow_error": report.shadow_error,
1721 "is_clean": report.is_clean(),
1722 }),
1723 )),
1724 Err(e) => Ok((
1725 500,
1726 serde_json::json!({ "error": format!("shadow execute: {}", e) }),
1727 )),
1728 }
1729 }
1730
1731 #[cfg(feature = "circuit-breaker")]
1750 async fn handle_circuit_status(state: &Arc<AdminState>) -> Result<(u16, serde_json::Value)> {
1751 let mgr = match state.circuit_breaker.read().await.clone() {
1752 Some(m) => m,
1753 None => {
1754 return Ok((
1755 503,
1756 serde_json::json!({ "error": "circuit breaker not attached" }),
1757 ))
1758 }
1759 };
1760 let nodes = state.config_snapshot.read().await.nodes.clone();
1761 let circuits: Vec<serde_json::Value> = nodes
1762 .iter()
1763 .map(|n| {
1764 let st = mgr.get_breaker(&n.address).get_state();
1765 serde_json::json!({
1766 "node": n.address,
1767 "state": format!("{:?}", st).to_lowercase(),
1768 })
1769 })
1770 .collect();
1771 Ok((200, serde_json::json!({ "circuits": circuits })))
1772 }
1773
1774 async fn handle_chaos_request(
1775 body: Option<&str>,
1776 state: &Arc<AdminState>,
1777 ) -> Result<(u16, serde_json::Value)> {
1778 let raw =
1779 body.ok_or_else(|| ProxyError::Internal("chaos: empty request body".to_string()))?;
1780 let action: ChaosAction = match serde_json::from_str(raw) {
1781 Ok(a) => a,
1782 Err(e) => {
1783 return Ok((
1784 400,
1785 serde_json::json!({ "error": format!("invalid body: {}", e) }),
1786 ));
1787 }
1788 };
1789 match action {
1790 ChaosAction::ForceUnhealthy { target_node } => {
1791 if let Err(e) = Self::set_node_enabled(state, &target_node, false).await {
1792 return Ok((404, serde_json::json!({ "error": e.to_string() })));
1793 }
1794 state.chaos_overrides.write().await.insert(
1795 target_node.clone(),
1796 ChaosOverride {
1797 since: chrono::Utc::now().to_rfc3339(),
1798 kind: "force_unhealthy".to_string(),
1799 note: "forced unhealthy via chaos endpoint".to_string(),
1800 },
1801 );
1802 Ok((
1803 200,
1804 serde_json::json!({
1805 "applied": "force_unhealthy",
1806 "target_node": target_node,
1807 }),
1808 ))
1809 }
1810 ChaosAction::Restore { target_node } => {
1811 if let Err(e) = Self::set_node_enabled(state, &target_node, true).await {
1812 return Ok((404, serde_json::json!({ "error": e.to_string() })));
1813 }
1814 state.chaos_overrides.write().await.remove(&target_node);
1815 Ok((
1816 200,
1817 serde_json::json!({
1818 "restored": target_node,
1819 }),
1820 ))
1821 }
1822 ChaosAction::Reset => {
1823 let overrides: Vec<String> =
1824 state.chaos_overrides.read().await.keys().cloned().collect();
1825 let mut restored = Vec::with_capacity(overrides.len());
1826 for addr in overrides {
1827 let _ = Self::set_node_enabled(state, &addr, true).await;
1828 restored.push(addr);
1829 }
1830 state.chaos_overrides.write().await.clear();
1831 Ok((
1832 200,
1833 serde_json::json!({
1834 "reset": true,
1835 "restored": restored,
1836 }),
1837 ))
1838 }
1839 }
1840 }
1841
1842 #[cfg(feature = "wasm-plugins")]
1848 async fn handle_plugins_list(state: &Arc<AdminState>) -> Result<(u16, serde_json::Value)> {
1849 let pm = match state.plugin_manager.read().await.clone() {
1850 Some(p) => p,
1851 None => {
1852 return Ok((
1853 503,
1854 serde_json::json!({ "error": "plugin manager not attached" }),
1855 ));
1856 }
1857 };
1858 let plugins: Vec<PluginListEntry> = pm
1859 .list_plugins()
1860 .into_iter()
1861 .map(|info| PluginListEntry {
1862 name: info.name,
1863 version: info.version,
1864 description: info.description,
1865 hooks: info
1866 .hooks
1867 .iter()
1868 .map(|h| h.export_name().to_string())
1869 .collect(),
1870 state: format!("{:?}", info.state),
1871 invocations: info.stats.total_calls,
1872 errors: info.stats.error_count,
1873 })
1874 .collect();
1875 Ok((200, serde_json::to_value(plugins)?))
1876 }
1877
1878 #[cfg(not(feature = "wasm-plugins"))]
1879 async fn handle_plugins_list(_state: &Arc<AdminState>) -> Result<(u16, serde_json::Value)> {
1880 Ok((
1881 503,
1882 serde_json::json!({ "error": "wasm-plugins feature not compiled in" }),
1883 ))
1884 }
1885
1886 #[cfg(feature = "wasm-plugins")]
1913 async fn handle_kv(
1914 method: &str,
1915 path: &str,
1916 body: Option<&str>,
1917 state: &Arc<AdminState>,
1918 ) -> Result<(u16, serde_json::Value)> {
1919 use serde_json::json;
1920 let list_prefix = parse_query_params(path)
1930 .remove("prefix")
1931 .unwrap_or_default();
1932 let path = path.split('?').next().unwrap_or(path);
1933 let rest = &path["/admin/kv/".len()..];
1934 let Some((plugin, key)) = rest.split_once('/') else {
1935 return Ok((400, json!({ "error": "expected /admin/kv/<plugin>/<key>" })));
1936 };
1937 if plugin.is_empty() {
1941 return Ok((400, json!({ "error": "empty plugin name" })));
1942 }
1943 let pm = state.plugin_manager.read().await.clone();
1944 let Some(pm) = pm else {
1945 return Ok((503, json!({ "error": "plugin runtime not enabled" })));
1946 };
1947 let kv = pm.kv();
1948 match (method, key.is_empty()) {
1949 ("GET", true) => Ok((
1950 200,
1951 json!({ "plugin": plugin, "keys": kv.list_keys(plugin, list_prefix.as_bytes()) }),
1952 )),
1953 ("GET", false) => match kv.get(plugin, key.as_bytes()) {
1954 Some(v) => Ok((
1955 200,
1956 json!({ "plugin": plugin, "key": key, "value": String::from_utf8_lossy(&v) }),
1957 )),
1958 None => Ok((404, json!({ "error": "key not found" }))),
1959 },
1960 ("PUT", false) => {
1961 let body_bytes = body.unwrap_or("").as_bytes();
1967 let cap = kv.max_value_bytes();
1968 if cap != 0 && body_bytes.len() > cap {
1969 return Ok((413, json!({ "error": "kv_max_value_bytes exceeded" })));
1970 }
1971 if kv.set(plugin, key.as_bytes().to_vec(), body_bytes.to_vec()) {
1972 Ok((200, json!({ "ok": true })))
1973 } else {
1974 Ok((
1975 413,
1976 json!({
1977 "error": "kv_max_value_bytes, kv_max_keys_per_plugin, kv_max_plugins, or kv_max_total_bytes exceeded"
1978 }),
1979 ))
1980 }
1981 }
1982 ("DELETE", false) => {
1983 kv.delete(plugin, key.as_bytes());
1984 Ok((200, json!({ "ok": true })))
1985 }
1986 _ => Ok((405, json!({ "error": "method not allowed" }))),
1987 }
1988 }
1989
1990 async fn compute_topology(state: &Arc<AdminState>) -> TopologyResponse {
1996 let health = state.node_health.read().await;
1997 let cfg = state.config_snapshot.read().await;
1998
1999 let mut current_primary: Option<String> = None;
2000 for n in &cfg.nodes {
2001 if n.role.eq_ignore_ascii_case("primary") {
2002 let healthy = health.get(&n.address).map(|h| h.healthy).unwrap_or(false);
2003 if healthy {
2004 current_primary = Some(n.address.clone());
2005 break;
2006 }
2007 }
2008 }
2009
2010 let healthy_nodes = health.values().filter(|h| h.healthy).count() as u32;
2011 let unhealthy_nodes = health.values().filter(|h| !h.healthy).count() as u32;
2012 let total_nodes = cfg.nodes.len() as u32;
2013
2014 TopologyResponse {
2015 current_primary,
2016 healthy_nodes,
2017 unhealthy_nodes,
2018 total_nodes,
2019 last_failover_at: None,
2020 }
2021 }
2022
2023 fn format_prometheus_metrics(metrics: &ServerMetricsSnapshot) -> String {
2025 let mut output = String::new();
2026
2027 output.push_str("# HELP heliosdb_proxy_connections_total Total connections accepted\n");
2028 output.push_str("# TYPE heliosdb_proxy_connections_total counter\n");
2029 output.push_str(&format!(
2030 "heliosdb_proxy_connections_total {}\n",
2031 metrics.connections_accepted
2032 ));
2033
2034 output.push_str("# HELP heliosdb_proxy_connections_closed Total connections closed\n");
2035 output.push_str("# TYPE heliosdb_proxy_connections_closed counter\n");
2036 output.push_str(&format!(
2037 "heliosdb_proxy_connections_closed {}\n",
2038 metrics.connections_closed
2039 ));
2040
2041 output.push_str("# HELP heliosdb_proxy_queries_total Total queries processed\n");
2042 output.push_str("# TYPE heliosdb_proxy_queries_total counter\n");
2043 output.push_str(&format!(
2044 "heliosdb_proxy_queries_total {}\n",
2045 metrics.queries_processed
2046 ));
2047
2048 output.push_str("# HELP heliosdb_proxy_bytes_received_total Total bytes received\n");
2049 output.push_str("# TYPE heliosdb_proxy_bytes_received_total counter\n");
2050 output.push_str(&format!(
2051 "heliosdb_proxy_bytes_received_total {}\n",
2052 metrics.bytes_received
2053 ));
2054
2055 output.push_str("# HELP heliosdb_proxy_bytes_sent_total Total bytes sent\n");
2056 output.push_str("# TYPE heliosdb_proxy_bytes_sent_total counter\n");
2057 output.push_str(&format!(
2058 "heliosdb_proxy_bytes_sent_total {}\n",
2059 metrics.bytes_sent
2060 ));
2061
2062 output.push_str("# HELP heliosdb_proxy_failovers_total Total failovers\n");
2063 output.push_str("# TYPE heliosdb_proxy_failovers_total counter\n");
2064 output.push_str(&format!(
2065 "heliosdb_proxy_failovers_total {}\n",
2066 metrics.failovers
2067 ));
2068
2069 output
2070 }
2071
2072 async fn send_response(
2074 writer: &mut tokio::net::tcp::WriteHalf<'_>,
2075 status: u16,
2076 status_text: &str,
2077 body: &str,
2078 ) -> Result<()> {
2079 let response = format!(
2080 "HTTP/1.1 {} {}\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
2081 status,
2082 status_text,
2083 body.len(),
2084 body
2085 );
2086
2087 writer
2088 .write_all(response.as_bytes())
2089 .await
2090 .map_err(|e| ProxyError::Network(format!("Write error: {}", e)))?;
2091
2092 Ok(())
2093 }
2094
2095 async fn send_json_response<T: Serialize>(
2097 writer: &mut tokio::net::tcp::WriteHalf<'_>,
2098 status: u16,
2099 body: &T,
2100 ) -> Result<()> {
2101 let json = serde_json::to_string(body)
2102 .map_err(|e| ProxyError::Internal(format!("JSON error: {}", e)))?;
2103
2104 let status_text = match status {
2105 200 => "OK",
2106 400 => "Bad Request",
2107 404 => "Not Found",
2108 500 => "Internal Server Error",
2109 503 => "Service Unavailable",
2110 _ => "Unknown",
2111 };
2112
2113 let response = format!(
2114 "HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
2115 status,
2116 status_text,
2117 json.len(),
2118 json
2119 );
2120
2121 writer
2122 .write_all(response.as_bytes())
2123 .await
2124 .map_err(|e| ProxyError::Network(format!("Write error: {}", e)))?;
2125
2126 Ok(())
2127 }
2128
2129 pub fn shutdown(&self) {
2131 let _ = self.shutdown_tx.send(());
2132 }
2133}
2134
2135impl AdminState {
2136 pub fn new() -> Self {
2138 Self {
2139 node_health: RwLock::new(HashMap::new()),
2140 metrics: RwLock::new(ServerMetricsSnapshot {
2141 connections_accepted: 0,
2142 connections_closed: 0,
2143 queries_processed: 0,
2144 bytes_received: 0,
2145 bytes_sent: 0,
2146 failovers: 0,
2147 }),
2148 active_sessions: RwLock::new(0),
2149 config_snapshot: RwLock::new(ConfigSnapshot {
2150 listen_address: String::new(),
2151 admin_address: String::new(),
2152 tr_enabled: false,
2153 tr_mode: String::new(),
2154 pool_min_connections: 0,
2155 pool_max_connections: 0,
2156 nodes: Vec::new(),
2157 }),
2158 proxy_config: RwLock::new(None),
2159 read_lb_counter: AtomicUsize::new(0),
2160 commands: RwLock::new(HashMap::new()),
2161 #[cfg(feature = "pool-modes")]
2162 pool_manager: RwLock::new(None),
2163 #[cfg(feature = "circuit-breaker")]
2164 circuit_breaker: RwLock::new(None),
2165 #[cfg(feature = "ha-tr")]
2166 replay_engine: RwLock::new(None),
2167 #[cfg(feature = "wasm-plugins")]
2168 plugin_manager: RwLock::new(None),
2169 chaos_overrides: RwLock::new(HashMap::new()),
2170 #[cfg(feature = "anomaly-detection")]
2171 anomaly_detector: RwLock::new(None),
2172 #[cfg(feature = "query-analytics")]
2173 analytics: RwLock::new(None),
2174 #[cfg(feature = "edge-proxy")]
2175 edge_cache: RwLock::new(None),
2176 #[cfg(feature = "edge-proxy")]
2177 edge_registry: RwLock::new(None),
2178 auth_token: RwLock::new(None),
2179 migration: RwLock::new(None),
2180 branch: RwLock::new(None),
2181 }
2182 }
2183
2184 pub async fn with_auth_token(&self, token: Option<String>) {
2186 *self.auth_token.write().await = token;
2187 }
2188
2189 pub async fn with_migration(&self, info: MigrationInfo) {
2191 *self.migration.write().await = Some(info);
2192 }
2193
2194 pub async fn with_branch(&self, cfg: crate::config::BranchConfig) {
2196 *self.branch.write().await = Some(cfg);
2197 }
2198
2199 #[cfg(feature = "pool-modes")]
2202 pub async fn with_pool_manager(&self, manager: Arc<crate::pool::ConnectionPoolManager>) {
2203 *self.pool_manager.write().await = Some(manager);
2204 }
2205
2206 #[cfg(feature = "circuit-breaker")]
2209 pub async fn with_circuit_breaker(
2210 &self,
2211 manager: Arc<crate::circuit_breaker::CircuitBreakerManager>,
2212 ) {
2213 *self.circuit_breaker.write().await = Some(manager);
2214 }
2215
2216 #[cfg(feature = "anomaly-detection")]
2219 pub async fn with_anomaly_detector(&self, detector: Arc<AnomalyDetector>) {
2220 *self.anomaly_detector.write().await = Some(detector);
2221 }
2222
2223 #[cfg(feature = "query-analytics")]
2225 pub async fn with_analytics(&self, analytics: Arc<crate::analytics::QueryAnalytics>) {
2226 *self.analytics.write().await = Some(analytics);
2227 }
2228
2229 #[cfg(feature = "edge-proxy")]
2232 pub async fn with_edge(&self, cache: Arc<EdgeCache>, registry: Arc<EdgeRegistry>) {
2233 *self.edge_cache.write().await = Some(cache);
2234 *self.edge_registry.write().await = Some(registry);
2235 }
2236
2237 #[cfg(feature = "ha-tr")]
2242 pub async fn with_replay_engine(&self, engine: Arc<ReplayEngine>) {
2243 *self.replay_engine.write().await = Some(engine);
2244 }
2245
2246 #[cfg(feature = "wasm-plugins")]
2250 pub async fn with_plugin_manager(&self, manager: Arc<PluginManager>) {
2251 *self.plugin_manager.write().await = Some(manager);
2252 }
2253
2254 pub async fn set_proxy_config(&self, config: ProxyConfig) {
2256 let mut proxy_config = self.proxy_config.write().await;
2257 *proxy_config = Some(config);
2258 }
2259
2260 pub async fn register_command<F>(&self, name: &str, handler: F)
2262 where
2263 F: Fn(&[&str]) -> Result<String> + Send + Sync + 'static,
2264 {
2265 let mut commands = self.commands.write().await;
2266 commands.insert(name.to_string(), Arc::new(handler));
2267 }
2268
2269 pub async fn execute_command(&self, name: &str, args: &[&str]) -> Result<String> {
2271 let commands = self.commands.read().await;
2272 match commands.get(name) {
2273 Some(handler) => handler(args),
2274 None => Err(ProxyError::Internal(format!("Unknown command: {}", name))),
2275 }
2276 }
2277}
2278
2279impl Default for AdminState {
2280 fn default() -> Self {
2281 Self::new()
2282 }
2283}
2284
2285#[derive(Debug, Deserialize)]
2289struct SqlRequest {
2290 query: String,
2292}
2293
2294#[derive(Debug, Serialize)]
2296struct SqlResponse {
2297 query_type: String,
2299 routed_to: String,
2301 node_role: String,
2303 result: serde_json::Value,
2305}
2306
2307#[derive(Serialize)]
2308struct HealthResponse {
2309 status: &'static str,
2310}
2311
2312#[derive(Serialize)]
2313struct ReadinessResponse {
2314 ready: bool,
2315 message: &'static str,
2316}
2317
2318#[derive(Serialize)]
2319struct LivenessResponse {
2320 alive: bool,
2321}
2322
2323#[derive(Serialize)]
2324struct ErrorResponse {
2325 error: String,
2326}
2327
2328#[derive(Serialize)]
2329struct MetricsResponse {
2330 connections_accepted: u64,
2331 connections_closed: u64,
2332 connections_active: u64,
2333 queries_processed: u64,
2334 bytes_received: u64,
2335 bytes_sent: u64,
2336 failovers: u64,
2337}
2338
2339impl From<ServerMetricsSnapshot> for MetricsResponse {
2340 fn from(m: ServerMetricsSnapshot) -> Self {
2341 Self {
2342 connections_accepted: m.connections_accepted,
2343 connections_closed: m.connections_closed,
2344 connections_active: m.connections_accepted.saturating_sub(m.connections_closed),
2345 queries_processed: m.queries_processed,
2346 bytes_received: m.bytes_received,
2347 bytes_sent: m.bytes_sent,
2348 failovers: m.failovers,
2349 }
2350 }
2351}
2352
2353#[derive(Serialize)]
2354struct NodeHealthResponse {
2355 address: String,
2356 healthy: bool,
2357 last_check: String,
2358 failure_count: u32,
2359 last_error: Option<String>,
2360 latency_ms: f64,
2361 replication_lag_bytes: Option<u64>,
2362}
2363
2364impl From<NodeHealth> for NodeHealthResponse {
2365 fn from(h: NodeHealth) -> Self {
2366 Self {
2367 address: h.address,
2368 healthy: h.healthy,
2369 last_check: h.last_check.to_rfc3339(),
2370 failure_count: h.failure_count,
2371 last_error: h.last_error,
2372 latency_ms: h.latency_ms,
2373 replication_lag_bytes: h.replication_lag_bytes,
2374 }
2375 }
2376}
2377
2378#[derive(Serialize)]
2379struct SessionsResponse {
2380 active_sessions: u64,
2381}
2382
2383#[cfg(feature = "edge-proxy")]
2385#[derive(Debug, Deserialize)]
2386struct EdgeRegisterBody {
2387 edge_id: String,
2388 region: String,
2389 base_url: String,
2390}
2391
2392#[cfg(feature = "edge-proxy")]
2396#[derive(Debug, Deserialize)]
2397struct EdgeInvalidateBody {
2398 #[serde(default)]
2399 tables: Vec<String>,
2400 #[serde(default)]
2401 up_to_version: Option<u64>,
2402}
2403
2404#[cfg(any(feature = "anomaly-detection", feature = "query-analytics"))]
2411fn parse_limit_query(path: &str, default: usize, max: usize) -> usize {
2412 let q = match path.find('?') {
2413 Some(i) => &path[i + 1..],
2414 None => return default,
2415 };
2416 for kv in q.split('&') {
2417 let mut it = kv.splitn(2, '=');
2418 if let (Some(k), Some(v)) = (it.next(), it.next()) {
2419 if k == "limit" {
2420 if let Ok(n) = v.parse::<usize>() {
2421 return n.min(max);
2422 }
2423 }
2424 }
2425 }
2426 default
2427}
2428
2429#[cfg(any(feature = "edge-proxy", feature = "wasm-plugins"))]
2435fn percent_decode(s: &str) -> String {
2436 let bytes = s.as_bytes();
2437 let mut out = Vec::with_capacity(bytes.len());
2438 let mut i = 0;
2439 while i < bytes.len() {
2440 if bytes[i] == b'%' && i + 2 < bytes.len() {
2441 let hi = (bytes[i + 1] as char).to_digit(16);
2442 let lo = (bytes[i + 2] as char).to_digit(16);
2443 if let (Some(hi), Some(lo)) = (hi, lo) {
2444 out.push(((hi as u8) << 4) | lo as u8);
2445 i += 3;
2446 continue;
2447 }
2448 }
2449 out.push(bytes[i]);
2450 i += 1;
2451 }
2452 String::from_utf8_lossy(&out).into_owned()
2453}
2454
2455#[cfg(any(feature = "edge-proxy", feature = "wasm-plugins"))]
2461fn parse_query_params(path: &str) -> HashMap<String, String> {
2462 let mut out = HashMap::new();
2463 let q = match path.find('?') {
2464 Some(i) => &path[i + 1..],
2465 None => return out,
2466 };
2467 for kv in q.split('&') {
2468 let mut it = kv.splitn(2, '=');
2469 if let (Some(k), Some(v)) = (it.next(), it.next()) {
2470 if !k.is_empty() {
2471 out.insert(k.to_string(), percent_decode(v));
2472 }
2473 }
2474 }
2475 out
2476}
2477
2478#[cfg(feature = "ha-tr")]
2480#[derive(Debug, Deserialize)]
2481struct ShadowRequestBody {
2482 sql: String,
2484 #[serde(default)]
2487 params: Option<Vec<String>>,
2488
2489 source_host: String,
2492 source_port: u16,
2493 #[serde(default)]
2494 source_user: Option<String>,
2495 #[serde(default)]
2496 source_password: Option<String>,
2497 #[serde(default)]
2498 source_database: Option<String>,
2499
2500 shadow_host: String,
2503 shadow_port: u16,
2504 #[serde(default)]
2505 shadow_user: Option<String>,
2506 #[serde(default)]
2507 shadow_password: Option<String>,
2508 #[serde(default)]
2509 shadow_database: Option<String>,
2510}
2511
2512#[derive(Debug, Deserialize)]
2517#[serde(tag = "action", rename_all = "snake_case")]
2518enum ChaosAction {
2519 ForceUnhealthy { target_node: String },
2523 Restore { target_node: String },
2525 Reset,
2527}
2528
2529#[cfg(feature = "wasm-plugins")]
2532#[derive(Serialize)]
2533struct PluginListEntry {
2534 name: String,
2535 version: String,
2536 description: String,
2537 hooks: Vec<String>,
2539 state: String,
2541 invocations: u64,
2542 errors: u64,
2543}
2544
2545#[cfg(feature = "ha-tr")]
2547#[derive(Debug, Deserialize)]
2548struct ReplayRequestBody {
2549 from: DateTime<Utc>,
2551 to: DateTime<Utc>,
2553 target_host: String,
2555 target_port: u16,
2557 #[serde(default)]
2563 target_user: Option<String>,
2564 #[serde(default)]
2565 target_password: Option<String>,
2566 #[serde(default)]
2567 target_database: Option<String>,
2568}
2569
2570#[derive(Serialize)]
2574struct TopologyResponse {
2575 #[serde(rename = "currentPrimary")]
2576 current_primary: Option<String>,
2577 #[serde(rename = "healthyNodes")]
2578 healthy_nodes: u32,
2579 #[serde(rename = "unhealthyNodes")]
2580 unhealthy_nodes: u32,
2581 #[serde(rename = "totalNodes")]
2582 total_nodes: u32,
2583 #[serde(rename = "lastFailoverAt")]
2586 last_failover_at: Option<String>,
2587}
2588
2589#[derive(Serialize)]
2590struct PoolStatsResponse {
2591 node: String,
2592 active_connections: u64,
2593 idle_connections: u64,
2594 pending_requests: u64,
2595 total_connections_created: u64,
2596 total_connections_closed: u64,
2597}
2598
2599#[derive(Serialize)]
2600struct VersionResponse {
2601 version: String,
2602 build_time: String,
2603}
2604
2605#[cfg(test)]
2606mod tests {
2607 use super::*;
2608
2609 #[tokio::test]
2610 async fn test_admin_state_creation() {
2611 let state = AdminState::new();
2612 let sessions = state.active_sessions.read().await;
2613 assert_eq!(*sessions, 0);
2614 }
2615
2616 #[test]
2617 fn test_admin_authorized() {
2618 let h = |s: &str| vec!["GET /topology HTTP/1.1".to_string(), s.to_string()];
2619 assert!(AdminServer::admin_authorized(
2620 &h("Authorization: Bearer s3cret"),
2621 "s3cret"
2622 ));
2623 assert!(AdminServer::admin_authorized(
2625 &h("authorization: Bearer s3cret"),
2626 "s3cret"
2627 ));
2628 assert!(!AdminServer::admin_authorized(
2630 &h("Authorization: Bearer nope"),
2631 "s3cret"
2632 ));
2633 assert!(!AdminServer::admin_authorized(
2634 &h("Authorization: Basic s3cret"),
2635 "s3cret"
2636 ));
2637 assert!(!AdminServer::admin_authorized(
2638 &["GET / HTTP/1.1".to_string()],
2639 "s3cret"
2640 ));
2641 }
2642
2643 #[test]
2644 fn test_constant_time_eq_str() {
2645 assert!(constant_time_eq_str("abc", "abc"));
2646 assert!(!constant_time_eq_str("abc", "abd"));
2647 assert!(!constant_time_eq_str("abc", "abcd"));
2648 }
2649
2650 #[tokio::test]
2651 async fn test_readiness_check_no_nodes() {
2652 let state = Arc::new(AdminState::new());
2653 let ready = AdminServer::check_readiness(&state).await;
2654 assert!(!ready);
2655 }
2656
2657 #[tokio::test]
2658 async fn test_readiness_check_with_healthy_node() {
2659 let state = Arc::new(AdminState::new());
2660
2661 {
2662 let mut health = state.node_health.write().await;
2663 health.insert(
2664 "localhost:5432".to_string(),
2665 NodeHealth {
2666 address: "localhost:5432".to_string(),
2667 healthy: true,
2668 last_check: chrono::Utc::now(),
2669 failure_count: 0,
2670 last_error: None,
2671 latency_ms: 1.0,
2672 replication_lag_bytes: None,
2673 },
2674 );
2675 }
2676
2677 let ready = AdminServer::check_readiness(&state).await;
2678 assert!(ready);
2679 }
2680
2681 #[tokio::test]
2682 async fn test_healthz_alias_matches_health() {
2683 let state = Arc::new(AdminState::new());
2684 let (status, body) = AdminServer::route_request("GET", "/health", None, &state)
2685 .await
2686 .unwrap();
2687 let (alias_status, alias_body) =
2688 AdminServer::route_request("GET", "/healthz", None, &state)
2689 .await
2690 .unwrap();
2691 assert_eq!(status, 200);
2692 assert_eq!(alias_status, 200);
2693 assert_eq!(body, alias_body);
2694 assert_eq!(alias_body, serde_json::json!({ "status": "ok" }));
2695 }
2696
2697 #[tokio::test]
2698 async fn test_livez_alias_matches_health_live() {
2699 let state = Arc::new(AdminState::new());
2700 let (status, body) = AdminServer::route_request("GET", "/health/live", None, &state)
2701 .await
2702 .unwrap();
2703 let (alias_status, alias_body) = AdminServer::route_request("GET", "/livez", None, &state)
2704 .await
2705 .unwrap();
2706 assert_eq!(status, 200);
2707 assert_eq!(alias_status, 200);
2708 assert_eq!(body, alias_body);
2709 assert_eq!(alias_body, serde_json::json!({ "alive": true }));
2710 }
2711
2712 #[tokio::test]
2713 async fn test_readyz_alias_matches_health_ready() {
2714 let state = Arc::new(AdminState::new());
2716 let (status, body) = AdminServer::route_request("GET", "/health/ready", None, &state)
2717 .await
2718 .unwrap();
2719 let (alias_status, alias_body) = AdminServer::route_request("GET", "/readyz", None, &state)
2720 .await
2721 .unwrap();
2722 assert_eq!(status, 503);
2723 assert_eq!(alias_status, status);
2724 assert_eq!(body, alias_body);
2725
2726 {
2728 let mut health = state.node_health.write().await;
2729 health.insert(
2730 "localhost:5432".to_string(),
2731 NodeHealth {
2732 address: "localhost:5432".to_string(),
2733 healthy: true,
2734 last_check: chrono::Utc::now(),
2735 failure_count: 0,
2736 last_error: None,
2737 latency_ms: 1.0,
2738 replication_lag_bytes: None,
2739 },
2740 );
2741 }
2742 let (status, body) = AdminServer::route_request("GET", "/health/ready", None, &state)
2743 .await
2744 .unwrap();
2745 let (alias_status, alias_body) = AdminServer::route_request("GET", "/readyz", None, &state)
2746 .await
2747 .unwrap();
2748 assert_eq!(status, 200);
2749 assert_eq!(alias_status, status);
2750 assert_eq!(body, alias_body);
2751 }
2752
2753 #[tokio::test]
2754 async fn test_command_registration() {
2755 let state = AdminState::new();
2756
2757 state
2758 .register_command("test", |args| {
2759 Ok(format!("Test command with {} args", args.len()))
2760 })
2761 .await;
2762
2763 let result = state.execute_command("test", &["arg1", "arg2"]).await;
2764 assert!(result.is_ok());
2765 assert_eq!(result.unwrap(), "Test command with 2 args");
2766 }
2767
2768 #[tokio::test]
2769 async fn test_unknown_command() {
2770 let state = AdminState::new();
2771 let result = state.execute_command("unknown", &[]).await;
2772 assert!(result.is_err());
2773 }
2774
2775 #[test]
2776 fn test_prometheus_metrics_format() {
2777 let metrics = ServerMetricsSnapshot {
2778 connections_accepted: 100,
2779 connections_closed: 50,
2780 queries_processed: 1000,
2781 bytes_received: 50000,
2782 bytes_sent: 100000,
2783 failovers: 2,
2784 };
2785
2786 let output = AdminServer::format_prometheus_metrics(&metrics);
2787 assert!(output.contains("heliosdb_proxy_connections_total 100"));
2788 assert!(output.contains("heliosdb_proxy_queries_total 1000"));
2789 assert!(output.contains("heliosdb_proxy_failovers_total 2"));
2790 }
2791
2792 #[test]
2793 fn test_metrics_response_active_connections() {
2794 let snapshot = ServerMetricsSnapshot {
2795 connections_accepted: 100,
2796 connections_closed: 30,
2797 queries_processed: 500,
2798 bytes_received: 10000,
2799 bytes_sent: 20000,
2800 failovers: 1,
2801 };
2802
2803 let response = MetricsResponse::from(snapshot);
2804 assert_eq!(response.connections_active, 70);
2805 }
2806
2807 async fn topology_state(nodes: &[(&str, &str, bool)]) -> Arc<AdminState> {
2810 let state = Arc::new(AdminState::new());
2811 {
2812 let mut cfg = state.config_snapshot.write().await;
2813 cfg.nodes = nodes
2814 .iter()
2815 .map(|(addr, role, _)| NodeSnapshot {
2816 address: (*addr).to_string(),
2817 role: (*role).to_string(),
2818 weight: 100,
2819 enabled: true,
2820 })
2821 .collect();
2822 }
2823 {
2824 let mut health = state.node_health.write().await;
2825 for (addr, _, healthy) in nodes {
2826 health.insert(
2827 (*addr).to_string(),
2828 NodeHealth {
2829 address: (*addr).to_string(),
2830 healthy: *healthy,
2831 last_check: chrono::Utc::now(),
2832 failure_count: 0,
2833 last_error: None,
2834 latency_ms: 1.0,
2835 replication_lag_bytes: None,
2836 },
2837 );
2838 }
2839 }
2840 state
2841 }
2842
2843 #[tokio::test]
2844 async fn test_topology_returns_healthy_primary() {
2845 let state = topology_state(&[
2846 ("primary.svc:5432", "primary", true),
2847 ("standby-a.svc:5432", "standby", true),
2848 ("standby-b.svc:5432", "standby", false),
2849 ])
2850 .await;
2851
2852 let topo = AdminServer::compute_topology(&state).await;
2853 assert_eq!(topo.current_primary.as_deref(), Some("primary.svc:5432"));
2854 assert_eq!(topo.healthy_nodes, 2);
2855 assert_eq!(topo.unhealthy_nodes, 1);
2856 assert_eq!(topo.total_nodes, 3);
2857 }
2858
2859 #[tokio::test]
2860 async fn test_topology_no_primary_when_primary_unhealthy() {
2861 let state = topology_state(&[
2864 ("primary.svc:5432", "primary", false),
2865 ("standby.svc:5432", "standby", true),
2866 ])
2867 .await;
2868
2869 let topo = AdminServer::compute_topology(&state).await;
2870 assert_eq!(topo.current_primary, None);
2871 assert_eq!(topo.healthy_nodes, 1);
2872 assert_eq!(topo.unhealthy_nodes, 1);
2873 }
2874
2875 #[tokio::test]
2876 async fn test_topology_handles_empty_cluster() {
2877 let state = Arc::new(AdminState::new());
2878 let topo = AdminServer::compute_topology(&state).await;
2879 assert_eq!(topo.current_primary, None);
2880 assert_eq!(topo.healthy_nodes, 0);
2881 assert_eq!(topo.unhealthy_nodes, 0);
2882 assert_eq!(topo.total_nodes, 0);
2883 }
2884
2885 #[tokio::test]
2886 async fn test_topology_role_match_is_case_insensitive() {
2887 let state = topology_state(&[("primary.svc:5432", "PRIMARY", true)]).await;
2888 let topo = AdminServer::compute_topology(&state).await;
2889 assert_eq!(topo.current_primary.as_deref(), Some("primary.svc:5432"));
2890 }
2891
2892 #[cfg(feature = "ha-tr")]
2893 #[tokio::test]
2894 async fn test_replay_returns_503_when_engine_unattached() {
2895 let state = Arc::new(AdminState::new());
2896 let body = r#"{
2897 "from": "2026-04-25T10:00:00Z",
2898 "to": "2026-04-25T11:00:00Z",
2899 "target_host": "127.0.0.1",
2900 "target_port": 5432
2901 }"#;
2902 let (status, value) = AdminServer::handle_replay_request(Some(body), &state)
2903 .await
2904 .expect("handler returns Ok with status code");
2905 assert_eq!(status, 503);
2906 assert_eq!(value["error"], "replay engine not attached");
2907 }
2908
2909 #[cfg(feature = "ha-tr")]
2910 #[tokio::test]
2911 async fn test_replay_400_on_malformed_body() {
2912 let state = Arc::new(AdminState::new());
2913 let (status, _) = AdminServer::handle_replay_request(Some("not json"), &state)
2914 .await
2915 .expect("handler returns Ok with status code");
2916 assert_eq!(status, 400);
2917 }
2918
2919 #[cfg(feature = "ha-tr")]
2920 #[tokio::test]
2921 async fn test_replay_errors_on_empty_body() {
2922 let state = Arc::new(AdminState::new());
2923 let err = AdminServer::handle_replay_request(None, &state).await;
2924 assert!(err.is_err(), "empty body must surface as Err");
2925 }
2926
2927 #[cfg(feature = "wasm-plugins")]
2928 #[tokio::test]
2929 async fn test_plugins_list_returns_503_when_manager_unattached() {
2930 let state = Arc::new(AdminState::new());
2931 let (status, value) = AdminServer::handle_plugins_list(&state)
2932 .await
2933 .expect("handler returns Ok with status code");
2934 assert_eq!(status, 503);
2935 assert_eq!(value["error"], "plugin manager not attached");
2936 }
2937
2938 #[cfg(not(feature = "wasm-plugins"))]
2939 #[tokio::test]
2940 async fn test_plugins_list_503_without_feature() {
2941 let state = Arc::new(AdminState::new());
2942 let (status, _) = AdminServer::handle_plugins_list(&state)
2943 .await
2944 .expect("handler returns Ok");
2945 assert_eq!(status, 503);
2946 }
2947
2948 #[cfg(feature = "wasm-plugins")]
2953 async fn kv_state_with_manager() -> Arc<AdminState> {
2954 use crate::plugins::{PluginManager, PluginRuntimeConfig};
2955 let state = Arc::new(AdminState::new());
2956 let pm = Arc::new(PluginManager::new(PluginRuntimeConfig::default()).unwrap());
2957 state.with_plugin_manager(pm).await;
2958 state
2959 }
2960
2961 #[cfg(feature = "wasm-plugins")]
2962 #[tokio::test]
2963 async fn test_kv_put_then_get_roundtrip() {
2964 let state = kv_state_with_manager().await;
2965 let (status, _) =
2966 AdminServer::route_request("PUT", "/admin/kv/demo/greeting", Some("hello"), &state)
2967 .await
2968 .expect("PUT returns Ok");
2969 assert_eq!(status, 200);
2970
2971 let (status, body) =
2972 AdminServer::route_request("GET", "/admin/kv/demo/greeting", None, &state)
2973 .await
2974 .expect("GET returns Ok");
2975 assert_eq!(status, 200);
2976 assert_eq!(body["plugin"], "demo");
2977 assert_eq!(body["key"], "greeting");
2978 assert_eq!(body["value"], "hello");
2979 }
2980
2981 #[cfg(feature = "wasm-plugins")]
2982 #[tokio::test]
2983 async fn test_kv_get_missing_returns_404() {
2984 let state = kv_state_with_manager().await;
2985 let (status, body) = AdminServer::route_request("GET", "/admin/kv/demo/nope", None, &state)
2986 .await
2987 .expect("GET returns Ok");
2988 assert_eq!(status, 404);
2989 assert_eq!(body["error"], "key not found");
2990 }
2991
2992 #[cfg(feature = "wasm-plugins")]
2993 #[tokio::test]
2994 async fn test_kv_delete_is_idempotent_200() {
2995 let state = kv_state_with_manager().await;
2996 let (status, _) = AdminServer::route_request("DELETE", "/admin/kv/demo/x", None, &state)
2998 .await
2999 .expect("DELETE returns Ok");
3000 assert_eq!(status, 200);
3001
3002 AdminServer::route_request("PUT", "/admin/kv/demo/x", Some("v"), &state)
3004 .await
3005 .expect("PUT returns Ok");
3006 let (status, _) = AdminServer::route_request("DELETE", "/admin/kv/demo/x", None, &state)
3007 .await
3008 .expect("DELETE returns Ok");
3009 assert_eq!(status, 200);
3010 let (status, _) = AdminServer::route_request("GET", "/admin/kv/demo/x", None, &state)
3011 .await
3012 .expect("GET returns Ok");
3013 assert_eq!(status, 404);
3014 }
3015
3016 #[cfg(feature = "wasm-plugins")]
3017 #[tokio::test]
3018 async fn test_kv_list_with_trailing_slash() {
3019 let state = kv_state_with_manager().await;
3020 AdminServer::route_request("PUT", "/admin/kv/demo/a", Some("1"), &state)
3021 .await
3022 .unwrap();
3023 AdminServer::route_request("PUT", "/admin/kv/demo/b", Some("2"), &state)
3024 .await
3025 .unwrap();
3026 let (status, body) = AdminServer::route_request("GET", "/admin/kv/demo/", None, &state)
3027 .await
3028 .expect("list returns Ok");
3029 assert_eq!(status, 200);
3030 assert_eq!(body["plugin"], "demo");
3031 let mut names: Vec<String> = body["keys"]
3032 .as_array()
3033 .expect("keys is an array")
3034 .iter()
3035 .map(|v| v.as_str().unwrap().to_string())
3036 .collect();
3037 names.sort();
3038 assert_eq!(names, vec!["a".to_string(), "b".to_string()]);
3039 }
3040
3041 #[cfg(feature = "wasm-plugins")]
3042 #[tokio::test]
3043 async fn test_kv_503_when_no_manager() {
3044 let state = Arc::new(AdminState::new());
3046 let (status, body) = AdminServer::route_request("GET", "/admin/kv/demo/x", None, &state)
3047 .await
3048 .expect("handler returns Ok");
3049 assert_eq!(status, 503);
3050 assert_eq!(body["error"], "plugin runtime not enabled");
3051 }
3052
3053 #[cfg(feature = "wasm-plugins")]
3054 #[tokio::test]
3055 async fn test_kv_malformed_path_returns_400() {
3056 let state = kv_state_with_manager().await;
3057 let (status, _) = AdminServer::route_request("GET", "/admin/kv/demo", None, &state)
3060 .await
3061 .expect("handler returns Ok");
3062 assert_eq!(status, 400);
3063 }
3064
3065 #[cfg(feature = "wasm-plugins")]
3066 #[tokio::test]
3067 async fn test_kv_wrong_method_returns_405() {
3068 let state = kv_state_with_manager().await;
3069 let (status, _) = AdminServer::route_request("POST", "/admin/kv/demo/x", Some("v"), &state)
3070 .await
3071 .expect("handler returns Ok");
3072 assert_eq!(status, 405);
3073 }
3074
3075 #[cfg(feature = "wasm-plugins")]
3076 #[tokio::test]
3077 async fn test_kv_put_413_on_value_cap() {
3078 use crate::plugins::{PluginManager, PluginRuntimeConfig};
3079 let cfg = PluginRuntimeConfig {
3080 kv_max_value_bytes: 3,
3081 ..PluginRuntimeConfig::default()
3082 };
3083 let state = Arc::new(AdminState::new());
3084 state
3085 .with_plugin_manager(Arc::new(PluginManager::new(cfg).unwrap()))
3086 .await;
3087 let (status, _) =
3089 AdminServer::route_request("PUT", "/admin/kv/demo/k", Some("toolong"), &state)
3090 .await
3091 .expect("handler returns Ok");
3092 assert_eq!(status, 413);
3093 }
3094
3095 #[cfg(feature = "wasm-plugins")]
3096 #[tokio::test]
3097 async fn test_kv_put_413_on_namespace_cap() {
3098 use crate::plugins::{PluginManager, PluginRuntimeConfig};
3099 let cfg = PluginRuntimeConfig {
3101 kv_max_plugins: 1,
3102 ..PluginRuntimeConfig::default()
3103 };
3104 let state = Arc::new(AdminState::new());
3105 state
3106 .with_plugin_manager(Arc::new(PluginManager::new(cfg).unwrap()))
3107 .await;
3108 let (status, _) = AdminServer::route_request("PUT", "/admin/kv/a/k", Some("1"), &state)
3110 .await
3111 .expect("handler returns Ok");
3112 assert_eq!(status, 200);
3113 let (status, body) = AdminServer::route_request("PUT", "/admin/kv/b/k", Some("2"), &state)
3116 .await
3117 .expect("handler returns Ok");
3118 assert_eq!(status, 413);
3119 assert!(body["error"].as_str().unwrap().contains("kv_max_plugins"));
3120 let (status, _) = AdminServer::route_request("PUT", "/admin/kv/a/k2", Some("3"), &state)
3122 .await
3123 .expect("handler returns Ok");
3124 assert_eq!(status, 200);
3125 }
3126
3127 #[cfg(feature = "wasm-plugins")]
3128 #[tokio::test]
3129 async fn test_kv_put_413_on_total_bytes_cap() {
3130 use crate::plugins::{PluginManager, PluginRuntimeConfig};
3131 let cfg = PluginRuntimeConfig {
3135 kv_max_total_bytes: 12,
3136 ..PluginRuntimeConfig::default()
3137 };
3138 let state = Arc::new(AdminState::new());
3139 state
3140 .with_plugin_manager(Arc::new(PluginManager::new(cfg).unwrap()))
3141 .await;
3142 let (status, _) = AdminServer::route_request("PUT", "/admin/kv/demo/a", Some("1"), &state)
3144 .await
3145 .expect("handler returns Ok");
3146 assert_eq!(status, 200);
3147 let (status, _) = AdminServer::route_request("PUT", "/admin/kv/demo/b", Some("1"), &state)
3149 .await
3150 .expect("handler returns Ok");
3151 assert_eq!(status, 200);
3152 let (status, body) =
3155 AdminServer::route_request("PUT", "/admin/kv/demo/c", Some("way-too-long"), &state)
3156 .await
3157 .expect("handler returns Ok");
3158 assert_eq!(status, 413);
3159 assert!(body["error"]
3160 .as_str()
3161 .unwrap()
3162 .contains("kv_max_total_bytes"));
3163 }
3164
3165 #[cfg(feature = "wasm-plugins")]
3166 #[tokio::test]
3167 async fn test_kv_empty_plugin_returns_400() {
3168 let state = kv_state_with_manager().await;
3169 let (status, body) = AdminServer::route_request("PUT", "/admin/kv//k", Some("v"), &state)
3171 .await
3172 .expect("handler returns Ok");
3173 assert_eq!(status, 400);
3174 assert_eq!(body["error"], "empty plugin name");
3175 let (status, body) = AdminServer::route_request("GET", "/admin/kv//", None, &state)
3177 .await
3178 .expect("handler returns Ok");
3179 assert_eq!(status, 400);
3180 assert_eq!(body["error"], "empty plugin name");
3181 }
3182
3183 #[cfg(feature = "wasm-plugins")]
3184 #[tokio::test]
3185 async fn test_kv_query_string_is_stripped() {
3186 let state = kv_state_with_manager().await;
3187 let (status, _) =
3189 AdminServer::route_request("PUT", "/admin/kv/demo/k?x=1", Some("v"), &state)
3190 .await
3191 .expect("handler returns Ok");
3192 assert_eq!(status, 200);
3193 let (status, body) = AdminServer::route_request("GET", "/admin/kv/demo/k", None, &state)
3195 .await
3196 .expect("handler returns Ok");
3197 assert_eq!(status, 200);
3198 assert_eq!(body["value"], "v");
3199 let (status, body) = AdminServer::route_request("GET", "/admin/kv/demo/", None, &state)
3201 .await
3202 .expect("handler returns Ok");
3203 assert_eq!(status, 200);
3204 let keys: Vec<String> = body["keys"]
3205 .as_array()
3206 .expect("keys array")
3207 .iter()
3208 .map(|v| v.as_str().unwrap().to_string())
3209 .collect();
3210 assert_eq!(keys, vec!["k".to_string()]);
3211 }
3212
3213 #[cfg(feature = "wasm-plugins")]
3214 #[tokio::test]
3215 async fn test_kv_list_prefix_filter() {
3216 let state = kv_state_with_manager().await;
3217 for (k, v) in [("ba", "1"), ("bb", "2"), ("cc", "3")] {
3218 AdminServer::route_request("PUT", &format!("/admin/kv/demo/{k}"), Some(v), &state)
3219 .await
3220 .unwrap();
3221 }
3222 let (status, body) =
3224 AdminServer::route_request("GET", "/admin/kv/demo/?prefix=b", None, &state)
3225 .await
3226 .expect("handler returns Ok");
3227 assert_eq!(status, 200);
3228 let mut keys: Vec<String> = body["keys"]
3229 .as_array()
3230 .expect("keys array")
3231 .iter()
3232 .map(|v| v.as_str().unwrap().to_string())
3233 .collect();
3234 keys.sort();
3235 assert_eq!(keys, vec!["ba".to_string(), "bb".to_string()]);
3236 }
3237
3238 #[cfg(not(feature = "wasm-plugins"))]
3239 #[tokio::test]
3240 async fn test_kv_501_without_feature() {
3241 let state = Arc::new(AdminState::new());
3242 let (status, body) = AdminServer::route_request("GET", "/admin/kv/demo/x", None, &state)
3243 .await
3244 .expect("handler returns Ok");
3245 assert_eq!(status, 501);
3246 assert!(body["error"].as_str().unwrap().contains("wasm-plugins"));
3247 }
3248
3249 async fn chaos_state_with_node(addr: &str) -> Arc<AdminState> {
3251 let state = Arc::new(AdminState::new());
3252 state.node_health.write().await.insert(
3253 addr.to_string(),
3254 NodeHealth {
3255 address: addr.to_string(),
3256 healthy: true,
3257 last_check: chrono::Utc::now(),
3258 failure_count: 0,
3259 last_error: None,
3260 latency_ms: 1.0,
3261 replication_lag_bytes: None,
3262 },
3263 );
3264 state
3265 }
3266
3267 #[tokio::test]
3268 async fn test_chaos_force_unhealthy_flips_node_and_records_override() {
3269 let state = chaos_state_with_node("primary.svc:5432").await;
3270 let body = r#"{"action":"force_unhealthy","target_node":"primary.svc:5432"}"#;
3271 let (status, value) = AdminServer::handle_chaos_request(Some(body), &state)
3272 .await
3273 .expect("handler returns Ok");
3274 assert_eq!(status, 200);
3275 assert_eq!(value["applied"], "force_unhealthy");
3276 assert!(!state.node_health.read().await["primary.svc:5432"].healthy);
3278 assert!(state
3280 .chaos_overrides
3281 .read()
3282 .await
3283 .contains_key("primary.svc:5432"));
3284 }
3285
3286 #[tokio::test]
3287 async fn test_chaos_restore_clears_override_and_flips_back() {
3288 let state = chaos_state_with_node("primary.svc:5432").await;
3289 let _ = AdminServer::handle_chaos_request(
3290 Some(r#"{"action":"force_unhealthy","target_node":"primary.svc:5432"}"#),
3291 &state,
3292 )
3293 .await
3294 .unwrap();
3295 let (status, _) = AdminServer::handle_chaos_request(
3296 Some(r#"{"action":"restore","target_node":"primary.svc:5432"}"#),
3297 &state,
3298 )
3299 .await
3300 .unwrap();
3301 assert_eq!(status, 200);
3302 assert!(state.node_health.read().await["primary.svc:5432"].healthy);
3303 assert!(state.chaos_overrides.read().await.is_empty());
3304 }
3305
3306 #[tokio::test]
3307 async fn test_chaos_reset_restores_all_overrides() {
3308 let state = chaos_state_with_node("a:5432").await;
3309 state.node_health.write().await.insert(
3310 "b:5432".to_string(),
3311 NodeHealth {
3312 address: "b:5432".to_string(),
3313 healthy: true,
3314 last_check: chrono::Utc::now(),
3315 failure_count: 0,
3316 last_error: None,
3317 latency_ms: 1.0,
3318 replication_lag_bytes: None,
3319 },
3320 );
3321 for addr in &["a:5432", "b:5432"] {
3322 let body = format!(r#"{{"action":"force_unhealthy","target_node":"{}"}}"#, addr);
3323 let _ = AdminServer::handle_chaos_request(Some(&body), &state)
3324 .await
3325 .unwrap();
3326 }
3327 let (status, value) =
3328 AdminServer::handle_chaos_request(Some(r#"{"action":"reset"}"#), &state)
3329 .await
3330 .unwrap();
3331 assert_eq!(status, 200);
3332 assert_eq!(value["reset"], true);
3333 let restored = value["restored"].as_array().unwrap();
3334 assert_eq!(restored.len(), 2);
3335 for addr in &["a:5432", "b:5432"] {
3337 assert!(state.node_health.read().await[*addr].healthy);
3338 }
3339 assert!(state.chaos_overrides.read().await.is_empty());
3340 }
3341
3342 #[tokio::test]
3343 async fn test_chaos_force_unhealthy_404s_when_node_unknown() {
3344 let state = Arc::new(AdminState::new());
3345 let body = r#"{"action":"force_unhealthy","target_node":"missing.svc:5432"}"#;
3346 let (status, _) = AdminServer::handle_chaos_request(Some(body), &state)
3347 .await
3348 .expect("handler returns Ok");
3349 assert_eq!(status, 404);
3350 }
3351
3352 #[tokio::test]
3353 async fn test_chaos_400_on_malformed_body() {
3354 let state = Arc::new(AdminState::new());
3355 let (status, _) = AdminServer::handle_chaos_request(Some("not json"), &state)
3356 .await
3357 .expect("handler returns Ok");
3358 assert_eq!(status, 400);
3359 }
3360
3361 #[tokio::test]
3362 async fn test_chaos_400_on_unknown_action() {
3363 let state = Arc::new(AdminState::new());
3364 let body = r#"{"action":"format_disk","target_node":"x"}"#;
3365 let (status, _) = AdminServer::handle_chaos_request(Some(body), &state)
3366 .await
3367 .expect("handler returns Ok");
3368 assert_eq!(status, 400);
3369 }
3370
3371 #[cfg(feature = "ha-tr")]
3372 #[tokio::test]
3373 async fn test_shadow_400_on_malformed_body() {
3374 let (status, _) = AdminServer::handle_shadow_request(Some("not json"))
3375 .await
3376 .expect("handler returns Ok");
3377 assert_eq!(status, 400);
3378 }
3379
3380 #[cfg(feature = "ha-tr")]
3381 #[tokio::test]
3382 async fn test_shadow_500_on_source_unreachable() {
3383 let body = r#"{
3386 "sql": "SELECT 1",
3387 "source_host": "127.0.0.1",
3388 "source_port": 1,
3389 "shadow_host": "127.0.0.1",
3390 "shadow_port": 1
3391 }"#;
3392 let (status, value) = AdminServer::handle_shadow_request(Some(body))
3393 .await
3394 .expect("handler returns Ok");
3395 assert_eq!(status, 500);
3396 let err = value["error"].as_str().expect("error field");
3397 assert!(
3398 err.contains("source connect"),
3399 "expected source connect error, got {}",
3400 err
3401 );
3402 }
3403
3404 #[cfg(feature = "ha-tr")]
3405 #[tokio::test]
3406 async fn test_shadow_errors_on_empty_body() {
3407 let err = AdminServer::handle_shadow_request(None).await;
3408 assert!(err.is_err(), "empty body must surface as Err");
3409 }
3410
3411 #[cfg(feature = "anomaly-detection")]
3412 #[tokio::test]
3413 async fn test_anomalies_returns_503_when_detector_unattached() {
3414 let state = Arc::new(AdminState::new());
3415 let (status, value) = AdminServer::handle_anomalies_list("/anomalies", &state)
3416 .await
3417 .expect("handler returns Ok");
3418 assert_eq!(status, 503);
3419 assert_eq!(value["error"], "anomaly detector not attached");
3420 }
3421
3422 #[cfg(feature = "anomaly-detection")]
3423 #[tokio::test]
3424 async fn test_anomalies_returns_attached_detector_events() {
3425 use crate::anomaly::{AnomalyConfig, AnomalyDetector, QueryObservation};
3426 let state = Arc::new(AdminState::new());
3427 let det = Arc::new(AnomalyDetector::new(AnomalyConfig::default()));
3428 let _ = det.record_query(&QueryObservation {
3430 tenant: "test".into(),
3431 fingerprint: "fp".into(),
3432 sql: "SELECT * FROM users WHERE id = 1 OR 1=1 --".into(),
3433 timestamp: std::time::Instant::now(),
3434 });
3435 state.with_anomaly_detector(det.clone()).await;
3436
3437 let (status, value) = AdminServer::handle_anomalies_list("/anomalies", &state)
3438 .await
3439 .expect("handler returns Ok");
3440 assert_eq!(status, 200);
3441 let count = value["count"].as_u64().expect("count field");
3442 assert!(count > 0, "expected at least one event, got {}", count);
3443 }
3444
3445 #[cfg(feature = "anomaly-detection")]
3446 #[tokio::test]
3447 async fn test_anomalies_limit_query_string_respected() {
3448 use crate::anomaly::{AnomalyConfig, AnomalyDetector, QueryObservation};
3449 let state = Arc::new(AdminState::new());
3450 let det = Arc::new(AnomalyDetector::new(AnomalyConfig::default()));
3451 for i in 0..50 {
3452 let fp = format!("fp{}", i);
3453 let _ = det.record_query(&QueryObservation {
3454 tenant: "test".into(),
3455 fingerprint: fp,
3456 sql: "SELECT 1".into(),
3457 timestamp: std::time::Instant::now(),
3458 });
3459 }
3460 state.with_anomaly_detector(det).await;
3461
3462 let (status, value) = AdminServer::handle_anomalies_list("/anomalies?limit=5", &state)
3463 .await
3464 .expect("handler returns Ok");
3465 assert_eq!(status, 200);
3466 assert_eq!(value["limit"].as_u64().unwrap(), 5);
3467 assert_eq!(value["events"].as_array().unwrap().len(), 5);
3468 }
3469
3470 #[cfg(any(feature = "anomaly-detection", feature = "query-analytics"))]
3471 #[test]
3472 fn test_parse_limit_query_helper() {
3473 assert_eq!(parse_limit_query("/anomalies", 100, 1024), 100);
3474 assert_eq!(parse_limit_query("/anomalies?limit=42", 100, 1024), 42);
3475 assert_eq!(parse_limit_query("/anomalies?limit=99999", 100, 1024), 1024);
3476 assert_eq!(parse_limit_query("/anomalies?limit=abc", 100, 1024), 100);
3477 assert_eq!(
3478 parse_limit_query("/anomalies?other=x&limit=7", 100, 1024),
3479 7
3480 );
3481 }
3482
3483 #[test]
3488 fn test_admin_ui_html_contains_token_handling() {
3489 assert!(
3490 ADMIN_UI_HTML.contains("helios_admin_token"),
3491 "admin UI must inject the admin bearer token via the helios_admin_token key"
3492 );
3493 }
3494
3495 #[test]
3502 fn test_admin_ui_html_contains_esc_helper() {
3503 assert!(
3504 ADMIN_UI_HTML.contains("function esc("),
3505 "admin UI must define the esc() HTML-escape helper to prevent stored XSS in innerHTML sinks"
3506 );
3507 }
3508
3509 #[cfg(feature = "edge-proxy")]
3510 async fn edge_state() -> Arc<AdminState> {
3511 use crate::edge::{EdgeCache, EdgeRegistry};
3512 use std::time::Duration;
3513 let s = Arc::new(AdminState::new());
3514 let cache = Arc::new(EdgeCache::new(100));
3515 let registry = Arc::new(EdgeRegistry::new(8, Duration::from_secs(60)));
3516 s.with_edge(cache, registry).await;
3517 s
3518 }
3519
3520 #[cfg(feature = "edge-proxy")]
3521 #[tokio::test]
3522 async fn test_edge_status_returns_empty_lists_initially() {
3523 let s = edge_state().await;
3524 let (status, value) = AdminServer::handle_edge_status(&s)
3525 .await
3526 .expect("handler returns Ok");
3527 assert_eq!(status, 200);
3528 assert_eq!(value["edge_count"].as_u64().unwrap(), 0);
3529 assert_eq!(value["registered"].as_array().unwrap().len(), 0);
3530 assert!(value["cache"].is_object(), "cache stats present");
3531 }
3532
3533 #[cfg(feature = "edge-proxy")]
3534 #[tokio::test]
3535 async fn test_edge_register_then_status_lists_edge() {
3536 let s = edge_state().await;
3537 let body = r#"{"edge_id":"e1","region":"us-east","base_url":"https://e1.svc"}"#;
3538 let (status, _) = AdminServer::handle_edge_register(Some(body), &s)
3539 .await
3540 .expect("handler ok");
3541 assert_eq!(status, 201);
3542 let (status2, value2) = AdminServer::handle_edge_status(&s).await.unwrap();
3543 assert_eq!(status2, 200);
3544 assert_eq!(value2["edge_count"].as_u64().unwrap(), 1);
3545 assert_eq!(value2["registered"][0]["edge_id"].as_str().unwrap(), "e1");
3546 }
3547
3548 #[cfg(feature = "edge-proxy")]
3549 #[tokio::test]
3550 async fn test_edge_register_400_on_malformed_body() {
3551 let s = edge_state().await;
3552 let (status, _) = AdminServer::handle_edge_register(Some("not json"), &s)
3553 .await
3554 .expect("handler ok");
3555 assert_eq!(status, 400);
3556 }
3557
3558 #[cfg(feature = "edge-proxy")]
3559 #[tokio::test]
3560 async fn test_edge_invalidate_drops_local_cache_entries() {
3561 use crate::edge::{CacheEntry, CacheKey};
3562 use std::time::{Duration, Instant};
3563 let s = edge_state().await;
3564 let cache = s.edge_cache.read().await.clone().unwrap();
3566 cache.insert(
3567 CacheKey::new("fp1", "p1"),
3568 CacheEntry {
3569 version: 1,
3570 response_bytes: bytes::Bytes::from_static(b"row"),
3571 tables: vec!["users".into()],
3572 expires_at: Instant::now() + Duration::from_secs(60),
3573 },
3574 );
3575 assert!(cache.get(&CacheKey::new("fp1", "p1")).is_some());
3576
3577 let body = r#"{"tables":["users"]}"#;
3578 let (status, value) = AdminServer::handle_edge_invalidate(Some(body), &s)
3579 .await
3580 .expect("handler ok");
3581 assert_eq!(status, 200);
3582 assert_eq!(value["dropped_local"].as_u64().unwrap(), 1);
3583 assert!(cache.get(&CacheKey::new("fp1", "p1")).is_none());
3584 }
3585
3586 #[cfg(feature = "edge-proxy")]
3587 #[tokio::test]
3588 async fn test_edge_invalidate_503_when_cache_unattached() {
3589 let s = Arc::new(AdminState::new());
3590 let body = r#"{"tables":["users"]}"#;
3591 let (status, _) = AdminServer::handle_edge_invalidate(Some(body), &s)
3592 .await
3593 .expect("handler ok");
3594 assert_eq!(status, 503);
3595 }
3596
3597 #[cfg(feature = "edge-proxy")]
3600 #[test]
3601 fn test_parse_query_params_decodes_and_defaults() {
3602 let p = parse_query_params(
3603 "/api/edge/subscribe?edge_id=e1®ion=us-east&base_url=http%3A%2F%2Fe1.svc%3A9090",
3604 );
3605 assert_eq!(p.get("edge_id").unwrap(), "e1");
3606 assert_eq!(p.get("region").unwrap(), "us-east");
3607 assert_eq!(p.get("base_url").unwrap(), "http://e1.svc:9090");
3608 assert!(parse_query_params("/api/edge/subscribe").is_empty());
3610 let p2 = parse_query_params("/x?flag&edge_id=a%2&b=%zz");
3613 assert!(!p2.contains_key("flag"));
3614 assert_eq!(p2.get("edge_id").unwrap(), "a%2");
3615 assert_eq!(p2.get("b").unwrap(), "%zz");
3616 }
3617
3618 #[cfg(feature = "edge-proxy")]
3622 async fn connect_admin(state: Arc<AdminState>) -> tokio::net::TcpStream {
3623 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
3624 let addr = listener.local_addr().unwrap();
3625 tokio::spawn(async move {
3626 let (stream, peer) = listener.accept().await.unwrap();
3627 let _ = AdminServer::handle_connection(stream, peer, state).await;
3628 });
3629 tokio::net::TcpStream::connect(addr).await.unwrap()
3630 }
3631
3632 #[cfg(feature = "edge-proxy")]
3635 async fn read_until(stream: &mut tokio::net::TcpStream, needle: &str) -> String {
3636 let mut buf = Vec::new();
3637 let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
3638 loop {
3639 let mut chunk = [0u8; 1024];
3640 let n = tokio::time::timeout_at(deadline, stream.read(&mut chunk))
3641 .await
3642 .expect("read timed out")
3643 .expect("read failed");
3644 if n == 0 {
3645 break;
3646 }
3647 buf.extend_from_slice(&chunk[..n]);
3648 if String::from_utf8_lossy(&buf).contains(needle) {
3649 break;
3650 }
3651 }
3652 String::from_utf8_lossy(&buf).into_owned()
3653 }
3654
3655 #[cfg(feature = "edge-proxy")]
3656 #[tokio::test]
3657 async fn test_edge_subscribe_unauthenticated_gets_401() {
3658 let s = edge_state().await;
3659 *s.auth_token.write().await = Some("s3cret".to_string());
3660 let mut c = connect_admin(s).await;
3661 c.write_all(b"GET /api/edge/subscribe?edge_id=e1 HTTP/1.1\r\n\r\n")
3662 .await
3663 .unwrap();
3664 let resp = read_until(&mut c, "401").await;
3665 assert!(resp.starts_with("HTTP/1.1 401"), "got: {resp}");
3666 }
3667
3668 #[cfg(feature = "edge-proxy")]
3669 #[tokio::test]
3670 async fn test_edge_subscribe_400_without_edge_id() {
3671 let s = edge_state().await;
3672 let mut c = connect_admin(s).await;
3673 c.write_all(b"GET /api/edge/subscribe?region=us-east HTTP/1.1\r\n\r\n")
3674 .await
3675 .unwrap();
3676 let resp = read_until(&mut c, "edge_id").await;
3677 assert!(resp.starts_with("HTTP/1.1 400"), "got: {resp}");
3678 assert!(resp.contains("edge_id query parameter is required"));
3679 }
3680
3681 #[cfg(feature = "edge-proxy")]
3682 #[tokio::test]
3683 async fn test_edge_subscribe_503_when_registry_full() {
3684 use crate::edge::{EdgeCache, EdgeRegistry};
3685 let s = Arc::new(AdminState::new());
3686 s.with_edge(
3687 Arc::new(EdgeCache::new(16)),
3688 Arc::new(EdgeRegistry::new(1, std::time::Duration::from_secs(60))),
3689 )
3690 .await;
3691 let registry = s.edge_registry.read().await.clone().unwrap();
3692 let _held = registry.register("occupant", "r", "u", "ts").unwrap();
3695 let mut c = connect_admin(s).await;
3696 c.write_all(b"GET /api/edge/subscribe?edge_id=e2 HTTP/1.1\r\n\r\n")
3697 .await
3698 .unwrap();
3699 let resp = read_until(&mut c, "full").await;
3700 assert!(resp.starts_with("HTTP/1.1 503"), "got: {resp}");
3701 }
3702
3703 #[cfg(feature = "edge-proxy")]
3704 #[tokio::test]
3705 async fn test_edge_subscribe_streams_preamble_and_invalidations() {
3706 let s = edge_state().await;
3707 *s.auth_token.write().await = Some("s3cret".to_string());
3708 let registry = s.edge_registry.read().await.clone().unwrap();
3709 let mut c = connect_admin(s).await;
3710 c.write_all(
3711 b"GET /api/edge/subscribe?edge_id=e1®ion=eu&base_url=http%3A%2F%2Fe1 HTTP/1.1\r\nAuthorization: Bearer s3cret\r\n\r\n",
3712 )
3713 .await
3714 .unwrap();
3715 let preamble = read_until(&mut c, "\r\n\r\n").await;
3716 assert!(preamble.starts_with("HTTP/1.1 200 OK"), "got: {preamble}");
3717 assert!(preamble.contains("Content-Type: text/event-stream"));
3718
3719 let nodes = registry.list();
3723 assert_eq!(nodes.len(), 1);
3724 assert_eq!(nodes[0].edge_id, "e1");
3725 assert_eq!(nodes[0].region, "eu");
3726 assert_eq!(nodes[0].base_url, "http://e1");
3727
3728 let mut hello = preamble;
3734 if !hello.contains("}\n\n") {
3735 hello.push_str(&read_until(&mut c, "}\n\n").await);
3736 }
3737 assert!(hello.contains("event: invalidate"), "got: {hello}");
3738 assert!(hello.contains("\"epoch\":"), "got: {hello}");
3739 assert!(
3740 !hello.contains("\"epoch\":0}"),
3741 "hello epoch must be non-zero: {hello}"
3742 );
3743
3744 let (sent, _) = registry
3746 .broadcast(InvalidationEvent {
3747 up_to_version: 7,
3748 tables: vec!["users".into()],
3749 committed_at: "ts".into(),
3750 epoch: 0,
3751 })
3752 .await;
3753 assert_eq!(sent, 1);
3754 let frame = read_until(&mut c, "\"up_to_version\":7").await;
3755 assert!(frame.contains("event: invalidate"), "got: {frame}");
3756 assert!(frame.contains("\"up_to_version\":7"), "got: {frame}");
3757 }
3758}