diff --git a/src/ast.rs b/src/ast.rs index 7eeece2..9bc53ae 100644 --- a/src/ast.rs +++ b/src/ast.rs @@ -127,7 +127,7 @@ pub enum Expression { // Vmap expressions (value maps) Vmap { - expr: Box, + expr: Option>, map: Vec<(Expression, Expression)>, }, diff --git a/src/parser.rs b/src/parser.rs index d250d4d..d28fffc 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -311,9 +311,8 @@ impl Parser { self.advance(); // consume 'policy' let policy = self.parse_policy()?; chain = chain.with_policy(policy); + self.consume(TokenKind::Semicolon, "Expected ';' after policy")?; } - - self.consume(TokenKind::Semicolon, "Expected ';' after policy")?; } Some(TokenKind::CommentLine(_)) => { self.advance(); @@ -511,7 +510,7 @@ impl Parser { // Return a vmap expression with the previous expression as the mapping target return Ok(Expression::Vmap { - expr: Box::new(expr), + expr: Some(Box::new(expr)), map, }); } @@ -828,8 +827,8 @@ impl Parser { self.consume(TokenKind::RightBrace, "Expected '}' to close vmap")?; - // The expression that came before "vmap" is the expr being mapped - let expr = Box::new(Expression::Identifier("dummy".to_string())); // This will be replaced in post-processing + // No expression available at parse time, will be filled by post-processing if needed + let expr = None; Ok(Expression::Vmap { expr, map }) } @@ -854,6 +853,23 @@ impl Parser { self.consume(TokenKind::RightBrace, "Expected '}' to close set")?; Ok(Expression::Set(elements)) } + Some(TokenKind::Accept) => { + self.advance(); + Ok(Expression::Identifier("accept".to_string())) + } + Some(TokenKind::Drop) => { + self.advance(); + Ok(Expression::Identifier("drop".to_string())) + } + Some(TokenKind::Reject) => { + self.advance(); + Ok(Expression::Identifier("reject".to_string())) + } + Some(TokenKind::Protocol) => { + self.advance(); // consume 'protocol' + let protocol = self.parse_identifier_or_keyword()?; + Ok(Expression::Protocol(protocol)) + } _ => Err(ParseError::InvalidExpression { message: format!( "Unexpected token in expression: {}", diff --git a/src/syntax.rs b/src/syntax.rs index 37bcb65..1b57ab3 100644 --- a/src/syntax.rs +++ b/src/syntax.rs @@ -175,10 +175,11 @@ impl NftablesFormatter { // Add policy on the same line if present if let Some(policy) = &chain.policy { write!(output, " policy {}", policy).unwrap(); + output.push_str(";\n"); + } else { + output.push_str("\n"); } - output.push_str(";\n"); - if !chain.rules.is_empty() && !self.config.optimize { output.push('\n'); } @@ -294,8 +295,11 @@ impl NftablesFormatter { } Expression::Vmap { expr, map } => { - self.format_expression(output, expr); - output.push_str(" vmap { "); + if let Some(expr) = expr { + self.format_expression(output, expr); + output.push(' '); + } + output.push_str("vmap { "); for (i, (key, value)) in map.iter().enumerate() { if i > 0 { output.push_str(", ");