plugin/rewrite: Add EDNS0 Unset Action (#7380)

* plugin/rewrite: EDNS0 unset action

Signed-off-by: Dennis Simmons <215134900+dennis-ix@users.noreply.github.com>

* plugin/rewrite: EDNS0 unset tests

Signed-off-by: Dennis Simmons <215134900+dennis-ix@users.noreply.github.com>

* plugin/rewrite: EDNS0 unset documentation

Signed-off-by: Dennis Simmons <215134900+dennis-ix@users.noreply.github.com>

* gofmt whitespace fixes

Signed-off-by: Dennis Simmons <215134900+dennis-ix@users.noreply.github.com>

* plugin/rewrite: improve edns0 test coverage

Signed-off-by: Dennis Simmons <215134900+dennis-ix@users.noreply.github.com>

---------

Signed-off-by: Dennis Simmons <215134900+dennis-ix@users.noreply.github.com>

plugin/rewrite/README.md

@@ -394,11 +394,12 @@ The values of FROM and TO can be any of the following, text value or numeric:
 
 ## EDNS0 Options
 
-Using the FIELD edns0, you can set, append, or replace specific EDNS0 options in the request.
+Using the FIELD edns0, you can set, append, replace, or unset specific EDNS0 options in the request.
 
 * `replace` will modify any "matching" option with the specified option. The criteria for "matching" varies based on EDNS0 type.
 * `append` will add the option only if no matching option exists
 * `set` will modify a matching option or add one if none is found
+* `unset` will remove the matching option if one exists
 
 Currently supported are `EDNS0_LOCAL`, `EDNS0_NSID` and `EDNS0_SUBNET`.
 
@@ -444,10 +445,17 @@ some-plugin
 rewrite edns0 local set 0xffee {some-plugin/some-label}
 ~~~
 
+A local option may be removed by unsetting its code. Example:
+
+~~~
+rewrite edns0 local unset 0xffee
+~~~
+
 ### EDNS0_NSID
 
 This has no fields; it will add an NSID option with an empty string for the NSID. If the option already exists
 and the action is `replace` or `set`, then the NSID in the option will be set to the empty string.  
+The option can be removed with the `unset` action.
 
 ### EDNS0_SUBNET
 
@@ -463,6 +471,12 @@ rewrite edns0 subnet set 24 56
 * If the query's source IP address is an IPv4 address, the first 24 bits in the IP will be the network subnet.
 * If the query's source IP address is an IPv6 address, the first 56 bits in the IP will be the network subnet.
 
+This option can be removed by using `unset`:
+
+~~~
+rewrite edns0 subnet unset
+~~~
+
 ### EDNS0 Revert
 
 Using the `revert` flag, you can revert the changes made by this rewrite call, so the response will not contain this option.

plugin/rewrite/edns0.go

@@ -80,10 +80,25 @@ func setupEdns0Opt(r *dns.Msg) *dns.OPT {
 	return o
 }
 
+func unsetEdns0Option(opt *dns.OPT, code uint16) {
+	var newOpts []dns.EDNS0
+	for _, o := range opt.Option {
+		if o.Option() != code {
+			newOpts = append(newOpts, o)
+		}
+	}
+	opt.Option = newOpts
+}
+
 // Rewrite will alter the request EDNS0 NSID option
 func (rule *edns0NsidRule) Rewrite(ctx context.Context, state request.Request) (ResponseRules, Result) {
 	o := setupEdns0Opt(state.Req)
 
+	if rule.action == Unset {
+		unsetEdns0Option(o, dns.EDNS0NSID)
+		return nil, RewriteDone
+	}
+
 	var resp ResponseRules
 
 	for _, s := range o.Option {
@@ -118,6 +133,11 @@ func (rule *edns0NsidRule) Mode() string { return rule.mode }
 func (rule *edns0LocalRule) Rewrite(ctx context.Context, state request.Request) (ResponseRules, Result) {
 	o := setupEdns0Opt(state.Req)
 
+	if rule.action == Unset {
+		unsetEdns0Option(o, rule.code)
+		return nil, RewriteDone
+	}
+
 	var resp ResponseRules
 
 	for _, s := range o.Option {
@@ -162,6 +182,8 @@ func newEdns0Rule(mode string, args ...string) (Rule, error) {
 	case Append:
 	case Replace:
 	case Set:
+	case Unset:
+		return newEdns0UnsetRule(mode, action, ruleType, args...)
 	default:
 		return nil, fmt.Errorf("invalid action: %q", action)
 	}
@@ -198,6 +220,28 @@ func newEdns0Rule(mode string, args ...string) (Rule, error) {
 	}
 }
 
+func newEdns0UnsetRule(mode string, action string, ruleType string, args ...string) (Rule, error) {
+	switch ruleType {
+	case "local":
+		if len(args) != 3 {
+			return nil, fmt.Errorf("local unset action requires exactly two arguments")
+		}
+		return newEdns0LocalRule(mode, action, args[2], "", false)
+	case "nsid":
+		if len(args) != 2 {
+			return nil, fmt.Errorf("nsid unset action requires exactly one argument")
+		}
+		return &edns0NsidRule{mode, action, false}, nil
+	case "subnet":
+		if len(args) != 2 {
+			return nil, fmt.Errorf("subnet unset action requires exactly one argument")
+		}
+		return &edns0SubnetRule{mode, 0, 0, action, false}, nil
+	default:
+		return nil, fmt.Errorf("invalid rule type %q", ruleType)
+	}
+}
+
 func newEdns0LocalRule(mode, action, code, data string, revert bool) (*edns0LocalRule, error) {
 	c, err := strconv.ParseUint(code, 0, 16)
 	if err != nil {
@@ -393,6 +437,11 @@ func (rule *edns0SubnetRule) fillEcsData(state request.Request, ecs *dns.EDNS0_S
 func (rule *edns0SubnetRule) Rewrite(ctx context.Context, state request.Request) (ResponseRules, Result) {
 	o := setupEdns0Opt(state.Req)
 
+	if rule.action == Unset {
+		unsetEdns0Option(o, dns.EDNS0SUBNET)
+		return nil, RewriteDone
+	}
+
 	var resp ResponseRules
 
 	for _, s := range o.Option {
@@ -433,6 +482,7 @@ const (
 	Replace = "replace"
 	Set     = "set"
 	Append  = "append"
+	Unset   = "unset"
 )
 
 // Supported local EDNS0 variables

plugin/rewrite/rewrite_test.go

@@ -65,9 +65,12 @@ func TestNewRule(t *testing.T) {
 		{[]string{"class", "XY", "WV"}, true, nil},
 		{[]string{"class", "IN", "WV"}, true, nil},
 		{[]string{"edns0"}, true, nil},
+		{[]string{"edns0", "unknown-rule-type", "set"}, true, nil},
+		{[]string{"edns0", "unknown-rule-type", "unset"}, true, nil},
 		{[]string{"edns0", "local"}, true, nil},
 		{[]string{"edns0", "local", "set"}, true, nil},
 		{[]string{"edns0", "local", "set", "0xffee"}, true, nil},
+		{[]string{"edns0", "local", "set", "invalid-uint", "abcdefg"}, true, nil},
 		{[]string{"edns0", "local", "set", "65518", "abcdefg"}, false, reflect.TypeOf(&edns0LocalRule{})},
 		{[]string{"edns0", "local", "set", "0xffee", "abcdefg"}, false, reflect.TypeOf(&edns0LocalRule{})},
 		{[]string{"edns0", "local", "set", "0xffee", "abcdefg", "revert"}, false, reflect.TypeOf(&edns0LocalRule{})},
@@ -75,6 +78,9 @@ func TestNewRule(t *testing.T) {
 		{[]string{"edns0", "local", "append", "0xffee", "abcdefg", "revert"}, false, reflect.TypeOf(&edns0LocalRule{})},
 		{[]string{"edns0", "local", "replace", "0xffee", "abcdefg"}, false, reflect.TypeOf(&edns0LocalRule{})},
 		{[]string{"edns0", "local", "replace", "0xffee", "abcdefg", "revert"}, false, reflect.TypeOf(&edns0LocalRule{})},
+		{[]string{"edns0", "local", "unset", "0xffee"}, false, reflect.TypeOf(&edns0LocalRule{})},
+		{[]string{"edns0", "local", "unset", "0xffee", "abcdefg"}, true, reflect.TypeOf(&edns0LocalRule{})},
+		{[]string{"edns0", "local", "unset", "0xffee", "revert"}, true, reflect.TypeOf(&edns0LocalRule{})},
 		{[]string{"edns0", "local", "foo", "0xffee", "abcdefg"}, true, nil},
 		{[]string{"edns0", "local", "set", "0xffee", "0xabcdefg"}, true, nil},
 		{[]string{"edns0", "nsid", "set", "junk"}, true, nil},
@@ -84,7 +90,10 @@ func TestNewRule(t *testing.T) {
 		{[]string{"edns0", "nsid", "append", "revert"}, false, reflect.TypeOf(&edns0NsidRule{})},
 		{[]string{"edns0", "nsid", "replace"}, false, reflect.TypeOf(&edns0NsidRule{})},
 		{[]string{"edns0", "nsid", "replace", "revert"}, false, reflect.TypeOf(&edns0NsidRule{})},
+		{[]string{"edns0", "nsid", "unset"}, false, reflect.TypeOf(&edns0NsidRule{})},
+		{[]string{"edns0", "nsid", "unset", "revert"}, true, reflect.TypeOf(&edns0NsidRule{})},
 		{[]string{"edns0", "nsid", "foo"}, true, nil},
+		{[]string{"edns0", "local", "set", "invalid-uint", "{qname}"}, true, nil},
 		{[]string{"edns0", "local", "set", "0xffee", "{dummy}"}, true, nil},
 		{[]string{"edns0", "local", "set", "0xffee", "{qname}"}, false, reflect.TypeOf(&edns0VariableRule{})},
 		{[]string{"edns0", "local", "set", "0xffee", "{qtype}"}, false, reflect.TypeOf(&edns0VariableRule{})},
@@ -123,6 +132,9 @@ func TestNewRule(t *testing.T) {
 		{[]string{"edns0", "subnet", "append", "24", "56", "revert"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"edns0", "subnet", "replace", "24", "56"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"edns0", "subnet", "replace", "24", "56", "revert"}, false, reflect.TypeOf(&edns0SubnetRule{})},
+		{[]string{"edns0", "subnet", "unset"}, false, reflect.TypeOf(&edns0SubnetRule{})},
+		{[]string{"edns0", "subnet", "unset", "24", "56"}, true, reflect.TypeOf(&edns0SubnetRule{})},
+		{[]string{"edns0", "subnet", "unset", "revert"}, true, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"unknown-action", "name", "a.com", "b.com"}, true, nil},
 		{[]string{"stop", "name", "a.com", "b.com"}, false, reflect.TypeOf(&exactNameRule{})},
 		{[]string{"continue", "name", "a.com", "b.com"}, false, reflect.TypeOf(&exactNameRule{})},
@@ -157,9 +169,11 @@ func TestNewRule(t *testing.T) {
 		{[]string{"stop", "edns0", "subnet", "set", "24", "56"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"stop", "edns0", "subnet", "append", "24", "56"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"stop", "edns0", "subnet", "replace", "24", "56"}, false, reflect.TypeOf(&edns0SubnetRule{})},
+		{[]string{"stop", "edns0", "subnet", "unset"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"continue", "edns0", "subnet", "set", "24", "56"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"continue", "edns0", "subnet", "append", "24", "56"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 		{[]string{"continue", "edns0", "subnet", "replace", "24", "56"}, false, reflect.TypeOf(&edns0SubnetRule{})},
+		{[]string{"continue", "edns0", "subnet", "unset"}, false, reflect.TypeOf(&edns0SubnetRule{})},
 	}
 
 	for i, tc := range tests {
@@ -1046,3 +1060,101 @@ func TestRewriteEDNS0Revert(t *testing.T) {
 		}
 	}
 }
+
+func TestRewriteEDNS0Unset(t *testing.T) {
+	rw := Rewrite{
+		Next:         plugin.HandlerFunc(msgPrinter),
+		RevertPolicy: NewRevertPolicy(false, false),
+	}
+
+	tests := []struct {
+		fromOpts []dns.EDNS0
+		args     []string
+		toOpts   []dns.EDNS0
+	}{
+		{
+			[]dns.EDNS0{},
+			[]string{"local", "unset", "0xffee"},
+			[]dns.EDNS0{},
+		},
+		{
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte{0xab, 0xcd, 0xef}}},
+			[]string{"local", "unset", "0xffee"},
+			[]dns.EDNS0{},
+		},
+		{
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte{0xab, 0xcd, 0xef}}},
+			[]string{"local", "unset", "0xffed"},
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte{0xab, 0xcd, 0xef}}},
+		},
+		{
+			[]dns.EDNS0{&dns.EDNS0_NSID{Code: dns.EDNS0NSID, Nsid: ""}},
+			[]string{"nsid", "unset"},
+			[]dns.EDNS0{},
+		},
+		{
+			[]dns.EDNS0{},
+			[]string{"nsid", "unset"},
+			[]dns.EDNS0{},
+		},
+		{
+			[]dns.EDNS0{&dns.EDNS0_SUBNET{Code: 0x8,
+				Family:        0x1,
+				SourceNetmask: 0x0,
+				SourceScope:   0x0,
+				Address:       []byte{0x00, 0x00, 0x00, 0x00},
+			}},
+			[]string{"subnet", "unset"},
+			[]dns.EDNS0{},
+		},
+		{
+			[]dns.EDNS0{},
+			[]string{"subnet", "unset"},
+			[]dns.EDNS0{},
+		},
+		{
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte("foobar")}, &dns.EDNS0_NSID{Code: dns.EDNS0NSID, Nsid: ""}},
+			[]string{"nsid", "unset"},
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte("foobar")}},
+		},
+		{
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte("foobar")}, &dns.EDNS0_NSID{Code: dns.EDNS0NSID, Nsid: ""}},
+			[]string{"local", "unset", "0xffee"},
+			[]dns.EDNS0{&dns.EDNS0_NSID{Code: dns.EDNS0NSID, Nsid: ""}},
+		},
+		{
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte("foobar")}, &dns.EDNS0_NSID{Code: dns.EDNS0NSID, Nsid: ""}},
+			[]string{"subnet", "unset"},
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte("foobar")}, &dns.EDNS0_NSID{Code: dns.EDNS0NSID, Nsid: ""}},
+		},
+	}
+
+	ctx := context.TODO()
+	for i, tc := range tests {
+		m := new(dns.Msg)
+		m.SetQuestion("example.com.", dns.TypeA)
+		m.Question[0].Qclass = dns.ClassINET
+		o := m.IsEdns0()
+		if tc.fromOpts != nil {
+			if o == nil {
+				m.SetEdns0(4096, true)
+				o = m.IsEdns0()
+			}
+			o.Option = append(o.Option, tc.fromOpts...)
+		}
+
+		r, err := newEdns0Rule("stop", tc.args...)
+		if err != nil {
+			t.Errorf("Error creating test rule: %s", err)
+			continue
+		}
+		rw.Rules = []Rule{r}
+
+		rec := dnstest.NewRecorder(&test.ResponseWriter{})
+		rw.ServeDNS(ctx, rec, m)
+
+		if !optsEqual(o.Option, tc.toOpts) {
+			t.Errorf("Test %d: Expected %v but got %v", i, tc.toOpts, o)
+		}
+	}
+}