From aaa3717227e47c22195f2cf7e95d96bc01179cd8 Mon Sep 17 00:00:00 2001
From: Nacho Barrientos <nacho.barrientos@cern.ch>
Date: Thu, 3 Dec 2020 17:09:32 +0100
Subject: [PATCH] Implement nftables:;simplerule::daddr

---
 manifests/simplerule.pp         | 14 +++++--
 spec/defines/simplerule_spec.rb | 67 ++++++++++++++++++++++++++++++++-
 templates/simplerule.epp        | 24 ++++++++++--
 3 files changed, 96 insertions(+), 9 deletions(-)

diff --git a/manifests/simplerule.pp b/manifests/simplerule.pp
index a3d06f73..77a8035f 100644
--- a/manifests/simplerule.pp
+++ b/manifests/simplerule.pp
@@ -20,6 +20,10 @@
     $dport  = undef,
   Optional[Enum['tcp', 'tcp4', 'tcp6', 'udp', 'udp4', 'udp6']]
     $proto  = undef,
+  Optional[Variant[Stdlib::IP::Address::V6, Stdlib::IP::Address::V4, Pattern[/^@[-a-zA-Z0-9_]+$/]]]
+    $daddr = undef,
+  Enum['ip', 'ip6']
+    $set_type = 'ip6',
 ){
 
   if $dport and !$proto {
@@ -30,10 +34,12 @@
     nftables::rule{"${chain}-${rulename}":
       content => epp('nftables/simplerule.epp',
         {
-          'action'  => $action,
-          'comment' => $comment,
-          'dport'   => $dport,
-          'proto'   => $proto,
+          'action'   => $action,
+          'comment'  => $comment,
+          'dport'    => $dport,
+          'proto'    => $proto,
+          'daddr'    => $daddr,
+          'set_type' => $set_type,
         }
       ),
       order   => $order,
diff --git a/spec/defines/simplerule_spec.rb b/spec/defines/simplerule_spec.rb
index 3e6d3da9..828d9dc9 100644
--- a/spec/defines/simplerule_spec.rb
+++ b/spec/defines/simplerule_spec.rb
@@ -38,13 +38,14 @@
             dport: 333,
             proto: 'udp',
             chain: 'default_out',
+            daddr: '2001:1458::/32',
           }
         end
 
         it { is_expected.to compile }
         it {
           is_expected.to contain_nftables__rule('default_out-my_big_rule').with(
-            content: 'udp dport 333 comment "this is my rule" accept',
+            content: 'udp dport 333 ip6 daddr 2001:1458::/32 accept comment "this is my rule"',
             order: '50',
           )
         }
@@ -113,6 +114,70 @@
           )
         }
       end
+
+      describe 'with an IPv4 CIDR as daddr' do
+        let(:params) do
+          {
+            daddr: '192.168.0.1/24',
+            dport: 33,
+            proto: 'tcp',
+          }
+        end
+
+        it { is_expected.to compile }
+        it {
+          is_expected.to contain_nftables__rule('default_in-my_default_rule_name').with(
+            content: 'tcp dport 33 ip daddr 192.168.0.1/24 accept',
+          )
+        }
+      end
+
+      describe 'with an IPv6 address as daddr' do
+        let(:params) do
+          {
+            daddr: '2001:1458::1',
+          }
+        end
+
+        it { is_expected.to compile }
+        it {
+          is_expected.to contain_nftables__rule('default_in-my_default_rule_name').with(
+            content: 'ip6 daddr 2001:1458::1 accept',
+          )
+        }
+      end
+
+      describe 'with an IPv6 set as daddr, default set_type' do
+        let(:params) do
+          {
+            daddr: '@my6_set',
+          }
+        end
+
+        it { is_expected.to compile }
+        it {
+          is_expected.to contain_nftables__rule('default_in-my_default_rule_name').with(
+            content: 'ip6 daddr @my6_set accept',
+          )
+        }
+      end
+
+      describe 'with a IPv4 set as daddr' do
+        let(:params) do
+          {
+            daddr: '@my4_set',
+            set_type: 'ip',
+          }
+        end
+
+        it { is_expected.to compile }
+        it {
+          is_expected.to contain_nftables__rule('default_in-my_default_rule_name').with(
+            content: 'ip daddr @my4_set accept',
+          )
+        }
+      end
+
     end
   end
 end
diff --git a/templates/simplerule.epp b/templates/simplerule.epp
index 5b797388..a66c2806 100644
--- a/templates/simplerule.epp
+++ b/templates/simplerule.epp
@@ -2,6 +2,8 @@
       Optional[String]        $comment,
       Optional[Variant[Array[Stdlib::Port, 1], Stdlib::Port, String]] $dport,
       Optional[String]        $proto,
+      Optional[Variant[Stdlib::IP::Address::V6, Stdlib::IP::Address::V4, Pattern[/^@[-a-zA-Z0-9_]+$/]]] $daddr,
+      Enum['ip', 'ip6']       $set_type,
 | -%>
 <%- if $proto {
   $_proto = $proto ? {
@@ -16,18 +18,32 @@
 } else {
   $_ip_version_filter = undef
 } -%>
+<%- if $daddr {
+  if $daddr =~ Stdlib::IP::Address::V6 {
+    $_dst_hosts = "ip6 daddr ${daddr}"
+  } elsif $daddr =~ Stdlib::IP::Address::V4 {
+    $_dst_hosts = "ip daddr ${daddr}"
+  } else {
+    $_dst_hosts = $set_type ? {
+      'ip'  => "ip daddr ${daddr}",
+      'ip6' => "ip6 daddr ${daddr}",
+    }
+  }
+} else {
+  $_dst_hosts = undef
+} -%>
 <%- if $proto and $dport {
   if $dport =~ Array {
-    $_destination = "${_proto} dport {${dport.join(', ')}}"
+    $_dst_port = "${_proto} dport {${dport.join(', ')}}"
   } else {
-    $_destination = "${_proto} dport $dport"
+    $_dst_port = "${_proto} dport $dport"
   }
 } else {
-  $_destination = undef
+  $_dst_port = undef
 } -%>
 <%- if $comment {
   $_comment = "comment \"${comment}\""
 } else {
   $_comment = undef
 } -%>
-<%= regsubst(strip([$_ip_version_filter, $_destination, $_comment, $action].join(' ')), '\s+', ' ', 'G') -%>
+<%= regsubst(strip([$_ip_version_filter, $_dst_port, $_dst_hosts, $action, $_comment].join(' ')), '\s+', ' ', 'G') -%>