commit 3a469f2b35f19d7e91fe2f745a124faf8271b45b
Author: Nacho Barrientos <nacho.barrientos@cern.ch>
Date:   Wed Dec 9 13:02:25 2020 +0100

    Implement nftables::simplerule::saddr

diff --git a/manifests/simplerule.pp b/manifests/simplerule.pp
index 2f889a9..6ef963b 100644
--- a/manifests/simplerule.pp
+++ b/manifests/simplerule.pp
@@ -47,6 +47,9 @@
 # @param sport
 #   The source port, ports or port range.
 #
+# @param saddr
+#   The source address, CIDR or set to match.
+#
 # @param counter
 #   Enable traffic counters for the matched traffic.
 
@@ -63,6 +66,7 @@ define nftables::simplerule (
   Optional[Variant[Stdlib::IP::Address::V6, Stdlib::IP::Address::V4, Pattern[/^@[-a-zA-Z0-9_]+$/]]] $daddr = undef,
   Enum['ip', 'ip6'] $set_type = 'ip6',
   Optional[Variant[Array[Stdlib::Port, 1], Stdlib::Port, Pattern[/\d+-\d+/]]] $sport = undef,
+  Optional[Variant[Stdlib::IP::Address::V6, Stdlib::IP::Address::V4, Pattern[/^@[-a-zA-Z0-9_]+$/]]] $saddr = undef,
   Boolean $counter = false,
 ) {
   if $dport and !$proto {
@@ -84,6 +88,7 @@ define nftables::simplerule (
           'proto'    => $proto,
           'daddr'    => $daddr,
           'set_type' => $set_type,
+          'saddr'    => $saddr,
           'sport'    => $sport,
         }
       ),
diff --git a/spec/defines/simplerule_spec.rb b/spec/defines/simplerule_spec.rb
index 42f8d2f..413e8a7 100644
--- a/spec/defines/simplerule_spec.rb
+++ b/spec/defines/simplerule_spec.rb
@@ -51,13 +51,14 @@ describe 'nftables::simplerule' do
             proto: 'udp',
             chain: 'default_out',
             daddr: '2001:1458::/32',
+            saddr: '2001:145c::/32',
           }
         end
 
         it { is_expected.to compile }
         it {
           is_expected.to contain_nftables__rule('default_out-my_big_rule').with(
-            content: 'udp sport {444} udp dport {333} ip6 daddr 2001:1458::/32 counter accept comment "this is my rule"',
+            content: 'udp sport {444} udp dport {333} ip6 saddr 2001:145c::/32 ip6 daddr 2001:1458::/32 counter accept comment "this is my rule"',
             order: '50',
           )
         }
@@ -177,6 +178,21 @@ describe 'nftables::simplerule' do
         }
       end
 
+      describe 'with an IPv6 address as saddr' do
+        let(:params) do
+          {
+            saddr: '2001:1458:0000:0000:0000:0000:0000:0003',
+          }
+        end
+
+        it { is_expected.to compile }
+        it {
+          is_expected.to contain_nftables__rule('default_in-my_default_rule_name').with(
+            content: 'ip6 saddr 2001:1458:0000:0000:0000:0000:0000:0003 accept',
+          )
+        }
+      end
+
       describe 'with an IPv6 set as daddr, default set_type' do
         let(:params) do
           {
@@ -208,6 +224,22 @@ describe 'nftables::simplerule' do
         }
       end
 
+      describe 'with a IPv6 set as saddr' do
+        let(:params) do
+          {
+            saddr: '@my6_set',
+            set_type: 'ip6',
+          }
+        end
+
+        it { is_expected.to compile }
+        it {
+          is_expected.to contain_nftables__rule('default_in-my_default_rule_name').with(
+            content: 'ip6 saddr @my6_set accept',
+          )
+        }
+      end
+
       describe 'with counter enabled' do
         let(:params) do
           {
diff --git a/templates/simplerule.epp b/templates/simplerule.epp
index 9d43453..46cf928 100644
--- a/templates/simplerule.epp
+++ b/templates/simplerule.epp
@@ -5,6 +5,7 @@
       Optional[Variant[Stdlib::IP::Address::V6, Stdlib::IP::Address::V4, Pattern[/^@[-a-zA-Z0-9_]+$/]]] $daddr,
       Enum['ip', 'ip6']       $set_type,
       Optional[Variant[Array[Stdlib::Port, 1], Stdlib::Port, String]] $sport,
+      Optional[Variant[Stdlib::IP::Address::V6, Stdlib::IP::Address::V4, Pattern[/^@[-a-zA-Z0-9_]+$/]]] $saddr,
       Boolean                 $counter,
 | -%>
 <%- if $proto {
@@ -34,6 +35,20 @@
 } else {
   $_dst_hosts = undef
 } -%>
+<%- if $saddr {
+  if $saddr =~ Stdlib::IP::Address::V6 {
+    $_src_hosts = "ip6 saddr ${saddr}"
+  } elsif $daddr =~ Stdlib::IP::Address::V4 {
+    $_src_hosts = "ip saddr ${saddr}"
+  } else {
+    $_src_hosts = $set_type ? {
+      'ip'  => "ip saddr ${saddr}",
+      'ip6' => "ip6 saddr ${saddr}",
+    }
+  }
+} else {
+  $_src_hosts = undef
+} -%>
 <%- if $proto and $dport {
   $_dst_port = "${_proto} dport {${Array($dport, true).join(', ')}}"
 } else {
@@ -54,4 +69,4 @@
 } else {
   $_counter = undef
 } -%>
-<%= regsubst(strip([$_ip_version_filter, $_src_port, $_dst_port, $_dst_hosts, $_counter, $action, $_comment].join(' ')), '\s+', ' ', 'G') -%>
+<%= regsubst(strip([$_ip_version_filter, $_src_port, $_dst_port, $_src_hosts, $_dst_hosts, $_counter, $action, $_comment].join(' ')), '\s+', ' ', 'G') -%>
