Nftables ensure comments are attached to expected position in NFT syntax (after verdict or counter rather than standalone line).

PiperOrigin-RevId: 504134275

Author: TheLinuxGuy

capirca/lib/nftables.py

@@ -349,7 +349,7 @@ def _OptionsHandler(self, term):
     else:
       return ''
 
-  def GroupExpressions(self, address_expr, pp_expr, options, verdict):
+  def GroupExpressions(self, address_expr, pp_expr, options, verdict, comment):
     """Combines all expressions with a verdict (decision).
 
     The inputs are already pre-sanitized by RulesetGenerator. NFTables processes
@@ -362,11 +362,13 @@ def GroupExpressions(self, address_expr, pp_expr, options, verdict):
       pp_expr: pre-processed list of nftables protocols and ports.
       options: string value to append before verdict for NFT special options.
       verdict: action to take on resulting final statement (allow/deny).
+      comment: term.comment string adhering to NFT limits.
 
     Returns:
       list of strings representing valid nftables statements.
     """
     statement = []
+    statement_with_comment = []
     if address_expr:
       for addr in address_expr:
         if pp_expr:
@@ -388,8 +390,13 @@ def GroupExpressions(self, address_expr, pp_expr, options, verdict):
         statement.append(pstat + Add(options) + Add(verdict))
     else:
       # If no addresses or ports & protocol. Verdict only statement.
-      statement.append((Add(options) + verdict))
-    return statement
+      statement.append((Add(options) + Add(verdict)))
+    # Handling of comments should always be done after verdict statement.
+    if comment != 'comment ':
+      statement_with_comment.append(statement[0] + Add(comment))
+      return statement_with_comment
+    else:
+      return statement
 
   def _AddrStatement(self, address_family, src_addr, dst_addr):
     """Builds an NFTables address statement.
@@ -455,12 +462,12 @@ def RulesetGenerator(self, term):
       list of strings. Representing a ruleset for later formatting.
     """
     term_ruleset = []
+    comment = 'comment '
 
     # COMMENT handling.
     if self.verbose:
-      term_ruleset.append(
-          'comment ' + aclgenerator.TruncateWords(
-              self.term.comment, Nftables.COMMENT_CHAR_LIMIT))
+      comment += aclgenerator.TruncateWords(
+          self.term.comment, Nftables.COMMENT_CHAR_LIMIT)
     # OPTIONS / LOGGING / COUNTERS
     opt = self._OptionsHandler(term)
     # STATEMENT VERDICT / ACTION.
@@ -483,7 +490,7 @@ def RulesetGenerator(self, term):
 
       # TODO: If verdict is not supported, drop nftable_rule for it.
       nftable_rule = self.GroupExpressions(address_list, proto_and_ports, opt,
-                                           verdict)
+                                           verdict, comment)
       term_ruleset.extend(nftable_rule)
     return term_ruleset
 
@@ -795,20 +802,22 @@ def __str__(self):
           # First time we comment it out so .nft file is human-readable.
           nft_config.append(
               TabSpacer(8, '#' + ' '.join(base_chain_dict[item]['comment'])))
-          # Second time so 'nft list ruleset' keeps the comment in memory.
-          nft_config.append(
-              TabSpacer(
-                  8, 'comment ' +
-                  aclgenerator.TruncateWords(
-              base_chain_dict[item]['comment'], self.COMMENT_CHAR_LIMIT)))
         nft_config.append(
             TabSpacer(
                 8, 'type filter hook %s priority %s; policy %s;' %
                 (base_chain_dict[item]['hook'],
                  base_chain_dict[item]['priority'],
                  base_chain_dict[item]['policy'])))
-        # stateful firewall: allow reply traffic.
-        nft_config.append(TabSpacer(8, 'ct state established,related accept'))
+        # Add policy header comment after stateful firewall rule.
+        if base_chain_dict[item]['comment']:
+          nft_config.append(TabSpacer(8, 'ct state established,related accept'
+                                      + Add('comment') +
+                                      Add(aclgenerator.TruncateWords(
+                                          base_chain_dict[item]['comment'],
+                                          self.COMMENT_CHAR_LIMIT))))
+        else:
+          # stateful firewall: allows reply traffic.
+          nft_config.append(TabSpacer(8, 'ct state established,related accept'))
         # Reference the child chains with jump.
         for child_chain in base_chain_dict[item]['rules'][item].keys():
           nft_config.append(TabSpacer(8, 'jump %s' % child_chain))

