Merge pull request #144 from telekom-mms/feature/add-set-allowed-hosts

Add commands for setting allowed hosts to TrafPol
telekom-mms · Jan 29, 2025 · 611c956 · 611c956
2 parents a677aee + f56fc37
commit 611c956
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 61 deletions.
diff --git a/internal/cmdtmpl/command.go b/internal/cmdtmpl/command.go
@@ -225,29 +225,21 @@ add element inet oc-daemon-filter allowdevs { {{.}} }
 			},
 			defaultTemplate: TrafPolDefaultTemplate,
 		}
-	case "TrafPolFlushAllowedHosts":
-		// Flush Allowed Hosts
-		cl = &CommandList{
-			Name: name,
-			Commands: []*Command{
-				{Line: "{{.Executables.Nft}} -f - flush set inet oc-daemon-filter allowhosts4"},
-				{Line: "{{.Executables.Nft}} -f - flush set inet oc-daemon-filter allowhosts6"},
-			},
-			defaultTemplate: TrafPolDefaultTemplate,
-		}
-	case "TrafPolAddAllowedHost":
-		// Add Allowed Host
+	case "TrafPolSetAllowedHosts":
+		// Set Allowed Hosts
 		cl = &CommandList{
 			Name: name,
 			Commands: []*Command{
 				{Line: "{{.Executables.Nft}} -f -",
-					Stdin: `
-				{{if .AllowedIP.Addr.Is4}}
-				add element inet oc-daemon-filter allowhosts4 { {{.AllowedIP}} }
-				{{else}}
-				add element inet oc-daemon-filter allowhosts6 { {{.AllowedIP}} }
-				{{end}}
-				`,
+					Stdin: `flush set inet oc-daemon-filter allowhosts4
+flush set inet oc-daemon-filter allowhosts6
+{{range .AllowedIPs -}}
+{{if .Addr.Is4 -}}
+add element inet oc-daemon-filter allowhosts4 { {{.}} }
+{{else -}}
+add element inet oc-daemon-filter allowhosts6 { {{.}} }
+{{end -}}
+{{end}}`,
 				},
 			},
 			defaultTemplate: TrafPolDefaultTemplate,

diff --git a/internal/cmdtmpl/command_test.go b/internal/cmdtmpl/command_test.go
@@ -38,8 +38,7 @@ func TestGetCommandList(t *testing.T) {
 		"TrafPolSetFilterRules",
 		"TrafPolUnsetFilterRules",
 		"TrafPolSetAllowedDevices",
-		"TrafPolFlushAllowedHosts",
-		"TrafPolAddAllowedHost",
+		"TrafPolSetAllowedHosts",
 		"TrafPolSetAllowedPorts",
 		"TrafPolCleanup",
 
@@ -86,8 +85,7 @@ func TestGetCmds(t *testing.T) {
 		"TrafPolSetFilterRules",
 		"TrafPolUnsetFilterRules",
 		// TrafPolSetAllowedDevices", // skip, requires devices
-		"TrafPolFlushAllowedHosts",
-		// "TrafPolAddAllowedHost", // skip, requires host
+		// "TrafPolSetAllowedHosts", // skip, requires hosts
 		//"TrafPolSetAllowedPorts", // skip, requires ports
 		"TrafPolCleanup",
 
@@ -110,7 +108,7 @@ func TestGetCmds(t *testing.T) {
 	for _, name := range []string{
 		// Traffic Policing
 		"TrafPolSetAllowedDevices",
-		"TrafPolAddAllowedHost",
+		"TrafPolSetAllowedHosts",
 		"TrafPolSetAllowedPorts",
 
 		// VPN Setup

diff --git a/internal/trafpol/filter.go b/internal/trafpol/filter.go
@@ -80,54 +80,29 @@ func setAllowedDevices(ctx context.Context, conf *daemoncfg.Config, devices []st
 
 // setAllowedIPs set the allowed hosts.
 func setAllowedIPs(ctx context.Context, conf *daemoncfg.Config, ips []netip.Prefix) {
-	// we perform all nft commands separately here and not as one atomic
-	// operation to avoid issues where the whole update fails because nft
-	// runs into "file exists" errors even though we remove duplicates from
-	// ips before calling this function and we flush the existing entries
-
-	// flush allowed hosts
-	cmds, err := cmdtmpl.GetCmds("TrafPolFlushAllowedHosts", conf)
+	// set allowed hosts
+	data := &struct {
+		daemoncfg.Config
+		AllowedIPs []netip.Prefix
+	}{
+		Config:     *conf,
+		AllowedIPs: ips,
+	}
+	cmds, err := cmdtmpl.GetCmds("TrafPolSetAllowedHosts", data)
 	if err != nil {
-		log.WithError(err).Error("TrafPol could not get flush allowed hosts commands")
+		log.WithError(err).Error("TrafPol could not get set allowed hosts commands")
 	}
 	for _, c := range cmds {
 		if stdout, stderr, err := c.Run(ctx); err != nil && !errors.Is(err, context.Canceled) {
 			log.WithFields(log.Fields{
+				"hosts":   ips,
 				"command": c.Cmd,
 				"args":    c.Args,
 				"stdin":   c.Stdin,
 				"stdout":  string(stdout),
 				"stderr":  string(stderr),
 				"error":   err,
-			}).Error("TrafPol could not run flush allowed hosts command")
-		}
-	}
-
-	// add allowed hosts
-	for _, ip := range ips {
-		data := &struct {
-			daemoncfg.Config
-			AllowedIP netip.Prefix
-		}{
-			Config:    *conf,
-			AllowedIP: ip,
-		}
-		cmds, err := cmdtmpl.GetCmds("TrafPolAddAllowedHost", data)
-		if err != nil {
-			log.WithError(err).Error("TrafPol could not get add allowed host commands")
-		}
-		for _, c := range cmds {
-			if stdout, stderr, err := c.Run(ctx); err != nil && !errors.Is(err, context.Canceled) {
-				log.WithFields(log.Fields{
-					"host":    ip,
-					"command": c.Cmd,
-					"args":    c.Args,
-					"stdin":   c.Stdin,
-					"stdout":  string(stdout),
-					"stderr":  string(stderr),
-					"error":   err,
-				}).Error("TrafPol could not run add allowed host command")
-			}
+			}).Error("TrafPol could not run set allowed hosts command")
 		}
 	}
 }