Add support for DNS SRV lookups

Author: Matt Hamann

glide.lock

@@ -5,3 +5,7 @@ import:
 - package: golang.org/x/net
   subpackages:
   - websocket
+- package: github.com/benschw/srv-lb
+  subpackages:
+  - dns
+  - lb

glide.yaml

@@ -0,0 +1,41 @@
+package resolver
+
+import (
+	"strings"
+
+	"github.com/benschw/srv-lb/lb"
+)
+
+// DNSConfig accepts an address and optional load balancer config
+type DNSConfig struct {
+	Addr  string
+	LbCfg *lb.Config
+}
+
+// ResolveSrvAddr returns a load-balanced host:port based on DNS SRV lookup (or `addr` if not a hostname)
+func ResolveSrvAddr(dnsCfg DNSConfig) (string, error) {
+	addr := dnsCfg.Addr
+
+	if v := strings.Split(addr, ":"); len(v) < 2 {
+		var cfg *lb.Config
+		var err error
+		if dnsCfg.LbCfg == nil {
+			cfg, err = lb.DefaultConfig()
+
+			if err != nil {
+				return addr, err
+			}
+		} else {
+			cfg = dnsCfg.LbCfg
+		}
+
+		l := lb.New(cfg, addr)
+		resolvAddr, err := l.Next()
+
+		if err == nil {
+			addr = resolvAddr.String()
+		}
+	}
+
+	return addr, nil
+}

resolver/dns.go

@@ -0,0 +1,29 @@
+package resolver
+
+import (
+	"testing"
+
+	"github.com/benschw/srv-lb/lb"
+)
+
+func TestSrvLookup(t *testing.T) {
+	lbCfg, _ := lb.DefaultConfig()
+	lbCfg.Strategy = MockStrategy
+	dnsCfg := DNSConfig{Addr: "foo.example.com", LbCfg: lbCfg}
+
+	addr, _ := ResolveSrvAddr(dnsCfg)
+
+	if addr != "1.2.3.4:1234" {
+		t.Error("expected address string of 1.2.3.4:1234, got:", addr)
+	}
+}
+
+func TestIpPassthrough(t *testing.T) {
+	dnsCfg := DNSConfig{Addr: "10.0.0.1:1234"}
+
+	addr, _ := ResolveSrvAddr(dnsCfg)
+
+	if addr != "10.0.0.1:1234" {
+		t.Error("expected output to equal input (10.0.0.1:1234), got:", addr)
+	}
+}

resolver/dns_test.go

@@ -0,0 +1,30 @@
+package resolver
+
+import (
+	"github.com/benschw/srv-lb/dns"
+	"github.com/benschw/srv-lb/lb"
+)
+
+// MockStrategy is used for testing DNS load balancing
+const MockStrategy lb.StrategyType = "mock"
+
+// New creates a new instance of the load balancer
+func New(lib dns.Lookup) lb.GenericLoadBalancer {
+	lb := new(MockClb)
+	lb.dnsLib = lib
+	return lb
+}
+
+// MockClb contains the dnslib
+type MockClb struct {
+	dnsLib dns.Lookup
+}
+
+// Next gets the next server in the available nodes
+func (lb *MockClb) Next(name string) (dns.Address, error) {
+	return dns.Address{Address: "1.2.3.4", Port: 1234}, nil
+}
+
+func init() {
+	lb.RegisterStrategy(MockStrategy, New)
+}

resolver/lb_mock_strategy.go

@@ -4,6 +4,7 @@ import (
 	"net"
 
 	"github.com/gliderlabs/logspout/adapters/raw"
+	"github.com/gliderlabs/logspout/resolver"
 	"github.com/gliderlabs/logspout/router"
 )
 
@@ -21,7 +22,12 @@ func rawTCPAdapter(route *router.Route) (router.LogAdapter, error) {
 type tcpTransport int
 
 func (t *tcpTransport) Dial(addr string, options map[string]string) (net.Conn, error) {
-	raddr, err := net.ResolveTCPAddr("tcp", addr)
+	daddr, err := resolver.ResolveSrvAddr(resolver.DNSConfig{Addr: addr})
+	if err != nil {
+		return nil, err
+	}
+
+	raddr, err := net.ResolveTCPAddr("tcp", daddr)
 	if err != nil {
 		return nil, err
 	}

transports/tcp/tcp.go

@@ -5,6 +5,7 @@ import (
 	"net"
 
 	"github.com/gliderlabs/logspout/adapters/raw"
+	"github.com/gliderlabs/logspout/resolver"
 	"github.com/gliderlabs/logspout/router"
 )
 
@@ -22,7 +23,12 @@ func rawTLSAdapter(route *router.Route) (router.LogAdapter, error) {
 type tlsTransport int
 
 func (t *tlsTransport) Dial(addr string, options map[string]string) (net.Conn, error) {
-	conn, err := tls.Dial("tcp", addr, nil)
+	daddr, err := resolver.ResolveSrvAddr(resolver.DNSConfig{Addr: addr})
+	if err != nil {
+		return nil, err
+	}
+
+	conn, err := tls.Dial("tcp", daddr, nil)
 	if err != nil {
 		return nil, err
 	}

transports/tls/tls.go

@@ -4,6 +4,7 @@ import (
 	"net"
 
 	"github.com/gliderlabs/logspout/adapters/raw"
+	"github.com/gliderlabs/logspout/resolver"
 	"github.com/gliderlabs/logspout/router"
 )
 
@@ -26,7 +27,12 @@ func rawUDPAdapter(route *router.Route) (router.LogAdapter, error) {
 type udpTransport int
 
 func (t *udpTransport) Dial(addr string, options map[string]string) (net.Conn, error) {
-	raddr, err := net.ResolveUDPAddr("udp", addr)
+	daddr, err := resolver.ResolveSrvAddr(resolver.DNSConfig{Addr: addr})
+	if err != nil {
+		return nil, err
+	}
+
+	raddr, err := net.ResolveUDPAddr("udp", daddr)
 	if err != nil {
 		return nil, err
 	}