tests/lib/nftables_test.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Google Inc. All Rights Reserved.
+# Copyright 2023 Google Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -193,6 +193,21 @@ def __init__(self, in_dict: dict):
 }
 """
 
+COMMENT_TERM = """
+term good-icmpv6-type {
+  comment:: "This term has a comment"
+  protocol:: tcp
+  action:: accept
+}
+"""
+
+NOCOMMENT_TERM = """
+term good-icmpv6-type {
+  protocol:: tcp
+  action:: accept
+}
+"""
+
 GOOD_TERM_1 = """
 term good-term-1 {
   action:: accept
@@ -294,17 +309,16 @@ def testCreateAnonymousSet(self, input_data, expected):
     self.assertEqual(result, expected)
 
   @parameterized.parameters(
-      ([
-          'ip6 saddr 2606:4700:4700::1111/128 ip6 daddr { 2001:4860:4860::8844/128, 2001:4860:4860::8888/128'
-      ], ['tcp sport 80 tcp dport 80'],
-       'ct state { ESTABLISHED, RELATED } log prefix "combo_cnt_log_established" counter',
-       'accept', [
-           'ip6 saddr 2606:4700:4700::1111/128 ip6 daddr { 2001:4860:4860::8844/128, 2001:4860:4860::8888/128 tcp sport 80 tcp dport 80 ct state { ESTABLISHED, RELATED } log prefix "combo_cnt_log_established" counter accept'
-       ]),)
-  def testGroupExpressions(self, address_expr, porst_proto_expr, opt, verdict,
-                           expected_output):
+      (['ip6 saddr 2606:4700:4700::1111/128 ip6 daddr { 2001:4860:4860::8844/128, 2001:4860:4860::8888/128 }'], ['tcp sport 80 tcp dport 80'],'ct state { ESTABLISHED, RELATED } log prefix "combo_cnt_log_established" counter',
+       'accept', 'comment ', ['ip6 saddr 2606:4700:4700::1111/128 ip6 daddr { 2001:4860:4860::8844/128, 2001:4860:4860::8888/128 } tcp sport 80 tcp dport 80 ct state { ESTABLISHED, RELATED } log prefix "combo_cnt_log_established" counter accept'
+       ]),
+      (['ip daddr 8.8.8.8/32'], ['tcp sport 53 tcp dport 53'],'ct state new','accept', 'comment "this is a term with a comment"', ['ip daddr 8.8.8.8/32 tcp sport 53 tcp dport 53 ct state new accept comment "this is a term with a comment"'])
+      )
+  def testGroupExpressions(self, address_expr, porst_proto_expr, opt,
+                           verdict, comment, expected_output):
     result = self.dummyterm.GroupExpressions(address_expr, porst_proto_expr,
-                                             opt, verdict)
+                                             opt, verdict, comment)
+
     self.assertEqual(result, expected_output)
 
   def testDuplicateTerm(self):
@@ -530,11 +544,12 @@ def testSkippedTerm(self, termdata, messagetxt):
   @parameterized.parameters(
       (GOOD_HEADER_1 + GOOD_TERM_2, 'inet6'),
       (GOOD_HEADER_1 + ICMPV6_TERM, 'inet6'),
+      (GOOD_HEADER_1 + COMMENT_TERM, 'mixed'),
       (GOOD_HEADER_2 + GOOD_TERM_2, 'mixed'),
       (GOOD_HEADER_3 + GOOD_TERM_2, 'inet'),
       (GOOD_HEADER_3 + ICMP_TERM, 'inet'),
   )
-  def testRulesetGenerator(self, policy_data: str, expected_inet: str):
+  def testRulesetGeneratorAF(self, policy_data: str, expected_inet: str):
     self.naming.GetNetAddr.return_value = TEST_IPS
     self.naming.GetServiceByProto.return_value = ['22']