diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index d815620..10c138e 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -13,6 +13,7 @@ AuditBridge uses environment variables. A direct secret variable and its | `RETRY_MAX_ATTEMPTS` | `5` | Attempts per fetch or delivery cycle | | `RETRY_BASE_DELAY_MS` | `500` | Full-jitter retry backoff base in milliseconds | | `RETRY_MAX_DELAY_MS` | `30000` | Retry backoff maximum in milliseconds | +| `BATCH_SIZE` | `500` | Maximum events per delivery request; `0` disables chunking | | `CURSOR_FILE` | unset | Persistent per-sink delivery watermark path | | `METRICS_PORT` | `9090` | Health and Prometheus HTTP server port | | `RUST_LOG` | `info` | Rust log filter | @@ -24,4 +25,11 @@ AuditBridge uses environment variables. A direct secret variable and its > NetBird's audit endpoint returns the full history and exposes no > server-side paging or cursor parameters. +> [!NOTE] +> A fresh install, a lost cursor, or a restored old cursor delivers the whole +> audit history in one poll. `BATCH_SIZE` splits that replay into +> `BATCH_SIZE`-event requests so payloads stay under intake limits (Loki's +> push API and most HTTP/SIEM endpoints cap request size); the per-sink +> watermark only advances when every chunk has been delivered. + See [Sinks](SINKS.md) for sink-specific variables. diff --git a/src/config.rs b/src/config.rs index b92f760..26c0810 100644 --- a/src/config.rs +++ b/src/config.rs @@ -51,6 +51,8 @@ pub struct Config { // just mean a silently-discarded file every restart. pub cursor_file: Option, pub retry: RetryConfig, + // 0 means "no chunking": send the whole pending batch in one request. + pub batch_size: usize, pub metrics_port: u16, } @@ -74,6 +76,10 @@ impl Config { sinks, cursor_file: env::var("CURSOR_FILE").ok(), retry: RetryConfig::from_env(), + batch_size: env::var("BATCH_SIZE") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(500), metrics_port: env::var("METRICS_PORT") .ok() .and_then(|v| v.parse().ok()) diff --git a/src/main.rs b/src/main.rs index 6c75d6f..09574d5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -119,7 +119,15 @@ async fn run( return; } - process_cycle(nb_client, sinks, cursors, &config.retry, metrics).await; + process_cycle( + nb_client, + sinks, + cursors, + &config.retry, + config.batch_size, + metrics, + ) + .await; if let Some(path) = &config.cursor_file { let to_save: HashMap> = cursors @@ -187,6 +195,7 @@ async fn process_cycle( sinks: &[Box], cursors: &mut HashMap>>, retry_cfg: &RetryConfig, + batch_size: usize, metrics: &Metrics, ) { let mut events = match with_retry("netbird fetch", retry_cfg, || nb_client.fetch_events()).await @@ -222,20 +231,44 @@ async fn process_cycle( let count = pending.len(); let op_name = format!("sink '{}' send", sink.name()); - match with_retry(&op_name, retry_cfg, || sink.send(&pending)).await { - Ok(_) => { - if let Some(last_event) = pending.last() { - if let Ok(ts) = DateTime::parse_from_rfc3339(&last_event.timestamp) { - *cursor = Some(ts.with_timezone(&Utc)); - } + // A full-history replay (fresh install, lost cursor) can produce a + // very large batch in one poll; most intake endpoints cap payload + // size, so split into chunks that each get the usual retry/backoff. + // 0 = no chunking. The watermark advances only when every chunk + // delivers, so a partial failure retries the whole set next cycle. + let chunk_size = if batch_size == 0 { + count.max(1) + } else { + batch_size + }; + + let mut all_delivered = true; + for chunk in pending.chunks(chunk_size) { + match with_retry(&op_name, retry_cfg, || sink.send(chunk)).await { + Ok(_) => {} + Err(e) => { + all_delivered = false; + metrics.record_sink_error(sink.name()); + error!( + "Failed to send {} of {} events to {}: {}", + chunk.len(), + count, + sink.name(), + e + ); + break; } - metrics.record_sink_success(sink.name(), count); - info!("Delivered {} events to {}", count, sink.name()); } - Err(e) => { - metrics.record_sink_error(sink.name()); - error!("Failed to send {} events to {}: {}", count, sink.name(), e); + } + + if all_delivered { + if let Some(last_event) = pending.last() { + if let Ok(ts) = DateTime::parse_from_rfc3339(&last_event.timestamp) { + *cursor = Some(ts.with_timezone(&Utc)); + } } + metrics.record_sink_success(sink.name(), count); + info!("Delivered {} events to {}", count, sink.name()); } } } diff --git a/src/tests.rs b/src/tests.rs index bdcc869..8d3d797 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -35,6 +35,7 @@ fn test_config() -> Config { sinks: vec![], cursor_file: None, retry: no_retry(), + batch_size: 500, metrics_port: 0, } } @@ -188,6 +189,7 @@ async fn test_process_cycle_does_not_advance_watermark_on_send_failure() { &sinks, &mut cursors, &no_retry(), + 500, &Metrics::default(), ) .await; @@ -231,6 +233,7 @@ async fn test_process_cycle_advances_watermark_on_send_success() { &sinks, &mut cursors, &no_retry(), + 500, &Metrics::default(), ) .await; @@ -289,6 +292,7 @@ async fn test_process_cycle_one_failing_sink_does_not_block_the_other() { &sinks, &mut cursors, &no_retry(), + 500, &Metrics::default(), ) .await; @@ -812,3 +816,166 @@ async fn test_run_lets_in_flight_cycle_finish_before_exiting() { "the in-flight cycle must finish and its result apply even though shutdown fired mid-fetch" ); } + +fn sink_for(mock: &MockServer) -> Box { + Box::new(HttpSink::new( + "sink".to_string(), + format!("{}/ingest", mock.uri()), + Method::POST, + vec![], + Encoding::Json, + )) +} + +async fn mount_events(mock: &MockServer, events: Vec) { + Mock::given(method("GET")) + .and(path("/api/events/audit")) + .respond_with(ResponseTemplate::new(200).set_body_json(events)) + .mount(mock) + .await; +} + +async fn received_batch_sizes(mock: &MockServer) -> Vec { + mock.received_requests() + .await + .expect("no requests received") + .iter() + .map(|req| { + serde_json::from_slice::>(&req.body) + .unwrap() + .len() + }) + .collect() +} + +#[tokio::test] +async fn test_process_cycle_splits_large_batches_into_chunks() { + let nb_mock = MockServer::start().await; + let sink_mock = MockServer::start().await; + mount_events( + &nb_mock, + vec![sample_event("1"), sample_event("2"), sample_event("3")], + ) + .await; + Mock::given(method("POST")) + .and(path("/ingest")) + .respond_with(ResponseTemplate::new(200)) + .mount(&sink_mock) + .await; + + let nb_client = NetbirdClient::new(nb_mock.uri(), "token".to_string()); + let sinks: Vec> = vec![sink_for(&sink_mock)]; + let mut cursors: HashMap>> = HashMap::new(); + + process_cycle( + &nb_client, + &sinks, + &mut cursors, + &no_retry(), + 2, + &Metrics::default(), + ) + .await; + + assert_eq!( + received_batch_sizes(&sink_mock).await, + vec![2, 1], + "3 events with batch_size 2 must produce two requests of 2 and 1" + ); + assert!( + cursors.get("sink").copied().flatten().is_some(), + "watermark must advance once every chunk delivered" + ); +} + +#[tokio::test] +async fn test_process_cycle_batch_size_zero_disables_chunking() { + let nb_mock = MockServer::start().await; + let sink_mock = MockServer::start().await; + mount_events( + &nb_mock, + vec![sample_event("1"), sample_event("2"), sample_event("3")], + ) + .await; + Mock::given(method("POST")) + .and(path("/ingest")) + .respond_with(ResponseTemplate::new(200)) + .mount(&sink_mock) + .await; + + let nb_client = NetbirdClient::new(nb_mock.uri(), "token".to_string()); + let sinks: Vec> = vec![sink_for(&sink_mock)]; + let mut cursors: HashMap>> = HashMap::new(); + + process_cycle( + &nb_client, + &sinks, + &mut cursors, + &no_retry(), + 0, + &Metrics::default(), + ) + .await; + + assert_eq!(received_batch_sizes(&sink_mock).await, vec![3]); +} + +#[tokio::test] +async fn test_process_cycle_does_not_advance_watermark_when_a_chunk_fails() { + let nb_mock = MockServer::start().await; + let sink_mock = MockServer::start().await; + mount_events(&nb_mock, vec![sample_event("1"), sample_event("2")]).await; + + // First chunk succeeds, second chunk is rejected: mocks match in mount + // order, and the 200 response is exhausted after one use. + Mock::given(method("POST")) + .and(path("/ingest")) + .respond_with(ResponseTemplate::new(200)) + .up_to_n_times(1) + .mount(&sink_mock) + .await; + Mock::given(method("POST")) + .and(path("/ingest")) + .respond_with(ResponseTemplate::new(500)) + .mount(&sink_mock) + .await; + + let nb_client = NetbirdClient::new(nb_mock.uri(), "token".to_string()); + let sinks: Vec> = vec![sink_for(&sink_mock)]; + let mut cursors: HashMap>> = HashMap::new(); + + process_cycle( + &nb_client, + &sinks, + &mut cursors, + &no_retry(), + 1, + &Metrics::default(), + ) + .await; + + assert_eq!( + cursors.get("sink").copied().flatten(), + None, + "watermark must not advance when any chunk fails, so nothing is skipped" + ); +} + +#[test] +fn test_config_batch_size_defaults_to_500_and_parses_env() { + temp_env::with_vars( + [ + ("NETBIRD_API_TOKEN", Some("test_token")), + ("BATCH_SIZE", None), + ], + || assert_eq!(Config::from_env().unwrap().batch_size, 500), + ); + + temp_env::with_vars( + [ + ("NETBIRD_API_TOKEN", Some("test_token")), + ("BATCH_SIZE", Some("7")), + ], + || assert_eq!(Config::from_env().unwrap().batch_size, 7), + ); +}