Add dns wildcarding

dns_server.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net"
 	"strconv"
+	"strings"
 	"sync"
 
 	"github.com/miekg/dns"
@@ -20,23 +21,39 @@ var dnsAddr string
 
 type dnsRecords struct {
 	sync.RWMutex
-	dnsMap  map[string]string
-	hostMap *HostMap
+	dnsMap             map[string]string
+	hostMap            *HostMap
+	dnsWildcardEnabled bool
+	dnsWildcardLimit   int
 }
 
 func newDnsRecords(hostMap *HostMap) *dnsRecords {
 	return &dnsRecords{
-		dnsMap:  make(map[string]string),
-		hostMap: hostMap,
+		dnsMap:             make(map[string]string),
+		hostMap:            hostMap,
+		dnsWildcardEnabled: false,
+		dnsWildcardLimit:   5,
 	}
 }
 
 func (d *dnsRecords) Query(data string) string {
 	d.RLock()
-	if r, ok := d.dnsMap[data]; ok {
-		d.RUnlock()
-		return r
+
+	if d.dnsWildcardEnabled {
+		hostParts := strings.SplitN(data, ".", d.dnsWildcardLimit+1)
+		for i := range hostParts {
+			if r, ok := d.dnsMap[strings.Join(hostParts[i:], ".")]; ok {
+				d.RUnlock()
+				return r
+			}
+		}
+	} else {
+		if r, ok := d.dnsMap[data]; ok {
+			d.RUnlock()
+			return r
+		}
 	}
+
 	d.RUnlock()
 	return ""
 }
@@ -113,6 +130,8 @@ func handleDnsRequest(l *logrus.Logger, w dns.ResponseWriter, r *dns.Msg) {
 
 func dnsMain(l *logrus.Logger, hostMap *HostMap, c *config.C) func() {
 	dnsR = newDnsRecords(hostMap)
+	dnsR.dnsWildcardEnabled = c.GetBool("dns.wildcard", false)
+	dnsR.dnsWildcardLimit = c.GetInt("dns.wildcard_limit", 5)
 
 	// attach request handler func
 	dns.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {

dns_server_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/miekg/dns"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestParsequery(t *testing.T) {
@@ -17,3 +18,45 @@ func TestParsequery(t *testing.T) {
 
 	//parseQuery(m)
 }
+
+func TestDnsWildcardLookup(t *testing.T) {
+	hostMap := &HostMap{}
+	ds := newDnsRecords(hostMap)
+	ds.Add("test.com.com", "1.2.3.4")
+	ds.dnsWildcardEnabled = true
+
+	result := ds.Query("test.com.com")
+	assert.Equal(t, "1.2.3.4", result)
+
+	result = ds.Query("foo.test.com.com")
+	assert.Equal(t, "1.2.3.4", result)
+}
+
+func TestDnsWildcardDisabled(t *testing.T) {
+	hostMap := &HostMap{}
+	ds := newDnsRecords(hostMap)
+	ds.Add("test.com.com", "1.2.3.4")
+
+	result := ds.Query("test.com.com")
+	assert.Equal(t, "1.2.3.4", result)
+
+	result = ds.Query("foo.test.com.com")
+	assert.Equal(t, "", result)
+}
+
+func TestDnsWildcardLookupLimit(t *testing.T) {
+	hostMap := &HostMap{}
+	ds := newDnsRecords(hostMap)
+	ds.Add("test.com.com", "1.2.3.4")
+	ds.dnsWildcardEnabled = true
+	ds.dnsWildcardLimit = 1
+
+	result := ds.Query("test.com.com")
+	assert.Equal(t, "1.2.3.4", result)
+
+	result = ds.Query("foo.test.com.com")
+	assert.Equal(t, "1.2.3.4", result)
+
+	result = ds.Query("foo.bar.test.com.com")
+	assert.Equal(t, "", result)
+}