Update config loading to support multiple YAML files

buty4649 · Mar 30, 2024 · 74d013a · 74d013a
1 parent de33ae1
commit 74d013a
Show file tree

Hide file tree

Showing 9 changed files with 141 additions and 64 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -35,7 +35,7 @@ import (
 )
 
 type Flags struct {
-	ConfigPath   string
+	ConfigDir    string
 	IpCmdPath    string
 	Debug, Quiet bool
 }
@@ -70,7 +70,7 @@ var rootCmd = &cobra.Command{
 		}))
 		slog.SetDefault(logger)
 
-		cfg, err = config.LoadConfig(flags.ConfigPath)
+		cfg, err = config.LoadYamlFiles(flags.ConfigDir)
 		if err != nil {
 			return err
 		}
@@ -95,7 +95,7 @@ func Execute() {
 }
 
 func init() {
-	rootCmd.PersistentFlags().StringVarP(&flags.ConfigPath, "config", "c", "./netnsplan.yaml", "config file")
+	rootCmd.PersistentFlags().StringVarP(&flags.ConfigDir, "config-dir", "d", "/etc/netnsplan", "config file directory")
 	rootCmd.PersistentFlags().StringVar(&flags.IpCmdPath, "cmd", "/bin/ip", "ip command path")
 
 	rootCmd.PersistentFlags().BoolVar(&flags.Debug, "debug", false, "debug mode")

diff --git a/config/config.go b/config/config.go
@@ -23,7 +23,9 @@ package config
 
 import (
 	"os"
+	"path/filepath"
 
+	"github.com/TwiN/deepmerge"
 	yaml "gopkg.in/yaml.v3"
 )
 
@@ -61,14 +63,45 @@ type Route struct {
 	Via string `yaml:"via"`
 }
 
-func LoadConfig(path string) (*Config, error) {
-	data, err := os.ReadFile(path)
+func mergeMaps(dst, src map[string]interface{}) {
+	for key, valueSrc := range src {
+		if valueDst, ok := dst[key]; ok {
+			if mapValueDst, ok := valueDst.(map[string]interface{}); ok {
+				if mapValueSrc, ok := valueSrc.(map[string]interface{}); ok {
+					mergeMaps(mapValueDst, mapValueSrc)
+					continue
+				}
+			}
+		}
+		dst[key] = valueSrc
+	}
+}
+
+func LoadYamlFiles(dirPath string) (*Config, error) {
+	files, err := filepath.Glob(filepath.Join(dirPath, "*.yaml"))
 	if err != nil {
 		return nil, err
 	}
 
+	var mergedYaml []byte
+	for _, file := range files {
+		bytes, err := os.ReadFile(file)
+		if err != nil {
+			return nil, err
+		}
+
+		if mergedYaml == nil {
+			mergedYaml = bytes
+		} else {
+			mergedYaml, err = deepmerge.YAML(mergedYaml, bytes)
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+
 	var config Config
-	err = yaml.Unmarshal(data, &config)
+	err = yaml.Unmarshal(mergedYaml, &config)
 	if err != nil {
 		return nil, err
 	}

diff --git a/config/config_test.go b/config/config_test.go
@@ -23,89 +23,81 @@ package config
 
 import (
 	"os"
+	"path/filepath"
 	"reflect"
 	"testing"
 )
 
-const testYAML = `
-netns:
-  netns1:
-    ethernets:
-      eth0:
-        addresses:
-          - "192.168.1.1/24"
-        routes:
-          - to: "0.0.0.0/0"
-            via: "192.168.1.254"
-    dummy-devices:
-      dummy0:
-        addresses:
-          - "10.0.0.1/8"
-    veth-devices:
-      veth0:
-        addresses:
-          - "10.1.0.1/24"
-        peer:
-          name: "veth0-peer"
-          netns: "netns2"
-          addresses:
-            - "10.1.0.2/24"
-    post-script: |
-        echo 'Hello'
-        echo 'World!'
-`
-
-func TestLoadConfig(t *testing.T) {
-	tmpfile, err := os.CreateTemp("", "test*.yaml")
-	if err != nil {
-		t.Fatalf("Failed to create temp file: %v", err)
-	}
-	defer os.Remove(tmpfile.Name())
-
-	if _, err := tmpfile.Write([]byte(testYAML)); err != nil {
-		t.Fatalf("Failed to write to temp file: %v", err)
-	}
-	if err := tmpfile.Close(); err != nil {
-		t.Fatalf("Failed to close temp file: %v", err)
-	}
-
-	config, err := LoadConfig(tmpfile.Name())
-	if err != nil {
-		t.Fatalf("LoadConfig returned an error: %v", err)
-	}
+func TestLoadYamlFiles(t *testing.T) {
+	wd, _ := os.Getwd()
+	testdataDir := filepath.Join(wd, "..", "testdata", "config")
 
 	expected := &Config{
 		Netns: map[string]Netns{
-			"netns1": {
+			"sample1": {
 				Ethernets: map[string]Ethernet{
 					"eth0": {
-						Addresses: []string{"192.168.1.1/24"},
-						Routes: []Route{
-							{To: "0.0.0.0/0", Via: "192.168.1.254"},
+						Addresses: []string{
+							"192.168.0.1/24",
+							"2001:db8:beaf:cafe::1/112",
+						},
+						Routes: []Route{{
+							To:  "default",
+							Via: "192.168.0.254",
+						}},
+					},
+					"eth1": {
+						Addresses: []string{
+							"192.168.1.1/24",
+							"10.0.0.1/24",
 						},
 					},
 				},
 				DummyDevices: map[string]Ethernet{
 					"dummy0": {
-						Addresses: []string{"10.0.0.1/8"},
+						Addresses: []string{"192.168.10.1/24"},
+						Routes: []Route{{
+							To:  "192.168.11.0/24",
+							Via: "192.168.10.254",
+						}},
 					},
 				},
 				VethDevices: map[string]VethDevice{
 					"veth0": {
-						Addresses: []string{"10.1.0.1/24"},
+						Addresses: []string{"192.168.20.1/24"},
+						Routes: []Route{{
+							To:  "192.168.21.0/24",
+							Via: "192.168.20.254",
+						}},
 						Peer: Peer{
 							Name:      "veth0-peer",
-							Netns:     "netns2",
-							Addresses: []string{"10.1.0.2/24"},
+							Netns:     "sample2",
+							Addresses: []string{"192.168.20.2/24"},
+							Routes: []Route{{
+								To:  "192.168.21.0/24",
+								Via: "192.168.20.2",
+							}},
 						},
 					},
 				},
-				PostScript: "echo 'Hello'\necho 'World!'\n",
+				PostScript: "echo 'Hello, World!'\n",
+			},
+			"sample2": {
+				Ethernets: map[string]Ethernet{
+					"eth2": {
+						Addresses: []string{"172.16.0.1/24"},
+					},
+				},
 			},
 		},
 	}
 
-	if !reflect.DeepEqual(config, expected) {
-		t.Errorf("Config does not match expected\nGot: %#v\nWant: %#v", config, expected)
+	result, err := LoadYamlFiles(testdataDir)
+	if err != nil {
+		t.Fatalf("LoadYamlFiles returned an error: %v", err)
+	}
+
+	if !reflect.DeepEqual(result, expected) {
+		t.Errorf("Expected %v, got %v", expected, result)
 	}
 }
diff --git a/go.mod b/go.mod
@@ -3,6 +3,7 @@ module netnsplan
 go 1.22.1
 
 require (
+	github.com/TwiN/deepmerge v0.2.1
 	github.com/spf13/cobra v1.8.0
 	gitlab.com/greyxor/slogor v1.2.6
 	gopkg.in/yaml.v3 v3.0.1

diff --git a/go.sum b/go.sum
@@ -1,3 +1,5 @@
+github.com/TwiN/deepmerge v0.2.1 h1:GowJr9O4THTVW4awX63x1BVg1hgr4q+35XKKCYbwsSs=
+github.com/TwiN/deepmerge v0.2.1/go.mod h1:LVBmCEBQvibYSF8Gyl/NqhHXH7yIiT7Ozqf9dHxGPW0=
 github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=

diff --git a/sample/netnsplan.yaml b/sample/netnsplan.yaml
@@ -9,7 +9,7 @@ netns:
             via: 192.168.20.254
     post-script: |
       sysctl --system
-        iptables-restore /etc/iptables/rules.v4
+      iptables-restore /etc/iptables/rules.v4
   sample2:
     ethernets:
       eth2:

diff --git a/testdata/config/sample1.yaml b/testdata/config/sample1.yaml
@@ -0,0 +1,37 @@
+netns:
+  sample1:
+    ethernets:
+      eth0:
+        addresses:
+          - 192.168.0.1/24
+          - 2001:db8:beaf:cafe::1/112
+        routes:
+          - to: default
+            via: 192.168.0.254
+      eth1:
+        addresses:
+          - 192.168.1.1/24
+    dummy-devices:
+      dummy0:
+        addresses:
+          - 192.168.10.1/24
+        routes:
+          - to: 192.168.11.0/24
+            via: 192.168.10.254
+    veth-devices:
+      veth0:
+        addresses:
+          - 192.168.20.1/24
+        routes:
+          - to: 192.168.21.0/24
+            via: 192.168.20.254
+        peer:
+          name: veth0-peer
+          netns: sample2
+          addresses:
+            - 192.168.20.2/24
+          routes:
+            - to: 192.168.21.0/24
+              via: 192.168.20.2
+    post-script: |
+      echo 'Hello, World!'
diff --git a/testdata/config/sample2.yaml b/testdata/config/sample2.yaml
@@ -0,0 +1,6 @@
+netns:
+  sample2:
+    ethernets:
+      eth2:
+        addresses:
+          - 172.16.0.1/24
diff --git a/testdata/config/zz_sampl3.yaml b/testdata/config/zz_sampl3.yaml
@@ -0,0 +1,6 @@
+netns:
+  sample1:
+    ethernets:
+      eth1:
+        addresses:
+          - 10.0.0.1/24