refine code && support scope and mode for chaos (#6)

Signed-off-by: xiang <[email protected]>

Author: WangXiangUSTC

chaos.go

@@ -4,14 +4,57 @@ import (
 	"context"
 	"math/rand"
 	"net"
+	"time"
 
+	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/request"
 	"github.com/miekg/dns"
 	api "k8s.io/api/core/v1"
 	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
-func (k Kubernetes) chaosDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg, state request.Request) (int, error) {
+const (
+	// ScopeInner means chaos only works on the inner host in Kubernetes cluster
+	ScopeInner = "INNER"
+	// ScopeOuter means chaos only works on the outer host of Kubernetes cluster
+	ScopeOuter = "OUTER"
+	// ScopeAll means chaos works on all host
+	ScopeAll = "ALL"
+
+	// ModeError means return error for DNS request
+	ModeError = "ERROR"
+	// ModeRandom means return random IP for DNS request
+	ModeRandom = "RANDOM"
+)
+
+// PodInfo saves some information for pod
+type PodInfo struct {
+	Namespace      string
+	Name           string
+	Mode           string
+	Scope          string
+	IP             string
+	LastUpdateTime time.Time
+}
+
+// IsOverdue ...
+func (p *PodInfo) IsOverdue() bool {
+	// if the pod's IP is not updated greater than 10 seconds, will treate it as overdue
+	// and need to update it
+	if time.Since(p.LastUpdateTime) > time.Duration(time.Second*10) {
+		return true
+	}
+
+	return false
+}
+
+func (k Kubernetes) chaosDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg, state request.Request, podInfo *PodInfo) (int, error) {
+	if podInfo.Mode == ModeError {
+		return dns.RcodeServerFailure, nil
+	}
+
+	// return random IP
+
 	answers := []dns.RR{}
 	qname := state.Name()
 
@@ -22,6 +65,7 @@ func (k Kubernetes) chaosDNS(ctx context.Context, w dns.ResponseWriter, r *dns.M
 		log.Infof("dns.TypeA %v", ips)
 		answers = a(qname, 10, ips)
 	case dns.TypeAAAA:
+		// TODO: return random IP
 		ips := []net.IP{net.IP{0x20, 0x1, 0xd, 0xb8, 0, 0, 0, 0, 0, 0, 0x1, 0x23, 0, 0x12, 0, 0x1}}
 		log.Infof("dns.TypeAAAA %v", ips)
 		answers = aaaa(qname, 10, ips)
@@ -77,35 +121,68 @@ func aaaa(zone string, ttl uint32, ips []net.IP) []dns.RR {
 	return answers
 }
 
-func (k Kubernetes) getChaosMode(pod *api.Pod) string {
+func (k Kubernetes) getChaosPod(ip string) (*PodInfo, error) {
 	k.RLock()
-	defer k.RUnlock()
 
-	if pod == nil {
-		return ""
+	podInfo := k.ipPodMap[ip]
+	if podInfo == nil {
+		k.RUnlock()
+		return nil, nil
 	}
 
-	if _, ok := k.podChaosMap[pod.Namespace]; ok {
-		return k.podChaosMap[pod.Namespace][pod.Name]
+	if podInfo.IsOverdue() {
+		k.RUnlock()
+
+		v1Pod, err := k.getPodFromCluster(podInfo.Namespace, podInfo.Name)
+		if err != nil {
+			return nil, err
+		}
+
+		if v1Pod.Status.PodIP != podInfo.IP {
+			k.Lock()
+			podInfo.IP = v1Pod.Status.PodIP
+			podInfo.LastUpdateTime = time.Now()
+
+			delete(k.ipPodMap, podInfo.IP)
+			k.ipPodMap[v1Pod.Status.PodIP] = podInfo
+			k.Unlock()
+		}
+
+		return podInfo, nil
 	}
 
-	return ""
+	k.RUnlock()
+	return podInfo, nil
 }
 
-func (k Kubernetes) getChaosPod() ([]api.Pod, error) {
-	k.RLock()
-	defer k.RUnlock()
+// needChaos judges weather should do chaos for the request
+func (k Kubernetes) needChaos(podInfo *PodInfo, state request.Request) bool {
+	if podInfo == nil {
+		return false
+	}
 
-	pods := make([]api.Pod, 0, 10)
-	for namespace := range k.podChaosMap {
-		podList, err := k.Client.Pods(namespace).List(context.Background(), meta.ListOptions{})
-		if err != nil {
-			return nil, err
+	if podInfo.Scope == ScopeAll {
+		return true
+	}
+
+	qname := state.QName()
+	zone := plugin.Zones(k.Zones).Matches(qname)
+
+	if zone == "" {
+		// is outer host
+		if podInfo.Scope == ScopeOuter {
+			return true
 		}
-		for _, pod := range podList.Items {
-			pods = append(pods, pod)
+	} else {
+		// is inner host
+		if podInfo.Scope == ScopeInner {
+			return true
 		}
 	}
 
-	return pods, nil
+	return false
+}
+
+func (k Kubernetes) getPodFromCluster(namespace, name string) (*api.Pod, error) {
+	return k.Client.Pods(namespace).Get(context.Background(), name, meta.GetOptions{})
 }

grpc_server.go

@@ -3,6 +3,7 @@ package kubernetes
 import (
 	"context"
 	"net"
+	"time"
 
 	"github.com/chaos-mesh/k8s_dns_chaos/pb"
 	"google.golang.org/grpc"
@@ -30,15 +31,39 @@ func (k Kubernetes) CreateGRPCServer(port string) error {
 // SetDNSChaos ...
 func (k Kubernetes) SetDNSChaos(ctx context.Context, req *pb.SetDNSChaosRequest) (*pb.DNSChaosResponse, error) {
 	log.Infof("receive SetDNSChaos request %v", req)
+
 	k.Lock()
 	defer k.Unlock()
 
 	k.chaosMap[req.Name] = req
+
 	for _, pod := range req.Pods {
-		if _, ok := k.podChaosMap[pod.Namespace]; !ok {
-			k.podChaosMap[pod.Namespace] = make(map[string]string)
+		v1Pod, err := k.getPodFromCluster(pod.Namespace, pod.Name)
+		if err != nil {
+			return nil, err
 		}
-		k.podChaosMap[pod.Namespace][pod.Name] = req.Mode
+
+		if _, ok := k.podMap[pod.Namespace]; !ok {
+			k.podMap[pod.Namespace] = make(map[string]*PodInfo)
+		}
+
+		if oldPod, ok := k.podMap[pod.Namespace][pod.Name]; ok {
+			// Pod's IP maybe changed, so delete the old pod info
+			delete(k.podMap[pod.Namespace], pod.Name)
+			delete(k.ipPodMap, oldPod.IP)
+		}
+
+		podInfo := &PodInfo{
+			Namespace:      pod.Namespace,
+			Name:           pod.Name,
+			Mode:           req.Mode,
+			Scope:          req.Scope,
+			IP:             v1Pod.Status.PodIP,
+			LastUpdateTime: time.Now(),
+		}
+
+		k.podMap[pod.Namespace][pod.Name] = podInfo
+		k.ipPodMap[v1Pod.Status.PodIP] = podInfo
 	}
 
 	return &pb.DNSChaosResponse{
@@ -52,19 +77,22 @@ func (k Kubernetes) CancelDNSChaos(ctx context.Context, req *pb.CancelDNSChaosRe
 	k.Lock()
 	defer k.Unlock()
 	for _, pod := range k.chaosMap[req.Name].Pods {
-		if _, ok := k.podChaosMap[pod.Namespace]; ok {
-			delete(k.podChaosMap[pod.Namespace], pod.Name)
+		if _, ok := k.podMap[pod.Namespace]; ok {
+			if podInfo, ok := k.podMap[pod.Namespace][pod.Name]; ok {
+				delete(k.podMap[pod.Namespace], pod.Name)
+				delete(k.ipPodMap, podInfo.IP)
+			}
 		}
 	}
 
 	shouldDeleteNs := make([]string, 0, 1)
-	for namespace, pods := range k.podChaosMap {
+	for namespace, pods := range k.podMap {
 		if len(pods) == 0 {
 			shouldDeleteNs = append(shouldDeleteNs, namespace)
 		}
 	}
 	for _, namespace := range shouldDeleteNs {
-		delete(k.podChaosMap, namespace)
+		delete(k.podMap, namespace)
 	}
 
 	delete(k.chaosMap, req.Name)

handler.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/request"
-	api "k8s.io/api/core/v1"
 
 	"github.com/miekg/dns"
 )
@@ -16,21 +15,12 @@ func (k Kubernetes) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.M
 	sourceIP := state.IP()
 	log.Infof("k8s ServeDNS, source IP: %s", sourceIP)
 
-	var sourcePod *api.Pod
-
-	pods, err := k.getChaosPod()
+	chaosPod, err := k.getChaosPod(sourceIP)
 	if err != nil {
-		log.Errorf("list pods, error %v", err)
-	}
-	for _, pod := range pods {
-		log.Infof("list pod name: %s, ip: %s", pod.Name, pod.Status.PodIP)
-		if pod.Status.PodIP == sourceIP {
-			sourcePod = &pod
-		}
+		log.Infof("fail to get pod information from cluster, IP: %s, error: %v", sourceIP, err)
 	}
-	mode := k.getChaosMode(sourcePod)
-	if len(mode) != 0 {
-		return k.chaosDNS(ctx, w, r, state)
+	if k.needChaos(chaosPod, state) {
+		return k.chaosDNS(ctx, w, r, state, chaosPod)
 	}
 
 	qname := state.QName()

kubernetes.go

@@ -57,12 +57,14 @@ type Kubernetes struct {
 	localIPs         []net.IP
 	autoPathSearch   []string // Local search path from /etc/resolv.conf. Needed for autopath.
 	TransferTo       []string
-	Client           typev1.CoreV1Interface //client.Client
+	Client           typev1.CoreV1Interface
 
 	sync.RWMutex
 	chaosMap map[string]*pb.SetDNSChaosRequest
-	// namespace -> pod_name -> chaos_mode
-	podChaosMap map[string]map[string]string
+	// namespace -> pod_name -> pod info
+	podMap map[string]map[string]*PodInfo
+
+	ipPodMap map[string]*PodInfo
 }
 
 // New returns a initialized Kubernetes. It default interfaceAddrFunc to return 127.0.0.1. All other
@@ -74,7 +76,8 @@ func New(zones []string) *Kubernetes {
 	k.podMode = podModeDisabled
 	k.ttl = defaultTTL
 	k.chaosMap = make(map[string]*pb.SetDNSChaosRequest)
-	k.podChaosMap = make(map[string]map[string]string)
+	k.podMap = make(map[string]map[string]*PodInfo)
+	k.ipPodMap = make(map[string]*PodInfo)
 	rand.Seed(time.Now().UnixNano())
 
 	return k