DNS: New Features

app/dns/nameserver.go

@@ -26,14 +26,17 @@ type Server interface {
 
 // Client is the interface for DNS client.
 type Client struct {
-	server             Server
-	skipFallback       bool
-	domains            []string
-	expectedIPs        []*router.GeoIPMatcher
-	allowUnexpectedIPs bool
-	tag                string
-	timeoutMs          time.Duration
-	ipOption           *dns.IPOption
+	server        Server
+	skipFallback  bool
+	domains       []string
+	expectedIPs   []*router.GeoIPMatcher
+	unexpectedIPs []*router.GeoIPMatcher
+	priorIPs      []*router.GeoIPMatcher
+	unpriorIPs    []*router.GeoIPMatcher
+	tag           string
+	timeoutMs     time.Duration
+	finalQuery    bool
+	ipOption      *dns.IPOption
 }
 
 // NewServer creates a name server object according to the network destination url.
@@ -150,13 +153,43 @@ func NewClient(
 		}
 
 		// Establish expected IPs
-		var matchers []*router.GeoIPMatcher
-		for _, geoip := range ns.Geoip {
+		var expectedMatchers []*router.GeoIPMatcher
+		for _, geoip := range ns.ExpectedGeoip {
 			matcher, err := router.GlobalGeoIPContainer.Add(geoip)
 			if err != nil {
-				return errors.New("failed to create ip matcher").Base(err).AtWarning()
+				return errors.New("failed to create expected ip matcher").Base(err).AtWarning()
 			}
-			matchers = append(matchers, matcher)
+			expectedMatchers = append(expectedMatchers, matcher)
+		}
+
+		// Establish unexpected IPs
+		var unexpectedMatchers []*router.GeoIPMatcher
+		for _, geoip := range ns.UnexpectedGeoip {
+			matcher, err := router.GlobalGeoIPContainer.Add(geoip)
+			if err != nil {
+				return errors.New("failed to create unexpected ip matcher").Base(err).AtWarning()
+			}
+			unexpectedMatchers = append(unexpectedMatchers, matcher)
+		}
+
+		// Establish prior IPs
+		var priorMatchers []*router.GeoIPMatcher
+		for _, geoip := range ns.PriorGeoip {
+			matcher, err := router.GlobalGeoIPContainer.Add(geoip)
+			if err != nil {
+				return errors.New("failed to create prior ip matcher").Base(err).AtWarning()
+			}
+			priorMatchers = append(priorMatchers, matcher)
+		}
+
+		// Establish unprior IPs
+		var unpriorMatchers []*router.GeoIPMatcher
+		for _, geoip := range ns.UnpriorGeoip {
+			matcher, err := router.GlobalGeoIPContainer.Add(geoip)
+			if err != nil {
+				return errors.New("failed to create unprior ip matcher").Base(err).AtWarning()
+			}
+			unpriorMatchers = append(unpriorMatchers, matcher)
 		}
 
 		if len(clientIP) > 0 {
@@ -176,10 +209,13 @@ func NewClient(
 		client.server = server
 		client.skipFallback = ns.SkipFallback
 		client.domains = rules
-		client.expectedIPs = matchers
-		client.allowUnexpectedIPs = ns.AllowUnexpectedIPs
+		client.expectedIPs = expectedMatchers
+		client.unexpectedIPs = unexpectedMatchers
+		client.priorIPs = priorMatchers
+		client.unpriorIPs = unpriorMatchers
 		client.tag = tag
 		client.timeoutMs = timeoutMs
+		client.finalQuery = ns.FinalQuery
 		client.ipOption = &ipOption
 		return nil
 	})
@@ -191,6 +227,10 @@ func (c *Client) Name() string {
 	return c.server.Name()
 }
 
+func (c *Client) IsFinalQuery() bool {
+	return c.finalQuery
+}
+
 // QueryIP sends DNS query to the name server with the client's IP.
 func (c *Client) QueryIP(ctx context.Context, domain string, option dns.IPOption) ([]net.IP, uint32, error) {
 	option.IPv4Enable = option.IPv4Enable && c.ipOption.IPv4Enable
@@ -213,34 +253,40 @@ func (c *Client) QueryIP(ctx context.Context, domain string, option dns.IPOption
 	}
 
 	if len(c.expectedIPs) > 0 {
-		newIps := c.MatchExpectedIPs(domain, ips)
-		if len(newIps) == 0 {
-			if !c.allowUnexpectedIPs {
-				return nil, 0, dns.ErrEmptyResponse
-			}
-		} else {
-			ips = newIps
+		ips = router.MatchIPs(c.expectedIPs, ips, false)
+		errors.LogDebug(context.Background(), "domain ", domain, " expectedIPs ", ips, " matched at server ", c.Name())
+		if len(ips) == 0 {
+			return nil, 0, dns.ErrEmptyResponse
+		}
+	}
+
+	if len(c.unexpectedIPs) > 0 {
+		ips = router.MatchIPs(c.unexpectedIPs, ips, true)
+		errors.LogDebug(context.Background(), "domain ", domain, " unexpectedIPs ", ips, " matched at server ", c.Name())
+		if len(ips) == 0 {
+			return nil, 0, dns.ErrEmptyResponse
+		}
+	}
+
+	if len(c.priorIPs) > 0 {
+		ipsNew := router.MatchIPs(c.priorIPs, ips, false)
+		if len(ipsNew) > 0 {
+			ips = ipsNew
+			errors.LogDebug(context.Background(), "domain ", domain, " priorIPs ", ips, " matched at server ", c.Name())
+		}
+	}
+
+	if len(c.unpriorIPs) > 0 {
+		ipsNew := router.MatchIPs(c.unpriorIPs, ips, true)
+		if len(ipsNew) > 0 {
+			ips = ipsNew
+			errors.LogDebug(context.Background(), "domain ", domain, " unpriorIPs ", ips, " matched at server ", c.Name())
 		}
 	}
 
 	return ips, ttl, nil
 }
 
-// MatchExpectedIPs matches queried domain IPs with expected IPs and returns matched ones.
-func (c *Client) MatchExpectedIPs(domain string, ips []net.IP) []net.IP {
-	var newIps []net.IP
-	for _, ip := range ips {
-		for _, matcher := range c.expectedIPs {
-			if matcher.Match(ip) {
-				newIps = append(newIps, ip)
-				break
-			}
-		}
-	}
-	errors.LogDebug(context.Background(), "domain ", domain, " expectedIPs ", newIps, " matched at server ", c.Name())
-	return newIps
-}
-
 func ResolveIpOptionOverride(queryStrategy QueryStrategy, ipOption dns.IPOption) dns.IPOption {
 	switch queryStrategy {
 	case QueryStrategy_USE_IP: