Merge pull request #7 from caseydavenport/casey-load-table

Support for loading all rules in one command

nftables.go

@@ -43,7 +43,8 @@ type Interface interface {
 	// list and no error.
 	List(ctx context.Context, objectType string) ([]string, error)
 
-	// ListRules returns a list of the rules in a chain, in order. Note that at the
+	// ListRules returns a list of the rules in a chain, in order. If no chain name is
+	// specified, then all rules within the table will be returned. Note that at the
 	// present time, the Rule objects will have their `Comment` and `Handle` fields
 	// filled in, but *not* the actual `Rule` field. So this can only be used to find
 	// the handles of rules if they have unique comments to recognize them by, or if
@@ -298,7 +299,13 @@ func (nft *realNFTables) List(ctx context.Context, objectType string) ([]string,
 
 // ListRules is part of Interface
 func (nft *realNFTables) ListRules(ctx context.Context, chain string) ([]*Rule, error) {
-	cmd := exec.CommandContext(ctx, nft.path, "--json", "list", "chain", string(nft.family), nft.table, chain)
+	// If no chain is given, return all rules from within the table.
+	var cmd *exec.Cmd
+	if chain == "" {
+		cmd = exec.CommandContext(ctx, nft.path, "--json", "list", "table", string(nft.family), nft.table)
+	} else {
+		cmd = exec.CommandContext(ctx, nft.path, "--json", "list", "chain", string(nft.family), nft.table, chain)
+	}
 	out, err := nft.exec.Run(cmd)
 	if err != nil {
 		return nil, fmt.Errorf("failed to run nft: %w", err)
@@ -311,8 +318,12 @@ func (nft *realNFTables) ListRules(ctx context.Context, chain string) ([]*Rule,
 
 	rules := make([]*Rule, 0, len(jsonRules))
 	for _, jsonRule := range jsonRules {
+		parentChain, ok := jsonVal[string](jsonRule, "chain")
+		if !ok {
+			return nil, fmt.Errorf("unexpected JSON output from nft (rule with no chain)")
+		}
 		rule := &Rule{
-			Chain: chain,
+			Chain: parentChain,
 		}
 
 		// handle is written as an integer in nft's output, but json.Unmarshal

nftables_test.go

@@ -209,21 +209,25 @@ func TestRun(t *testing.T) {
 func TestListRules(t *testing.T) {
 	for _, tc := range []struct {
 		name       string
+		chain      string
 		nftOutput  string
 		nftError   string
 		listOutput []*Rule
 	}{
 		{
 			name:     "no such chain",
+			chain:    "testchain",
 			nftError: "Error: No such file or directory\nlist chain ip testing testchain\n                      ^^^^^^^^^\n",
 		},
 		{
 			name:       "no rules",
+			chain:      "testchain",
 			nftOutput:  `{"nftables": [{"metainfo": {"version": "1.0.1", "release_name": "Fearless Fosdick #3", "json_schema_version": 1}}, {"chain": {"family": "ip", "table": "testing", "name": "testchain", "handle": 21}}]}`,
 			listOutput: []*Rule{},
 		},
 		{
 			name:      "normal output",
+			chain:     "testchain",
 			nftOutput: `{"nftables": [{"metainfo": {"version": "1.0.1", "release_name": "Fearless Fosdick #3", "json_schema_version": 1}}, {"chain": {"family": "ip", "table": "testing", "name": "testchain", "handle": 165}}, {"rule": {"family": "ip", "table": "testing", "chain": "testchain", "handle": 169, "expr": [{"match": {"op": "==", "left": {"ct": {"key": "state"}}, "right": {"set": ["established", "related"]}}}, {"accept": null}]}}, {"rule": {"family": "ip", "table": "testing", "chain": "testchain", "handle": 170, "comment": "This rule does something", "expr": [{"match": {"op": "in", "left": {"ct": {"key": "status"}}, "right": "dnat"}}, {"accept": null}]}}, {"rule": {"family": "ip", "table": "testing", "chain": "testchain", "handle": 171, "expr": [{"match": {"op": "==", "left": {"meta": {"key": "iifname"}}, "right": "lo"}}, {"accept": null}]}}]}`,
 			listOutput: []*Rule{
 				{
@@ -241,6 +245,20 @@ func TestListRules(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:      "all rules in table",
+			nftOutput: `{"nftables": [{"metainfo": {"version": "1.0.2", "release_name": "Lester Gooch", "json_schema_version": 1}}, {"table": {"family": "ip", "name": "testing", "handle": 3}}, {"chain": {"family": "ip", "table": "testing", "name": "chain1", "handle": 1}}, {"rule": {"family": "ip", "table": "testing", "chain": "chain1", "handle": 3, "expr": [{"match": {"op": "==", "left": {"payload": {"protocol": "ip", "field": "daddr"}}, "right": "8.8.8.8"}}, {"counter": {"packets": 0, "bytes": 0}}]}}, {"chain": {"family": "ip", "table": "testing", "name": "chain2", "handle": 2}}, {"rule": {"family": "ip", "table": "testing", "chain": "chain2", "handle": 4, "expr": [{"match": {"op": "==", "left": {"payload": {"protocol": "ip", "field": "daddr"}}, "right": "1.2.3.4"}}, {"counter": {"packets": 0, "bytes": 0}}]}}]}`,
+			listOutput: []*Rule{
+				{
+					Chain:  "chain1",
+					Handle: PtrTo(3),
+				},
+				{
+					Chain:  "chain2",
+					Handle: PtrTo(4),
+				},
+			},
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			nft, fexec, _ := newTestInterface(t, IPv4Family, "testing")
@@ -249,14 +267,24 @@ func TestListRules(t *testing.T) {
 			if tc.nftError != "" {
 				err = fmt.Errorf(tc.nftError)
 			}
-			fexec.expected = append(fexec.expected,
-				expectedCmd{
-					args:   []string{"/nft", "--json", "list", "chain", "ip", "testing", "testchain"},
-					stdout: strings.TrimSpace(dedent.Dedent(tc.nftOutput)),
-					err:    err,
-				},
-			)
-			result, err := nft.ListRules(context.Background(), "testchain")
+			if tc.chain == "" {
+				fexec.expected = append(fexec.expected,
+					expectedCmd{
+						args:   []string{"/nft", "--json", "list", "table", "ip", "testing"},
+						stdout: strings.TrimSpace(dedent.Dedent(tc.nftOutput)),
+						err:    err,
+					},
+				)
+			} else {
+				fexec.expected = append(fexec.expected,
+					expectedCmd{
+						args:   []string{"/nft", "--json", "list", "chain", "ip", "testing", "testchain"},
+						stdout: strings.TrimSpace(dedent.Dedent(tc.nftOutput)),
+						err:    err,
+					},
+				)
+			}
+			result, err := nft.ListRules(context.Background(), tc.chain)
 			if err != nil {
 				if tc.nftError == "" {
 					t.Errorf("unexpected error: %v", err)