Cache based on query and ECS subnet address (#73)

cache.go

@@ -83,7 +83,7 @@ func (r *Cache) Resolve(q *dns.Msg, ci ClientInfo) (*dns.Msg, error) {
 	}
 
 	// Put the upstream response into the cache and return it
-	r.storeInCache(a)
+	r.storeInCache(q, a)
 	return a, nil
 }
 
@@ -93,11 +93,10 @@ func (r *Cache) String() string {
 
 // Returns an answer from the cache with it's TTL updated or false in case of a cache-miss.
 func (r *Cache) answerFromCache(q *dns.Msg) (*dns.Msg, bool) {
-	question := q.Question[0]
 	var answer *dns.Msg
 	var timestamp time.Time
 	r.mu.Lock()
-	if a := r.lru.get(question); a != nil {
+	if a := r.lru.get(q); a != nil {
 		answer = a.Copy()
 		timestamp = a.timestamp
 	}
@@ -123,7 +122,7 @@ func (r *Cache) answerFromCache(q *dns.Msg) (*dns.Msg, bool) {
 			}
 			h := a.Header()
 			if age >= h.Ttl {
-				r.evictFromCache(question)
+				r.evictFromCache(q)
 				return nil, false
 			}
 			h.Ttl -= age
@@ -133,7 +132,7 @@ func (r *Cache) answerFromCache(q *dns.Msg) (*dns.Msg, bool) {
 	return answer, true
 }
 
-func (r *Cache) storeInCache(answer *dns.Msg) {
+func (r *Cache) storeInCache(query, answer *dns.Msg) {
 	now := time.Now()
 
 	// Prepare an item for the cache, without expiry for now
@@ -161,14 +160,14 @@ func (r *Cache) storeInCache(answer *dns.Msg) {
 
 	// Store it in the cache
 	r.mu.Lock()
-	r.lru.add(item)
+	r.lru.add(query, item)
 	r.mu.Unlock()
 }
 
-func (r *Cache) evictFromCache(questions ...dns.Question) {
+func (r *Cache) evictFromCache(queries ...*dns.Msg) {
 	r.mu.Lock()
-	for _, question := range questions {
-		r.lru.delete(question)
+	for _, query := range queries {
+		r.lru.delete(query)
 	}
 	r.mu.Unlock()
 }

cmd/routedns/example-config/cache.toml

@@ -14,8 +14,3 @@ cache-negative-ttl = 10         # Optional, TTL to apply to responses without a
 address = "127.0.0.1:53"
 protocol = "udp"
 resolver = "cloudflare-cached"
-
-[listeners.local-tcp]
-address = "127.0.0.1:53"
-protocol = "tcp"
-resolver = "cloudflare-cached"

doc/configuration.md

@@ -243,7 +243,7 @@ Example config files: [doq-listener.toml](../cmd/routedns/example-config/doq-lis
 
 ### Cache
 
-A cache will store the responses to queries in memory and respond to further identical queries with the same response. To determine how long an item is kept in memory, the cache uses the lowest TTL of the RRs in the response. Responses served from the cache have their TTL updated according to the time the records spent in memory.
+A cache will store the responses to queries in memory and respond to further identical queries with the same response. To determine how long an item is kept in memory, the cache uses the lowest TTL of the RRs in the response. Responses served from the cache have their TTL updated according to the time the records spent in memory. If a query has an [ECS Subnet](https://tools.ietf.org/html/rfc7871) option, the subnet address forms part of they key to support subnet-specific answers.
 
 Caches can be combined with a [TTL Modifier](#TTL-Modifier) to avoid too many cache-misses due to excessively low TTL values.
 

lru-cache.go

@@ -8,15 +8,21 @@ import (
 
 type lruCache struct {
 	maxItems   int
-	items      map[dns.Question]*cacheItem
+	items      map[lruKey]*cacheItem
 	head, tail *cacheItem
 }
 
 type cacheItem struct {
+	key lruKey
 	*cacheAnswer
 	prev, next *cacheItem
 }
 
+type lruKey struct {
+	question dns.Question
+	net      string
+}
+
 type cacheAnswer struct {
 	timestamp time.Time // Time the record was cached. Needed to adjust TTL
 	expiry    time.Time // Time the record expires and should be removed
@@ -31,33 +37,34 @@ func newLRUCache(capacity int) *lruCache {
 
 	return &lruCache{
 		maxItems: capacity,
-		items:    make(map[dns.Question]*cacheItem),
+		items:    make(map[lruKey]*cacheItem),
 		head:     head,
 		tail:     tail,
 	}
 }
 
-func (c *lruCache) add(answer *cacheAnswer) {
-	question := answer.Question[0]
-	item := c.touch(question)
+func (c *lruCache) add(query *dns.Msg, answer *cacheAnswer) {
+	key := lruKeyFromQuery(query)
+	item := c.touch(key)
 	if item != nil {
 		return
 	}
 	// Add new item to the top of the linked list
 	item = &cacheItem{
+		key:         key,
 		cacheAnswer: answer,
 		next:        c.head.next,
 		prev:        c.head,
 	}
 	c.head.next.prev = item
 	c.head.next = item
-	c.items[question] = item
+	c.items[key] = item
 	c.resize()
 }
 
 // Loads a cache item and puts it to the top of the queue (most recent).
-func (c *lruCache) touch(question dns.Question) *cacheItem {
-	item := c.items[question]
+func (c *lruCache) touch(key lruKey) *cacheItem {
+	item := c.items[key]
 	if item == nil {
 		return nil
 	}
@@ -71,25 +78,27 @@ func (c *lruCache) touch(question dns.Question) *cacheItem {
 	return item
 }
 
-func (c *lruCache) delete(question dns.Question) {
-	item := c.items[question]
+func (c *lruCache) delete(q *dns.Msg) {
+	key := lruKeyFromQuery(q)
+	item := c.items[key]
 	if item == nil {
 		return
 	}
 	item.prev.next = item.next
 	item.next.prev = item.prev
-	delete(c.items, item.Question[0])
+	delete(c.items, key)
 }
 
-func (c *lruCache) get(question dns.Question) *cacheAnswer {
-	item := c.touch(question)
+func (c *lruCache) get(query *dns.Msg) *cacheAnswer {
+	key := lruKeyFromQuery(query)
+	item := c.touch(key)
 	if item != nil {
 		return item.cacheAnswer
 	}
 	return nil
 }
 
-// Shrink the cache down to the maximum number of itmes.
+// Shrink the cache down to the maximum number of items.
 func (c *lruCache) resize() {
 	if c.maxItems <= 0 { // no size limit
 		return
@@ -99,7 +108,7 @@ func (c *lruCache) resize() {
 		item := c.tail.prev
 		item.prev.next = c.tail
 		c.tail.prev = item.prev
-		delete(c.items, item.Question[0])
+		delete(c.items, item.key)
 	}
 }
 
@@ -111,7 +120,7 @@ func (c *lruCache) deleteFunc(f func(*cacheAnswer) bool) {
 		if f(item.cacheAnswer) {
 			item.prev.next = item.next
 			item.next.prev = item.prev
-			delete(c.items, item.Question[0])
+			delete(c.items, item.key)
 		}
 		item = item.next
 	}
@@ -120,3 +129,18 @@ func (c *lruCache) deleteFunc(f func(*cacheAnswer) bool) {
 func (c *lruCache) size() int {
 	return len(c.items)
 }
+
+func lruKeyFromQuery(q *dns.Msg) lruKey {
+	key := lruKey{question: q.Question[0]}
+
+	edns0 := q.IsEdns0()
+	if edns0 != nil {
+		// See if we have a subnet option
+		for _, opt := range edns0.Option {
+			if subnet, ok := opt.(*dns.EDNS0_SUBNET); ok {
+				key.net = subnet.Address.String()
+			}
+		}
+	}
+	return key
+}

lru-cache_test.go

@@ -12,7 +12,12 @@ import (
 func TestLRUAddGet(t *testing.T) {
 	c := newLRUCache(5)
 
-	var answers []*cacheAnswer
+	type item struct {
+		query  *dns.Msg
+		answer *cacheAnswer
+	}
+	var items []item
+
 	for i := 0; i < 10; i++ {
 		msg := new(dns.Msg)
 		msg.SetQuestion(fmt.Sprintf("test%d.com.", i), dns.TypeA)
@@ -27,28 +32,31 @@ func TestLRUAddGet(t *testing.T) {
 				A: net.IP{127, 0, 0, 1},
 			},
 		}
-		item := &cacheAnswer{Msg: msg}
-		answers = append(answers, item)
+		answer := &cacheAnswer{Msg: msg}
+		items = append(items, item{
+			query:  msg,
+			answer: answer,
+		})
 		// Load into the cache
-		c.add(item)
+		c.add(msg, answer)
 	}
 
 	// Since the capacity is only 5 and we loaded 10, only the last 5 should be in there
 	require.Equal(t, 5, c.size())
 
 	// Check it's the right items in the cache
-	for _, item := range answers[:5] {
-		answer := c.get(item.Question[0])
+	for _, item := range items[:5] {
+		answer := c.get(item.query)
 		require.Nil(t, answer)
 	}
-	for _, item := range answers[5:] {
-		answer := c.get(item.Question[0])
+	for _, item := range items[5:] {
+		answer := c.get(item.query)
 		require.NotNil(t, answer)
-		require.Equal(t, item, answer)
+		require.Equal(t, item.answer, answer)
 	}
 
 	// Delete one of the items directly
-	c.delete(answers[5].Question[0])
+	c.delete(items[5].query)
 	require.Equal(t, 4, c.size())
 
 	// Use an iterator to delete two more