blackhole: block with ip

dispatcher/plugin/executable/blackhole/blackhole.go

@@ -19,28 +19,34 @@ package blackhole
 
 import (
 	"context"
+	"fmt"
 	"github.com/IrineSistiana/mosdns/dispatcher/handler"
 	"github.com/miekg/dns"
+	"net"
 )
 
 const PluginType = "blackhole"
 
 func init() {
 	handler.RegInitFunc(PluginType, Init)
 
-	handler.MustRegPlugin(handler.WrapExecutablePlugin("_drop_response", PluginType, &blackhole{rCode: -1}))
-	handler.MustRegPlugin(handler.WrapExecutablePlugin("_block_with_servfail", PluginType, &blackhole{rCode: dns.RcodeServerFailure}))
-	handler.MustRegPlugin(handler.WrapExecutablePlugin("_block_with_nxdomain", PluginType, &blackhole{rCode: dns.RcodeNameError}))
+	handler.MustRegPlugin(handler.WrapExecutablePlugin("_drop_response", PluginType, &blackhole{args: &Args{RCode: -1}}))
+	handler.MustRegPlugin(handler.WrapExecutablePlugin("_block_with_servfail", PluginType, &blackhole{args: &Args{RCode: dns.RcodeServerFailure}}))
+	handler.MustRegPlugin(handler.WrapExecutablePlugin("_block_with_nxdomain", PluginType, &blackhole{args: &Args{RCode: dns.RcodeNameError}}))
 }
 
 var _ handler.Executable = (*blackhole)(nil)
 
 type blackhole struct {
-	rCode int
+	args *Args
+	ipv4 net.IP
+	ipv6 net.IP
 }
 
 type Args struct {
-	RCode int `yaml:"rcode"`
+	IPv4  string `yaml:"ipv4"` // block by responding specific IP
+	IPv6  string `yaml:"ipv6"`
+	RCode int    `yaml:"rcode"` // block by responding specific RCode
 }
 
 func Init(tag string, argsMap map[string]interface{}) (p handler.Plugin, err error) {
@@ -50,25 +56,75 @@ func Init(tag string, argsMap map[string]interface{}) (p handler.Plugin, err err
 		return nil, handler.NewErrFromTemplate(handler.ETInvalidArgs, err)
 	}
 
-	b := new(blackhole)
-	b.rCode = args.RCode
-
+	b, err := newBlackhole(tag, args)
+	if err != nil {
+		return nil, err
+	}
 	return handler.WrapExecutablePlugin(tag, PluginType, b), nil
 }
 
