basic DNS cache

unit tests

nameserver/cache.go

@@ -0,0 +1,370 @@
+package nameserver
+
+import (
+	"container/heap"
+	"errors"
+	"github.com/miekg/dns"
+	. "github.com/zettio/weave/common"
+	"math"
+	"math/rand"
+	"sync"
+	"time"
+)
+
+var (
+	errInvalidCapacity = errors.New("Invalid cache capacity")
+	errCouldNotResolve = errors.New("Could not resolve")
+	errTimeout         = errors.New("Timeout while waiting for resolution")
+	errNoLocalReplies  = errors.New("No local replies")
+)
+
+const (
+	defPendingTimeout int = 5  // timeout for a resolution
+)
+
+type entryStatus uint8
+
+const (
+	stPending  entryStatus = iota // someone is waiting for the resolution
+	stResolved entryStatus = iota // resolved
+)
+
+const (
+	CacheNoLocalReplies uint8 = 1 << iota // not found in local network (stored in the cache so we skip another local lookup or some time)
+)
+
+// shuffleAnswers reorders answers for very basic load balancing
+func shuffleAnswers(answers []dns.RR) []dns.RR {
+	n := len(answers)
+	if n > 1 {
+		rand.Seed(time.Now().UTC().UnixNano())
+
+		for i := 0; i < n; i++ {
+			r := i + rand.Intn(n-i)
+			answers[r], answers[i] = answers[i], answers[r]
+		}
+	}
+
+	return answers
+}
+
+// a cache entry
+type cacheEntry struct {
+	Status   entryStatus // status of the entry
+	Flags    uint8       // some extra flags
+	ReplyLen int
+
+	question dns.Question
+	protocol dnsProtocol
+	reply    dns.Msg
+
+	validUntil time.Time // obtained from the reply and stored here for convenience/speed
+	putTime    time.Time
+
+	waitChan chan struct{}
+
+	index int // for fast lookups in the heap
+}
+
+func newCacheEntry(question *dns.Question, reply *dns.Msg, status entryStatus, flags uint8, now time.Time) *cacheEntry {
+	e := &cacheEntry{
+		Status:   status,
+		Flags:    flags,
+		question: *question,
+		index:    -1,
+	}
+
+	if e.Status == stPending {
+		e.validUntil = now.Add(time.Duration(defPendingTimeout) * time.Second)
+		e.waitChan = make(chan struct{})
+	} else {
+		e.setReply(reply, flags, now)
+	}
+
+	return e
+}
+
+// Get a copy of the reply stored in the entry, but with some values adjusted like the TTL
+func (e *cacheEntry) getReply(request *dns.Msg, maxLen int, now time.Time) (*dns.Msg, error) {
+	if e.Status != stResolved {
+		return nil, nil
+	}
+
+	// if the reply has expired or is invalid, force the caller to start a new resolution
+	if e.hasExpired(now) {
+		return nil, nil
+	}
+
+	if e.Flags&CacheNoLocalReplies != 0 {
+		return nil, errNoLocalReplies
+	}
+
+	if e.ReplyLen >= maxLen {
+		Debug.Printf("[cache] returning truncated reponse: %d > %d", e.ReplyLen, maxLen)
+		return makeTruncatedReply(request), nil
+	}
+
+	// create a copy of the reply, with values for this particular query
+	reply := e.reply
+	reply.SetReply(request)
+	reply.Rcode = e.reply.Rcode
+	reply.Authoritative = true
+
+	// adjust the TTLs
+	passedSecs := uint32(now.Sub(e.putTime).Seconds())
+	for _, rr := range reply.Answer {
+		ttl := rr.Header().Ttl
+		if passedSecs < ttl {
+			rr.Header().Ttl = ttl - passedSecs
+		} else {
+			rr.Header().Ttl = 0
+		}
+	}
+
+	// shuffle the values, etc...
+	reply.Answer = shuffleAnswers(reply.Answer)
+
+	return &reply, nil
+}
+
+func (e cacheEntry) hasExpired(now time.Time) bool {
+	return e.validUntil.Before(now) || e.validUntil == now
+}
+
+// set the reply for the entry
+// returns True if the entry has changed the validUntil time
+func (e *cacheEntry) setReply(reply *dns.Msg, flags uint8, now time.Time) bool {
+	shouldNotify := (e.Status == stPending)
+
+	var prevValidUntil time.Time
+	if e.Status == stResolved {
+		prevValidUntil = e.validUntil
+	}
+
+	e.Status = stResolved
+	e.Flags = flags
+	e.putTime = now
+
+	if e.Flags&CacheNoLocalReplies != 0 {
+		// use a fixed timeout for negative local resolutions
+		e.validUntil = now.Add(time.Second * time.Duration(negLocalTTL))
+	} else {
+		// calculate the validUntil from the reply TTL
+		var minTtl uint32 = math.MaxUint32
+		for _, rr := range reply.Answer {
+			ttl := rr.Header().Ttl
+			if ttl < minTtl {
+				minTtl = ttl // TODO: improve the minTTL calculation (maybe we should skip some RRs)
+			}
+		}
+		e.validUntil = now.Add(time.Second * time.Duration(minTtl))
+	}
+
+	if reply != nil {
+		e.reply = *reply
+		e.ReplyLen = reply.Len()
+	}
+
+	if shouldNotify {
+		close(e.waitChan) // notify all the waiters by closing the channel
+	}
+
+	return (prevValidUntil != e.validUntil)
+}
+
+// wait until a valid reply is set in the cache
+func (e *cacheEntry) waitReply(request *dns.Msg, timeout time.Duration, maxLen int, now time.Time) (*dns.Msg, error) {
+	if e.Status == stResolved {
+		return e.getReply(request, maxLen, now)
+	}
+
+	if timeout > 0 {
+		select {
+		case <-e.waitChan:
+			return e.getReply(request, maxLen, now)
+		case <-time.After(time.Second * timeout):
+			return nil, errTimeout
+		}
+	}
+
+	return nil, errCouldNotResolve
+}
+
+func (e *cacheEntry) close() {
+	if e.Status == stPending {
+		close(e.waitChan)
+	}
+}
+
+//////////////////////////////////////////////////////////////////////////////////////
+
+// An entriesPtrHeap is a min-heap of cache entries.
+type entriesPtrsHeap []*cacheEntry
+
+func (h entriesPtrsHeap) Len() int           { return len(h) }
+func (h entriesPtrsHeap) Less(i, j int) bool { return h[i].validUntil.Before(h[j].validUntil) }
+func (h entriesPtrsHeap) Swap(i, j int) {
+	h[i], h[j] = h[j], h[i]
+	h[i].index = i
+	h[j].index = j
+}
+
+func (h *entriesPtrsHeap) Push(x interface{}) {
+	// Push and Pop use pointer receivers because they modify the slice's length,
+	// not just its contents.
+	n := len(*h)
+	entry := x.(*cacheEntry)
+	entry.index = n
+	*h = append(*h, entry)
+}
+
+func (h *entriesPtrsHeap) Pop() interface{} {
+	old := *h
+	n := len(old)
+	item := old[n-1]
+	item.index = -1 // for safety
+	*h = old[0 : n-1]
+	return item
+}
+
+//////////////////////////////////////////////////////////////////////////////////////
+
+type cacheKey dns.Question
+type entries map[cacheKey]*cacheEntry
+
+// Cache is a thread-safe fixed capacity LRU cache.
+type Cache struct {
+	Capacity int
+
+	entries  entries
+	entriesH entriesPtrsHeap // len(entriesH) <= len(entries), as pending entries can be in entries but not in entriesH
+	lock     sync.RWMutex
+}
+
+// NewCache creates a cache of the given capacity
+func NewCache(capacity int) (*Cache, error) {
+	if capacity <= 0 {
+		return nil, errInvalidCapacity
+	}
+	c := &Cache{
+		Capacity: capacity,
+		entries:  make(entries, capacity),
+	}
+
+	heap.Init(&c.entriesH)
+	return c, nil
+}
+
+// Clear removes all the entries in the cache
+func (c *Cache) Clear() {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	c.entries = make(entries, c.Capacity)
+	heap.Init(&c.entriesH)
+}
+
+// Purge removes the old elements in the cache
+func (c *Cache) Purge(now time.Time) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	for i, entry := range c.entriesH {
+		if entry.hasExpired(now) {
+			heap.Remove(&c.entriesH, i)
+			delete(c.entries, cacheKey(entry.question))
+		} else {
+			return // all remaining entries must be still valid...
+		}
+	}
+}
+
+// Add adds a reply to the cache.
+func (c *Cache) Put(request *dns.Msg, reply *dns.Msg, flags uint8, now time.Time) int {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	question := request.Question[0]
+	key := cacheKey(question)
+	ent, found := c.entries[key]
+	if found {
+		Debug.Printf("[cache msgid %d] replacing response in cache", request.MsgHdr.Id)
+		updated := ent.setReply(reply, flags, now)
+		if updated {
+			heap.Fix(&c.entriesH, ent.index)
+		}
+	} else {
+		// If we will add a new item and the capacity has been exceeded, make some room...
+		if len(c.entriesH) >= c.Capacity {
+			lowestEntry := heap.Pop(&c.entriesH).(*cacheEntry)
+			lowestEntry.close()
+			delete(c.entries, cacheKey(lowestEntry.question))
+		}
+		ent = newCacheEntry(&question, reply, stResolved, flags, now)
+		heap.Push(&c.entriesH, ent)
+		c.entries[key] = ent
+	}
+	return ent.ReplyLen
+}
+
+// Look up for a question's reply from the cache.
+// If no reply is stored in the cache, it returns a `nil` reply and no error. The caller can then `Wait()`
+// for another goroutine `Put`ing a reply in the cache.
+func (c *Cache) Get(request *dns.Msg, maxLen int, now time.Time) (reply *dns.Msg, err error) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	question := request.Question[0]
+	key := cacheKey(question)
+	if ent, found := c.entries[key]; found {
+		reply, err = ent.getReply(request, maxLen, now)
+		if ent.hasExpired(now) {
+			Debug.Printf("[cache msgid %d] expired: removing", request.MsgHdr.Id)
+			if ent.index > 0 {
+				heap.Remove(&c.entriesH, ent.index)
+			}
+			delete(c.entries, key)
+			reply = nil
+		}
+	} else {
+		// we are the first asking for this name: create an entry with no reply... the caller must wait
+		Debug.Printf("[cache msgid %d] addind in pending state", request.MsgHdr.Id)
+		c.entries[key] = newCacheEntry(&question, nil, stPending, 0, now)
+	}
+	return
+}
+
+// Wait for a reply for a question in the cache
+// Notice that the caller could Get() and then Wait() for a question, but the corresponding cache
+// entry could have been removed in between. In that case, the caller should retry the query (and
+// the user should increase the cache size!)
+func (c *Cache) Wait(request *dns.Msg, timeout time.Duration, maxLen int, now time.Time) (reply *dns.Msg, err error) {
+	// do not try to lock the cache: otherwise, no one else could `Put()` the reply
+	question := request.Question[0]
+	if entry, found := c.entries[cacheKey(question)]; found {
+		reply, err = entry.waitReply(request, timeout, maxLen, now)
+	}
+	return
+}
+
+// Remove removes the provided question from the cache.
+func (c *Cache) Remove(question *dns.Question) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	key := cacheKey(*question)
+	if entry, found := c.entries[key]; found {
+		if entry.index > 0 {
+			heap.Remove(&c.entriesH, entry.index)
+		}
+		delete(c.entries, key)
+	}
+}
+
+// Len returns the number of entries in the cache.
+func (c *Cache) Len() int {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+
+	return len(c.entries)
+}

