From 49ba21a820160f0416bc39a66edd4656ffd6e673 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Sat, 22 Feb 2014 17:31:11 -0800
Subject: [PATCH] consul: Default recursor port 53. Fixes #4

---
 command/agent/dns.go      | 28 ++++++++++++++++++++++++++++
 command/agent/dns_test.go | 10 ++++++++++
 2 files changed, 38 insertions(+)

diff --git a/command/agent/dns.go b/command/agent/dns.go
index 076840405515..da5312d31de6 100644
--- a/command/agent/dns.go
+++ b/command/agent/dns.go
@@ -68,6 +68,11 @@ func NewDNSServer(agent *Agent, logOutput io.Writer, domain, bind, recursor stri
 		mux.HandleFunc(consulDomain, srv.handleTest)
 	}
 	if recursor != "" {
+		recursor, err := recursorAddr(recursor)
+		if err != nil {
+			return nil, fmt.Errorf("Invalid recursor address: %v", err)
+		}
+		srv.recursor = recursor
 		mux.HandleFunc(".", srv.handleRecurse)
 	}
 
@@ -124,6 +129,29 @@ func NewDNSServer(agent *Agent, logOutput io.Writer, domain, bind, recursor stri
 	return srv, nil
 }
 
+// recursorAddr is used to add a port to the recursor if omitted.
+func recursorAddr(recursor string) (string, error) {
+	// Add the port if none
+START:
+	_, _, err := net.SplitHostPort(recursor)
+	if ae, ok := err.(*net.AddrError); ok && ae.Err == "missing port in address" {
+		recursor = fmt.Sprintf("%s:%d", recursor, 53)
+		goto START
+	}
+	if err != nil {
+		return "", err
+	}
+
+	// Get the address
+	addr, err := net.ResolveTCPAddr("tcp", recursor)
+	if err != nil {
+		return "", err
+	}
+
+	// Return string
+	return addr.String(), nil
+}
+
 // handleQUery is used to handle DNS queries in the configured domain
 func (d *DNSServer) handleQuery(resp dns.ResponseWriter, req *dns.Msg) {
 	q := req.Question[0]
diff --git a/command/agent/dns_test.go b/command/agent/dns_test.go
index a1776cad640f..1ff478dd3a22 100644
--- a/command/agent/dns_test.go
+++ b/command/agent/dns_test.go
@@ -21,6 +21,16 @@ func makeDNSServer(t *testing.T) (string, *DNSServer) {
 	return dir, server
 }
 
+func TestRecursorAddr(t *testing.T) {
+	addr, err := recursorAddr("8.8.8.8")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if addr != "8.8.8.8:53" {
+		t.Fatalf("bad: %v", addr)
+	}
+}
+
 func TestDNS_IsAlive(t *testing.T) {
 	dir, srv := makeDNSServer(t)
 	defer os.RemoveAll(dir)