-// Do drops or replaces qCtx.R with a simple denial response.
-// It never returns an err.
-func (b *blackhole) Exec(_ context.Context, qCtx *handler.Context) (err error) {
-	if qCtx == nil || qCtx.Q == nil {
-		return nil
+func newBlackhole(tag string, args *Args) (*blackhole, error) {
+	b := &blackhole{args: args}
+	if len(args.IPv4) != 0 {
+		ip := net.ParseIP(args.IPv4)
+		if ip == nil {
+			return nil, fmt.Errorf("%s is an invalid ipv4 addr", args.IPv4)
+		}
+		b.ipv4 = ip
 	}
+	if len(args.IPv6) != 0 {
+		ip := net.ParseIP(args.IPv6)
+		if ip == nil {
+			return nil, fmt.Errorf("%s is an invalid ipv6 addr", args.IPv6)
+		}
+		b.ipv6 = ip
+	}
+	return b, nil
+}
 
+// Exec
+// sets qCtx.R with IP response if query type is A/AAAA and Args.IPv4 / Args.IPv6 is not empty.
+// sets qCtx.R with empty response with rcode = Args.RCode.
+// drops qCtx.R if Args.RCode < 0
+// It never returns an err.
+func (b *blackhole) Exec(_ context.Context, qCtx *handler.Context) (err error) {
 	switch {
-	case b.rCode >= 0:
+	case b.ipv4 != nil && len(qCtx.Q.Question) == 1 && qCtx.Q.Question[0].Qtype == dns.TypeA:
 		r := new(dns.Msg)
 		r.SetReply(qCtx.Q)
-		r.Rcode = b.rCode
+		rr := &dns.A{
+			Hdr: dns.RR_Header{
+				Name:   qCtx.Q.Question[0].Name,
+				Rrtype: dns.TypeA,
+				Class:  dns.ClassINET,
+				Ttl:    3600,
+			},
+			A: b.ipv4,
+		}
+		r.Answer = append(r.Answer, rr)
 		qCtx.SetResponse(r, handler.ContextStatusRejected)
+
+	case b.ipv6 != nil && len(qCtx.Q.Question) == 1 && qCtx.Q.Question[0].Qtype == dns.TypeAAAA:
+		r := new(dns.Msg)
+		r.SetReply(qCtx.Q)
+		rr := &dns.AAAA{
+			Hdr: dns.RR_Header{
+				Name:   qCtx.Q.Question[0].Name,
+				Rrtype: dns.TypeA,
+				Class:  dns.ClassINET,
+				Ttl:    3600,
+			},
+			AAAA: b.ipv6,
+		}
+		r.Answer = append(r.Answer, rr)
+		qCtx.SetResponse(r, handler.ContextStatusRejected)
+
+	case b.args.RCode >= 0:
+		r := new(dns.Msg)
+		r.SetReply(qCtx.Q)
+		r.Rcode = b.args.RCode
+		qCtx.SetResponse(r, handler.ContextStatusRejected)
+
 	default:
 		qCtx.SetResponse(nil, handler.ContextStatusDropped)
 	}

dispatcher/plugin/executable/blackhole/blackhole_test.go

@@ -4,36 +4,43 @@ import (
 	"context"
 	"github.com/IrineSistiana/mosdns/dispatcher/handler"
 	"github.com/miekg/dns"
+	"net"
 	"testing"
 )
 
-func Test_blackhole_Do(t *testing.T) {
+func Test_blackhole_Exec(t *testing.T) {
 	tests := []struct {
 		name         string
-		argsRcode    int
-		wantRcode    int
+		args         *Args
+		queryType    uint16
 		wantResponse bool
+		wantRcode    int
+		wantIP       string
 	}{
-		{"Drop response", -1, 0, false},
-		{"Respond with rcode 2", 2, 2, true},
-		{"Respond with rcode 3", 3, 3, true},
+		{"drop response1", &Args{RCode: -1}, dns.TypeA, false, 0, ""},
+		{"respond with rcode 2", &Args{RCode: 2}, dns.TypeA, true, 2, ""},
+		{"respond with ipv4 1", &Args{IPv4: "127.0.0.1"}, dns.TypeA, true, 0, "127.0.0.1"},
+		{"respond with ipv4 2", &Args{IPv4: "127.0.0.1", RCode: 2}, dns.TypeAAAA, true, 2, ""},
+		{"respond with ipv6", &Args{IPv6: "127.0.0.1"}, dns.TypeAAAA, true, 0, "127.0.0.1"},
 	}
 
 	ctx := context.Background()
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			b := &blackhole{
-				rCode: tt.argsRcode,
+			b, err := newBlackhole("test", tt.args)
+			if err != nil {
+				t.Fatal(err)
 			}
+
 			qCtx := new(handler.Context)
 			q := new(dns.Msg)
-			q.SetQuestion("example.com", dns.TypeA)
+			q.SetQuestion("example.com", tt.queryType)
 			r := new(dns.Msg)
 			r.SetReply(q)
 			qCtx.Q = q
 			qCtx.SetResponse(r, handler.ContextStatusResponded)
 
-			err := b.Exec(ctx, qCtx)
+			err = b.Exec(ctx, qCtx)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -43,8 +50,22 @@ func Test_blackhole_Do(t *testing.T) {
 			}
 
 			if tt.wantResponse {
+				if len(tt.wantIP) != 0 {
+					wantIP := net.ParseIP(tt.wantIP)
+					var gotIP net.IP
+					switch tt.queryType {
+					case dns.TypeA:
+						gotIP = qCtx.R.Answer[0].(*dns.A).A
+					case dns.TypeAAAA:
+						gotIP = qCtx.R.Answer[0].(*dns.AAAA).AAAA
+					}
+					if !wantIP.Equal(gotIP) {
+						t.Fatalf("ip mismatched, want %v, got %v", wantIP, gotIP)
+					}
+				}
+
 				if tt.wantRcode != qCtx.R.Rcode {
-					t.Errorf("response should have rcode %d, but got %d", tt.wantRcode, qCtx.R.Rcode)
+					t.Fatalf("response should have rcode %d, but got %d", tt.wantRcode, qCtx.R.Rcode)
 				}
 			}
 		})