use super::{
target::{
forward_target, join_optional_target, join_required_target, optional_forward_target,
response_from_target, selected_response,
},
TargetRequest,
};
use crate::AppState;
use axum::{
extract::{OriginalUri, State},
http::Method,
response::{IntoResponse, Response},
};
use bytes::Bytes;
use chrono::Utc;
use futures::future::{join3, BoxFuture};
use moonlight_core::{
compare::{capture_body_with_redaction_patterns, capture_headers, CompareConfig},
config::{AppConfig, ResponseTiming, ReturnFallback, ReturnTarget},
run::{build_comparison_run, CapturedTargets, RunMetadata},
target::CapturedTarget,
Adapter, ComparisonRun, RunInput,
};
use std::sync::Arc;
use tokio::task::JoinHandle;
use tracing::Instrument;
use uuid::Uuid;
pub async fn proxy_handler(
State(state): State<Arc<AppState>>,
OriginalUri(uri): OriginalUri,
method: Method,
headers: axum::http::HeaderMap,
body: Bytes,
) -> Response {
if body.len() > state.config.max_request_body_bytes {
return axum::http::StatusCode::PAYLOAD_TOO_LARGE.into_response();
}
let id = Uuid::new_v4();
let timestamp = Utc::now();
let method_label = method.to_string();
let path_label = uri.path().to_string();
let path = uri.path().to_string();
let query = uri
.query()
.map(|query| redact_query(query, &state.config.redact_query_params));
let path_and_query = uri
.path_and_query()
.map(|value| value.as_str().to_string())
.unwrap_or_else(|| path.clone());
let metadata = RunMetadata {
id,
timestamp,
adapter: Adapter::Http,
input: RunInput::Http {
method: method_label.clone(),
path,
query,
},
request_headers: capture_headers(&headers, &state.config.redact_headers),
request_body: capture_body_with_redaction_patterns(
&body,
state.config.max_body_capture_bytes,
&state.config.redact_json_paths,
&state.config.redact_json_path_patterns,
),
};
let target_request = TargetRequest {
method,
path_and_query,
headers,
body,
};
state.metrics.record_comparison_started();
let span = tracing::info_span!(
"moonlight_proxy_run",
run_id = %id,
method = %method_label,
path = %path_label,
return_target = ?state.config.return_target,
response_timing = ?state.config.response_timing,
);
async move {
let primary = forward_target(
state.clone(),
"primary",
state.config.primary_url.clone(),
target_request.clone(),
);
let candidate = forward_target(
state.clone(),
"candidate",
state.config.candidate_url.clone(),
target_request.clone(),
);
let secondary = optional_forward_target(
state.clone(),
"secondary",
state.config.secondary_url.clone(),
state.config.enable_secondary,
target_request,
);
match state.config.response_timing {
ResponseTiming::WaitAll => {
proxy_wait_all(state, metadata, primary, candidate, secondary).await
}
ResponseTiming::ReturnSelected => {
proxy_return_selected(state, metadata, primary, candidate, secondary).await
}
}
}
.instrument(span)
.await
}
async fn proxy_wait_all(
state: Arc<AppState>,
metadata: RunMetadata,
primary: BoxFuture<'static, CapturedTarget>,
candidate: BoxFuture<'static, CapturedTarget>,
secondary: BoxFuture<'static, Option<CapturedTarget>>,
) -> Response {
let (primary, candidate, secondary) = join3(primary, candidate, secondary).await;
let response = selected_response(&state, &primary, &candidate);
let id = metadata.id;
let run = build_run(metadata, primary, candidate, secondary, &state.config);
persist_run(state, id, run).await;
response
}
async fn proxy_return_selected(
state: Arc<AppState>,
metadata: RunMetadata,
primary: BoxFuture<'static, CapturedTarget>,
candidate: BoxFuture<'static, CapturedTarget>,
secondary: BoxFuture<'static, Option<CapturedTarget>>,
) -> Response {
match state.config.return_target {
ReturnTarget::Primary => {
let candidate = tokio::spawn(candidate);
let secondary = tokio::spawn(secondary);
let primary = primary.await;
let response = response_from_target(&primary);
spawn_persist_run(state, metadata, primary, candidate, secondary);
response
}
ReturnTarget::Candidate => {
let primary = tokio::spawn(primary);
let secondary = tokio::spawn(secondary);
let candidate = candidate.await;
if candidate.observation.error.is_some()
&& state.config.return_fallback == ReturnFallback::Primary
{
let primary = join_required_target(primary, "primary").await;
let response = response_from_target(&primary);
spawn_persist_run_with_primary(state, metadata, primary, candidate, secondary);
response
} else {
let response = response_from_target(&candidate);
spawn_persist_run_with_candidate(state, metadata, primary, candidate, secondary);
response
}
}
}
}
fn spawn_persist_run(
state: Arc<AppState>,
metadata: RunMetadata,
primary: CapturedTarget,
candidate: JoinHandle<CapturedTarget>,
secondary: JoinHandle<Option<CapturedTarget>>,
) {
tokio::spawn(async move {
let candidate = join_required_target(candidate, "candidate").await;
persist_run_with_targets(state, metadata, primary, candidate, secondary).await;
});
}
fn spawn_persist_run_with_primary(
state: Arc<AppState>,
metadata: RunMetadata,
primary: CapturedTarget,
candidate: CapturedTarget,
secondary: JoinHandle<Option<CapturedTarget>>,
) {
tokio::spawn(async move {
persist_run_with_targets(state, metadata, primary, candidate, secondary).await;
});
}
fn spawn_persist_run_with_candidate(
state: Arc<AppState>,
metadata: RunMetadata,
primary: JoinHandle<CapturedTarget>,
candidate: CapturedTarget,
secondary: JoinHandle<Option<CapturedTarget>>,
) {
tokio::spawn(async move {
let primary = join_required_target(primary, "primary").await;
persist_run_with_targets(state, metadata, primary, candidate, secondary).await;
});
}
async fn persist_run_with_targets(
state: Arc<AppState>,
metadata: RunMetadata,
primary: CapturedTarget,
candidate: CapturedTarget,
secondary: JoinHandle<Option<CapturedTarget>>,
) {
let secondary = join_optional_target(secondary, "secondary").await;
let id = metadata.id;
let run = build_run(metadata, primary, candidate, secondary, &state.config);
persist_run(state, id, run).await;
}
async fn persist_run(state: Arc<AppState>, id: Uuid, run: ComparisonRun) {
let classification = run.comparison.classification.clone();
let target_error_count = [
run.primary.error.as_ref(),
run.candidate.error.as_ref(),
run.secondary
.as_ref()
.and_then(|target| target.error.as_ref()),
]
.into_iter()
.flatten()
.count();
match state.storage.insert(run.clone()).await {
Ok(()) => {
state.metrics.record_persisted_run(&run);
tracing::info!(
run_id = %id,
classification = ?classification,
target_error_count,
"persisted moonlight comparison run"
);
}
Err(error) => {
state.metrics.record_persistence_failure();
tracing::error!(
run_id = %id,
error = %error,
"failed to persist moonlight comparison run"
);
}
}
}
fn build_run(
metadata: RunMetadata,
primary: CapturedTarget,
candidate: CapturedTarget,
secondary: Option<CapturedTarget>,
config: &AppConfig,
) -> ComparisonRun {
let compare_config = CompareConfig::new_with_patterns(
&config.ignore_json_paths,
&config.ignore_json_path_patterns,
&config.redact_json_paths,
&config.redact_json_path_patterns,
&config.ignore_headers,
config.ignore_stderr,
);
build_comparison_run(
metadata,
CapturedTargets {
primary,
candidate,
secondary,
},
&compare_config,
)
}
fn redact_query(query: &str, redact_query_params: &[String]) -> String {
query
.split('&')
.map(|part| {
let (key, separator, value) = match part.split_once('=') {
Some((key, value)) => (key, "=", value),
None => (part, "", ""),
};
if redact_query_params
.iter()
.any(|redacted| redacted.eq_ignore_ascii_case(key))
{
format!("{key}{separator}[redacted]")
} else if separator.is_empty() {
key.to_string()
} else {
format!("{key}={value}")
}
})
.collect::<Vec<_>>()
.join("&")
}