update domain name match logic (#14)

Author: WangXiangUSTC

chaos.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/coredns/coredns/request"
 	"github.com/miekg/dns"
+	selector "github.com/pingcap/tidb-tools/pkg/table-rule-selector"
 	api "k8s.io/api/core/v1"
 	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
@@ -33,6 +34,7 @@ type PodInfo struct {
 	Name           string
 	Action         string
 	Scope          string
+	Selector       selector.Selector
 	IP             string
 	LastUpdateTime time.Time
 }
@@ -156,7 +158,7 @@ func (k Kubernetes) getChaosPod(ip string) (*PodInfo, error) {
 }
 
 // needChaos judges weather should do chaos for the request
-func (k Kubernetes) needChaos(podInfo *PodInfo, records []dns.RR, err error) bool {
+func (k Kubernetes) needChaos(podInfo *PodInfo, records []dns.RR, name string) bool {
 	if podInfo == nil {
 		return false
 	}
@@ -165,33 +167,17 @@ func (k Kubernetes) needChaos(podInfo *PodInfo, records []dns.RR, err error) boo
 		return true
 	}
 
-	if err != nil {
-		// not found in cluster, is outer host
-		if k.IsNameError(err) {
-			if podInfo.Scope == ScopeOuter {
-				return true
-			}
-			return false
-		}
-
-		// can't judge the host is outer or inner host, ignore chaos
+	rules := podInfo.Selector.Match(name, "")
+	if len(rules) == 0 {
 		return false
 	}
 
-	if len(records) == 0 {
-		// not found in cluster, is outer host {
-		if podInfo.Scope == ScopeOuter {
-			return true
-		}
+	match, ok := rules[0].(bool)
+	if !ok {
 		return false
 	}
 
-	// is inner host
-	if podInfo.Scope == ScopeInner {
-		return true
-	}
-
-	return false
+	return match
 }
 
 func (k Kubernetes) getPodFromCluster(namespace, name string) (*api.Pod, error) {

grpc_server.go

@@ -4,9 +4,11 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"strings"
 	"time"
 
 	"github.com/chaos-mesh/k8s_dns_chaos/pb"
+	trieselector "github.com/pingcap/tidb-tools/pkg/table-rule-selector"
 	"google.golang.org/grpc"
 )
 
@@ -45,6 +47,30 @@ func (k Kubernetes) SetDNSChaos(ctx context.Context, req *pb.SetDNSChaosRequest)
 
 	k.chaosMap[req.Name] = req
 
+	var scope string
+	if len(req.Patterns) == 0 {
+		scope = ScopeAll
+	}
+
+	// build selector
+	selector := trieselector.NewTrieSelector()
+	for _, pattern := range req.Patterns {
+		err := selector.Insert(pattern, "", true, trieselector.Insert)
+		if err != nil {
+			log.Errorf("fail to build selector %v", err)
+			return nil, err
+		}
+
+		if !strings.Contains(pattern, "*") {
+			// when send dns request to the dns server, will add a '.' at the end of the domain name.
+			err := selector.Insert(fmt.Sprintf("%s.", pattern), "", 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)
 		if err != nil {
@@ -66,7 +92,8 @@ func (k Kubernetes) SetDNSChaos(ctx context.Context, req *pb.SetDNSChaosRequest)
 			Namespace:      pod.Namespace,
 			Name:           pod.Name,
 			Action:         req.Action,
-			Scope:          req.Scope,
+			Scope:          scope,
+			Selector:       selector,
 			IP:             v1Pod.Status.PodIP,
 			LastUpdateTime: time.Now(),
 		}

handler.go

@@ -23,7 +23,7 @@ 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, err) {
+	if k.needChaos(chaosPod, records, state.QName()) {
 		return k.chaosDNS(ctx, w, r, state, chaosPod)
 	}
 

pb/dns.pb.go

@@ -22,6 +22,7 @@ message SetDNSChaosRequest {
   //   "all":   chaos works on all host
   string scope = 4;
   string selector = 5;
+  repeated string patterns = 6;
 }
 
 message Pod {