nameserver/cache_test.go

@@ -0,0 +1,258 @@
+package nameserver
+
+import (
+	"fmt"
+	"github.com/miekg/dns"
+	. "github.com/zettio/weave/common"
+	wt "github.com/zettio/weave/testing"
+	"net"
+	"testing"
+	"time"
+)
+
+// Check that the cache keeps its intended capacity constant
+func TestCacheLength(t *testing.T) {
+	InitDefaultLogging(true)
+
+	const cacheLen = 128
+
+	l, err := NewCache(cacheLen)
+	wt.AssertNoErr(t, err)
+
+	insTime := time.Now()
+
+	t.Logf("Inserting 256 questions in the cache at '%s', with TTL from 0 to 255", insTime)
+	for i := 0; i < cacheLen*2; i++ {
+		questionMsg := new(dns.Msg)
+		questionMsg.SetQuestion(fmt.Sprintf("name%d", i), dns.TypeA)
+		questionMsg.RecursionDesired = true
+
+		question := &questionMsg.Question[0]
+
+		ip := net.ParseIP(fmt.Sprintf("10.0.1.%d", i))
+		ips := []net.IP{ip}
+		reply := makeAddressReply(questionMsg, question, ips)
+		reply.Answer[0].Header().Ttl = uint32(i)
+
+		l.Put(questionMsg, reply, 0, insTime)
+	}
+
+	wt.AssertEqualInt(t, l.Len(), cacheLen, "cache length")
+
+	minExpectedTime := insTime.Add(time.Duration(cacheLen) * time.Second)
+	t.Logf("Checking all remaining entries expire after insert_time + %d secs='%s'", cacheLen, minExpectedTime)
+	for _, entry := range l.entries {
+		if entry.validUntil.Before(minExpectedTime) {
+			t.Fatalf("Entry valid until %s", entry.validUntil)
+		}
+	}
+}
+
+// Check that the cache entries are ok
+func TestCacheEntries(t *testing.T) {
+	InitDefaultLogging(true)
+
+	Info.Println("Checking cache consistency")
+
+	const cacheLen = 128
+
+	l, err := NewCache(cacheLen)
+	wt.AssertNoErr(t, err)
+
+	questionMsg := new(dns.Msg)
+	questionMsg.SetQuestion("some.name", dns.TypeA)
+	questionMsg.RecursionDesired = true
+
+	question := &questionMsg.Question[0]
+
+	t.Logf("Trying to get a name")
+	resp, err := l.Get(questionMsg, minUdpSize, time.Now())
+	wt.AssertNoErr(t, err)
+	if resp != nil {
+		t.Logf("Got '%s'", resp)
+		t.Fatalf("ERROR: Did not expect a reponse from Get() yet")
+	}
+	t.Logf("Trying to get it again")
+	resp, err = l.Get(questionMsg, minUdpSize, time.Now())
+	wt.AssertNoErr(t, err)
+	if resp != nil {
+		t.Logf("Got '%s'", resp)
+		t.Fatalf("ERROR: Did not expect a reponse from Get() yet")
+	}
+
+	t.Logf("Inserting the reply")
+	reply1 := makeAddressReply(questionMsg, question, []net.IP{net.ParseIP("10.0.1.1")})
+	l.Put(questionMsg, reply1, 0, time.Now())
+
+	timeGet1 := time.Now()
+	t.Logf("Checking we can Get() the reply now")
+	resp, err = l.Get(questionMsg, minUdpSize, timeGet1)
+	wt.AssertNoErr(t, err)
+	wt.AssertTrue(t, resp != nil, "reponse from Get()")
+	t.Logf("Received '%s'", resp.Answer[0])
+	wt.AssertType(t, resp.Answer[0], (*dns.A)(nil), "DNS record")
+	ttlGet1 := resp.Answer[0].Header().Ttl
+
+	t.Logf("Checking a Wait() with timeout=0 gets the same result")
+	resp, err = l.Wait(questionMsg, time.Duration(0)*time.Second, minUdpSize, time.Now())
+	wt.AssertNoErr(t, err)
+	wt.AssertTrue(t, resp != nil, "reponse from a Wait(timeout=0)")
+	t.Logf("Received '%s'", resp.Answer[0])
+	wt.AssertType(t, resp.Answer[0], (*dns.A)(nil), "DNS record")
+
+	timeGet2 := timeGet1.Add(time.Duration(1) * time.Second)
+	t.Logf("Checking that a second Get(), after 1 second, gets the same result, but with reduced TTL")
+	resp, err = l.Get(questionMsg, minUdpSize, timeGet2)
+	wt.AssertNoErr(t, err)
+	wt.AssertTrue(t, resp != nil, "reponse from a second Get()")
+	t.Logf("Received '%s'", resp.Answer[0])
+	wt.AssertType(t, resp.Answer[0], (*dns.A)(nil), "DNS record")
+	ttlGet2 := resp.Answer[0].Header().Ttl
+	wt.AssertEqualInt(t, int(ttlGet1 - ttlGet2), 1, "TTL difference")
+
+	timeGet3 := timeGet1.Add(time.Duration(localTTL) * time.Second)
+	t.Logf("Checking that a third Get(), after %d second, gets no result", localTTL)
+	resp, err = l.Get(questionMsg, minUdpSize, timeGet3)
+	wt.AssertNoErr(t, err)
+	if resp != nil {
+		t.Logf("Got '%s'", resp)
+		t.Fatalf("ERROR: Did NOT expect a reponse from the second Get()")
+	}
+
+	t.Logf("Checking that an Remove() results in Get() returning nothing")
+	replyTemp := makeAddressReply(questionMsg, question, []net.IP{net.ParseIP("10.0.9.9")})
+	l.Put(questionMsg, replyTemp, 0, time.Now())
+	lenBefore := l.Len()
+	l.Remove(question)
+	wt.AssertEqualInt(t, l.Len(), lenBefore - 1, "cache length")
+	l.Remove(question) // do it again: should have no effect...
+	wt.AssertEqualInt(t, l.Len(), lenBefore - 1, "cache length")
+
+	resp, err = l.Get(questionMsg, minUdpSize, timeGet1)
+	wt.AssertNoErr(t, err)
+	wt.AssertTrue(t, resp == nil, "reponse from the Get() after a Remove()")
+
+	t.Logf("Inserting a two replies for the same query")
+	timePut2 := time.Now()
+	reply2 := makeAddressReply(questionMsg, question, []net.IP{net.ParseIP("10.0.1.2")})
+	l.Put(questionMsg, reply2, 0, timePut2)
+	timePut3 := timePut2.Add(time.Duration(1) * time.Second)
+	reply3 := makeAddressReply(questionMsg, question, []net.IP{net.ParseIP("10.0.1.3")})
+	l.Put(questionMsg, reply3, 0, timePut3)
+
+	t.Logf("Checking we get the last one...")
+	resp, err = l.Get(questionMsg, minUdpSize, timePut3)
+	wt.AssertNoErr(t, err)
+	wt.AssertTrue(t, resp != nil, "reponse from the Get()")
+	t.Logf("Received '%s'", resp.Answer[0])
+	wt.AssertType(t, resp.Answer[0], (*dns.A)(nil), "DNS record")
+	wt.AssertEqualString(t, resp.Answer[0].(*dns.A).A.String(), "10.0.1.3", "IP address")
+	wt.AssertEqualInt(t, int(resp.Answer[0].Header().Ttl), int(localTTL), "TTL")
+
+	resp, err = l.Get(questionMsg, minUdpSize, timePut3.Add(time.Duration(localTTL - 1) * time.Second))
+	wt.AssertNoErr(t, err)
+	wt.AssertTrue(t, resp != nil, "reponse from the Get()")
+	t.Logf("Received '%s'", resp.Answer[0])
+	wt.AssertType(t, resp.Answer[0], (*dns.A)(nil), "DNS record")
+	wt.AssertEqualString(t, resp.Answer[0].(*dns.A).A.String(), "10.0.1.3", "IP address")
+	wt.AssertEqualInt(t, int(resp.Answer[0].Header().Ttl), 1, "TTL")
+
+	t.Logf("Checking we get empty replies when they are expired...")
+	lenBefore = l.Len()
+	resp, err = l.Get(questionMsg, minUdpSize, timePut3.Add(time.Duration(localTTL) * time.Second))
+	wt.AssertNoErr(t, err)
+	if resp != nil {
+		t.Logf("Received '%s'", resp.Answer[0])
+		t.Fatalf("ERROR: Did NOT expect a reponse from the Get()")
+	}
+	wt.AssertEqualInt(t, l.Len(), lenBefore - 1, "cache length (after getting an expired entry)")
+
+	questionMsg2 := new(dns.Msg)
+	questionMsg2.SetQuestion("some.other.name", dns.TypeA)
+	questionMsg2.RecursionDesired = true
+	question2 := &questionMsg2.Question[0]
+
+	t.Logf("Trying to Get() a name")
+	resp, err = l.Get(questionMsg2, minUdpSize, time.Now())
+	wt.AssertNoErr(t, err)
+	wt.AssertNil(t, resp, "reponse from Get() yet")
+
+	t.Logf("Checking that an Remove() between Get() and Put() does not break things")
+	replyTemp2 := makeAddressReply(questionMsg2, question2, []net.IP{net.ParseIP("10.0.9.9")})
+	l.Remove(question2)
+	l.Put(questionMsg2, replyTemp2, 0, time.Now())
+	resp, err = l.Get(questionMsg2, minUdpSize, time.Now())
+	wt.AssertNoErr(t, err)
+	wt.AssertNotNil(t, resp, "reponse from Get()")
+	resp, err = l.Wait(questionMsg2, time.Duration(0)*time.Second, minUdpSize, time.Now())
+	wt.AssertNoErr(t, err)
+	wt.AssertNotNil(t, resp, "reponse from Get()")
+
+	questionMsg3 := new(dns.Msg)
+	questionMsg3.SetQuestion("some.other.name", dns.TypeA)
+	questionMsg3.RecursionDesired = true
+	question3 := &questionMsg3.Question[0]
+
+	t.Logf("Checking that a entry with CacheNoLocalReplies return an error")
+	timePut3 = time.Now()
+	l.Put(questionMsg3, nil, CacheNoLocalReplies, timePut3)
+	resp, err = l.Get(questionMsg3, minUdpSize, timePut3)
+	wt.AssertNil(t, resp, "Get() response with CacheNoLocalReplies")
+	wt.AssertNotNil(t, err, "Get() error with CacheNoLocalReplies")
+
+	timeExpiredGet3 := timePut3.Add(time.Second * time.Duration(negLocalTTL + 1))
+	t.Logf("Checking that we get an expired response after %f seconds", timeExpiredGet3.Sub(timePut3).Seconds())
+	resp, err = l.Get(questionMsg3, minUdpSize, timeExpiredGet3)
+	wt.AssertNil(t, resp, "expired Get() response with CacheNoLocalReplies")
+	wt.AssertNil(t, err, "expired Get() error with CacheNoLocalReplies")
+
+	l.Remove(question3)
+	t.Logf("Checking that Put&Get with CacheNoLocalReplies with a Remove in the middle returns nothing")
+	l.Put(questionMsg3, nil, CacheNoLocalReplies, time.Now())
+	l.Remove(question3)
+	resp, err = l.Get(questionMsg3, minUdpSize, time.Now())
+	wt.AssertNil(t, resp, "Get() reponse with CacheNoLocalReplies")
+	wt.AssertNil(t, err, "Get() error with CacheNoLocalReplies")
+}
+
+// Check that waiters are unblocked when the name they are waiting for is inserted
+func TestCacheBlockingOps(t *testing.T) {
+	InitDefaultLogging(true)
+
+	const cacheLen = 256
+
+	l, err := NewCache(cacheLen)
+	wt.AssertNoErr(t, err)
+
+	requests := []*dns.Msg{}
+
+	t.Logf("Starting 256 queries that will block...")
+	for i := 0; i < cacheLen; i++ {
+		questionName := fmt.Sprintf("name%d", i)
+		questionMsg := new(dns.Msg)
+		questionMsg.SetQuestion(questionName, dns.TypeA)
+		questionMsg.RecursionDesired = true
+
+		requests = append(requests, questionMsg)
+
+		go func(request *dns.Msg) {
+			t.Logf("Querying about %s...", request.Question[0].Name)
+			_, err := l.Get(request, minUdpSize, time.Now())
+			wt.AssertNoErr(t, err)
+			t.Logf("Waiting for %s...", request.Question[0].Name)
+			r, err := l.Wait(request, 1*time.Second, minUdpSize, time.Now())
+			t.Logf("Obtained response for %s:\n%s", request.Question[0].Name, r)
+			wt.AssertNoErr(t, err)
+		}(questionMsg)
+	}
+
+	// insert the IPs for those names
+	for i, requestMsg := range requests {
+		ip := net.ParseIP(fmt.Sprintf("10.0.1.%d", i))
+		ips := []net.IP{ip}
+		reply := makeAddressReply(requestMsg, &requestMsg.Question[0], ips)
+
+		t.Logf("Inserting response for %s...", requestMsg.Question[0].Name)
+		l.Put(requestMsg, reply, 0, time.Now())
+	}
+}

