Add support for static domain-to-IP mapping in k8s_dns_server for DNSChaos

Makefile

@@ -14,6 +14,9 @@ protoc: ## Generate the protobuf code
 	go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
 	protoc --proto_path=pb --go_out=pb --go_opt=paths=source_relative ./pb/dns.proto
 
+fmt:
+	find . -type f -name '*.go' -not -path './pb/**' -exec goimports -w -l {} +
+
 # The help will print out all targets with their descriptions organized bellow their categories. The categories are represented by `##@` and the target descriptions by `##`.
 # The awk commands is responsible to read the entire set of makefiles included in this invocation, looking for lines of the file as xyz: ## something, and then pretty-format the target and help. Then, if there's a line with ##@ something, that gets pretty-printed as a category.
 # More info over the usage of ANSI control characters for terminal formatting: https://en.wikipedia.org/wiki/ANSI_escape_code#SGR_parameters

chaos.go

@@ -26,6 +26,8 @@ const (
 	ActionError = "error"
 	// ActionRandom means return random IP for DNS request
 	ActionRandom = "random"
+	// ActionChaos means return chaos IP for DNS request
+	ActionStatic = "static"
 )
 
 // PodInfo saves some information for pod
@@ -55,8 +57,6 @@ func (k Kubernetes) chaosDNS(ctx context.Context, w dns.ResponseWriter, r *dns.M
 		return dns.RcodeServerFailure, fmt.Errorf("dns chaos error")
 	}
 
-	// return random IP
-
 	answers := []dns.RR{}
 	qname := state.Name()
 
@@ -166,6 +166,15 @@ func (k Kubernetes) needChaos(podInfo *PodInfo, records []dns.RR, name string) b
 	if podInfo.Scope == ScopeAll {
 		return true
 	}
+	if podInfo.Action == ActionStatic {
+		domainMap := k.domainIPMapByNamespacedName[podInfo.Namespace][podInfo.Name]
+		if domainMap != nil {
+			if _, ok := domainMap[name]; ok {
+				return true
+			}
+		}
+		return false
+	}
 
 	rules := podInfo.Selector.Match(name, "")
 	if len(rules) == 0 {
@@ -188,3 +197,39 @@ func (k Kubernetes) getPodFromCluster(namespace, name string) (*api.Pod, error)
 	}
 	return pods.Get(context.Background(), name, meta.GetOptions{})
 }
+
+func generateDNSRecords(state request.Request, domainIPMapByNamespacedName map[string]string, r *dns.Msg, w dns.ResponseWriter) (int, error) {
+	answers := []dns.RR{}
+	qname := state.Name()
+	if domainIPMapByNamespacedName == nil {
+		return dns.RcodeServerFailure, nil
+	}
+	ipStr, ok := domainIPMapByNamespacedName[qname]
+	if !ok {
+		return dns.RcodeServerFailure, fmt.Errorf("domain %s not found", qname)
+	}
+	ip := net.ParseIP(ipStr)
+	switch state.QType() {
+	case dns.TypeA:
+		ipv4 := ip.To4()
+		if ipv4 == nil {
+			return dns.RcodeServerFailure, fmt.Errorf("not a valid IPv4 address: %s", ipStr)
+		}
+		answers = a(qname, 10, []net.IP{ipv4})
+		log.Debugf("dns.TypeA %v", ipv4)
+	case dns.TypeAAAA:
+		ipv6 := ip.To16()
+		if ip.To4() != nil {
+			return dns.RcodeServerFailure, fmt.Errorf("not a valid IPv6 address: %s", ipStr)
+		}
+		log.Debugf("dns.TypeAAAA %v", ipv6)
+		answers = aaaa(qname, 10, []net.IP{ipv6})
+	}
+	m := new(dns.Msg)
+	m.SetReply(r)
+	m.Authoritative = true
+	m.Answer = answers
+
+	w.WriteMsg(m)
+	return dns.RcodeSuccess, nil
+}

go.mod

@@ -9,7 +9,6 @@ require (
 	github.com/miekg/dns v1.1.43
 	github.com/pingcap/tidb-tools v6.3.0+incompatible
 	github.com/prometheus/client_golang v1.11.0
-	golang.org/x/net v0.0.0-20210614182718-04defd469f4e
 	google.golang.org/grpc v1.41.0
 	k8s.io/api v0.22.2
 	k8s.io/apimachinery v0.22.2
@@ -45,6 +44,7 @@ require (
 	github.com/prometheus/common v0.31.1 // indirect
 	github.com/prometheus/procfs v0.6.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	golang.org/x/net v0.0.0-20210614182718-04defd469f4e // indirect
 	golang.org/x/oauth2 v0.0.0-20210819190943-2bc19b11175f // indirect
 	golang.org/x/sys v0.0.0-20210917161153-d61c044b1678 // indirect
 	golang.org/x/term v0.0.0-20210220032956-6a3ed077a48d // indirect

grpc_server.go

@@ -44,7 +44,6 @@ func (k Kubernetes) SetDNSChaos(ctx context.Context, req *pb.SetDNSChaosRequest)
 
 	k.Lock()
 	defer k.Unlock()
-
 	k.chaosMap[req.Name] = req
 
 	var scope string
@@ -70,6 +69,15 @@ func (k Kubernetes) SetDNSChaos(ctx context.Context, req *pb.SetDNSChaosRequest)
 			}
 		}
 	}
+	if req.Action == ActionStatic && req.IpDomainMaps != nil {
+		for _, domainIPMap := range req.IpDomainMaps {
+			err := selector.Insert(domainIPMap.Domain, "", true, trieselector.Insert)
+			if err != nil {
+				log.Errorf("fail to build selector %v", err)
+				return nil, err
+			}
+		}
+	}
 
 	for _, pod := range req.Pods {
 		v1Pod, err := k.getPodFromCluster(pod.Namespace, pod.Name)
@@ -100,8 +108,15 @@ func (k Kubernetes) SetDNSChaos(ctx context.Context, req *pb.SetDNSChaosRequest)
 
 		k.podMap[pod.Namespace][pod.Name] = podInfo
 		k.ipPodMap[v1Pod.Status.PodIP] = podInfo
-	}
+		domainIPMap := saveDomainAndIp(req.IpDomainMaps)
+		if domainIPMap != nil {
+			if _, ok := k.domainIPMapByNamespacedName[pod.Namespace]; !ok {
+				k.domainIPMapByNamespacedName[pod.Namespace] = make(map[string]map[string]string)
+			}
+			k.domainIPMapByNamespacedName[pod.Namespace][pod.Name] = domainIPMap
+		}
 
+	}
 	return &pb.DNSChaosResponse{
 		Result: true,
 	}, nil
@@ -125,6 +140,9 @@ func (k Kubernetes) CancelDNSChaos(ctx context.Context, req *pb.CancelDNSChaosRe
 				delete(k.podMap[pod.Namespace], pod.Name)
 				delete(k.ipPodMap, podInfo.IP)
 			}
+			if _, ok1 := k.domainIPMapByNamespacedName[pod.Namespace][pod.Name]; ok1 {
+				delete(k.domainIPMapByNamespacedName[pod.Namespace], pod.Name)
+			}
 		}
 	}
 
@@ -136,6 +154,7 @@ func (k Kubernetes) CancelDNSChaos(ctx context.Context, req *pb.CancelDNSChaosRe
 	}
 	for _, namespace := range shouldDeleteNs {
 		delete(k.podMap, namespace)
+		delete(k.domainIPMapByNamespacedName, namespace)
 	}
 
 	delete(k.chaosMap, req.Name)
@@ -144,3 +163,16 @@ func (k Kubernetes) CancelDNSChaos(ctx context.Context, req *pb.CancelDNSChaosRe
 		Result: true,
 	}, nil
 }
+
+// save domain and ip
+func saveDomainAndIp(domainMapList []*pb.IpDomainMap) map[string]string {
+	if len(domainMapList) == 0 {
+		return nil
+	}
+	domainIPMap := make(map[string]string)
+	for _, domainMap := range domainMapList {
+		key := fmt.Sprintf("%s.", domainMap.Domain)
+		domainIPMap[key] = domainMap.Ip
+	}
+	return domainIPMap
+}

handler.go

@@ -22,10 +22,23 @@ func (k Kubernetes) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.M
 
 	records, extra, zone, err := k.getRecords(ctx, state)
 	log.Debugf("records: %v, err: %v", records, err)
