feat(dns): support fixed domain ttl (#100)

* feat(dns): support fixed domain ttl

* docs

common/utils.go

@@ -466,3 +466,11 @@ func IsValidHttpMethod(method string) bool {
 		return false
 	}
 }
+
+func StringSet(list []string) map[string]struct{} {
+	m := make(map[string]struct{})
+	for _, s := range list {
+		m[s] = struct{}{}
+	}
+	return m
+}

config/config.go

@@ -86,6 +86,7 @@ type DnsRouting struct {
 type KeyableString string
 type Dns struct {
 	IpVersionPrefer int             `mapstructure:"ipversion_prefer"`
+	FixedDomainTtl  []KeyableString `mapstructure:"fixed_domain_ttl"`
 	Upstream        []KeyableString `mapstructure:"upstream"`
 	Routing         DnsRouting      `mapstructure:"routing"`
 }

config/desc.go

@@ -59,6 +59,7 @@ var GlobalDesc = Desc{
 
 var DnsDesc = Desc{
 	"ipversion_prefer": "For example, if ipversion_prefer is 4 and the domain name has both type A and type AAAA records, the dae will only respond to type A queries and response empty answer to type AAAA queries.",
+	"fixed_domain_ttl": "Give a fixed ttl for domains. Zero means that dae will request to upstream every time and not cache DNS results for these domains.",
 	"upstream":         "Value can be scheme://host:port, where the scheme can be tcp/udp/tcp+udp.\nIf host is a domain and has both IPv4 and IPv6 record, dae will automatically choose IPv4 or IPv6 to use according to group policy (such as min latency policy).\nPlease make sure DNS traffic will go through and be forwarded by dae, which is REQUIRED for domain routing.\nIf dial_mode is \"ip\", the upstream DNS answer SHOULD NOT be polluted, so domestic public DNS is not recommended.",
 	"request": `DNS requests will follow this routing.
 Built-in outbound: asis.

control/bpf_utils.go

@@ -142,6 +142,24 @@ func BpfMapBatchUpdate(m *ebpf.Map, keys interface{}, values interface{}, opts *
 	return vKeys.Len(), nil
 }
 
+// BpfMapBatchDelete deletes keys and ignores ErrKeyNotExist.
+func BpfMapBatchDelete(m *ebpf.Map, keys interface{}) (n int, err error) {
+	// Simulate
+	vKeys := reflect.ValueOf(keys)
+	if vKeys.Kind() != reflect.Slice {
+		return 0, fmt.Errorf("keys must be slice")
+	}
+	length := vKeys.Len()
+
+	for i := 0; i < length; i++ {
+		vKey := vKeys.Index(i)
+		if err = m.Delete(vKey.Interface()); err != nil && !errors.Is(err, ebpf.ErrKeyNotExist) {
+			return i, err
+		}
+	}
+	return vKeys.Len(), nil
+}
+
 // detectCgroupPath returns the first-found mount point of type cgroup2
 // and stores it in the cgroupPath global variable.
 // Copied from https://github.com/cilium/ebpf/blob/v0.10.0/examples/cgroup_skb/main.go

control/control_plane.go

@@ -371,6 +371,10 @@ func NewControlPlane(
 		return nil, err
 	}
 	/// Dns controller.
+	fixedDomainTtl, err := ParseFixedDomainTtl(dnsConfig.FixedDomainTtl)
+	if err != nil {
+		return nil, err
+	}
 	if plane.dnsController, err = NewDnsController(dnsUpstream, &DnsControllerOption{
 		Log: log,
 		CacheAccessCallback: func(cache *DnsCache) (err error) {
@@ -381,6 +385,14 @@ func NewControlPlane(
 			}
 			return nil
 		},
+		CacheRemoveCallback: func(cache *DnsCache) (err error) {
+			// Write mappings into eBPF map:
+			// IP record (from dns lookup) -> domain routing
+			if err = core.BatchRemoveDomainRouting(cache); err != nil {
+				return fmt.Errorf("BatchUpdateDomainRouting: %w", err)
+			}
+			return nil
+		},
 		NewCache: func(fqdn string, answers []dnsmessage.Resource, deadline time.Time) (cache *DnsCache, err error) {
 			return &DnsCache{
 				DomainBitmap: plane.routingMatcher.domainMatcher.MatchDomainBitmap(fqdn),
@@ -390,6 +402,7 @@ func NewControlPlane(
 		},
 		BestDialerChooser: plane.chooseBestDnsDialer,
 		IpVersionPrefer:   dnsConfig.IpVersionPrefer,
+		FixedDomainTtl:    fixedDomainTtl,
 	}); err != nil {
 		return nil, err
 	}
@@ -405,7 +418,7 @@ func NewControlPlane(
 			}
 			host := cacheKey[:lastDot]
 			typ := cacheKey[lastDot+1:]
-			_ = plane.dnsController.UpdateDnsCache(host, typ, cache.Answers, cache.Deadline)
+			_ = plane.dnsController.UpdateDnsCacheDeadline(host, typ, cache.Answers, cache.Deadline)
 		}
 	} else if _bpf != nil {
 		// Is reloading, and dnsCache == nil.
@@ -430,6 +443,19 @@ func NewControlPlane(
 	return plane, nil
 }
 
+func ParseFixedDomainTtl(ks []config.KeyableString) (map[string]int, error) {
+	m := make(map[string]int)
+	for _, k := range ks {
+		key, value, _ := strings.Cut(string(k), ":")
+		ttl, err := strconv.ParseInt(strings.TrimSpace(value), 0, strconv.IntSize)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse ttl: %v", err)
+		}
+		m[strings.TrimSpace(key)] = int(ttl)
+	}
+	return m, nil
+}
+
 // EjectBpf will resect bpf from destroying life-cycle of control plane.
 func (c *ControlPlane) EjectBpf() *bpfObjects {
 	return c.core.EjectBpf()
@@ -485,7 +511,7 @@ func (c *ControlPlane) dnsUpstreamReadyCallback(dnsUpstream *dns.Upstream) (err
 				A: dnsUpstream.Ip4.As4(),
 			},
 		}}
-		if err = c.dnsController.UpdateDnsCache(dnsUpstream.Hostname, typ.String(), answers, deadline); err != nil {
+		if err = c.dnsController.UpdateDnsCacheDeadline(dnsUpstream.Hostname, typ.String(), answers, deadline); err != nil {
 			return err
 		}
 	}
@@ -503,7 +529,7 @@ func (c *ControlPlane) dnsUpstreamReadyCallback(dnsUpstream *dns.Upstream) (err
 				AAAA: dnsUpstream.Ip6.As16(),
 			},
 		}}
-		if err = c.dnsController.UpdateDnsCache(dnsUpstream.Hostname, typ.String(), answers, deadline); err != nil {
+		if err = c.dnsController.UpdateDnsCacheDeadline(dnsUpstream.Hostname, typ.String(), answers, deadline); err != nil {
 			return err
 		}
 	}

control/control_plane_core.go

@@ -561,6 +561,40 @@ func (c *controlPlaneCore) BatchUpdateDomainRouting(cache *DnsCache) error {
 	return nil
 }
 
+// BatchRemoveDomainRouting remove bpf map domain_routing.
+func (c *controlPlaneCore) BatchRemoveDomainRouting(cache *DnsCache) error {
+	// Parse ips from DNS resp answers.
+	var ips []netip.Addr
+	for _, ans := range cache.Answers {
+		var ip netip.Addr
+		switch ans.Header.Type {
+		case dnsmessage.TypeA:
+			ip = netip.AddrFrom4(ans.Body.(*dnsmessage.AResource).A)
+		case dnsmessage.TypeAAAA:
+			ip = netip.AddrFrom16(ans.Body.(*dnsmessage.AAAAResource).AAAA)
+		}
+		if ip.IsUnspecified() {
+			continue
+		}
+		ips = append(ips, ip)
+	}
+	if len(ips) == 0 {
+		return nil
+	}
+
+	// Update bpf map.
+	// Construct keys and vals, and BpfMapBatchUpdate.
+	var keys [][4]uint32
+	for _, ip := range ips {
+		ip6 := ip.As16()
+		keys = append(keys, common.Ipv6ByteSliceToUint32Array(ip6[:]))
+	}
+	if _, err := BpfMapBatchDelete(c.bpf.DomainRoutingMap, keys); err != nil {
+		return err
+	}
+	return nil
+}
+
 // EjectBpf will resect bpf from destroying life-cycle of control plane core.
 func (c *controlPlaneCore) EjectBpf() *bpfObjects {
 	if !c.bpfEjected && !c.isReload {

control/dns_control.go

@@ -32,9 +32,8 @@ import (
 )
 
 const (
-	MaxDnsLookupDepth      = 3
-	minFirefoxCacheTtl     = 120
-	minFirefoxCacheTimeout = minFirefoxCacheTtl * time.Second
+	MaxDnsLookupDepth  = 3
+	minFirefoxCacheTtl = 120
 )
 
 type IpVersionPrefer int
@@ -58,9 +57,11 @@ var (
 type DnsControllerOption struct {
 	Log                 *logrus.Logger
 	CacheAccessCallback func(cache *DnsCache) (err error)
+	CacheRemoveCallback func(cache *DnsCache) (err error)
 	NewCache            func(fqdn string, answers []dnsmessage.Resource, deadline time.Time) (cache *DnsCache, err error)
 	BestDialerChooser   func(req *udpRequest, upstream *dns.Upstream) (*dialArgument, error)
 	IpVersionPrefer     int
+	FixedDomainTtl      map[string]int
 }
 
 type DnsController struct {
@@ -71,9 +72,11 @@ type DnsController struct {
 
 	log                 *logrus.Logger
 	cacheAccessCallback func(cache *DnsCache) (err error)
+	cacheRemoveCallback func(cache *DnsCache) (err error)
 	newCache            func(fqdn string, answers []dnsmessage.Resource, deadline time.Time) (cache *DnsCache, err error)
 	bestDialerChooser   func(req *udpRequest, upstream *dns.Upstream) (*dialArgument, error)
 
+	fixedDomainTtl map[string]int
 	// mutex protects the dnsCache.
 	dnsCacheMu sync.Mutex
 	dnsCache   map[string]*DnsCache
@@ -105,11 +108,13 @@ func NewDnsController(routing *dns.Dns, option *DnsControllerOption) (c *DnsCont
 
 		log:                 option.Log,
 		cacheAccessCallback: option.CacheAccessCallback,
+		cacheRemoveCallback: option.CacheRemoveCallback,
 		newCache:            option.NewCache,
 		bestDialerChooser:   option.BestDialerChooser,
 
-		dnsCacheMu: sync.Mutex{},
-		dnsCache:   make(map[string]*DnsCache),
+		fixedDomainTtl: option.FixedDomainTtl,
+		dnsCacheMu:     sync.Mutex{},
+		dnsCache:       make(map[string]*DnsCache),
 	}, nil
 }
 
@@ -276,19 +281,14 @@ func (c *DnsController) updateDnsCache(msg *dnsmessage.Message, ttl uint32, q *d
 			"addition": FormatDnsRsc(msg.Additionals),
 		}).Tracef("Update DNS record cache")
 	}
-	cacheTimeout := time.Duration(ttl) * time.Second // TTL.
-	if cacheTimeout < minFirefoxCacheTimeout {
-		cacheTimeout = minFirefoxCacheTimeout
-	}
-	cacheTimeout += 5 * time.Second // DNS lookup timeout.
 
-	if err := c.UpdateDnsCache(q.Name.String(), q.Type.String(), msg.Answers, time.Now().Add(cacheTimeout)); err != nil {
+	if err := c.UpdateDnsCacheTtl(q.Name.String(), q.Type.String(), msg.Answers, int(ttl)); err != nil {
 		return err
 	}
 	return nil
 }
 
-func (c *DnsController) UpdateDnsCache(host string, dnsTyp string, answers []dnsmessage.Resource, deadline time.Time) (err error) {
+func (c *DnsController) __updateDnsCacheDeadline(host string, dnsTyp string, answers []dnsmessage.Resource, deadlineFunc func(now time.Time, host string) time.Time) (err error) {
 	var fqdn string
 	if strings.HasSuffix(host, ".") {
 		fqdn = host
@@ -300,15 +300,16 @@ func (c *DnsController) UpdateDnsCache(host string, dnsTyp string, answers []dns
 	if _, err = netip.ParseAddr(host); err == nil {
 		return nil
 	}
+
+	now := time.Now()
+	deadline := deadlineFunc(now, host)
+
 	cacheKey := fqdn + dnsTyp
 	c.dnsCacheMu.Lock()
 	cache, ok := c.dnsCache[cacheKey]
 	if ok {
-		// To avoid overwriting DNS upstream resolution.
-		if deadline.After(cache.Deadline) {
-			cache.Deadline = deadline
-		}
 		cache.Answers = answers
+		cache.Deadline = deadline
 		c.dnsCacheMu.Unlock()
 	} else {
 		cache, err = c.newCache(fqdn, answers, deadline)
@@ -322,9 +323,32 @@ func (c *DnsController) UpdateDnsCache(host string, dnsTyp string, answers []dns
 	if err = c.cacheAccessCallback(cache); err != nil {
 		return err
 	}
+
 	return nil
 }
 
+func (c *DnsController) UpdateDnsCacheDeadline(host string, dnsTyp string, answers []dnsmessage.Resource, deadline time.Time) (err error) {
+	return c.__updateDnsCacheDeadline(host, dnsTyp, answers, func(now time.Time, host string) time.Time {
+		if fixedTtl, ok := c.fixedDomainTtl[host]; ok {
+			/// NOTICE: Cannot set TTL accurately.
+			if now.Sub(deadline).Seconds() > float64(fixedTtl) {
+				return now.Add(time.Duration(fixedTtl) * time.Second)
+			}
+		}
+		return deadline
+	})
+}
+
+func (c *DnsController) UpdateDnsCacheTtl(host string, dnsTyp string, answers []dnsmessage.Resource, ttl int) (err error) {
+	return c.__updateDnsCacheDeadline(host, dnsTyp, answers, func(now time.Time, host string) time.Time {
+		if fixedTtl, ok := c.fixedDomainTtl[host]; ok {
+			return now.Add(time.Duration(fixedTtl) * time.Second)
+		} else {
+			return now.Add(time.Duration(ttl) * time.Second)
+		}
+	})
+}
+
 func (c *DnsController) DnsRespHandlerFactory(validateRushAnsFunc func(from netip.AddrPort) bool) func(data []byte, from netip.AddrPort) (msg *dnsmessage.Message, err error) {
 	return func(data []byte, from netip.AddrPort) (msg *dnsmessage.Message, err error) {
 		// Do not return conn-unrelated err in this func.

docs/dns.md

@@ -6,9 +6,17 @@ dae will intercept all UDP traffic to port 53 and sniff DNS. Here gives some exa
 
 ```shell
 dns {
-    # For example, if ipversion_prefer is 4 and the domain name has both type A and type AAAA records, the dae will only respond to type A queries and response empty answer to type AAAA queries.
+    # For example, if ipversion_prefer is 4 and the domain name has both type A and type AAAA records, the dae will only
+    # respond to type A queries and response empty answer to type AAAA queries.
     ipversion_prefer: 4
 
+    # Give a fixed ttl for domains. Zero means that dae will request to upstream every time and not cache DNS results
+    # for these domains.
+    fixed_domain_ttl {
+        ddns.example.org: 10
+        test.example.org: 3600
+    }
+
     upstream {
         # Value can be scheme://host:port.
         # Scheme list: tcp, udp, tcp+udp. Ongoing: https, tls, quic.

example.dae

@@ -108,9 +108,17 @@ node {
 
 # See https://github.com/daeuniverse/dae/blob/main/docs/dns.md for full examples.
 dns {
-    # For example, if ipversion_prefer is 4 and the domain name has both type A and type AAAA records, the dae will only respond to type A queries and response empty answer to type AAAA queries.
+    # For example, if ipversion_prefer is 4 and the domain name has both type A and type AAAA records, the dae will only
+    # respond to type A queries and response empty answer to type AAAA queries.
     #ipversion_prefer: 4
 
+    # Give a fixed ttl for domains. Zero means that dae will request to upstream every time and not cache DNS results
+    # for these domains.
+    #fixed_domain_ttl {
+    #    ddns.example.org: 10
+    #    test.example.org: 3600
+    #}
+
     upstream {
         # Value can be scheme://host:port, where the scheme can be tcp/udp/tcp+udp.
         # If host is a domain and has both IPv4 and IPv6 record, dae will automatically choose