nameserver/dns.go

@@ -6,8 +6,10 @@ import (
 )
 
 const (
-	localTTL uint32 = 300 // somewhat arbitrary; we don't expect anyone
-	// downstream to cache results
+	localTTL   uint32 = 30 // somewhat arbitrary; we don't expect anyone downstream to cache results
+	negLocalTTL       = 30 // TTL for negative local resolutions
+	minUdpSize        = 512
+	maxUdpSize        = 65535
 )
 
 func makeHeader(r *dns.Msg, q *dns.Question) *dns.RR_Header {
@@ -24,6 +26,14 @@ func makeReply(r *dns.Msg, as []dns.RR) *dns.Msg {
 	return m
 }
 
+func makeTruncatedReply(r *dns.Msg) *dns.Msg {
+	// for truncated response, we create a minimal reply with the Truncated bit set
+	reply := new(dns.Msg)
+	reply.SetReply(r)
+	reply.Truncated = true
+	return reply
+}
+
 func makeAddressReply(r *dns.Msg, q *dns.Question, addrs []net.IP) *dns.Msg {
 	answers := make([]dns.RR, len(addrs))
 	header := makeHeader(r, q)
@@ -53,3 +63,30 @@ func makePTRReply(r *dns.Msg, q *dns.Question, names []string) *dns.Msg {
 	}
 	return makeReply(r, answers)
 }
+
+func makeDNSFailResponse(r *dns.Msg) *dns.Msg {
+	m := new(dns.Msg)
+	m.SetReply(r)
+	m.RecursionAvailable = true
+	m.Rcode = dns.RcodeNameError
+	return m
+}
+
+func makeDNSNotImplResponse(r *dns.Msg) *dns.Msg {
+	m := new(dns.Msg)
+	m.SetReply(r)
+	m.RecursionAvailable = true
+	m.Rcode = dns.RcodeNotImplemented
+	return m
+}
+
+// get the maximum UDP-reply length
+func getMaxReplyLen(r *dns.Msg, proto dnsProtocol) int {
+	maxLen := minUdpSize
+	if proto == protTcp {
+		maxLen = maxUdpSize
+	} else if opt := r.IsEdns0(); opt != nil {
+		maxLen = int(opt.UDPSize())
+	}
+	return maxLen
+}

nameserver/mdns_server.go

@@ -9,6 +9,7 @@ import (
 type MDNSServer struct {
 	localAddrs []net.Addr
 	sendconn   *net.UDPConn
+	srv        *dns.Server
 	zone       Zone
 }
 
@@ -72,10 +73,19 @@ func (s *MDNSServer) Start(ifi *net.Interface, localDomain string) error {
 	mux.HandleFunc(localDomain, handleLocal)
 	mux.HandleFunc(RDNS_DOMAIN, handleReverse)
 
-	go dns.ActivateAndServe(nil, conn, mux)
+	s.srv = &dns.Server{
+		Listener:   nil,
+		PacketConn: conn,
+		Handler:    mux,
+	}
+	go s.srv.ActivateAndServe()
 	return err
 }
 
+func (s *MDNSServer) Stop() error {
+	return s.srv.Shutdown()
+}
+
 type LookupFunc func(Lookup, *dns.Msg, *dns.Question) *dns.Msg
 
 func (s *MDNSServer) makeHandler(qtype uint16, lookup LookupFunc) dns.HandlerFunc {

nameserver/server.go

@@ -1,30 +1,27 @@
 package nameserver
 
 import (
+	"bytes"
 	"fmt"
 	"github.com/miekg/dns"
 	. "github.com/zettio/weave/common"
 	"net"
 	"sync"
-	"bytes"
+	"time"
 )
 
 const (
-	DEFAULT_LOCAL_DOMAIN = "weave.local."     // The default name used for the local domain
-	DEFAULT_SERVER_PORT  = 53                 // The default server port
-	DEFAULT_CLI_CFG_FILE = "/etc/resolv.conf" // default "resolv.conf" file to try to load
-	DEFAULT_UDP_BUFLEN   = 4096               // bigger than the default 512
-	DEFAULT_IFACE_NAME   = "default interface"
+	DEFAULT_LOCAL_DOMAIN   = "weave.local."     // The default name used for the local domain
+	DEFAULT_SERVER_PORT    = 53                 // The default server port
+	DEFAULT_CLI_CFG_FILE   = "/etc/resolv.conf" // default "resolv.conf" file to try to load
+	DEFAULT_UDP_BUFLEN     = 4096               // bigger than the default 512
+	DEFAULT_CACHE_LEN      = 8192               // default cache capacity
+	DEFAULT_RESOLV_WORKERS = 8                  // default number of resolution workers
+	DEFAULT_TIMEOUT        = 5                  // default timeout for DNS resolutions
+	DEFAULT_IFACE_NAME     = "default interface"
+	DEFAULT_RESOLV_TRIES   = 3 // max # of times a worker tries to resolve a query
 )
 
-func makeDNSFailResponse(r *dns.Msg) *dns.Msg {
-	m := new(dns.Msg)
-	m.SetReply(r)
-	m.RecursionAvailable = true
-	m.Rcode = dns.RcodeNameError
-	return m
-}
-
 type DNSServerConfig struct {
 	// (Optional) client config file for resolving upstream servers
 	UpstreamCfgFile string
@@ -34,16 +31,54 @@ type DNSServerConfig struct {
 	Port int
 	// (Optional) local domain (ie, "weave.local.")
 	LocalDomain string
+	// (Optional) cache size
+	CacheLen int
+	// (Optional) timeout for DNS queries
+	Timeout int
+	// (Optional) UDP buffer length
+	UdpBufLen int
+}
+
+type dnsProtocol uint8
+
+const (
+	protUdp dnsProtocol = iota // UDP protocol
+	protTcp dnsProtocol = iota // TCP protocol
+)
+
+func (proto dnsProtocol) String() string {
+	switch proto {
+	case protUdp:
+		return "UDP"
+	case protTcp:
+		return "TCP"
+	}
+	return "unknown"
+}
+
+// get a new dns.Client for a protocol
+func (proto dnsProtocol) GetNewClient(bufsize int) *dns.Client {
+	switch proto {
+	case protTcp:
+		return &dns.Client{Net: "tcp"}
+	case protUdp:
+		return &dns.Client{Net: "udp", UDPSize: uint16(bufsize)}
+	}
+	return nil
 }
 
 type DNSServer struct {
-	udpSrv   *dns.Server
-	tcpSrv   *dns.Server
-	mdnsCli  *MDNSClient
-	mdnsSrv  *MDNSServer
-	zone     Zone
-	iface    *net.Interface
-	upstream *dns.ClientConfig
+	udpSrv      *dns.Server
+	tcpSrv      *dns.Server
+	mdnsCli     *MDNSClient
+	mdnsSrv     *MDNSServer
+	cache       *Cache
+	zone        Zone
+	iface       *net.Interface
+	upstream    *dns.ClientConfig
+	timeout     int
+	udpBuf      int
+	listenersWg *sync.WaitGroup
 
 	Domain     string // the local domain
 	IfaceName  string // the interface where mDNS is working on
@@ -53,8 +88,9 @@ type DNSServer struct {
 // Creates a new DNS server
 func NewDNSServer(config DNSServerConfig, zone Zone, iface *net.Interface) (s *DNSServer, err error) {
 	s = &DNSServer{
-		zone:  zone,
-		iface: iface,
+		zone:        zone,
+		iface:       iface,
+		listenersWg: new(sync.WaitGroup),
 
 		Domain:     DEFAULT_LOCAL_DOMAIN,
 		IfaceName:  DEFAULT_IFACE_NAME,
@@ -79,6 +115,12 @@ func NewDNSServer(config DNSServerConfig, zone Zone, iface *net.Interface) (s *D
 			return nil, err
 		}
 	}
+	if config.Timeout > 0 {
+		s.timeout = config.Timeout
+	}
+	if config.UdpBufLen > 0 {
+		s.udpBuf = config.UdpBufLen
+	}
 	s.mdnsCli, err = NewMDNSClient()
 	if err != nil {
 		return
@@ -90,19 +132,28 @@ func NewDNSServer(config DNSServerConfig, zone Zone, iface *net.Interface) (s *D
 	if iface != nil {
 		s.IfaceName = iface.Name
 	}
+	cacheLen := DEFAULT_CACHE_LEN
+	if config.CacheLen > 0 {
+		cacheLen = config.CacheLen
+	}
+	Debug.Printf("[dns] Initializing cache: %d entries", config.CacheLen)
+	s.cache, err = NewCache(cacheLen)
+	if err != nil {
+		return
+	}
 
 	// create two DNS request multiplexerers, depending on the protocol used by clients
 	// (we use the same protocol for asking upstream servers)
-	mux := func(client *dns.Client) *dns.ServeMux {
+	mux := func(proto dnsProtocol) *dns.ServeMux {
 		m := dns.NewServeMux()
-		m.HandleFunc(s.Domain, s.queryHandler([]Lookup{s.zone, s.mdnsCli}))
-		m.HandleFunc(RDNS_DOMAIN, s.rdnsHandler([]Lookup{s.zone, s.mdnsCli}, client))
-		m.HandleFunc(".", s.notUsHandler(client))
+		m.HandleFunc(s.Domain, s.queryHandler([]Lookup{s.zone, s.mdnsCli}, proto))
+		m.HandleFunc(RDNS_DOMAIN, s.rdnsHandler([]Lookup{s.zone, s.mdnsCli}, proto))
+		m.HandleFunc(".", s.notUsHandler(proto))
 		return m
 	}
 
-	s.udpSrv = &dns.Server{Addr: s.ListenAddr, Net: "udp", Handler: mux(&dns.Client{Net: "udp", UDPSize: DEFAULT_UDP_BUFLEN})}
-	s.tcpSrv = &dns.Server{Addr: s.ListenAddr, Net: "tcp", Handler: mux(&dns.Client{Net: "tcp"})}
+	s.udpSrv = &dns.Server{Addr: s.ListenAddr, Net: "udp", Handler: mux(protUdp)}
+	s.tcpSrv = &dns.Server{Addr: s.ListenAddr, Net: "tcp", Handler: mux(protTcp)}
 
 	return
 }
@@ -115,11 +166,10 @@ func (s *DNSServer) Start() error {
 	err = s.mdnsSrv.Start(s.iface, s.Domain)
 	CheckFatal(err)
 
-	wg := new(sync.WaitGroup)
-	wg.Add(2)
+	s.listenersWg.Add(2)
 
 	go func() {
-		defer wg.Done()
+		defer s.listenersWg.Done()
 
 		Debug.Printf("Listening for DNS on %s (UDP)", s.ListenAddr)
 		err = s.udpSrv.ListenAndServe()
@@ -128,7 +178,7 @@ func (s *DNSServer) Start() error {
 	}()
 
 	go func() {
-		defer wg.Done()
+		defer s.listenersWg.Done()
 
 		Debug.Printf("Listening for DNS on %s (TCP)", s.ListenAddr)
 		err = s.tcpSrv.ListenAndServe()
@@ -137,7 +187,7 @@ func (s *DNSServer) Start() error {
 	}()
 
 	// Waiting for all goroutines to finish (otherwise they die as main routine dies)
-	wg.Wait()
+	s.listenersWg.Wait()
 
 	Info.Printf("WeaveDNS server exiting...")
 	return nil
@@ -156,64 +206,136 @@ func (s *DNSServer) Status() string {
 
 // Perform a graceful shutdown
 func (s *DNSServer) Stop() error {
+	// Stop the listeners/handlers
 	if err := s.tcpSrv.Shutdown(); err != nil {
 		return err
 	}
 	if err := s.udpSrv.Shutdown(); err != nil {
 		return err
 	}
-	// TODO: shutdown the mDNS client/server
+	s.listenersWg.Wait()
+
+	// shutdown the mDNS server
+	s.mdnsSrv.Stop()
+
 	return nil
 }
 
-func (s *DNSServer) queryHandler(lookups []Lookup) dns.HandlerFunc {
+func (s *DNSServer) queryHandler(lookups []Lookup, proto dnsProtocol) dns.HandlerFunc {
 	return func(w dns.ResponseWriter, r *dns.Msg) {
+		now := time.Now()
 		q := r.Question[0]
+		maxLen := getMaxReplyLen(r, proto)
 		Debug.Printf("Query: %+v", q)
-		if q.Qtype == dns.TypeA {
-			for _, lookup := range lookups {
-				if ip, err := lookup.LookupName(q.Name); err == nil {
-					m := makeAddressReply(r, &q, []net.IP{ip})
-					w.WriteMsg(m)
-					return
-				}
-			}
+
+		if q.Qtype != dns.TypeA {
+			Debug.Printf("[dns msgid %d] Unsuported query type %s", r.MsgHdr.Id, dns.TypeToString[q.Qtype])
+			m := makeDNSNotImplResponse(r)
+			s.cache.Put(r, m, 0, now)
+			w.WriteMsg(m)
+			return
 		}
+
+		reply, err := s.cache.Get(r, maxLen, time.Now())
+		if err != nil {
+			Debug.Printf("[dns msgid %d] Error from cache: %s", r.MsgHdr.Id, err)
+			w.WriteMsg(makeDNSFailResponse(r))
+			return
+		}
+		if reply != nil {
+			Debug.Printf("[dns msgid %d] Returning reply from cache: %s/%d answers",
+				r.MsgHdr.Id, dns.RcodeToString[reply.MsgHdr.Rcode], len(reply.Answer))
+			w.WriteMsg(reply)
+			return
+		}
+		for _, lookup := range lookups {
+			if ip, err := lookup.LookupName(q.Name); err == nil {
+				m := makeAddressReply(r, &q, []net.IP{ip})
+
+				Debug.Printf("[dns msgid %d] Caching response for %s-query for \"%s\": %s [code:%s]",
+					m.MsgHdr.Id, dns.TypeToString[q.Qtype], q.Name, ip, dns.RcodeToString[m.Rcode])
+				s.cache.Put(r, m, 0, now)
+				w.WriteMsg(m)
+				return
+			}
+			now = time.Now()
+		}
+
 		Info.Printf("[dns msgid %d] No results for type %s query %s",
 			r.MsgHdr.Id, dns.TypeToString[q.Qtype], q.Name)
-		w.WriteMsg(makeDNSFailResponse(r))
+		m := makeDNSFailResponse(r)
+		s.cache.Put(r, m, 0, now)
+		w.WriteMsg(m)
 	}
 }
 
-func (s *DNSServer) rdnsHandler(lookups []Lookup, client *dns.Client) dns.HandlerFunc {
-	fallback := s.notUsHandler(client)
+func (s *DNSServer) rdnsHandler(lookups []Lookup, proto dnsProtocol) dns.HandlerFunc {
+	fallback := s.notUsHandler(proto)
 	return func(w dns.ResponseWriter, r *dns.Msg) {
+		now := time.Now()
 		q := r.Question[0]
+		maxLen := getMaxReplyLen(r, proto)
 		Debug.Printf("Reverse query: %+v", q)
-		if q.Qtype == dns.TypePTR {
-			for _, lookup := range lookups {
-				if name, err := lookup.LookupInaddr(q.Name); err == nil {
-					m := makePTRReply(r, &q, []string{name})
-					w.WriteMsg(m)
-					return
-				}
-			}
-			fallback(w, r)
+
+		if q.Qtype != dns.TypePTR {
+			Warning.Printf("[dns msgid %d] Unexpected reverse query type %s: %+v",
+				r.MsgHdr.Id, dns.TypeToString[q.Qtype], q)
+			m := makeDNSNotImplResponse(r)
+			s.cache.Put(r, m, 0, now)
+			w.WriteMsg(m)
 			return
 		}
-		Warning.Printf("[dns msgid %d] Unexpected reverse query type %s: %+v",
-			r.MsgHdr.Id, dns.TypeToString[q.Qtype], q)
+
+		reply, err := s.cache.Get(r, maxLen, time.Now())
+		if err != nil {
+			if err == errNoLocalReplies {
+				Debug.Printf("[dns msgid %d] Cached 'no local replies' - skipping local lookup", r.MsgHdr.Id)
+				fallback(w, r)
+			} else {
+				Debug.Printf("[dns msgid %d] Error from cache: %s", r.MsgHdr.Id, err)
+				w.WriteMsg(makeDNSFailResponse(r))
+			}
+			return
+		}
+		if reply != nil {
+			Debug.Printf("[dns msgid %d] Returning reply from cache: %s/%d answers",
+				r.MsgHdr.Id, dns.RcodeToString[reply.MsgHdr.Rcode], len(reply.Answer))
+			w.WriteMsg(reply)
+			return
+		}
+		for _, lookup := range lookups {
+			if name, err := lookup.LookupInaddr(q.Name); err == nil {
+				m := makePTRReply(r, &q, []string{name})
+				Debug.Printf("[dns msgid %d] Caching response for %s-query for \"%s\": %s [code:%s]",
+					m.MsgHdr.Id, dns.TypeToString[q.Qtype], q.Name, name, dns.RcodeToString[m.Rcode])
+				s.cache.Put(r, m, 0, now)
+				w.WriteMsg(m)
+				return
+			}
+			now = time.Now()
+		}
+
+		s.cache.Put(r, nil, CacheNoLocalReplies, now)
+		fallback(w, r)
 	}
 }
 
 // When we receive a request for a name outside of our '.weave.local.'
 // domain, ask the configured DNS server as a fallback.
-func (s *DNSServer) notUsHandler(dnsClient *dns.Client) dns.HandlerFunc {
+func (s *DNSServer) notUsHandler(proto dnsProtocol) dns.HandlerFunc {
+	dnsClient := proto.GetNewClient(DEFAULT_UDP_BUFLEN)
+
 	return func(w dns.ResponseWriter, r *dns.Msg) {
+		maxLen := getMaxReplyLen(r, proto)
 		q := r.Question[0]
-		Debug.Printf("[dns msgid %d] Fallback query: %+v", r.MsgHdr.Id, q)
+
+		// create a request where we announce our max payload size
+		rcopy := r
+		rcopy.SetEdns0(uint16(maxLen), false)
+
+		Debug.Printf("[dns msgid %d] Fallback query: %+v [%s, max:%d bytes]", rcopy.MsgHdr.Id, q, proto, maxLen)
 		for _, server := range s.upstream.Servers {
-			reply, _, err := dnsClient.Exchange(r, fmt.Sprintf("%s:%s", server, s.upstream.Port))
+			reply, _, err := dnsClient.Exchange(rcopy, fmt.Sprintf("%s:%s", server, s.upstream.Port))
 			if err != nil {
 				Debug.Printf("[dns msgid %d] Network error trying %s (%s)",
 					r.MsgHdr.Id, server, err)

nameserver/server_test.go

@@ -15,6 +15,7 @@ const (
 	testRDNSsuccess  = "1.2.2.10.in-addr.arpa."
 	testRDNSfail     = "4.3.2.1.in-addr.arpa."
 	testRDNSnonlocal = "8.8.8.8.in-addr.arpa."
+	testUdpBufSize   = 16384
 )
 
 var (
@@ -64,7 +65,7 @@ func TestUDPDNSServer(t *testing.T) {
 	}
 
 	// Run another DNS server for fallback
-	s, fallbackAddr, err := RunLocalUDPServer(t, "127.0.0.1:0", fallbackHandler)
+	s, fallbackAddr, err := runLocalUDPServer(t, "127.0.0.1:0", fallbackHandler)
 	wt.AssertNoErr(t, err)
 	defer s.Shutdown()
 
@@ -80,57 +81,34 @@ func TestUDPDNSServer(t *testing.T) {
 
 	var r *dns.Msg
 
-	r = testExchange(t, successTestName, dns.TypeA, 1, 1)
+	r = assertExchange(t, successTestName, dns.TypeA, 1, 1, 0)
 	wt.AssertType(t, r.Answer[0], (*dns.A)(nil), "DNS record")
 	wt.AssertEqualString(t, r.Answer[0].(*dns.A).A.String(), testAddr1, "IP address")
 
-	testExchange(t, failTestName, dns.TypeA, 0, 0)
+	assertExchange(t, failTestName, dns.TypeA, 0, 0, dns.RcodeNameError)
 
-	r = testExchange(t, testRDNSsuccess, dns.TypePTR, 1, 1)
+	r = assertExchange(t, testRDNSsuccess, dns.TypePTR, 1, 1, 0)
 	wt.AssertType(t, r.Answer[0], (*dns.PTR)(nil), "DNS record")
 	wt.AssertEqualString(t, r.Answer[0].(*dns.PTR).Ptr, successTestName, "IP address")
 
-	testExchange(t, testRDNSfail, dns.TypePTR, 0, 0)
+	assertExchange(t, testRDNSfail, dns.TypePTR, 0, 0, dns.RcodeNameError)
 
 	// This should fail because we don't handle MX records
-	testExchange(t, successTestName, dns.TypeMX, 0, 0)
+	assertExchange(t, successTestName, dns.TypeMX, 0, 0, dns.RcodeNotImplemented)
 
 	// This non-local query for an MX record should succeed by being
 	// passed on to the fallback server
-	testExchange(t, nonLocalName, dns.TypeMX, 1, -1)
+	assertExchange(t, nonLocalName, dns.TypeMX, 1, -1, 0)
 
 	// Now ask a query that we expect to return a lot of data.
-	testExchange(t, nonLocalName, dns.TypeANY, 5, -1)
+	assertExchange(t, nonLocalName, dns.TypeANY, 5, -1, 0)
 
-	testExchange(t, testRDNSnonlocal, dns.TypePTR, 1, -1)
+	assertExchange(t, testRDNSnonlocal, dns.TypePTR, 1, -1, 0)
 
 	// Not testing MDNS functionality of server here (yet), since it
 	// needs two servers, each listening on its own address
 }
 
-func testExchange(t *testing.T, z string, ty uint16, minAnswers int, maxAnswers int) *dns.Msg {
-	c := new(dns.Client)
-	c.UDPSize = DEFAULT_UDP_BUFLEN
-	m := new(dns.Msg)
-	m.RecursionDesired = true
-	m.SetQuestion(z, ty)
-	r, _, err := c.Exchange(m, fmt.Sprintf("127.0.0.1:%d", testPort))
-	wt.AssertNoErr(t, err)
-	if minAnswers == 0 && maxAnswers == 0 {
-		wt.AssertStatus(t, r.Rcode, dns.RcodeNameError, "DNS response code")
-	} else {
-		wt.AssertStatus(t, r.Rcode, dns.RcodeSuccess, "DNS response code")
-	}
-	answers := len(r.Answer)
-	if minAnswers >= 0 && answers < minAnswers {
-		wt.Fatalf(t, "Number of answers >= %d", minAnswers)
-	}
-	if maxAnswers >= 0 && answers > maxAnswers {
-		wt.Fatalf(t, "Number of answers <= %d", maxAnswers)
-	}
-	return r
-}
-
 func TestTCPDNSServer(t *testing.T) {
 	setupForTest()
 	const (
@@ -142,6 +120,7 @@ func TestTCPDNSServer(t *testing.T) {
 	InitDefaultLogging(true)
 	var zone = new(ZoneDb)
 
+	// generate a list of `numAnswers` IP addresses
 	var addrs []net.IP
 	bs := make([]byte, 4)
 	for i := 0; i < numAnswers; i++ {
@@ -151,10 +130,18 @@ func TestTCPDNSServer(t *testing.T) {
 
 	// handler for the fallback server: it will just return a very long response
 	fallbackUDPHandler := func(w dns.ResponseWriter, req *dns.Msg) {
+		maxLen := getMaxReplyLen(req, protUdp)
+
 		t.Logf("Fallback UDP server got asked: returning %d answers", numAnswers)
 		q := req.Question[0]
 		m := makeAddressReply(req, &q, addrs)
-		m.Truncated = true
+		mLen := m.Len()
+		m.SetEdns0(uint16(maxLen), false)
+
+		if mLen > maxLen {
+			t.Logf("... truncated response (%d > %d)", mLen, maxLen)
+			m.Truncated = true
+		}
 		w.WriteMsg(m)
 	}
 	fallbackTCPHandler := func(w dns.ResponseWriter, req *dns.Msg) {
@@ -164,17 +151,17 @@ func TestTCPDNSServer(t *testing.T) {
 		w.WriteMsg(m)
 	}
 
-	// Run another DNS server for fallback
-	us, fallbackUdpAddr, err := RunLocalUDPServer(t, "127.0.0.1:0", fallbackUDPHandler)
+	t.Logf("Running a DNS fallback server with UDP")
+	us, fallbackUdpAddr, err := runLocalUDPServer(t, "127.0.0.1:0", fallbackUDPHandler)
 	wt.AssertNoErr(t, err)
 	defer us.Shutdown()
 
 	_, fallbackPort, err := net.SplitHostPort(fallbackUdpAddr)
 	wt.AssertNoErr(t, err)
 
-	// start the TCP server on the same port as the UDP server
+	t.Logf("Starting another fallback server, with TCP, on the same port as the UDP server")
 	fallbackTcpAddr := fmt.Sprintf("127.0.0.1:%s", fallbackPort)
-	ts, fallbackTcpAddr, err := RunLocalTCPServer(t, fallbackTcpAddr, fallbackTCPHandler)
+	ts, fallbackTcpAddr, err := runLocalTCPServer(t, fallbackTcpAddr, fallbackTCPHandler)
 	wt.AssertNoErr(t, err)
 	defer ts.Shutdown()
 
@@ -188,7 +175,7 @@ func TestTCPDNSServer(t *testing.T) {
 
 	t.Logf("Creating a UDP and a TCP client")
 	uc := new(dns.Client)
-	uc.UDPSize = DEFAULT_UDP_BUFLEN
+	uc.UDPSize = minUdpSize
 	tc := new(dns.Client)
 	tc.Net = "tcp"
 
@@ -199,24 +186,73 @@ func TestTCPDNSServer(t *testing.T) {
 
 	t.Logf("Checking the fallback server at %s returns a truncated response with UDP", fallbackUdpAddr)
 	r, _, err := uc.Exchange(m, fallbackUdpAddr)
+	t.Logf("Got response from fallback server (UDP) with %d answers", len(r.Answer))
+	t.Logf("Response:\n%+v\n", r)
 	wt.AssertNoErr(t, err)
 	wt.AssertTrue(t, r.MsgHdr.Truncated, "DNS truncated reponse flag")
 	wt.AssertNotEqualInt(t, len(r.Answer), numAnswers, "number of answers (UDP)")
 
 	t.Logf("Checking the WeaveDNS server at %s returns a truncated reponse with UDP", dnsAddr)
 	r, _, err = uc.Exchange(m, dnsAddr)
+	t.Logf("UDP Response:\n%+v\n", r)
 	wt.AssertNoErr(t, err)
+	wt.AssertNotNil(t, r, "response")
+	t.Logf("%d answers", len(r.Answer))
 	wt.AssertTrue(t, r.MsgHdr.Truncated, "DNS truncated reponse flag")
 	wt.AssertNotEqualInt(t, len(r.Answer), numAnswers, "number of answers (UDP)")
 
 	t.Logf("Checking the WeaveDNS server at %s does not return a truncated reponse with TCP", dnsAddr)
 	r, _, err = tc.Exchange(m, dnsAddr)
+	t.Logf("TCP Response:\n%+v\n", r)
 	wt.AssertNoErr(t, err)
+	wt.AssertNotNil(t, r, "response")
+	t.Logf("%d answers", len(r.Answer))
 	wt.AssertFalse(t, r.MsgHdr.Truncated, "DNS truncated response flag")
 	wt.AssertEqualInt(t, len(r.Answer), numAnswers, "number of answers (TCP)")
+
+	t.Logf("Checking the WeaveDNS server at %s does not return a truncated reponse with UDP with a bigger buffer", dnsAddr)
+	m.SetEdns0(testUdpBufSize, false)
+	r, _, err = uc.Exchange(m, dnsAddr)
+	t.Logf("UDP-large Response:\n%+v\n", r)
+	wt.AssertNoErr(t, err)
+	wt.AssertNotNil(t, r, "response")
+	t.Logf("%d answers", len(r.Answer))
+	wt.AssertNoErr(t, err)
+	wt.AssertFalse(t, r.MsgHdr.Truncated, "DNS truncated response flag")
+	wt.AssertEqualInt(t, len(r.Answer), numAnswers, "number of answers (UDP-long)")
 }
 
-func RunLocalUDPServer(t *testing.T, laddr string, handler dns.HandlerFunc) (*dns.Server, string, error) {
+//////////////////////////////////////////////////////////////////////////////
+
+// perform a DNS query and assert the reply code, number or answers, etc
+func assertExchange(t *testing.T, z string, ty uint16, minAnswers int, maxAnswers int, expErr int) *dns.Msg {
+	c := new(dns.Client)
+	c.UDPSize = testUdpBufSize
+	m := new(dns.Msg)
+	m.RecursionDesired = true
+	m.SetQuestion(z, ty)
+	m.SetEdns0(testUdpBufSize, false) // we don't want to play with truncation here...
+
+	r, _, err := c.Exchange(m, fmt.Sprintf("127.0.0.1:%d", testPort))
+	t.Logf("Response:\n%+v\n", r)
+	wt.AssertNoErr(t, err)
+	if minAnswers == 0 && maxAnswers == 0 {
+		wt.AssertStatus(t, r.Rcode, expErr, "DNS response code")
+	} else {
+		wt.AssertStatus(t, r.Rcode, dns.RcodeSuccess, "DNS response code")
+	}
+	answers := len(r.Answer)
+	if minAnswers >= 0 && answers < minAnswers {
+		wt.Fatalf(t, "Number of answers >= %d", minAnswers)
+	}
+	if maxAnswers >= 0 && answers > maxAnswers {
+		wt.Fatalf(t, "Number of answers <= %d", maxAnswers)
+	}
+	return r
+}
+
+// run a UDP fallback server
+func runLocalUDPServer(t *testing.T, laddr string, handler dns.HandlerFunc) (*dns.Server, string, error) {
 	t.Logf("Starting fallback UDP server at %s", laddr)
 	pc, err := net.ListenPacket("udp", laddr)
 	if err != nil {
@@ -233,7 +269,8 @@ func RunLocalUDPServer(t *testing.T, laddr string, handler dns.HandlerFunc) (*dn
 	return server, pc.LocalAddr().String(), nil
 }
 
-func RunLocalTCPServer(t *testing.T, laddr string, handler dns.HandlerFunc) (*dns.Server, string, error) {
+// run a TCP fallback server
+func runLocalTCPServer(t *testing.T, laddr string, handler dns.HandlerFunc) (*dns.Server, string, error) {
 	t.Logf("Starting fallback TCP server at %s", laddr)
 	laddrTcp, err := net.ResolveTCPAddr("tcp", laddr)
 	if err != nil {

testing/util.go

@@ -9,13 +9,28 @@ import (
 
 func AssertTrue(t *testing.T, cond bool, desc string) {
 	if !cond {
-		Fatalf(t, "Expected %s to be true", desc)
+		Fatalf(t, "Expected %s", desc)
 	}
 }
 
 func AssertFalse(t *testing.T, cond bool, desc string) {
 	if cond {
-		Fatalf(t, "Expected %s to be false", desc)
+		Fatalf(t, "Unexpected %s", desc)
+	}
+}
+
+func AssertNil(t *testing.T, p interface{}, desc string) {
+	val := reflect.ValueOf(p)
+	if val.IsValid() && !val.IsNil() {
+		Fatalf(t, "Expected nil pointer for %s but got a \"%s\": \"%s\"",
+			desc, reflect.TypeOf(p), val.Interface())
+	}
+}
+
+func AssertNotNil(t *testing.T, p interface{}, desc string) {
+	val := reflect.ValueOf(p)
+	if !val.IsValid() || val.IsNil() {
+		Fatalf(t, "Unexpected nil pointer for %s", desc)
 	}
 }
 

weavedns/main.go

@@ -10,6 +10,9 @@ import (
 	"io"
 	"net"
 	"os"
+	"os/signal"
+	"runtime"
+	"syscall"
 )
 
 var version = "(unreleased version)"
@@ -24,6 +27,9 @@ func main() {
 		dnsPort     int
 		httpPort    int
 		wait        int
+		timeout     int
+		udpbuf      int
+		cacheLen    int
 		watch       bool
 		debug       bool
 		err         error
@@ -31,12 +37,15 @@ func main() {
 
 	flag.BoolVar(&justVersion, "version", false, "print version and exit")
 	flag.StringVar(&ifaceName, "iface", "", "name of interface to use for multicast")
-	flag.StringVar(&apiPath, "api", "unix:///var/run/docker.sock", "Path to Docker API socket")
-	flag.StringVar(&localDomain, "localDomain", weavedns.DEFAULT_LOCAL_DOMAIN, "local domain (e.g., 'weave.local.')")
+	flag.StringVar(&apiPath, "api", "unix:///var/run/docker.sock", "path to Docker API socket")
+	flag.StringVar(&localDomain, "localDomain", weavedns.DEFAULT_LOCAL_DOMAIN, "local domain (ie, 'weave.local.')")
 	flag.IntVar(&wait, "wait", 0, "number of seconds to wait for interface to be created and come up")
-	flag.IntVar(&dnsPort, "dnsport", 53, "port to listen to DNS requests")
+	flag.IntVar(&dnsPort, "dnsport", weavedns.DEFAULT_SERVER_PORT, "port to listen to DNS requests")
 	flag.IntVar(&httpPort, "httpport", 6785, "port to listen to HTTP requests")
-	flag.StringVar(&fallback, "fallback", "", "fallback server and port (ie, '8.8.8.8:53')")
+	flag.StringVar(&fallback, "fallback", "", "force a fallback (ie, '8.8.8.8:53') instead of /etc/resolv.conf values")
+	flag.IntVar(&timeout, "timeout", weavedns.DEFAULT_TIMEOUT, "timeout for resolutions")
+	flag.IntVar(&udpbuf, "udpbuf", weavedns.DEFAULT_UDP_BUFLEN, "UDP buffer length")
+	flag.IntVar(&cacheLen, "cache", weavedns.DEFAULT_CACHE_LEN, "cache length")
 	flag.BoolVar(&watch, "watch", true, "watch the docker socket for container events")
 	flag.BoolVar(&debug, "debug", false, "output debugging info to stderr")
 	flag.Parse()
@@ -70,8 +79,11 @@ func main() {
 	}
 
 	srvConfig := weavedns.DNSServerConfig{
-		Port:        dnsPort,
-		LocalDomain: localDomain,
+		Port:                dnsPort,
+		CacheLen:            cacheLen,
+		LocalDomain:         localDomain,
+		Timeout:             timeout,
+		UdpBufLen:           udpbuf,
 	}
 
 	if len(fallback) > 0 {
@@ -82,14 +94,36 @@ func main() {
 		srvConfig.UpstreamCfg = &dns.ClientConfig{Servers: []string{fallbackHost}, Port: fallbackPort}
 		Debug.Printf("DNS fallback at %s:%s", fallbackHost, fallbackPort)
 	}
-
+	
 	srv, err := weavedns.NewDNSServer(srvConfig, zone, iface)
 	if err != nil {
 		Error.Fatal("Failed to initialize the WeaveDNS server", err)
 	}
+
+	Debug.Printf("Starting the signals handler")
+	go handleSignals(srv)
+
 	go weavedns.ListenHttp(version, srv, localDomain, zone, httpPort)
 	err = srv.Start()
 	if err != nil {
 		Error.Fatal("Failed to start the WeaveDNS server", err)
 	}
 }
+
+func handleSignals(s *weavedns.DNSServer) {
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGQUIT)
+	buf := make([]byte, 1<<20)
+	for {
+		sig := <-sigs
+		switch sig {
+		case syscall.SIGINT:
+			Info.Printf("=== received SIGINT ===\n*** exiting\n")
+			s.Stop()
+			os.Exit(0)
+		case syscall.SIGQUIT:
+			stacklen := runtime.Stack(buf, true)
+			Info.Printf("=== received SIGQUIT ===\n*** goroutine dump...\n%s\n*** end\n", buf[:stacklen])
+		}
+	}
+}