-
-	if k.needChaos(chaosPod, records, state.QName()) {
+	if k.needChaos(chaosPod, records, state.QName()) && chaosPod.Action != ActionStatic {
 		return k.chaosDNS(ctx, w, r, state, chaosPod)
 	}
+	// Check if chaos testing is needed and the action type is static IP.
+	if k.needChaos(chaosPod, records, state.QName()) && chaosPod.Action == ActionStatic {
+		log.Debugf("need chaos, but action is static")
+		// Get the domain-IP mapping for the specific namespace and pod name.
+		domainMap := k.domainIPMapByNamespacedName[chaosPod.Namespace][chaosPod.Name]
+		// Check if the domain-IP mapping exists.
+		if domainMap != nil {
+			// Check if the requested domain exists in the mapping.
+			if _, ok := domainMap[state.Name()]; ok {
+				// Generate DNS records using the domain-IP mapping and return the result.
+				return generateDNSRecords(state, domainMap, r, w)
+			}
+		}
+	}
 
 	if k.IsNameError(err) {
 		if len(zone) == 0 {

handler_test.go

@@ -41,41 +41,41 @@ var dnsTestCases = []test.Case{
 	},
 	{
 		Qname: "svc1.testns.svc.cluster.local.", Qtype: dns.TypeSRV,
-		Rcode: dns.RcodeSuccess,
+		Rcode:  dns.RcodeSuccess,
 		Answer: []dns.RR{test.SRV("svc1.testns.svc.cluster.local.	5	IN	SRV	0 100 80 svc1.testns.svc.cluster.local.")},
-		Extra: []dns.RR{test.A("svc1.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
+		Extra:  []dns.RR{test.A("svc1.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
 	},
 	{
 		Qname: "svcempty.testns.svc.cluster.local.", Qtype: dns.TypeSRV,
-		Rcode: dns.RcodeSuccess,
+		Rcode:  dns.RcodeSuccess,
 		Answer: []dns.RR{test.SRV("svcempty.testns.svc.cluster.local.	5	IN	SRV	0 100 80 svcempty.testns.svc.cluster.local.")},
-		Extra: []dns.RR{test.A("svcempty.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
+		Extra:  []dns.RR{test.A("svcempty.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
 	},
 	{
 		Qname: "svc6.testns.svc.cluster.local.", Qtype: dns.TypeSRV,
-		Rcode: dns.RcodeSuccess,
+		Rcode:  dns.RcodeSuccess,
 		Answer: []dns.RR{test.SRV("svc6.testns.svc.cluster.local.	5	IN	SRV	0 100 80 svc6.testns.svc.cluster.local.")},
-		Extra: []dns.RR{test.AAAA("svc6.testns.svc.cluster.local.  5       IN      AAAA       1234:abcd::1")},
+		Extra:  []dns.RR{test.AAAA("svc6.testns.svc.cluster.local.  5       IN      AAAA       1234:abcd::1")},
 	},
 	// SRV Service (wildcard)
 	{
 		Qname: "svc1.*.svc.cluster.local.", Qtype: dns.TypeSRV,
-		Rcode: dns.RcodeSuccess,
+		Rcode:  dns.RcodeSuccess,
 		Answer: []dns.RR{test.SRV("svc1.*.svc.cluster.local.	5	IN	SRV	0 100 80 svc1.testns.svc.cluster.local.")},
-		Extra: []dns.RR{test.A("svc1.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
+		Extra:  []dns.RR{test.A("svc1.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
 	},
 	{
 		Qname: "svcempty.*.svc.cluster.local.", Qtype: dns.TypeSRV,
-		Rcode: dns.RcodeSuccess,
+		Rcode:  dns.RcodeSuccess,
 		Answer: []dns.RR{test.SRV("svcempty.*.svc.cluster.local.	5	IN	SRV	0 100 80 svcempty.testns.svc.cluster.local.")},
-		Extra: []dns.RR{test.A("svcempty.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
+		Extra:  []dns.RR{test.A("svcempty.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
 	},
 	// SRV Service (wildcards)
 	{
 		Qname: "*.any.svc1.*.svc.cluster.local.", Qtype: dns.TypeSRV,
-		Rcode: dns.RcodeSuccess,
+		Rcode:  dns.RcodeSuccess,
 		Answer: []dns.RR{test.SRV("*.any.svc1.*.svc.cluster.local.	5	IN	SRV	0 100 80 svc1.testns.svc.cluster.local.")},
-		Extra: []dns.RR{test.A("svc1.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
+		Extra:  []dns.RR{test.A("svc1.testns.svc.cluster.local.  5       IN      A       10.0.0.1")},
 	},
 	// A Service (wildcards)
 	{
@@ -207,7 +207,7 @@ var dnsTestCases = []test.Case{
 	// AAAA
 	{
 		Qname: "5678-abcd--2.hdls1.testns.svc.cluster.local", Qtype: dns.TypeAAAA,
-		Rcode: dns.RcodeSuccess,
+		Rcode:  dns.RcodeSuccess,
 		Answer: []dns.RR{test.AAAA("5678-abcd--2.hdls1.testns.svc.cluster.local.	5	IN	AAAA	5678:abcd::2")},
 	},
 	// CNAME External

kubernetes.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/chaos-mesh/k8s_dns_chaos/pb"
+
 	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/plugin/etcd/msg"
 	"github.com/coredns/coredns/plugin/kubernetes/object"
@@ -68,6 +69,8 @@ type Kubernetes struct {
 	podMap map[string]map[string]*PodInfo
 
 	ipPodMap map[string]*PodInfo
+
+	domainIPMapByNamespacedName map[string]map[string]map[string]string
 }
 
 // New returns a initialized Kubernetes. It default interfaceAddrFunc to return 127.0.0.1. All other
@@ -81,6 +84,7 @@ func New(zones []string) *Kubernetes {
 	k.chaosMap = make(map[string]*pb.SetDNSChaosRequest)
 	k.podMap = make(map[string]map[string]*PodInfo)
 	k.ipPodMap = make(map[string]*PodInfo)
+	k.domainIPMapByNamespacedName = make(map[string]map[string]map[string]string)
 	rand.Seed(time.Now().UnixNano())
 
 	return k