diff --git a/Cargo.lock b/Cargo.lock index 213504f..5e5757a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -281,6 +281,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "aws-sdk-ec2" +version = "1.200.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "919df2a48b8d6b4c431fdd329d5ba582ebf458f82e9736c249917afb85fa89af" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-query", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-smithy-xml", + "aws-types", + "fastrand", + "http 0.2.12", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-s3" version = "1.120.0" @@ -2603,6 +2627,7 @@ dependencies = [ "async-stream", "aws-config", "aws-sdk-cloudwatchlogs", + "aws-sdk-ec2", "aws-sdk-s3", "aws_lambda_events", "chrono", diff --git a/Cargo.toml b/Cargo.toml index 2be8da5..021e439 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ readme = "README.md" aws_lambda_events = { version = "1.0.0", default-features = false, features = ["cloudwatch_logs"] } aws-config = { version = "1.8", features = ["behavior-version-latest"] } aws-sdk-cloudwatchlogs = "1.115" +aws-sdk-ec2 = "1.200" aws-sdk-s3 = "1.120" hex = "0.4" chrono = "0.4" diff --git a/Makefile b/Makefile index 80dc2ca..c1deadd 100644 --- a/Makefile +++ b/Makefile @@ -2,6 +2,8 @@ IAM_ROLE ?= "arn:aws:iam::123456789012:role/lambda-execution-role" +FUNC_NAME ?= "rotel-lambda-forwarder" + build: cargo lambda build --release @@ -9,4 +11,4 @@ clippy: cargo clippy deploy: build - cargo lambda deploy --role ${IAM_ROLE} rotel-lambda-forwarder + cargo lambda deploy --role ${IAM_ROLE} --binary-name "rotel-lambda-forwarder" ${FUNC_NAME} diff --git a/README.md b/README.md index cee87a8..128da5f 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ expand as we verify support for additional services. | CloudTrail Logs | CloudWatch | | EKS Control Plane Logs | CloudWatch | | Lambda Logs | CloudWatch | +| VPC Flow Logs | CloudWatch | ## Deploying @@ -101,19 +102,18 @@ aws iam attach-role-policy \ --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole ``` -If your function needs additional AWS service access (e.g., S3, DynamoDB), attach those policies as well. - -**Required Permissions for Log Group Tag Enrichment** +**Additional Required Permissions** The Lambda function needs permissions to: 1. List tags on CloudWatch Logs log groups -2. Read and write to the S3 bucket for cache persistence +2. Describe VPC flow logs +3. Read and write to the S3 bucket for cache persistence -Create and attach a custom policy for tag caching: +Create and attach a custom policy: ```bash aws iam create-policy \ - --policy-name rotel-lambda-forwarder-tags-policy \ + --policy-name rotel-lambda-forwarder-policy \ --policy-document '{ "Version": "2012-10-17", "Statement": [ @@ -124,13 +124,20 @@ aws iam create-policy \ ], "Resource": "arn:aws:logs:*:*:log-group:*" }, + { + "Effect": "Allow", + "Action": [ + "ec2:DescribeFlowLogs" + ], + "Resource": "*" + }, { "Effect": "Allow", "Action": [ "s3:GetObject", "s3:PutObject" ], - "Resource": "arn:aws:s3:::YOUR_BUCKET_NAME/rotel-lambda-forwarder/cache/log-groups/*" + "Resource": "arn:aws:s3:::YOUR_BUCKET_NAME/rotel-lambda-forwarder/*" }, { "Effect": "Allow", @@ -144,7 +151,7 @@ aws iam create-policy \ aws iam attach-role-policy \ --role-name rotel-lambda-forwarder-role \ - --policy-arn arn:aws:iam::YOUR_ACCOUNT_ID:policy/rotel-lambda-forwarder-tags-policy + --policy-arn arn:aws:iam::YOUR_ACCOUNT_ID:policy/rotel-lambda-forwarder-policy ``` Replace `YOUR_BUCKET_NAME` with your S3 bucket name and `YOUR_ACCOUNT_ID` with your AWS account ID. @@ -282,6 +289,32 @@ from. We plan to make this configurable in the future. * Strips `responseElements.credentials.sessionToken` +## VPC Flow Logs + +The forwarder includes comprehensive support for AWS VPC Flow Logs with the following features: + +### Automatic Format Detection + +- **Dynamic Configuration**: Automatically fetches and caches flow log configurations from the EC2 API +- **Custom Formats**: Supports custom log formats defined in your VPC Flow Log configuration +- **Format Caching**: Caches parsed format configurations with a 30-minute TTL to minimize API calls +- **S3 Persistence**: Flow log configurations are persisted to S3 for durability across Lambda cold starts + +### Typed Field Parsing + +VPC Flow Log fields are parsed according to their AWS-defined data types per the [docs](https://docs.aws.amazon.com/vpc/latest/userguide/flow-log-records.html). + +### Configuration + +VPC Flow Logs are automatically detected when the log group name matches the pattern for VPC Flow Logs. No additional configuration is required. + +**Note:** The Lambda function requires the following IAM permissions to fetch flow log configurations: `ec2:DescribeFlowLogs`. + +These permissions are included in the CloudFormation templates by default. + +**Note:** Flow logs should be sent to unique CloudWatch Log Groups. If multiple flow logs are exported to the +same log group then the Lambda Forwarder will be unable to parse the fields correctly. + ## Configuration The Lambda function is configured via the same environment variables supported by Rotel, diff --git a/cloudformation-stacks/forwarder-clickhouse.template.yaml b/cloudformation-stacks/forwarder-clickhouse.template.yaml index c38ec7a..62e850a 100644 --- a/cloudformation-stacks/forwarder-clickhouse.template.yaml +++ b/cloudformation-stacks/forwarder-clickhouse.template.yaml @@ -176,6 +176,11 @@ Resources: - logs:ListTagsForResource Effect: Allow Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:*" + # Allow querying EC2 Flow Log configurations for dynamic parsing and tag extraction + - Action: + - ec2:DescribeFlowLogs + Effect: Allow + Resource: "*" - Action: - s3:GetObject - s3:PutObject diff --git a/cloudformation-stacks/forwarder-otlp.template.yaml b/cloudformation-stacks/forwarder-otlp.template.yaml index f33c7eb..6a4b92e 100644 --- a/cloudformation-stacks/forwarder-otlp.template.yaml +++ b/cloudformation-stacks/forwarder-otlp.template.yaml @@ -148,6 +148,11 @@ Resources: - logs:ListTagsForResource Effect: Allow Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:*" + # Allow querying EC2 Flow Log configurations for dynamic parsing and tag extraction + - Action: + - ec2:DescribeFlowLogs + Effect: Allow + Resource: "*" - Action: - s3:GetObject - s3:PutObject diff --git a/src/flowlogs/cache.rs b/src/flowlogs/cache.rs new file mode 100644 index 0000000..ede9712 --- /dev/null +++ b/src/flowlogs/cache.rs @@ -0,0 +1,513 @@ +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::{Arc, LazyLock}; +use std::time::{SystemTime, UNIX_EPOCH}; +use tracing::{debug, trace}; + +const MAX_CACHE_AGE_SECS: u64 = 30 * 60; + +/// Data type for a flow log field, based on Parquet data types from AWS documentation +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ParsedFieldType { + String, + Int32, + Int64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ParsedField { + /// The field name (e.g., "version", "account-id") + pub field_name: String, + pub field_type: ParsedFieldType, +} + +impl ParsedField { + pub fn new(field_name: String, field_type: ParsedFieldType) -> Self { + Self { + field_name, + field_type, + } + } +} + +/// Static mapping of non-string field names to their data types based on AWS VPC Flow Logs documentation +/// Reference: https://docs.aws.amazon.com/vpc/latest/userguide/flow-log-records.html +/// NOTE: All fields not listed here are assumed to be a String, this reduces map size. +static FIELD_TYPE_MAP: LazyLock> = LazyLock::new(|| { + let mut map = HashMap::new(); + + // Version 2 fields + map.insert("version", ParsedFieldType::Int32); + map.insert("srcport", ParsedFieldType::Int32); + map.insert("dstport", ParsedFieldType::Int32); + map.insert("protocol", ParsedFieldType::Int32); + map.insert("packets", ParsedFieldType::Int64); + map.insert("bytes", ParsedFieldType::Int64); + map.insert("start", ParsedFieldType::Int64); + map.insert("end", ParsedFieldType::Int64); + + // Version 3 fields + map.insert("tcp-flags", ParsedFieldType::Int32); + + // Version 5 fields + map.insert("traffic-path", ParsedFieldType::Int32); + + // Version 10 fields + map.insert("encryption-status", ParsedFieldType::Int32); + + map +}); + +/// Get the field type for a given field name, defaulting to String if unknown +pub fn get_field_type(field_name: &str) -> ParsedFieldType { + FIELD_TYPE_MAP + .get(field_name) + .copied() + .unwrap_or(ParsedFieldType::String) +} + +/// Result of parsing flow log format fields +#[derive(Debug, Clone, PartialEq)] +pub enum ParsedFields { + /// Successfully parsed field names and types from the log format + Success(Vec), + /// Failed to parse the log format, with error message + Error(String), +} + +/// Flow log configuration for a specific log group +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct FlowLogConfig { + /// The log format string (e.g., "${version} ${account-id} ${interface-id} ...") + pub log_format: String, + pub flow_log_id: String, + pub tags: std::collections::HashMap, + + /// Parsed field names from the log format (lazily computed, not serialized) + #[serde(skip)] + pub parsed_fields: Option>, // Use an Arc to reduce clone costs +} + +/// In-memory cache for flow log configurations +/// +/// Cache is refreshed when configurations are fetched from EC2 API. +/// Reading from the cache does not extend the TTL - the cache expires 30 minutes +/// after it was last refreshed from the API. +/// +/// Since all flow logs are queried together with a single DescribeFlowLogs call, +/// the last_seen timestamp applies to the entire cache, not individual entries. +#[derive(Debug, Clone)] +pub struct FlowLogCache { + inner: HashMap, + /// Unix timestamp in seconds when the cache was last refreshed from EC2 API + last_refreshed_secs: u64, +} + +impl FlowLogCache { + pub fn new() -> Self { + Self { + inner: HashMap::new(), + last_refreshed_secs: 0, + } + } + + /// Get flow log configuration for a log group if it exists and cache is not expired + /// Returns None if not found or cache is expired + /// + /// Note: This does not update the timestamp - cache expires 30 minutes + /// after it was last refreshed from the API. + pub fn get(&self, log_group: &str) -> Option { + if self.is_expired() { + debug!("Cache expired"); + return None; + } + + if let Some(config) = self.inner.get(log_group) { + trace!(log_group = %log_group, "Cache hit"); + Some(config.clone()) + } else { + trace!(log_group = %log_group, "Cache miss"); + None + } + } + + /// Get mutable reference to flow log configuration for a log group + /// Returns None if not found or cache is expired + /// + /// This is used for lazy initialization of parsed fields. + pub fn get_mut(&mut self, log_group: &str) -> Option<&mut FlowLogConfig> { + if self.is_expired() { + debug!("Cache expired"); + return None; + } + + if let Some(config) = self.inner.get_mut(log_group) { + trace!(log_group = %log_group, "Cache hit (mutable)"); + Some(config) + } else { + trace!(log_group = %log_group, "Cache miss"); + None + } + } + + /// Check if the entire cache is expired (older than 30 minutes) + pub fn is_expired(&self) -> bool { + if self.last_refreshed_secs == 0 { + return true; // Never been refreshed + } + + let now_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + + let age_secs = now_secs.saturating_sub(self.last_refreshed_secs); + age_secs > MAX_CACHE_AGE_SECS + } + + /// Insert or update flow log configuration for a log group + /// + /// Should only be called when configuration is freshly fetched from EC2 API. + pub fn insert(&mut self, log_group: String, config: FlowLogConfig) { + debug!(log_group = %log_group, flow_log_id = %config.flow_log_id, "Inserting flow log config into cache"); + self.inner.insert(log_group, config); + } + + /// Mark the cache as refreshed with the current timestamp + /// + /// This should be called after successfully fetching flow logs from EC2 API. + /// It resets the 30-minute TTL for the entire cache. + pub fn mark_refreshed(&mut self) { + self.last_refreshed_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + debug!( + timestamp = self.last_refreshed_secs, + "Cache marked as refreshed" + ); + } + + /// Get a snapshot of the current cache with timestamp + pub fn get_snapshot(&self) -> CacheSnapshot { + CacheSnapshot { + flow_logs: self.inner.clone(), + last_refreshed_secs: self.last_refreshed_secs, + } + } + + /// Load entries from a snapshot (used when restoring from S3) + pub fn load_snapshot(&mut self, snapshot: CacheSnapshot) { + debug!( + entry_count = snapshot.flow_logs.len(), + "Loading snapshot into cache" + ); + + if !snapshot.is_expired() { + self.inner = snapshot.flow_logs; + self.last_refreshed_secs = snapshot.last_refreshed_secs; + } else { + debug!("Snapshot is expired, not loading"); + } + } + + /// Get the number of entries in the cache + pub fn len(&self) -> usize { + self.inner.len() + } + + /// Check if the cache is empty + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + /// Clear all entries from the cache and reset timestamp + #[cfg(test)] + pub fn clear(&mut self) { + self.inner.clear(); + self.last_refreshed_secs = 0; + } +} + +/// Snapshot of the flow log cache for serialization/deserialization +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CacheSnapshot { + /// Map of log group names to their flow log configurations + pub flow_logs: HashMap, + /// Unix timestamp in seconds when the cache was last refreshed + pub last_refreshed_secs: u64, +} + +impl CacheSnapshot { + /// Check if this snapshot is expired (older than 30 minutes) + pub fn is_expired(&self) -> bool { + if self.last_refreshed_secs == 0 { + return true; + } + + let now_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + + let age_secs = now_secs.saturating_sub(self.last_refreshed_secs); + age_secs > MAX_CACHE_AGE_SECS + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + #[test] + fn test_cache_insert_and_get() { + let mut cache = FlowLogCache::new(); + let config = FlowLogConfig { + log_format: "${version} ${account-id} ${interface-id}".to_string(), + flow_log_id: "fl-1234567890abcdef0".to_string(), + tags: std::collections::HashMap::new(), + parsed_fields: None, + }; + + cache.insert("/aws/ec2/flowlogs".to_string(), config.clone()); + cache.mark_refreshed(); + + let retrieved = cache.get("/aws/ec2/flowlogs"); + assert!(retrieved.is_some()); + assert_eq!(retrieved.unwrap(), config); + } + + #[test] + fn test_cache_miss() { + let cache = FlowLogCache::new(); + let retrieved = cache.get("non-existent"); + assert!(retrieved.is_none()); + } + + #[test] + fn test_cache_expiration() { + let mut cache = FlowLogCache::new(); + + // New cache should be expired (never refreshed) + assert!(cache.is_expired()); + + // Mark as refreshed - should not be expired + cache.mark_refreshed(); + assert!(!cache.is_expired()); + + // Manually set old timestamp - should be expired + cache.last_refreshed_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() + .saturating_sub(31 * 60); // 31 minutes ago + assert!(cache.is_expired()); + } + + #[test] + fn test_snapshot() { + let mut cache = FlowLogCache::new(); + + let config1 = FlowLogConfig { + log_format: "${version} ${account-id}".to_string(), + flow_log_id: "fl-111".to_string(), + tags: std::collections::HashMap::new(), + parsed_fields: None, + }; + cache.insert("/aws/ec2/flowlogs1".to_string(), config1.clone()); + + let config2 = FlowLogConfig { + log_format: "${version} ${interface-id}".to_string(), + flow_log_id: "fl-222".to_string(), + tags: std::collections::HashMap::new(), + parsed_fields: None, + }; + cache.insert("/aws/ec2/flowlogs2".to_string(), config2.clone()); + cache.mark_refreshed(); + + let snapshot = cache.get_snapshot(); + assert_eq!(snapshot.flow_logs.len(), 2); + assert_eq!( + snapshot.flow_logs.get("/aws/ec2/flowlogs1").unwrap(), + &config1 + ); + assert_eq!( + snapshot.flow_logs.get("/aws/ec2/flowlogs2").unwrap(), + &config2 + ); + } + + #[test] + fn test_load_snapshot() { + let mut cache = FlowLogCache::new(); + + let config = FlowLogConfig { + log_format: "${version} ${account-id}".to_string(), + flow_log_id: "fl-123".to_string(), + tags: std::collections::HashMap::new(), + parsed_fields: None, + }; + let snapshot = CacheSnapshot { + flow_logs: { + let mut map = HashMap::new(); + map.insert("/aws/ec2/flowlogs".to_string(), config.clone()); + map + }, + last_refreshed_secs: SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(), + }; + + cache.load_snapshot(snapshot); + + let retrieved = cache.get("/aws/ec2/flowlogs"); + assert!(retrieved.is_some()); + assert_eq!(retrieved.unwrap(), config); + } + + #[test] + fn test_flow_log_config_with_tags() { + let mut cache = FlowLogCache::new(); + + // Create config with tags + let mut tags = std::collections::HashMap::new(); + tags.insert("Environment".to_string(), "production".to_string()); + tags.insert("Team".to_string(), "platform".to_string()); + tags.insert("Application".to_string(), "vpc-monitoring".to_string()); + + let config = FlowLogConfig { + log_format: "${version} ${account-id} ${interface-id}".to_string(), + flow_log_id: "fl-1234567890abcdef0".to_string(), + tags: tags.clone(), + parsed_fields: None, + }; + + cache.insert("/aws/ec2/flowlogs".to_string(), config.clone()); + cache.mark_refreshed(); + + let retrieved = cache.get("/aws/ec2/flowlogs"); + assert!(retrieved.is_some()); + + let retrieved_config = retrieved.unwrap(); + assert_eq!(retrieved_config.tags.len(), 3); + assert_eq!( + retrieved_config.tags.get("Environment").unwrap(), + "production" + ); + assert_eq!(retrieved_config.tags.get("Team").unwrap(), "platform"); + assert_eq!( + retrieved_config.tags.get("Application").unwrap(), + "vpc-monitoring" + ); + } + + #[test] + fn test_parsed_field_type_mapping() { + // Test some known fields + assert_eq!(get_field_type("version"), ParsedFieldType::Int32); + assert_eq!(get_field_type("account-id"), ParsedFieldType::String); + assert_eq!(get_field_type("srcport"), ParsedFieldType::Int32); + assert_eq!(get_field_type("packets"), ParsedFieldType::Int64); + assert_eq!(get_field_type("bytes"), ParsedFieldType::Int64); + assert_eq!(get_field_type("action"), ParsedFieldType::String); + + // Test unknown field defaults to String + assert_eq!(get_field_type("unknown-field"), ParsedFieldType::String); + } + + #[test] + fn test_parsed_field_creation() { + let field = ParsedField::new("version".to_string(), ParsedFieldType::Int32); + assert_eq!(field.field_name, "version"); + assert_eq!(field.field_type, ParsedFieldType::Int32); + } + + #[test] + fn test_parsed_fields_success() { + let fields = vec![ + ParsedField::new("version".to_string(), ParsedFieldType::Int32), + ParsedField::new("account-id".to_string(), ParsedFieldType::String), + ]; + let parsed = ParsedFields::Success(fields.clone()); + + assert_eq!(parsed, ParsedFields::Success(fields)); + } + + #[test] + fn test_parsed_fields_error() { + let error_msg = "Invalid format string".to_string(); + let parsed = ParsedFields::Error(error_msg.clone()); + + assert_eq!(parsed, ParsedFields::Error(error_msg)); + } + + #[test] + fn test_flow_log_config_with_parsed_fields() { + let mut cache = FlowLogCache::new(); + + let config = FlowLogConfig { + log_format: "${version} ${account-id}".to_string(), + flow_log_id: "fl-123".to_string(), + tags: std::collections::HashMap::new(), + parsed_fields: Some(Arc::new(ParsedFields::Success(vec![ + ParsedField::new("version".to_string(), ParsedFieldType::Int32), + ParsedField::new("account-id".to_string(), ParsedFieldType::String), + ]))), + }; + + cache.insert("/aws/ec2/flowlogs".to_string(), config.clone()); + cache.mark_refreshed(); + + let retrieved = cache.get("/aws/ec2/flowlogs"); + assert!(retrieved.is_some()); + + let retrieved_config = retrieved.unwrap(); + assert!(retrieved_config.parsed_fields.is_some()); + if let Some(parsed_fields) = &retrieved_config.parsed_fields { + if let ParsedFields::Success(fields) = parsed_fields.as_ref() { + assert_eq!(fields.len(), 2); + assert_eq!(fields[0].field_name, "version"); + assert_eq!(fields[0].field_type, ParsedFieldType::Int32); + assert_eq!(fields[1].field_name, "account-id"); + assert_eq!(fields[1].field_type, ParsedFieldType::String); + } else { + panic!("Expected ParsedFields::Success"); + } + } else { + panic!("Expected Some(parsed_fields)"); + } + } + + #[test] + fn test_flow_log_config_with_parse_error() { + let mut cache = FlowLogCache::new(); + + let config = FlowLogConfig { + log_format: "invalid format".to_string(), + flow_log_id: "fl-123".to_string(), + tags: std::collections::HashMap::new(), + parsed_fields: Some(Arc::new(ParsedFields::Error("Parse failed".to_string()))), + }; + + cache.insert("/aws/ec2/flowlogs".to_string(), config.clone()); + cache.mark_refreshed(); + + let retrieved = cache.get("/aws/ec2/flowlogs"); + assert!(retrieved.is_some()); + + let retrieved_config = retrieved.unwrap(); + assert!(retrieved_config.parsed_fields.is_some()); + + if let Some(parsed_fields) = &retrieved_config.parsed_fields { + if let ParsedFields::Error(msg) = parsed_fields.as_ref() { + assert_eq!(msg, "Parse failed"); + } else { + panic!("Expected ParsedFields::Error") + } + } else { + panic!("Expected Some(parsed_fields)"); + } + } +} diff --git a/src/flowlogs/ec2.rs b/src/flowlogs/ec2.rs new file mode 100644 index 0000000..6ddcc32 --- /dev/null +++ b/src/flowlogs/ec2.rs @@ -0,0 +1,161 @@ +use aws_sdk_ec2::{Client as Ec2Client, types::LogDestinationType}; +use std::collections::HashMap; +use thiserror::Error; +use tracing::{debug, info, warn}; + +use super::cache::FlowLogConfig; + +/// Errors that can occur during EC2 operations +#[derive(Debug, Error)] +pub enum Ec2Error { + #[error("EC2 API error: {0}")] + ApiError(String), + + #[error("Access denied: {0}")] + AccessDenied(String), + + #[error("Invalid flow log format: {0}")] + InvalidFormat(String), +} + +/// Fetcher for EC2 Flow Log configurations from EC2 API +/// +/// This fetcher queries the EC2 DescribeFlowLogs API to retrieve EC2 Flow Log +/// configurations, including the log format and tags. Tags are extracted and +/// will be applied to logs as resource attributes with the prefix +/// `ec2.flow-logs.tags.`. +pub struct Ec2FlowLogFetcher { + client: Ec2Client, +} + +impl Ec2FlowLogFetcher { + /// Create a new EC2 flow log fetcher + pub fn new(client: Ec2Client) -> Self { + Self { client } + } + + /// Fetch all EC2 Flow Logs that are delivered to CloudWatch Logs + /// + /// Returns a map of log group name to flow log configuration. + /// Only includes flow logs with destination type "cloud-watch-logs". + /// + pub async fn fetch_all_flow_logs(&self) -> Result, Ec2Error> { + debug!("Fetching VPC flow logs from EC2 API"); + + let result = self.client.describe_flow_logs().send().await; + + let response = match result { + Ok(resp) => resp, + Err(e) => { + let err_msg = format!("{:?}", e); + if err_msg.contains("UnauthorizedOperation") || err_msg.contains("AccessDenied") { + return Err(Ec2Error::AccessDenied(err_msg)); + } else { + return Err(Ec2Error::ApiError(err_msg)); + } + } + }; + + let mut flow_log_configs = HashMap::new(); + + // flow_logs() returns a slice, not an Option + for flow_log in response.flow_logs() { + // Only process CloudWatch Logs destinations + match flow_log.log_destination_type() { + Some(LogDestinationType::CloudWatchLogs) => { + // Continue processing this flow log + } + Some(v) => { + debug!( + flow_log_id = ?flow_log.flow_log_id(), + log_destination_type = v.to_string(), + "Skipping flow log with non-CloudWatch destination" + ); + continue; + } + None => { + debug!( + flow_log_id = ?flow_log.flow_log_id(), + "Skipping flow log with unset log_destination_type" + ); + continue; + } + } + + // Get the log group name + let log_group_name = match flow_log.log_group_name() { + Some(name) => name, + None => { + warn!( + flow_log_id = ?flow_log.flow_log_id(), + "Flow log missing log group name, skipping" + ); + continue; + } + }; + + let log_format = match flow_log.log_format().map(|s| s.to_string()) { + Some(log_format) => log_format, + None => { + warn!( + flow_log_id = ?flow_log.flow_log_id(), + "Flow log missing log format, skipping" + ); + + continue; + } + }; + + let flow_log_id = flow_log.flow_log_id().unwrap_or("unknown").to_string(); + + // Extract tags from the flow log + // These tags will be applied to logs as resource attributes + // with the prefix "ec2.flow-logs.tags." + let mut tags = HashMap::new(); + for tag in flow_log.tags() { + if let (Some(key), Some(value)) = (tag.key(), tag.value()) { + tags.insert(key.to_string(), value.to_string()); + } + } + + let config = FlowLogConfig { + log_format, + flow_log_id: flow_log_id.clone(), + tags, + parsed_fields: None, + }; + + debug!( + log_group = %log_group_name, + flow_log_id = %flow_log_id, + tag_count = config.tags.len(), + "Found EC2 flow log configuration with tags" + ); + + flow_log_configs.insert(log_group_name.to_string(), config); + } + + info!( + count = flow_log_configs.len(), + "Fetched EC2 flow log configurations from EC2" + ); + + Ok(flow_log_configs) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aws_config::BehaviorVersion; + + #[tokio::test] + async fn test_ec2_fetcher_creation() { + let config = aws_config::defaults(BehaviorVersion::latest()).load().await; + let ec2_client = Ec2Client::new(&config); + + let fetcher = Ec2FlowLogFetcher::new(ec2_client); + // Just verify we can create the fetcher + assert!(std::mem::size_of_val(&fetcher) > 0); + } +} diff --git a/src/flowlogs/mod.rs b/src/flowlogs/mod.rs new file mode 100644 index 0000000..6aa88a0 --- /dev/null +++ b/src/flowlogs/mod.rs @@ -0,0 +1,582 @@ +//! EC2 Flow Log Manager Module +//! +//! This module provides comprehensive management of EC2 Flow Log configurations, including: +//! - Fetching flow log configurations from EC2 API (DescribeFlowLogs) +//! - Extracting flow log format strings for dynamic parsing +//! - Extracting and applying flow log tags to resource attributes +//! - Caching configurations in-memory with 30-minute TTL +//! - Persisting configurations to S3 for durability across Lambda cold starts +//! + +mod cache; +mod ec2; + +pub use cache::{ + CacheSnapshot, FlowLogCache, FlowLogConfig, ParsedField, ParsedFieldType, ParsedFields, +}; +pub use ec2::{Ec2Error, Ec2FlowLogFetcher}; + +use aws_sdk_ec2::Client as Ec2Client; +use aws_sdk_s3::Client as S3Client; +use serde::{Deserialize, Serialize}; +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; +use thiserror::Error; +use tracing::{debug, error, info, warn}; + +use crate::{ + flowlogs::cache::get_field_type, + s3_cache::{S3Cache, S3CacheError}, +}; + +/// S3 cache key for storing flow log configurations +const FLOW_LOG_CACHE_KEY: &str = "rotel-lambda-forwarder/cache/flow-logs/configs.json.gz"; + +/// Errors that can occur during flow log operations +#[derive(Debug, Error)] +pub enum FlowLogError { + #[error("EC2 error: {0}")] + Ec2(#[from] Ec2Error), + + #[error("S3 error: {0}")] + S3(#[from] S3CacheError), + + #[error("No S3 bucket configured")] + NoS3Bucket, +} + +/// Serializable format for the cache file +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FlowLogCacheFile { + pub version: u32, + pub snapshot: CacheSnapshot, +} + +impl FlowLogCacheFile { + pub fn new(snapshot: CacheSnapshot) -> Self { + Self { + version: 1, + snapshot, + } + } +} + +/// Main flow log manager that coordinates cache, S3, and EC2 operations +/// +/// The `FlowLogManager` provides automatic caching and persistence of EC2 Flow Log +/// configurations. It fetches configurations from EC2 API on startup and caches them +/// in memory with a 30-minute TTL. Optionally, configs can be persisted to S3 for +/// durability across Lambda cold starts. +/// +pub struct FlowLogManager { + cache: FlowLogCache, + s3_cache: Option>, + ec2_fetcher: Ec2FlowLogFetcher, + persist_enabled: bool, + fetch_disabled_until: Option, + cooldown_duration: Duration, +} + +impl FlowLogManager { + /// Create a new flow log manager + pub fn new( + ec2_client: Ec2Client, + s3_client: Option, + s3_bucket: Option, + ) -> Self { + let persist_enabled = s3_client.is_some() && s3_bucket.is_some(); + + let s3_cache = match (s3_client, s3_bucket) { + (Some(client), Some(bucket)) => { + Some(S3Cache::new(client, bucket, FLOW_LOG_CACHE_KEY.to_string())) + } + _ => None, + }; + + Self { + cache: FlowLogCache::new(), + s3_cache, + ec2_fetcher: Ec2FlowLogFetcher::new(ec2_client), + persist_enabled, + fetch_disabled_until: None, + cooldown_duration: Duration::from_secs(30 * 60), // 30 minutes + } + } + + /// Initialize the flow log manager by loading the cache from S3 and fetching from EC2 + pub async fn initialize(&mut self) -> Result<(), FlowLogError> { + match self.reload_cache_if_needed().await { + Ok(_) => { + info!("Flow log manager initialization complete"); + Ok(()) + } + Err(e) => { + error!(error = %e, "Failed to fetch flow logs from EC2 during initialization"); + Err(FlowLogError::Ec2(e)) + } + } + } + + /// Reload cache from S3 or EC2 if needed + /// + /// First attempts to load a valid (non-expired) cache from S3. + /// If S3 cache doesn't exist or is expired, fetches from EC2. + /// Returns Ok(true) if cache was successfully loaded/refreshed, Ok(false) if using existing cache. + async fn reload_cache_if_needed(&mut self) -> Result<(), Ec2Error> { + // Check if fetching is currently disabled due to AccessDenied + if let Some(disabled_until) = self.fetch_disabled_until { + let now = Instant::now(); + if now < disabled_until { + let remaining = disabled_until.duration_since(now); + debug!( + remaining_seconds = remaining.as_secs(), + "Flow log fetching is disabled due to AccessDenied" + ); + return Ok(()); + } else { + // Cooldown period has elapsed, re-enable fetching + info!("Cooldown period elapsed, attempting to fetch flow logs again"); + self.fetch_disabled_until = None; + } + } + + // Load from S3 if available + if let Some(s3_cache) = &mut self.s3_cache { + match s3_cache.load().await { + Ok(Some(cache_file)) => { + let is_expired = cache_file.snapshot.is_expired(); + info!( + entry_count = cache_file.snapshot.flow_logs.len(), + expired = is_expired, + "Loaded flow log cache from S3" + ); + + if !is_expired { + // Cache is still valid, use it and skip EC2 API call + self.cache.load_snapshot(cache_file.snapshot); + debug!("Using valid cached flow log configurations, skipping EC2 API call"); + return Ok(()); + } else { + debug!("Cached flow log configurations are expired, will fetch from EC2"); + } + } + Ok(None) => { + debug!("No existing flow log cache found in S3"); + } + Err(e) => { + error!(error = %e, "Failed to load flow log cache from S3, will refresh"); + // Don't fail, just continue without S3 cache + } + } + } + + // Only fetch from EC2 if we don't have a valid cache, will persist to S3 if loaded + match self.fetch_and_update_all().await { + Ok(_) => { + info!("Successfully fetched and cached flow log configurations from EC2"); + Ok(()) + } + Err(Ec2Error::AccessDenied(e)) => { + warn!(error = ?e, "Access denied when fetching flow logs from EC2"); + // Disable fetching for the next 30 minutes + let disabled_until = Instant::now() + self.cooldown_duration; + self.fetch_disabled_until = Some(disabled_until); + + warn!( + cooldown_minutes = self.cooldown_duration.as_secs() / 60, + "AccessDenied error, disabling flow log fetching" + ); + + Ok(()) // Don't fail, just use cached data (if any) + } + Err(e) => { + // For other errors, propagate them + Err(e) + } + } + } + + /// Get flow log configuration for a log group + /// First checks the in-memory cache. Returns None if not found or expired. + /// Lazily parses the log format fields on first access and caches the result. + pub async fn get_config(&mut self, log_group_name: &str) -> Option { + // Check if cache is expired + if self.cache.is_expired() { + debug!( + log_group = %log_group_name, + "Cache expired, attempting to reload" + ); + + // Attempt to reload cache from S3 or EC2 + match self.reload_cache_if_needed().await { + Ok(_) => { + debug!("Successfully reloaded flow log cache"); + } + Err(e) => { + warn!( + error = %e, + log_group = %log_group_name, + "Cache expired and could not be reloaded" + ); + return None; + } + } + } + + // Get a mutable reference to the config so we can parse fields if needed + let config = self.cache.get_mut(log_group_name)?; + + // Parse fields if not already attempted + if config.parsed_fields.is_none() { + let fields = parse_log_format(&config.log_format); + + if fields.is_empty() { + // Parsing failed - cache the error + config.parsed_fields = Some(Arc::new(ParsedFields::Error( + "Failed to parse log format or no fields found".to_string(), + ))); + warn!( + log_group = %log_group_name, + log_format = %config.log_format, + "Failed to parse flow log format fields" + ); + } else { + // Parsing succeeded - cache the result + let field_count = fields.len(); + config.parsed_fields = Some(Arc::new(ParsedFields::Success(fields))); + debug!( + log_group = %log_group_name, + field_count = field_count, + "Parsed log format fields with types for flow log" + ); + } + } + + Some(config.clone()) + } + + /// Fetch all flow logs from EC2 and update the cache + async fn fetch_and_update_all(&mut self) -> Result<(), Ec2Error> { + debug!("Fetching flow logs from EC2"); + + let flow_log_configs = match self.ec2_fetcher.fetch_all_flow_logs().await { + Ok(configs) => configs, + Err(e) => { + return Err(e); + } + }; + + // Update cache with fetched configurations + for (log_group, config) in flow_log_configs { + self.cache.insert(log_group, config); + } + + // Mark the cache as refreshed + self.cache.mark_refreshed(); + + // Persist to S3 if enabled + if self.persist_enabled { + if let Err(e) = self.persist_cache().await { + error!(error = %e, "Failed to persist flow log cache to S3"); + // Don't fail the request, just log the error + } + } + + Ok(()) + } + + /// Persist the current cache to S3 + async fn persist_cache(&mut self) -> Result<(), FlowLogError> { + if let Some(s3_cache) = &mut self.s3_cache { + let snapshot = self.cache.get_snapshot(); + let cache_file = FlowLogCacheFile::new(snapshot); + + match s3_cache.save(cache_file).await { + Ok(_) => { + debug!("Successfully persisted flow log cache to S3"); + Ok(()) + } + Err(S3CacheError::ConditionalWriteFailed) => { + warn!("Conditional write failed, reloading and merging flow log cache"); + + // Reload from S3 and merge with our current snapshot + let local_snapshot = self.cache.get_snapshot(); + let local_cache_file = FlowLogCacheFile::new(local_snapshot); + + let merged = s3_cache + .reload_and_merge(local_cache_file, |from_s3, local| { + // Merge the snapshots, keeping the most recent one + let merged_snapshot = if local.snapshot.last_refreshed_secs + > from_s3.snapshot.last_refreshed_secs + { + debug!("Local snapshot is more recent, using it"); + local.snapshot + } else { + debug!("S3 snapshot is more recent, using it"); + from_s3.snapshot + }; + + FlowLogCacheFile::new(merged_snapshot) + }) + .await?; + + // Replace out local cache with the most recent snapshot + self.cache.load_snapshot(merged.snapshot.clone()); + + // Try to save again with the merged data. It is possible this + // could fail again, but we ignore it in the caller. + s3_cache.save(merged).await?; + + info!("Successfully merged and persisted flow log cache after conflict"); + Ok(()) + } + Err(e) => Err(FlowLogError::S3(e)), + } + } else { + Ok(()) + } + } + + /// Check if flow log fetching is currently disabled due to the circuit breaker + pub fn is_fetch_disabled(&self) -> bool { + self.fetch_disabled_until + .is_some_and(|disabled_until| Instant::now() < disabled_until) + } + + /// Get the remaining cooldown time if fetching is disabled + /// + /// Returns `Some(Duration)` with the remaining time if the circuit breaker is active, + /// or `None` if fetching is currently enabled. + pub fn remaining_cooldown(&self) -> Option { + self.fetch_disabled_until.and_then(|disabled_until| { + let now = Instant::now(); + if now < disabled_until { + Some(disabled_until.duration_since(now)) + } else { + None + } + }) + } + + /// Get cache statistics + /// + /// Returns information about the current state of the cache including + /// the number of cached entries and whether fetching is currently disabled. + #[cfg(test)] + pub fn cache_stats(&self) -> CacheStats { + CacheStats { + entry_count: self.cache.len(), + persist_enabled: self.persist_enabled, + fetch_disabled: self.is_fetch_disabled(), + } + } +} + +/// Parse the LogFormat string to extract field names and their types +/// +/// LogFormat strings look like: "${version} ${account-id} ${interface-id} ..." +/// This function extracts the field names between ${ and } and assigns types based on +/// the AWS VPC Flow Logs documentation. +pub fn parse_log_format(log_format: &str) -> Vec { + let mut fields = Vec::new(); + let mut chars = log_format.chars().peekable(); + + while let Some(ch) = chars.next() { + if ch == '$' { + if chars.peek() == Some(&'{') { + chars.next(); // consume '{' + let mut field_name = String::new(); + + // Read until '}' + while let Some(&ch) = chars.peek() { + if ch == '}' { + chars.next(); // consume '}' + break; + } + field_name.push(chars.next().unwrap()); + } + + if !field_name.is_empty() { + let field_type = get_field_type(&field_name); + fields.push(ParsedField::new(field_name, field_type)); + } + } + } + } + + debug!( + field_count = fields.len(), + "Parsed log format fields with types" + ); + fields +} + +/// Cache statistics +#[derive(Debug, Clone)] +#[cfg(test)] +pub struct CacheStats { + /// Number of entries currently in the cache + pub entry_count: usize, + /// Whether S3 persistence is enabled + pub persist_enabled: bool, + /// Whether flow log fetching is currently disabled due to AccessDenied circuit breaker + pub fetch_disabled: bool, +} + +#[cfg(test)] +mod tests { + use super::*; + use aws_config::BehaviorVersion; + + #[tokio::test] + async fn test_circuit_breaker_cooldown_duration() { + let config = aws_config::defaults(BehaviorVersion::latest()).load().await; + let ec2_client = Ec2Client::new(&config); + + let mut manager = FlowLogManager::new(ec2_client, None, None); + + // Set a short cooldown for testing + manager.cooldown_duration = Duration::from_millis(100); + + // Simulate AccessDenied by setting the disabled time + manager.fetch_disabled_until = Some(Instant::now() + manager.cooldown_duration); + + // Initially should be disabled + let stats = manager.cache_stats(); + assert!(stats.fetch_disabled); + + // Wait for cooldown to elapse + tokio::time::sleep(Duration::from_millis(150)).await; + + // Should now be enabled again + let stats = manager.cache_stats(); + assert!(!stats.fetch_disabled); + } + + #[tokio::test] + async fn test_cache_reuse_optimization() { + use crate::flowlogs::cache::{CacheSnapshot, FlowLogConfig}; + use std::collections::HashMap; + + let config = aws_config::defaults(BehaviorVersion::latest()).load().await; + let ec2_client = Ec2Client::new(&config); + + let mut manager = FlowLogManager::new(ec2_client, None, None); + + // Manually insert a valid cache entry + let mut flow_logs = HashMap::new(); + flow_logs.insert( + "/aws/ec2/test-flowlogs".to_string(), + FlowLogConfig { + log_format: "${version} ${account-id}".to_string(), + flow_log_id: "fl-test123".to_string(), + tags: HashMap::new(), + parsed_fields: None, + }, + ); + + let snapshot = CacheSnapshot { + flow_logs, + last_refreshed_secs: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(), + }; + + manager.cache.load_snapshot(snapshot); + + // Verify cache has the entry and is not expired + assert!(!manager.cache.is_expired()); + assert_eq!(manager.cache.len(), 1); + + let config = manager.get_config("/aws/ec2/test-flowlogs").await; + assert!(config.is_some()); + assert_eq!(config.unwrap().flow_log_id, "fl-test123"); + } + + #[test] + fn test_parse_log_format_default() { + let log_format = "${version} ${account-id} ${interface-id} ${srcaddr} ${dstaddr} ${srcport} ${dstport} ${protocol} ${packets} ${bytes} ${start} ${end} ${action} ${log-status}"; + let fields = parse_log_format(log_format); + + assert_eq!(fields.len(), 14); + assert_eq!(fields[0].field_name, "version"); + assert_eq!(fields[1].field_name, "account-id"); + assert_eq!(fields[2].field_name, "interface-id"); + assert_eq!(fields[13].field_name, "log-status"); + } + + #[test] + fn test_parse_log_format_custom() { + let log_format = "${version} ${vpc-id} ${subnet-id} ${instance-id} ${srcaddr} ${dstaddr}"; + let fields = parse_log_format(log_format); + + assert_eq!(fields.len(), 6); + assert_eq!(fields[0].field_name, "version"); + assert_eq!(fields[1].field_name, "vpc-id"); + assert_eq!(fields[2].field_name, "subnet-id"); + assert_eq!(fields[3].field_name, "instance-id"); + assert_eq!(fields[4].field_name, "srcaddr"); + assert_eq!(fields[5].field_name, "dstaddr"); + } + + #[test] + fn test_parse_log_format_with_extra_spaces() { + let log_format = "${version} ${account-id} ${interface-id}"; + let fields = parse_log_format(log_format); + + assert_eq!(fields.len(), 3); + assert_eq!(fields[0].field_name, "version"); + assert_eq!(fields[1].field_name, "account-id"); + assert_eq!(fields[2].field_name, "interface-id"); + } + + #[tokio::test] + async fn test_cache_expiration_and_reload() { + use crate::flowlogs::cache::{CacheSnapshot, FlowLogConfig}; + use std::collections::HashMap; + + let config = aws_config::defaults(BehaviorVersion::latest()).load().await; + let ec2_client = Ec2Client::new(&config); + + let mut manager = FlowLogManager::new(ec2_client, None, None); + + // Manually insert a cache entry with an expired timestamp + let mut flow_logs = HashMap::new(); + flow_logs.insert( + "/aws/ec2/expired-flowlogs".to_string(), + FlowLogConfig { + log_format: "${version} ${account-id}".to_string(), + flow_log_id: "fl-expired123".to_string(), + tags: HashMap::new(), + parsed_fields: None, + }, + ); + + // Create an expired snapshot (older than 30 minutes) + let expired_time = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + - 1900; // 31+ minutes ago + + let snapshot = CacheSnapshot { + flow_logs, + last_refreshed_secs: expired_time, + }; + + manager.cache.load_snapshot(snapshot); + + // Verify cache is expired but entries are still present + assert!(manager.cache.is_expired()); + // Note: cache.len() still returns 1 because expired cache doesn't clear entries, + // it just refuses to serve them + + // Attempt to get config - should return None since cache is expired + // and we can't reload from EC2 (no permissions) or S3 (not configured) + let config = manager.get_config("/aws/ec2/expired-flowlogs").await; + assert!(config.is_none()); + } +} diff --git a/src/forward/forwarder.rs b/src/forward/forwarder.rs index 51a73b6..ce03073 100644 --- a/src/forward/forwarder.rs +++ b/src/forward/forwarder.rs @@ -8,6 +8,7 @@ use tracing::{debug, error}; use crate::aws_attributes::AwsAttributes; use crate::events::{LambdaEvent, LambdaPayload}; +use crate::flowlogs::FlowLogManager; use crate::forward::{AckerBuilder, AckerWaiter}; use crate::parse::cwlogs; use crate::tags::TagManager; @@ -15,13 +16,19 @@ use crate::tags::TagManager; pub struct Forwarder { logs_tx: BoundedSender>, tag_manager: TagManager, + flow_log_manager: FlowLogManager, } impl Forwarder { - pub fn new(logs_tx: BoundedSender>, tag_manager: TagManager) -> Self { + pub fn new( + logs_tx: BoundedSender>, + tag_manager: TagManager, + flow_log_manager: FlowLogManager, + ) -> Self { Self { logs_tx, tag_manager, + flow_log_manager, } } } @@ -51,8 +58,12 @@ impl Forwarder { "Handling CloudWatch Logs event" ); - let mut parser = - cwlogs::Parser::new(aws_attributes, &context.request_id, &mut self.tag_manager); + let mut parser = cwlogs::Parser::new( + aws_attributes, + &context.request_id, + &mut self.tag_manager, + &mut self.flow_log_manager, + ); // Parse the logs let resource_logs = match parser.parse(logs_event).await { @@ -113,7 +124,10 @@ mod tests { let cw_client = aws_sdk_cloudwatchlogs::Client::new(&config); let tag_manager = TagManager::new(cw_client, None, None); - let mut forwarder = Forwarder::new(logs_tx, tag_manager); + let ec2_client = aws_sdk_ec2::Client::new(&config); + let flow_log_manager = FlowLogManager::new(ec2_client, None, None); + + let mut forwarder = Forwarder::new(logs_tx, tag_manager, flow_log_manager); let logs_event = LogsEvent::default(); let context = lambda_runtime::Context::default(); let aws_attributes = AwsAttributes::new(&context); diff --git a/src/lib.rs b/src/lib.rs index c9e6ca0..5f3258a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,8 @@ pub mod aws_attributes; pub mod events; +pub mod flowlogs; pub mod forward; pub mod init; pub mod parse; +pub mod s3_cache; pub mod tags; diff --git a/src/main.rs b/src/main.rs index 8791a16..25c6bc2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -162,12 +162,19 @@ async fn run_forwarder( let aws_config = aws_config::load_from_env().await; let cw_client = aws_sdk_cloudwatchlogs::Client::new(&aws_config); - let (s3_client, s3_bucket_name) = if let Some(bucket) = s3_bucket { - info!(bucket = %bucket, "Initializing tag manager with S3 cache"); - (Some(aws_sdk_s3::Client::new(&aws_config)), Some(bucket)) - } else { - (None, None) - }; + let (s3_client, s3_client_clone, s3_bucket_name, s3_bucket_clone) = + if let Some(bucket) = s3_bucket { + info!(bucket = %bucket, "Initializing managers with S3 cache"); + let s3_client = aws_sdk_s3::Client::new(&aws_config); + ( + Some(s3_client.clone()), + Some(s3_client), + Some(bucket.clone()), + Some(bucket), + ) + } else { + (None, None, None, None) + }; let mut tag_manager = TagManager::new(cw_client, s3_client, s3_bucket_name); @@ -177,7 +184,21 @@ async fn run_forwarder( // Don't fail startup, just log the error } - let mut forwarder = forward::Forwarder::new(logs_tx, tag_manager); + // Initialize FlowLogManager + let ec2_client = aws_sdk_ec2::Client::new(&aws_config); + let mut flow_log_manager = rotel_lambda_forwarder::flowlogs::FlowLogManager::new( + ec2_client, + s3_client_clone, + s3_bucket_clone, + ); + + // Initialize flow log manager (fetch from EC2 and load from S3) + if let Err(e) = flow_log_manager.initialize().await { + error!(error = %e, "Failed to initialize flow log manager"); + // Don't fail startup, just log the error + } + + let mut forwarder = forward::Forwarder::new(logs_tx, tag_manager, flow_log_manager); init_wait_rx .await diff --git a/src/parse/cwlogs.rs b/src/parse/cwlogs.rs index 36a8886..c798471 100644 --- a/src/parse/cwlogs.rs +++ b/src/parse/cwlogs.rs @@ -1,5 +1,6 @@ use regex::Regex; use std::sync::LazyLock; +use std::{collections::HashMap, sync::Arc}; use aws_lambda_events::cloudwatch_logs::LogsEvent; use opentelemetry_proto::tonic::{ @@ -9,13 +10,19 @@ use opentelemetry_proto::tonic::{ }; use tracing::debug; -use crate::{aws_attributes::AwsAttributes, parse::record_parser::RecordParser, tags::TagManager}; +use crate::{ + aws_attributes::AwsAttributes, + flowlogs::{FlowLogManager, ParsedFields}, + parse::record_parser::RecordParser, + tags::TagManager, +}; /// Parser handles the conversion of AWS CloudWatch Logs events into OpenTelemetry ResourceLogs pub struct Parser<'a> { aws_attributes: &'a AwsAttributes, request_id: &'a String, tag_manager: &'a mut TagManager, + flow_log_manager: &'a mut FlowLogManager, } impl<'a> Parser<'a> { @@ -23,11 +30,13 @@ impl<'a> Parser<'a> { aws_attributes: &'a AwsAttributes, request_id: &'a String, tag_manager: &'a mut TagManager, + flow_log_manager: &'a mut FlowLogManager, ) -> Self { Self { aws_attributes, request_id, tag_manager, + flow_log_manager, } } @@ -71,8 +80,9 @@ impl<'a> Parser<'a> { .as_nanos() as u64; // Detect platform and parser type - let (log_platform, parser_type) = - detect_log_type(&log_data.log_group, &log_data.log_stream); + let (log_platform, parser_type, flow_log_parsed_fields, flow_log_tags) = self + .detect_log_type(&log_data.log_group, &log_data.log_stream) + .await; // Fetch tags for the log group let log_group_tags = self @@ -137,7 +147,7 @@ impl<'a> Parser<'a> { }); } - // Add tags as resource attributes + // Add CloudWatch log group tags as resource attributes for (tag_key, tag_value) in log_group_tags { attributes.push(KeyValue { key: format!("cloudwatch.log.tags.{}", tag_key), @@ -147,7 +157,17 @@ impl<'a> Parser<'a> { }); } - let rec_parser = RecordParser::new(log_platform, parser_type); + // Add EC2 Flow Log tags as resource attributes + for (tag_key, tag_value) in flow_log_tags { + attributes.push(KeyValue { + key: format!("ec2.flow-logs.tags.{}", tag_key), + value: Some(AnyValue { + value: Some(Value::StringValue(tag_value)), + }), + }); + } + + let rec_parser = RecordParser::new(log_platform, parser_type, flow_log_parsed_fields); let log_records = log_data .log_events .into_iter() @@ -196,6 +216,7 @@ pub enum LogPlatform { Lambda, Codebuild, Cloudtrail, + VpcFlowLog, Unknown, } @@ -209,6 +230,7 @@ impl LogPlatform { LogPlatform::Lambda => "aws_lambda", LogPlatform::Codebuild => "aws_codebuild", LogPlatform::Cloudtrail => "aws_cloudtrail", + LogPlatform::VpcFlowLog => "aws_vpc_flow_log", LogPlatform::Unknown => "aws_unknown", } } @@ -219,29 +241,61 @@ impl LogPlatform { pub enum ParserType { Json, KeyValue, + VpcLog, #[default] Unknown, } -/// Detects the log platform and parser type based on log group and stream names -fn detect_log_type(log_group_name: &str, log_stream_name: &str) -> (LogPlatform, ParserType) { - // First, detect the platform - let platform = detect_log_platform(log_group_name, log_stream_name); - - // Then, determine the parser type based on platform and log stream name - let parser_type = match platform { - LogPlatform::Eks => { - if log_stream_name.starts_with("authenticator-") { - ParserType::KeyValue - } else { - ParserType::Unknown - } +impl<'a> Parser<'a> { + /// Detects the log platform and parser type based on log group and stream names + /// Returns (platform, parser_type, optional_flow_log_parsed_fields, flow_log_tags) + async fn detect_log_type( + &mut self, + log_group_name: &str, + log_stream_name: &str, + ) -> ( + LogPlatform, + ParserType, + Option>, + HashMap, + ) { + // First check if this is an EC2 Flow Log + if let Some(flow_log_config) = self.flow_log_manager.get_config(log_group_name).await { + debug!( + log_group = %log_group_name, + flow_log_id = %flow_log_config.flow_log_id, + "Detected EC2 Flow Log" + ); + + // Extract parsed fields if successful, None if parsing failed or not attempted + let parsed_fields = flow_log_config.parsed_fields.as_ref().cloned(); + + return ( + LogPlatform::VpcFlowLog, + ParserType::VpcLog, + parsed_fields, + flow_log_config.tags, + ); } - LogPlatform::Cloudtrail => ParserType::Json, - _ => ParserType::Unknown, - }; - (platform, parser_type) + // Otherwise, detect the platform normally + let platform = detect_log_platform(log_group_name, log_stream_name); + + // Then, determine the parser type based on platform and log stream name + let parser_type = match platform { + LogPlatform::Eks => { + if log_stream_name.starts_with("authenticator-") { + ParserType::KeyValue + } else { + ParserType::Unknown + } + } + LogPlatform::Cloudtrail => ParserType::Json, + _ => ParserType::Unknown, + }; + + (platform, parser_type, None, HashMap::new()) + } } static CLOUDTRAIL_REGEX: LazyLock = @@ -284,7 +338,10 @@ pub enum ParserError { JsonParseError(String), #[error("Invalid log format: {0}")] - InvalidFormat(String), + FormatParseError(String), + + #[error("Unable to parse EC2 flow log format")] + FlowLogFormatError, } #[cfg(test)] @@ -304,7 +361,15 @@ mod tests { let cw_client = aws_sdk_cloudwatchlogs::Client::new(&config); let mut tag_manager = crate::tags::TagManager::new(cw_client, None, None); - let mut parser = Parser::new(&aws_attributes, &request_id, &mut tag_manager); + let ec2_client = aws_sdk_ec2::Client::new(&config); + let mut flow_log_manager = crate::flowlogs::FlowLogManager::new(ec2_client, None, None); + + let mut parser = Parser::new( + &aws_attributes, + &request_id, + &mut tag_manager, + &mut flow_log_manager, + ); let result = parser.parse(logs_event).await; assert!(result.is_ok()); @@ -317,19 +382,13 @@ mod tests { async fn test_parse_eks_authenticator_log() { let log_msg = r#"time="2025-12-24T19:48:32Z" level=info msg="access granted" arn="arn:aws:iam::927209226484:role/AWSWesleyClusterManagerLambda-Add-AddonManagerRole-1CRTQUJF13T5U" client="127.0.0.1:54812" groups="[]" method=POST path=/authenticate stsendpoint=sts.us-east-1.amazonaws.com uid="aws-iam-authenticator:927209226484:AROA5PYP2AD2FVXU23CA6" username="eks:addon-manager""#; - // Test that EKS authenticator logs are detected as KeyValue parser type - let (platform, parser_type) = - detect_log_type("/aws/eks/cluster-name", "authenticator-12345"); - assert_eq!(platform, LogPlatform::Eks); - assert_eq!(parser_type, ParserType::KeyValue); - // Test parsing the log entry let mut log_entry = LogEntry::default(); log_entry.id = "test-id".to_string(); log_entry.timestamp = 1000; log_entry.message = log_msg.to_string(); - let rec_parser = RecordParser::new(platform, parser_type); + let rec_parser = RecordParser::new(LogPlatform::Eks, ParserType::KeyValue, None); let log_record = rec_parser.parse(123456789, log_entry); // Verify the log was parsed correctly @@ -358,54 +417,52 @@ mod tests { } #[test] - fn test_detect_log_type_eks() { - let (platform, parser_type) = detect_log_type("/aws/eks/cluster-name", "application-pod"); + fn test_detect_log_platform_eks() { + let platform = detect_log_platform("/aws/eks/cluster-name", "authenticator-12345"); assert_eq!(platform, LogPlatform::Eks); - assert_eq!(parser_type, ParserType::Unknown); } #[test] - fn test_detect_log_type_eks_authenticator() { - let (platform, parser_type) = - detect_log_type("/aws/eks/cluster-name", "authenticator-12345"); - assert_eq!(platform, LogPlatform::Eks); - assert_eq!(parser_type, ParserType::KeyValue); - } - - #[test] - fn test_detect_log_type_ecs() { - let (platform, parser_type) = detect_log_type("/aws/ecs/cluster-name", "task/service"); + fn test_detect_log_platform_ecs() { + let platform = detect_log_platform("/aws/ecs/cluster-name", "task/service"); assert_eq!(platform, LogPlatform::Ecs); - assert_eq!(parser_type, ParserType::Unknown); } #[test] - fn test_detect_log_type_rds() { - let (platform, parser_type) = detect_log_type("/aws/rds/instance-name", "error"); + fn test_detect_log_platform_rds() { + let platform = detect_log_platform("/aws/rds/instance-name", "error"); assert_eq!(platform, LogPlatform::Rds); - assert_eq!(parser_type, ParserType::Unknown); } #[test] - fn test_detect_log_type_lambda() { - let (platform, parser_type) = - detect_log_type("/aws/lambda/function-name", "2024/01/01/[$LATEST]abc123"); + fn test_detect_log_platform_lambda() { + let platform = + detect_log_platform("/aws/lambda/function-name", "2024/01/01/[$LATEST]abc123"); assert_eq!(platform, LogPlatform::Lambda); - assert_eq!(parser_type, ParserType::Unknown); } #[test] - fn test_detect_log_type_no_match() { - let (platform, parser_type) = detect_log_type("/aws/route53/function-name", "stream"); + fn test_detect_log_platform_no_match() { + let platform = detect_log_platform("/aws/route53/function-name", "stream"); assert_eq!(platform, LogPlatform::Unknown); - assert_eq!(parser_type, ParserType::Unknown); - let (platform, parser_type) = detect_log_type("/custom/log/group", "stream"); + let platform = detect_log_platform("/custom/log/group", "stream"); assert_eq!(platform, LogPlatform::Unknown); - assert_eq!(parser_type, ParserType::Unknown); - let (platform, parser_type) = detect_log_type("no-prefix", "stream"); + let platform = detect_log_platform("no-prefix", "stream"); assert_eq!(platform, LogPlatform::Unknown); - assert_eq!(parser_type, ParserType::Unknown); + } + + #[test] + fn test_detect_log_platform_cloudtrail() { + let platform = + detect_log_platform("/aws/cloudtrail/logs", "123456789012_CloudTrail_us-east-1"); + assert_eq!(platform, LogPlatform::Cloudtrail); + } + + #[test] + fn test_detect_log_platform_codebuild() { + let platform = detect_log_platform("/aws/codebuild/project-name", "stream"); + assert_eq!(platform, LogPlatform::Codebuild); } } diff --git a/src/parse/json.rs b/src/parse/json.rs index b903142..b56f101 100644 --- a/src/parse/json.rs +++ b/src/parse/json.rs @@ -21,7 +21,7 @@ pub(crate) fn parse_json_to_map( Ok(_) => { // Not an object return Err(RecordParserError( - ParserError::InvalidFormat("JSON log entry is not an object".to_string()), + ParserError::FormatParseError("JSON log entry is not an object".to_string()), msg, )); } @@ -107,7 +107,7 @@ mod tests { /// Test utility: Parse a log message and return the LogRecord fn parse_log_msg(message: &str, platform: LogPlatform) -> LogRecord { let log_entry = create_log_entry(message); - let parser = RecordParser::new(platform, ParserType::Json); + let parser = RecordParser::new(platform, ParserType::Json, None); parser.parse(123456789, log_entry) } diff --git a/src/parse/keyvalue.rs b/src/parse/keyvalue.rs index 7010f3e..5f581a5 100644 --- a/src/parse/keyvalue.rs +++ b/src/parse/keyvalue.rs @@ -20,7 +20,7 @@ pub(crate) fn parse_keyvalue_to_map( if pairs.is_empty() { return Err(RecordParserError( - ParserError::InvalidFormat("No valid key-value pairs found".to_string()), + ParserError::FormatParseError("No valid key-value pairs found".to_string()), input, )); } @@ -166,7 +166,7 @@ mod tests { /// Test utility: Parse a log message and return the LogRecord fn parse_log_msg(message: &str, platform: LogPlatform) -> LogRecord { let log_entry = create_log_entry(message); - let parser = RecordParser::new(platform, ParserType::KeyValue); + let parser = RecordParser::new(platform, ParserType::KeyValue, None); parser.parse(123456789, log_entry) } diff --git a/src/parse/mod.rs b/src/parse/mod.rs index e09b88b..9000f55 100644 --- a/src/parse/mod.rs +++ b/src/parse/mod.rs @@ -3,3 +3,4 @@ pub mod field_stripper; pub mod json; pub mod keyvalue; mod record_parser; +pub mod vpclog; diff --git a/src/parse/record_parser.rs b/src/parse/record_parser.rs index 3df8ab7..3d40576 100644 --- a/src/parse/record_parser.rs +++ b/src/parse/record_parser.rs @@ -1,4 +1,4 @@ -use std::sync::OnceLock; +use std::sync::{Arc, OnceLock}; use aws_lambda_events::cloudwatch_logs::LogEntry; use opentelemetry_proto::tonic::{ @@ -9,11 +9,13 @@ use regex::Regex; use serde_json::Value as JsonValue; use tracing::{debug, warn}; +use crate::flowlogs::ParsedFields; use crate::parse::{ cwlogs::{LogPlatform, ParserError, ParserType}, field_stripper::FieldStripper, json::{json_value_to_any_value, parse_json_to_map}, keyvalue::parse_keyvalue_to_map, + vpclog::parse_vpclog_to_map, }; /// Parser for CloudWatch log entries that converts them into OpenTelemetry LogRecords. @@ -24,17 +26,24 @@ pub(crate) struct RecordParser { platform: LogPlatform, parser_type: ParserType, field_stripper: FieldStripper, + flow_log_parsed_fields: Option>, } +#[derive(Debug)] pub(crate) struct RecordParserError(pub(crate) ParserError, pub(crate) String); impl RecordParser { /// Create a new RecordParser for a specific platform and parser type. - pub(crate) fn new(platform: LogPlatform, parser_type: ParserType) -> Self { + pub(crate) fn new( + platform: LogPlatform, + parser_type: ParserType, + flow_log_parsed_fields: Option>, + ) -> Self { Self { platform, parser_type, field_stripper: FieldStripper::new(platform), + flow_log_parsed_fields, } } @@ -60,6 +69,19 @@ impl RecordParser { let json_map = match self.parser_type { ParserType::Json => parse_json_to_map(message).map(|r| Some(r)), ParserType::KeyValue => parse_keyvalue_to_map(message).map(|r| Some(r)), + ParserType::VpcLog => { + // Always keep log in Body + lr.body = Some(AnyValue { + value: Some(Value::StringValue(message.clone())), + }); + + match self.flow_log_parsed_fields.as_ref() { + Some(parsed_fields) => { + parse_vpclog_to_map(message, parsed_fields.clone()).map(|r| Some(r)) + } + None => Ok(None), + } + } ParserType::Unknown => { // Auto-detect: try JSON first, otherwise try keyvalue, otherwise plain text if message.len() > 2 && message.starts_with("{") { @@ -304,7 +326,7 @@ mod tests { /// Test utility: Parse a log message and return the LogRecord fn parse_log_msg(message: &str, platform: LogPlatform, parser_type: ParserType) -> LogRecord { let log_entry = create_log_entry(message); - let parser = RecordParser::new(platform, parser_type); + let parser = RecordParser::new(platform, parser_type, None); parser.parse(123456789, log_entry) } diff --git a/src/parse/vpclog.rs b/src/parse/vpclog.rs new file mode 100644 index 0000000..c3b68f1 --- /dev/null +++ b/src/parse/vpclog.rs @@ -0,0 +1,653 @@ +//! EC2 Flow Log Parser Module +//! +//! This module provides parsing functionality for AWS EC2 Flow Logs. +//! It converts EC2 Flow Log records into `serde_json::Map` structures for unified processing. +//! +//! EC2 Flow Logs are space-separated records that capture network traffic flow information. +//! The format can be customized, so this parser dynamically determines the fields based on +//! the LogFormat string retrieved from the EC2 DescribeFlowLogs API. +//! +//! Reference: https://docs.aws.amazon.com/vpc/latest/userguide/flow-log-records.html + +use std::sync::Arc; + +use serde_json::Value as JsonValue; + +use crate::flowlogs::{ParsedFieldType, ParsedFields}; +use crate::parse::{cwlogs::ParserError, record_parser::RecordParserError}; + +/// Parse an EC2 Flow Log record from a string using pre-parsed field names +/// Fields with value "-" are excluded from the result +/// Returns an error if the record doesn't have the expected number of fields +pub(crate) fn parse_vpclog_to_map( + input: String, + parsed_fields: Arc, +) -> Result, RecordParserError> { + match parsed_fields.as_ref() { + ParsedFields::Success(parsed_fields) => { + let field_values = parse_vpclog_fields(&input); + + if field_values.len() != parsed_fields.len() { + return Err(RecordParserError( + ParserError::FormatParseError(format!( + "Expected {} fields but found {}", + parsed_fields.len(), + field_values.len() + )), + input, + )); + } + + let mut map = serde_json::Map::new(); + for (parsed_field, value) in parsed_fields.iter().zip(field_values.iter()) { + // Skip fields with "-" value (not present) + if value == "-" { + continue; + } + + match parsed_field.field_type { + ParsedFieldType::String => { + map.insert( + parsed_field.field_name.clone(), + JsonValue::String(value.clone()), + ); + } + ParsedFieldType::Int32 | ParsedFieldType::Int64 => match value.parse::() { + Ok(num) => { + map.insert( + parsed_field.field_name.clone(), + JsonValue::Number(serde_json::Number::from(num)), + ); + } + Err(e) => { + return Err(RecordParserError( + ParserError::FormatParseError(format!( + "Field {} unable to be parsed to integer: {}", + parsed_field.field_name, e + )), + value.clone(), + )); + } + }, + } + } + + Ok(map) + } + ParsedFields::Error(msg) => Err(RecordParserError( + ParserError::FlowLogFormatError, + msg.clone(), + )), + } +} + +/// Parse EC2 Flow Log fields from a string +/// Fields are space-separated +fn parse_vpclog_fields(input: &str) -> Vec { + input.split_whitespace().map(|s| s.to_string()).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::flowlogs::parse_log_format; + use crate::parse::cwlogs::{LogPlatform, ParserType}; + use crate::parse::record_parser::RecordParser; + use aws_lambda_events::cloudwatch_logs::LogEntry; + use opentelemetry_proto::tonic::{common::v1::any_value::Value, logs::v1::LogRecord}; + + const DEFAULT_FORMAT: &'static str = "${version} ${account-id} ${interface-id} ${srcaddr} ${dstaddr} ${srcport} ${dstport} ${protocol} ${packets} ${bytes} ${start} ${end} ${action} ${log-status}"; + + /// Test utility: Create a LogEntry from a message string + fn create_log_entry(message: &str) -> LogEntry { + let mut log_entry = LogEntry::default(); + log_entry.id = "test-id".to_string(); + log_entry.timestamp = 1000; + log_entry.message = message.to_string(); + log_entry + } + + /// Test utility: Parse a log message and return the LogRecord + fn parse_log_msg(message: &str) -> LogRecord { + let log_entry = create_log_entry(message); + let dflt_fields = parse_log_format(DEFAULT_FORMAT); + let parser = RecordParser::new( + LogPlatform::Unknown, + ParserType::VpcLog, + Some(Arc::new(ParsedFields::Success(dflt_fields))), + ); + parser.parse(123456789, log_entry) + } + + fn parse_vpclog_to_map_default( + input: String, + ) -> Result, RecordParserError> { + let dflt_fields = parse_log_format(DEFAULT_FORMAT); + parse_vpclog_to_map(input, Arc::new(ParsedFields::Success(dflt_fields))) + } + + #[test] + fn test_parse_vpclog_fields_basic() { + let input = "2 123456789010 eni-1235b8ca123456789 172.31.16.139 172.31.16.21 20641 22 6 20 4249 1418530010 1418530070 ACCEPT OK"; + let fields = parse_vpclog_fields(input); + + assert_eq!(fields.len(), 14); + assert_eq!(fields[0], "2"); + assert_eq!(fields[1], "123456789010"); + assert_eq!(fields[2], "eni-1235b8ca123456789"); + assert_eq!(fields[3], "172.31.16.139"); + assert_eq!(fields[12], "ACCEPT"); + assert_eq!(fields[13], "OK"); + } + + #[test] + fn test_parse_vpclog_fields_with_dashes() { + let input = + "2 123456789010 eni-1235b8ca123456789 - - - - - - - 1431280876 1431280934 - NODATA"; + let fields = parse_vpclog_fields(input); + + assert_eq!(fields.len(), 14); + assert_eq!(fields[3], "-"); + assert_eq!(fields[4], "-"); + assert_eq!(fields[12], "-"); + assert_eq!(fields[13], "NODATA"); + } + + #[test] + fn test_parse_vpclog_to_map_basic_tcp() { + let input = "2 123456789010 eni-1235b8ca123456789 172.31.16.139 172.31.16.21 20641 22 6 20 4249 1418530010 1418530070 ACCEPT OK"; + let result = parse_vpclog_to_map_default(input.to_string()); + + assert!(result.is_ok()); + let map = result.unwrap(); + + assert_eq!(map.len(), 14); // All fields present + assert_eq!( + map.get("version").unwrap(), + &JsonValue::Number(serde_json::Number::from(2)) + ); + assert_eq!( + map.get("account-id").unwrap(), + &JsonValue::String("123456789010".to_string()) + ); + assert_eq!( + map.get("interface-id").unwrap(), + &JsonValue::String("eni-1235b8ca123456789".to_string()) + ); + assert_eq!( + map.get("srcaddr").unwrap(), + &JsonValue::String("172.31.16.139".to_string()) + ); + assert_eq!( + map.get("dstaddr").unwrap(), + &JsonValue::String("172.31.16.21".to_string()) + ); + assert_eq!( + map.get("srcport").unwrap(), + &JsonValue::Number(serde_json::Number::from(20641)) + ); + assert_eq!( + map.get("dstport").unwrap(), + &JsonValue::Number(serde_json::Number::from(22)) + ); + assert_eq!( + map.get("protocol").unwrap(), + &JsonValue::Number(serde_json::Number::from(6)) + ); + assert_eq!( + map.get("packets").unwrap(), + &JsonValue::Number(serde_json::Number::from(20)) + ); + assert_eq!( + map.get("bytes").unwrap(), + &JsonValue::Number(serde_json::Number::from(4249)) + ); + assert_eq!( + map.get("start").unwrap(), + &JsonValue::Number(serde_json::Number::from(1418530010i64)) + ); + assert_eq!( + map.get("end").unwrap(), + &JsonValue::Number(serde_json::Number::from(1418530070i64)) + ); + assert_eq!( + map.get("action").unwrap(), + &JsonValue::String("ACCEPT".to_string()) + ); + assert_eq!( + map.get("log-status").unwrap(), + &JsonValue::String("OK".to_string()) + ); + } + + #[test] + fn test_parse_vpclog_to_map_with_format_custom() { + let log_format = "${version} ${account-id} ${srcaddr} ${dstaddr}"; + let input = "2 123456789010 172.31.16.139 172.31.16.21"; + let parsed_fields = parse_log_format(log_format); + let result = parse_vpclog_to_map( + input.to_string(), + Arc::new(ParsedFields::Success(parsed_fields)), + ); + + assert!(result.is_ok()); + let map = result.unwrap(); + + assert_eq!(map.len(), 4); + assert_eq!( + map.get("version").unwrap(), + &JsonValue::Number(serde_json::Number::from(2)) + ); + assert_eq!( + map.get("account-id").unwrap(), + &JsonValue::String("123456789010".to_string()) + ); + assert_eq!( + map.get("srcaddr").unwrap(), + &JsonValue::String("172.31.16.139".to_string()) + ); + assert_eq!( + map.get("dstaddr").unwrap(), + &JsonValue::String("172.31.16.21".to_string()) + ); + } + + #[test] + fn test_parse_vpclog_to_map_with_format_excludes_dashes() { + let log_format = "${version} ${account-id} ${srcaddr} ${dstaddr}"; + let input = "2 123456789010 - 172.31.16.21"; + let parsed_fields = parse_log_format(log_format); + let result = parse_vpclog_to_map( + input.to_string(), + Arc::new(ParsedFields::Success(parsed_fields)), + ); + + assert!(result.is_ok()); + let map = result.unwrap(); + + assert_eq!(map.len(), 3); // srcaddr should be excluded + assert_eq!( + map.get("version").unwrap(), + &JsonValue::Number(serde_json::Number::from(2)) + ); + assert_eq!( + map.get("account-id").unwrap(), + &JsonValue::String("123456789010".to_string()) + ); + assert!(map.get("srcaddr").is_none()); // Should be excluded + assert_eq!( + map.get("dstaddr").unwrap(), + &JsonValue::String("172.31.16.21".to_string()) + ); + } + + #[test] + fn test_parse_vpclog_to_map_reject() { + let input = "2 123456789010 eni-1235b8ca123456789 172.31.9.69 172.31.9.12 49761 3389 6 20 4249 1418530010 1418530070 REJECT OK"; + let result = parse_vpclog_to_map_default(input.to_string()); + + assert!(result.is_ok()); + let map = result.unwrap(); + + assert_eq!( + map.get("action").unwrap(), + &JsonValue::String("REJECT".to_string()) + ); + assert_eq!( + map.get("srcport").unwrap(), + &JsonValue::Number(serde_json::Number::from(49761)) + ); + assert_eq!( + map.get("dstport").unwrap(), + &JsonValue::Number(serde_json::Number::from(3389)) + ); + } + + #[test] + fn test_parse_vpclog_to_map_with_dashes() { + let input = + "2 123456789010 eni-1235b8ca123456789 - - - - - - - 1431280876 1431280934 - NODATA"; + let result = parse_vpclog_to_map_default(input.to_string()); + + assert!(result.is_ok()); + let map = result.unwrap(); + + // Only non-dash fields should be present + assert_eq!(map.len(), 6); + assert_eq!( + map.get("version").unwrap(), + &JsonValue::Number(serde_json::Number::from(2)) + ); + assert_eq!( + map.get("account-id").unwrap(), + &JsonValue::String("123456789010".to_string()) + ); + assert_eq!( + map.get("interface-id").unwrap(), + &JsonValue::String("eni-1235b8ca123456789".to_string()) + ); + assert_eq!( + map.get("start").unwrap(), + &JsonValue::Number(serde_json::Number::from(1431280876i64)) + ); + assert_eq!( + map.get("end").unwrap(), + &JsonValue::Number(serde_json::Number::from(1431280934i64)) + ); + + // Dash fields should not be present + assert!(map.get("srcaddr").is_none()); + assert!(map.get("dstaddr").is_none()); + assert!(map.get("action").is_none()); + assert_eq!( + map.get("log-status").unwrap(), + &JsonValue::String("NODATA".to_string()) + ); + } + + #[test] + fn test_parse_vpclog_to_map_icmp() { + let input = "2 123456789010 eni-1235b8ca123456789 203.0.113.12 172.31.16.139 0 0 1 4 336 1432917027 1432917142 ACCEPT OK"; + let result = parse_vpclog_to_map_default(input.to_string()); + + assert!(result.is_ok()); + let map = result.unwrap(); + + assert_eq!( + map.get("protocol").unwrap(), + &JsonValue::Number(serde_json::Number::from(1)) + ); // ICMP + assert_eq!( + map.get("srcport").unwrap(), + &JsonValue::Number(serde_json::Number::from(0)) + ); + assert_eq!( + map.get("dstport").unwrap(), + &JsonValue::Number(serde_json::Number::from(0)) + ); + assert_eq!( + map.get("srcaddr").unwrap(), + &JsonValue::String("203.0.113.12".to_string()) + ); + } + + #[test] + fn test_parse_vpclog_to_map_ipv6() { + let input = "2 123456789010 eni-1235b8ca123456789 2001:db8:1234:a100:8d6e:3477:df66:f105 2001:db8:1234:a102:3304:8879:34cf:4071 34892 22 6 54 8855 1477913708 1477913820 ACCEPT OK"; + let result = parse_vpclog_to_map_default(input.to_string()); + + assert!(result.is_ok()); + let map = result.unwrap(); + + assert_eq!( + map.get("srcaddr").unwrap(), + &JsonValue::String("2001:db8:1234:a100:8d6e:3477:df66:f105".to_string()) + ); + assert_eq!( + map.get("dstaddr").unwrap(), + &JsonValue::String("2001:db8:1234:a102:3304:8879:34cf:4071".to_string()) + ); + assert_eq!( + map.get("srcport").unwrap(), + &JsonValue::Number(serde_json::Number::from(34892)) + ); + assert_eq!( + map.get("dstport").unwrap(), + &JsonValue::Number(serde_json::Number::from(22)) + ); + assert_eq!( + map.get("protocol").unwrap(), + &JsonValue::Number(serde_json::Number::from(6)) + ); + } + + #[test] + fn test_parse_vpclog_to_map_invalid_field_count() { + let input = "2 123456789010 eni-1235b8ca123456789 172.31.16.139"; + let result = parse_vpclog_to_map_default(input.to_string()); + + assert!(result.is_err()); + } + + #[test] + fn test_parse_vpclog_to_map_empty() { + let input = ""; + let result = parse_vpclog_to_map_default(input.to_string()); + + assert!(result.is_err()); + } + + #[test] + fn test_parse_log_record_basic_tcp() { + let input = "2 123456789010 eni-1235b8ca123456789 172.31.16.139 172.31.16.21 20641 22 6 20 4249 1418530010 1418530070 ACCEPT OK"; + let log_record = parse_log_msg(input); + + // Check attributes exist + let attrs: std::collections::HashMap<_, _> = log_record + .attributes + .iter() + .map(|kv| (kv.key.as_str(), kv.value.as_ref())) + .collect(); + + assert!(attrs.contains_key("version")); + assert!(attrs.contains_key("account-id")); + assert!(attrs.contains_key("interface-id")); + assert!(attrs.contains_key("srcaddr")); + assert!(attrs.contains_key("dstaddr")); + assert!(attrs.contains_key("srcport")); + assert!(attrs.contains_key("dstport")); + assert!(attrs.contains_key("protocol")); + assert!(attrs.contains_key("packets")); + assert!(attrs.contains_key("bytes")); + assert!(attrs.contains_key("start")); + assert!(attrs.contains_key("end")); + assert!(attrs.contains_key("action")); + assert!(attrs.contains_key("log-status")); + + // Verify specific values + if let Some(Some(any_val)) = attrs.get("srcaddr") { + if let Some(Value::StringValue(s)) = &any_val.value { + assert_eq!(s, "172.31.16.139"); + } + } + + if let Some(Some(any_val)) = attrs.get("action") { + if let Some(Value::StringValue(s)) = &any_val.value { + assert_eq!(s, "ACCEPT"); + } + } + + if let Some(Some(any_val)) = attrs.get("protocol") { + if let Some(Value::IntValue(i)) = &any_val.value { + assert_eq!(*i, 6); // TCP + } + } + } + + #[test] + fn test_parse_log_record_reject() { + let input = "2 123456789010 eni-1235b8ca123456789 172.31.9.69 172.31.9.12 49761 3389 6 20 4249 1418530010 1418530070 REJECT OK"; + let log_record = parse_log_msg(input); + + let attrs: std::collections::HashMap<_, _> = log_record + .attributes + .iter() + .map(|kv| (kv.key.as_str(), kv.value.as_ref())) + .collect(); + + if let Some(Some(any_val)) = attrs.get("action") { + if let Some(Value::StringValue(s)) = &any_val.value { + assert_eq!(s, "REJECT"); + } + } + } + + #[test] + fn test_parse_log_record_with_dashes() { + let input = + "2 123456789010 eni-1235b8ca123456789 - - - - - - - 1431280876 1431280934 - NODATA"; + let log_record = parse_log_msg(input); + + let attrs: std::collections::HashMap<_, _> = log_record + .attributes + .iter() + .map(|kv| (kv.key.as_str(), kv.value.as_ref())) + .collect(); + + // Should have only non-dash fields plus cloudwatch.id + assert!(attrs.contains_key("version")); + assert!(attrs.contains_key("account-id")); + assert!(attrs.contains_key("interface-id")); + assert!(attrs.contains_key("start")); + assert!(attrs.contains_key("end")); + assert!(attrs.contains_key("log-status")); + assert!(attrs.contains_key("cloudwatch.id")); + + // Dash fields should not be present + assert!(!attrs.contains_key("srcaddr")); + assert!(!attrs.contains_key("dstaddr")); + assert!(!attrs.contains_key("srcport")); + assert!(!attrs.contains_key("dstport")); + assert!(!attrs.contains_key("protocol")); + assert!(!attrs.contains_key("packets")); + assert!(!attrs.contains_key("bytes")); + assert!(!attrs.contains_key("action")); + } + + #[test] + fn test_parse_log_record_icmp() { + let input = "2 123456789010 eni-1235b8ca123456789 203.0.113.12 172.31.16.139 0 0 1 4 336 1432917027 1432917142 ACCEPT OK"; + let log_record = parse_log_msg(input); + + let attrs: std::collections::HashMap<_, _> = log_record + .attributes + .iter() + .map(|kv| (kv.key.as_str(), kv.value.as_ref())) + .collect(); + + if let Some(Some(any_val)) = attrs.get("protocol") { + if let Some(Value::IntValue(i)) = &any_val.value { + assert_eq!(*i, 1); // ICMP + } + } + + if let Some(Some(any_val)) = attrs.get("srcport") { + if let Some(Value::IntValue(i)) = &any_val.value { + assert_eq!(*i, 0); + } + } + } + + #[test] + fn test_parse_log_record_ipv6() { + let input = "2 123456789010 eni-1235b8ca123456789 2001:db8:1234:a100:8d6e:3477:df66:f105 2001:db8:1234:a102:3304:8879:34cf:4071 34892 22 6 54 8855 1477913708 1477913820 ACCEPT OK"; + let log_record = parse_log_msg(input); + + let attrs: std::collections::HashMap<_, _> = log_record + .attributes + .iter() + .map(|kv| (kv.key.as_str(), kv.value.as_ref())) + .collect(); + + if let Some(Some(any_val)) = attrs.get("srcaddr") { + if let Some(Value::StringValue(s)) = &any_val.value { + assert_eq!(s, "2001:db8:1234:a100:8d6e:3477:df66:f105"); + } + } + + if let Some(Some(any_val)) = attrs.get("dstaddr") { + if let Some(Value::StringValue(s)) = &any_val.value { + assert_eq!(s, "2001:db8:1234:a102:3304:8879:34cf:4071"); + } + } + } + + #[test] + fn test_parse_log_record_timestamps() { + let input = "2 123456789010 eni-1235b8ca123456789 172.31.16.139 172.31.16.21 20641 22 6 20 4249 1418530010 1418530070 ACCEPT OK"; + let log_record = parse_log_msg(input); + + let attrs: std::collections::HashMap<_, _> = log_record + .attributes + .iter() + .map(|kv| (kv.key.as_str(), kv.value.as_ref())) + .collect(); + + // Check that start and end timestamps are present (Int64) + if let Some(Some(any_val)) = attrs.get("start") { + if let Some(Value::IntValue(i)) = &any_val.value { + assert_eq!(*i, 1418530010); + } + } + + if let Some(Some(any_val)) = attrs.get("end") { + if let Some(Value::IntValue(i)) = &any_val.value { + assert_eq!(*i, 1418530070); + } + } + } + + #[test] + fn test_all_example_logs() { + // Test all examples from the user's request + let examples = vec![ + "2 123456789010 eni-1235b8ca123456789 172.31.16.139 172.31.16.21 20641 22 6 20 4249 1418530010 1418530070 ACCEPT OK", + "2 123456789010 eni-1235b8ca123456789 172.31.9.69 172.31.9.12 49761 3389 6 20 4249 1418530010 1418530070 REJECT OK", + "2 123456789010 eni-1235b8ca123456789 - - - - - - - 1431280876 1431280934 - NODATA", + "2 123456789010 eni-1235b8ca123456789 203.0.113.12 172.31.16.139 0 0 1 4 336 1432917027 1432917142 ACCEPT OK", + "2 123456789010 eni-1235b8ca123456789 2001:db8:1234:a100:8d6e:3477:df66:f105 2001:db8:1234:a102:3304:8879:34cf:4071 34892 22 6 54 8855 1477913708 1477913820 ACCEPT OK", + ]; + + for example in examples { + let result = parse_vpclog_to_map_default(example.to_string()); + assert!(result.is_ok(), "Failed to parse: {}", example); + + let log_record = parse_log_msg(example); + assert!( + !log_record.attributes.is_empty(), + "No attributes parsed for: {}", + example + ); + } + } + + #[test] + fn test_parse_vpclog_to_map_invalid_int32() { + // Test parsing a field that should be Int32 but contains non-integer data + let log_format = "${version} ${account-id} ${srcport} ${dstport}"; + let input = "2 123456789010 invalid-port 22"; + let parsed_fields = parse_log_format(log_format); + let result = parse_vpclog_to_map( + input.to_string(), + Arc::new(ParsedFields::Success(parsed_fields)), + ); + + // Should return an error because srcport (Int32) cannot parse "invalid-port" + assert!(result.is_err()); + if let Err(e) = result { + let error_msg = format!("{:?}", e); + assert!(error_msg.contains("srcport")); + assert!(error_msg.contains("integer")); + } + } + + #[test] + fn test_parse_vpclog_to_map_invalid_int64() { + // Test parsing an Int64 field with non-integer data + let log_format = "${version} ${account-id} ${bytes} ${packets}"; + let input = "2 123456789010 not-a-number 100"; + let parsed_fields = parse_log_format(log_format); + let result = parse_vpclog_to_map( + input.to_string(), + Arc::new(ParsedFields::Success(parsed_fields)), + ); + + // Should return an error because bytes (Int64) cannot parse "not-a-number" + assert!(result.is_err()); + if let Err(e) = result { + let error_msg = format!("{:?}", e); + assert!(error_msg.contains("bytes")); + assert!(error_msg.contains("integer")); + } + } +} diff --git a/src/s3_cache.rs b/src/s3_cache.rs new file mode 100644 index 0000000..ee27b35 --- /dev/null +++ b/src/s3_cache.rs @@ -0,0 +1,338 @@ +//! Generic S3 Cache Module +//! +//! This module provides a reusable S3 persistence layer that can be used for different +//! types of cache data (tags, flow logs, etc.). It handles compression, serialization, +//! conditional writes, and conflict resolution. + +use aws_sdk_s3::Client as S3Client; +use aws_sdk_s3::primitives::ByteStream; +use flate2::Compression; +use flate2::read::GzDecoder; +use flate2::write::GzEncoder; +use serde::{Deserialize, Serialize}; +use std::io::{Read, Write}; +use thiserror::Error; +use tracing::{debug, error, info}; + +/// Errors that can occur during S3 operations +#[derive(Debug, Error)] +pub enum S3CacheError { + #[error("S3 operation failed: {0}")] + S3Operation(String), + + #[error("Compression failed: {0}")] + Compression(String), + + #[error("Decompression failed: {0}")] + Decompression(String), + + #[error("Serialization failed: {0}")] + Serialization(String), + + #[error("Deserialization failed: {0}")] + Deserialization(String), + + #[error("Cache file not found")] + NotFound, + + #[error("Conditional write failed - cache was modified")] + ConditionalWriteFailed, + + #[error("Invalid configuration: {0}")] + InvalidConfiguration(String), +} + +/// Generic S3 persistence layer for cache data +/// +/// This provides a reusable S3 caching mechanism that can store any serializable data. +/// It handles compression, ETags for conditional writes, and conflict resolution. +pub struct S3Cache +where + T: Serialize + for<'de> Deserialize<'de>, +{ + client: S3Client, + bucket: String, + key: String, + /// Current ETag of the cache file, if known + current_etag: Option, + _phantom: std::marker::PhantomData, +} + +impl S3Cache +where + T: Serialize + for<'de> Deserialize<'de>, +{ + /// Create a new S3 cache + pub fn new(client: S3Client, bucket: String, key: String) -> Self { + Self { + client, + bucket, + key, + current_etag: None, + _phantom: std::marker::PhantomData, + } + } + + /// Load the cache from S3 + /// Returns None if the cache file doesn't exist + pub async fn load(&mut self) -> Result, S3CacheError> { + debug!(bucket = %self.bucket, key = %self.key, "Loading cache from S3"); + + let result = self + .client + .get_object() + .bucket(&self.bucket) + .key(&self.key) + .send() + .await; + + match result { + Ok(output) => { + // Store the ETag for later conditional writes + self.current_etag = output.e_tag().map(|s| s.to_string()); + + let body = output.body.collect().await.map_err(|e| { + S3CacheError::S3Operation(format!("Failed to read body: {}", e)) + })?; + + let bytes = body.into_bytes(); + + // Decompress and deserialize the data + let data = decompress_and_deserialize(&bytes)?; + + info!( + bucket = %self.bucket, + key = %self.key, + etag = ?self.current_etag, + "Successfully loaded cache from S3" + ); + + Ok(Some(data)) + } + Err(e) => { + let err_msg = format!("{:?}", e); + if err_msg.contains("NoSuchKey") || err_msg.contains("NotFound") { + debug!(bucket = %self.bucket, key = %self.key, "Cache file not found in S3"); + Ok(None) + } else { + error!(bucket = %self.bucket, key = %self.key, error = %err_msg, "Failed to load cache from S3"); + Err(S3CacheError::S3Operation(err_msg)) + } + } + } + } + + /// Save the cache to S3 with conditional write + /// If the write fails due to ETag mismatch, returns ConditionalWriteFailed error + pub async fn save(&mut self, data: T) -> Result<(), S3CacheError> { + debug!( + bucket = %self.bucket, + key = %self.key, + "Saving cache to S3" + ); + + // Serialize and compress the data + let compressed_data = serialize_and_compress(&data)?; + + // Prepare the put request with conditional write if we have an ETag + let mut put_request = self + .client + .put_object() + .bucket(&self.bucket) + .key(&self.key) + .body(ByteStream::from(compressed_data)) + .content_type("application/gzip"); + + // Add conditional write based on ETag if we have one, otherwise + // do a conditional write on the non-existence of the file + if let Some(ref etag) = self.current_etag { + debug!(etag = %etag, "Using conditional write with ETag"); + put_request = put_request.if_match(etag); + } else { + debug!("Using conditional write with if-none-match"); + put_request = put_request.if_none_match("*"); + } + + match put_request.send().await { + Ok(output) => { + // Update our stored ETag + self.current_etag = output.e_tag().map(|s| s.to_string()); + + debug!( + bucket = %self.bucket, + key = %self.key, + new_etag = ?self.current_etag, + "Successfully saved cache to S3" + ); + Ok(()) + } + Err(e) => { + let err_msg = format!("{:?}", e); + + // Check if this is a conditional write failure (ETag mismatch) + if err_msg.contains("PreconditionFailed") || err_msg.contains("IfMatch") { + info!( + bucket = %self.bucket, + key = %self.key, + "Conditional write conflict to S3 cache file, will retry" + ); + Err(S3CacheError::ConditionalWriteFailed) + } else { + error!( + bucket = %self.bucket, + key = %self.key, + error = %err_msg, + "Failed to save cache to S3" + ); + Err(S3CacheError::S3Operation(err_msg)) + } + } + } + } + + /// Reload the cache from S3 and merge with the provided data using a merge function + /// This is used when a conditional write fails + pub async fn reload_and_merge( + &mut self, + local_data: T, + merge_fn: F, + ) -> Result + where + F: FnOnce(T, T) -> T, + { + debug!("Reloading cache from S3 to merge with new data"); + + // Load the current cache from S3 + let from_s3 = self.load().await?; + + // If nothing in S3, just use local data + let from_s3 = match from_s3 { + Some(data) => data, + None => { + debug!("No data in S3, using local data only"); + return Ok(local_data); + } + }; + + // Merge the data using the provided function + let merged = merge_fn(from_s3, local_data); + + debug!("Merged cache data"); + Ok(merged) + } + + /// Get the current ETag if known + pub fn etag(&self) -> Option<&str> { + self.current_etag.as_deref() + } + + /// Get the S3 bucket name + pub fn bucket(&self) -> &str { + &self.bucket + } + + /// Get the S3 key + pub fn key(&self) -> &str { + &self.key + } +} + +/// Serialize data and compress with gzip +fn serialize_and_compress(data: &T) -> Result, S3CacheError> +where + T: Serialize, +{ + // Serialize to JSON + let json = serde_json::to_vec(data).map_err(|e| S3CacheError::Serialization(e.to_string()))?; + + // Compress with gzip + let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); + encoder + .write_all(&json) + .map_err(|e| S3CacheError::Compression(e.to_string()))?; + + encoder + .finish() + .map_err(|e| S3CacheError::Compression(e.to_string())) +} + +/// Decompress gzip data and deserialize +fn decompress_and_deserialize(compressed_data: &[u8]) -> Result +where + T: for<'de> Deserialize<'de>, +{ + // Decompress + let mut decoder = GzDecoder::new(compressed_data); + let mut json = Vec::new(); + decoder + .read_to_end(&mut json) + .map_err(|e| S3CacheError::Decompression(e.to_string()))?; + + // Deserialize + let data: T = + serde_json::from_slice(&json).map_err(|e| S3CacheError::Deserialization(e.to_string()))?; + + Ok(data) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde::{Deserialize, Serialize}; + use std::collections::HashMap; + + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + struct TestData { + items: HashMap, + count: u32, + } + + #[test] + fn test_serialize_and_compress() { + let mut items = HashMap::new(); + items.insert("key1".to_string(), "value1".to_string()); + items.insert("key2".to_string(), "value2".to_string()); + + let data = TestData { items, count: 42 }; + + let result = serialize_and_compress(&data); + assert!(result.is_ok()); + + let compressed = result.unwrap(); + assert!(!compressed.is_empty()); + } + + #[test] + fn test_decompress_and_deserialize() { + let mut items = HashMap::new(); + items.insert("key1".to_string(), "value1".to_string()); + items.insert("key2".to_string(), "value2".to_string()); + + let data = TestData { + items: items.clone(), + count: 42, + }; + + let compressed = serialize_and_compress(&data).unwrap(); + let result: Result = decompress_and_deserialize(&compressed); + assert!(result.is_ok()); + + let decompressed = result.unwrap(); + assert_eq!(decompressed.items, items); + assert_eq!(decompressed.count, 42); + } + + #[test] + fn test_round_trip() { + let mut items = HashMap::new(); + items.insert("env".to_string(), "prod".to_string()); + items.insert("team".to_string(), "platform".to_string()); + + let data = TestData { items, count: 100 }; + + let compressed = serialize_and_compress(&data).unwrap(); + let decompressed: TestData = decompress_and_deserialize(&compressed).unwrap(); + + assert_eq!(decompressed, data); + } +}