commit 77abc10b1080ee450cc6080bdcee28f0ddb6b622
Author: Nacho Barrientos <nacho.barrientos@cern.ch>
Date:   Wed Dec 9 12:27:48 2020 +0100

    Implement nftables::simplerule::sport

diff --git a/manifests/simplerule.pp b/manifests/simplerule.pp
index 95c4d4c..2784fc3 100644
--- a/manifests/simplerule.pp
+++ b/manifests/simplerule.pp
@@ -1,6 +1,6 @@
 # @summary Provides a simplified interface to nftables::rule for basic use cases
 #
-# @example allow incoming traffic on port 543 TCP to a given IP range and count packets
+# @example allow incoming traffic from port 541 on port 543 TCP to a given IP range and count packets
 #   nftables::simplerule{'my_service_in':
 #     action  => 'accept',
 #     comment => 'allow traffic to port 543',
@@ -8,6 +8,7 @@
 #     proto   => 'tcp',
 #     dport   => 543,
 #     daddr   => '2001:1458::/32',
+#     sport   => 541,
 #   }
 
 define nftables::simplerule (
@@ -22,12 +23,17 @@ define nftables::simplerule (
   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',
+  Optional[Variant[Array[Stdlib::Port, 1], Stdlib::Port, Pattern[/\d+-\d+/]]] $sport = undef,
   Boolean $counter = false,
 ) {
   if $dport and !$proto {
     fail('Specifying a transport protocol via $proto is mandatory when passing a $dport')
   }
 
+  if $sport and !$proto {
+    fail('Specifying a transport protocol via $proto is mandatory when passing a $sport')
+  }
+
   if $ensure == 'present' {
     nftables::rule { "${chain}-${rulename}":
       content => epp('nftables/simplerule.epp',
@@ -39,6 +45,7 @@ define nftables::simplerule (
           'proto'    => $proto,
           'daddr'    => $daddr,
           'set_type' => $set_type,
+          'sport'    => $sport,
         }
       ),
       order   => $order,
diff --git a/spec/defines/simplerule_spec.rb b/spec/defines/simplerule_spec.rb
index b848ed9..4ef9600 100644
--- a/spec/defines/simplerule_spec.rb
+++ b/spec/defines/simplerule_spec.rb
@@ -19,7 +19,7 @@ describe 'nftables::simplerule' do
         }
       end
 
-      describe 'port without protocol' do
+      describe 'dport without protocol' do
         let(:params) do
           {
             dport: 333,
@@ -29,6 +29,16 @@ describe 'nftables::simplerule' do
         it { is_expected.not_to compile }
       end
 
+      describe 'sport without protocol' do
+        let(:params) do
+          {
+            sport: 333,
+          }
+        end
+
+        it { is_expected.not_to compile }
+      end
+
       describe 'all parameters provided' do
         let(:title) { 'my_big_rule' }
         let(:params) do
@@ -37,6 +47,7 @@ describe 'nftables::simplerule' do
             comment: 'this is my rule',
             counter: true,
             dport: 333,
+            sport: 444,
             proto: 'udp',
             chain: 'default_out',
             daddr: '2001:1458::/32',
@@ -46,7 +57,7 @@ describe 'nftables::simplerule' do
         it { is_expected.to compile }
         it {
           is_expected.to contain_nftables__rule('default_out-my_big_rule').with(
-            content: 'udp dport {333} ip6 daddr 2001:1458::/32 counter accept comment "this is my rule"',
+            content: 'udp sport {444} udp dport {333} ip6 daddr 2001:1458::/32 counter accept comment "this is my rule"',
             order: '50',
           )
         }
@@ -56,6 +67,7 @@ describe 'nftables::simplerule' do
         let(:params) do
           {
             dport: '333-334',
+            sport: '1-2',
             proto: 'tcp',
           }
         end
@@ -63,7 +75,7 @@ describe 'nftables::simplerule' do
         it { is_expected.to compile }
         it {
           is_expected.to contain_nftables__rule('default_in-my_default_rule_name').with(
-            content: 'tcp dport {333-334} accept',
+            content: 'tcp sport {1-2} tcp dport {333-334} accept',
           )
         }
       end
@@ -72,6 +84,23 @@ describe 'nftables::simplerule' do
         let(:params) do
           {
             dport: [333, 335],
+            sport: [433, 435],
+            proto: 'tcp',
+          }
+        end
+
+        it { is_expected.to compile }
+        it {
+          is_expected.to contain_nftables__rule('default_in-my_default_rule_name').with(
+            content: 'tcp sport {433, 435} tcp dport {333, 335} accept',
+          )
+        }
+      end
+
+      describe 'only sport TCP traffic' do
+        let(:params) do
+          {
+            sport: 555,
             proto: 'tcp',
           }
         end
@@ -79,7 +108,7 @@ describe 'nftables::simplerule' do
         it { is_expected.to compile }
         it {
           is_expected.to contain_nftables__rule('default_in-my_default_rule_name').with(
-            content: 'tcp dport {333, 335} accept',
+            content: 'tcp sport {555} accept',
           )
         }
       end
diff --git a/templates/simplerule.epp b/templates/simplerule.epp
index 0358492..9d43453 100644
--- a/templates/simplerule.epp
+++ b/templates/simplerule.epp
@@ -4,6 +4,7 @@
       Optional[String]        $proto,
       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,
       Boolean                 $counter,
 | -%>
 <%- if $proto {
@@ -43,9 +44,14 @@
 } else {
   $_comment = undef
 } -%>
+<%- if $proto and $sport {
+  $_src_port = "${_proto} sport {${Array($sport, true).join(', ')}}"
+} else {
+  $_src_port = undef
+} -%>
 <%- if $counter {
   $_counter = "counter"
 } else {
   $_counter = undef
 } -%>
-<%= regsubst(strip([$_ip_version_filter, $_dst_port, $_dst_hosts, $_counter, $action, $_comment].join(' ')), '\s+', ' ', 'G') -%>
+<%= regsubst(strip([$_ip_version_filter, $_src_port, $_dst_port, $_dst_hosts, $_counter, $action, $_comment].join(' ')), '\s+', ' ', 'G') -%>
