Establish connections to multiple LDAP servers and spread requests across them

Author: quentinmit

server/common/oursrc/scripts-proxy/ldap/conn.go

@@ -0,0 +1,62 @@
+package ldap
+
+import (
+	"fmt"
+	"log"
+	"sync"
+	"time"
+
+	ldap "gopkg.in/ldap.v3"
+)
+
+type conn struct {
+	server, baseDn string
+	// mu protects conn during reconnect cycles
+	// TODO: The ldap package supports multiple in-flight queries;
+	// by using a Mutex we are only going to issue one at a
+	// time. We should figure out how to do retry/reconnect
+	// behavior with parallel queries.
+	mu   sync.Mutex
+	conn *ldap.Conn
+}
+
+func (c *conn) reconnect() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.conn != nil {
+		c.conn.Close()
+	}
+	var err error
+	for {
+		log.Printf("connecting to %s", c.server)
+		c.conn, err = ldap.Dial("tcp", c.server)
+		if err == nil {
+			return
+		}
+		log.Printf("connecting to %s: %v", c.server, err)
+		time.Sleep(100 * time.Millisecond)
+	}
+}
+
+func (c *conn) resolvePool(hostname string) (string, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	escapedHostname := ldap.EscapeFilter(hostname)
+	req := ldap.NewSearchRequest(
+		c.baseDn,
+		ldap.ScopeWholeSubtree, ldap.NeverDerefAliases, 0, 0, false,
+		fmt.Sprintf("(|(scriptsVhostName=%s)(scriptsVhostAlias=%s))", escapedHostname, escapedHostname),
+		[]string{"scriptsVhostPoolIPv4"},
+		nil,
+	)
+	sr, err := c.conn.Search(req)
+	if err != nil {
+		return "", err
+	}
+	for _, entry := range sr.Entries {
+		return entry.GetAttributeValue("scriptsVhostPoolIPv4"), nil
+	}
+	// Not found is not an error
+	return "", nil
+}

server/common/oursrc/scripts-proxy/ldap/pool.go

@@ -0,0 +1,45 @@
+package ldap
+
+import "log"
+
+type Pool struct {
+	retries int
+	// connCh holds open connections to servers.
+	connCh chan *conn
+}
+
+func NewPool(servers []string, baseDn string, retries int) *Pool {
+	p := &Pool{
+		retries: retries,
+		connCh:  make(chan *conn, len(servers)),
+	}
+	for _, s := range servers {
+		c := &conn{
+			server: s,
+			baseDn: baseDn,
+		}
+		go p.reconnect(c)
+	}
+	return p
+}
+
+func (p *Pool) reconnect(c *conn) {
+	c.reconnect()
+	p.connCh <- c
+}
+
+func (p *Pool) ResolvePool(hostname string) (string, error) {
+	var ip string
+	var err error
+	for i := 0; i < p.retries; i++ {
+		c := <-p.connCh
+		ip, err = c.resolvePool(hostname)
+		if err == nil {
+			p.connCh <- c
+			return ip, err
+		}
+		log.Printf("resolving %q on %s: %v", hostname, c.server, err)
+		go p.reconnect(c)
+	}
+	return ip, err
+}

server/common/oursrc/scripts-proxy/main.go

@@ -8,39 +8,41 @@ import (
 	"net"
 	"strings"
 
-	ldap "gopkg.in/ldap.v3"
+	"github.com/mit-scripts/scripts/server/common/oursrc/scripts-proxy/ldap"
 	"inet.af/tcpproxy"
 )
 
 var (
 	httpAddrs   = flag.String("http_addrs", ":80", "comma-separated addresses to listen for HTTP traffic on")
 	sniAddrs    = flag.String("sni_addrs", ":443,:444", "comma-separated addresses to listen for SNI traffic on")
-	ldapServer  = flag.String("ldap_server", "scripts-ldap.mit.edu:389", "LDAP server to query")
+	ldapServers = flag.String("ldap_servers", "scripts-ldap.mit.edu:389", "comma-spearated LDAP servers to query")
 	defaultHost = flag.String("default_host", "scripts.mit.edu", "default host to route traffic to if SNI/Host header cannot be parsed or cannot be found in LDAP")
 	baseDn      = flag.String("base_dn", "ou=VirtualHosts,dc=scripts,dc=mit,dc=edu", "base DN to query for hosts")
 	localRange  = flag.String("local_range", "18.4.86.0/24", "IP block for client IP spoofing")
 )
 
+const ldapRetries = 3
+
 func always(context.Context, string) bool {
 	return true
 }
 
 type ldapTarget struct {
 	localPoolRange *net.IPNet
-	ldap           *ldap.Conn
+	ldap           *ldap.Pool
 }
 
 func (l *ldapTarget) HandleConn(netConn net.Conn) {
 	var pool string
 	var err error
 	if conn, ok := netConn.(*tcpproxy.Conn); ok {
-		pool, err = l.resolvePool(conn.HostName)
+		pool, err = l.ldap.ResolvePool(conn.HostName)
 		if err != nil {
 			log.Printf("resolving %q: %v", conn.HostName, err)
 		}
 	}
 	if pool == "" {
-		pool, err = l.resolvePool(*defaultHost)
+		pool, err = l.ldap.ResolvePool(*defaultHost)
 		if err != nil {
 			log.Printf("resolving default pool: %v", err)
 		}
@@ -72,44 +74,20 @@ func (l *ldapTarget) HandleConn(netConn net.Conn) {
 	dp.HandleConn(netConn)
 }
 
-func (l *ldapTarget) resolvePool(hostname string) (string, error) {
-	escapedHostname := ldap.EscapeFilter(hostname)
-	req := ldap.NewSearchRequest(
-		*baseDn,
-		ldap.ScopeWholeSubtree, ldap.NeverDerefAliases, 0, 0, false,
-		fmt.Sprintf("(|(scriptsVhostName=%s)(scriptsVhostAlias=%s))", escapedHostname, escapedHostname),
-		[]string{"scriptsVhostPoolIPv4"},
-		nil,
-	)
-	sr, err := l.ldap.Search(req)
-	if err != nil {
-		return "", err
-	}
-	for _, entry := range sr.Entries {
-		return entry.GetAttributeValue("scriptsVhostPoolIPv4"), nil
-	}
-	// Not found is not an error
-	return "", nil
-}
-
 func main() {
 	flag.Parse()
 
-	l, err := ldap.Dial("tcp", *ldapServer)
-	if err != nil {
-		log.Fatal(err)
-	}
-	defer l.Close()
-
 	_, ipnet, err := net.ParseCIDR(*localRange)
 	if err != nil {
 		log.Fatal(err)
 	}
 
+	ldapPool := ldap.NewPool(strings.Split(*ldapServers, ","), *baseDn, ldapRetries)
+
 	var p tcpproxy.Proxy
 	t := &ldapTarget{
 		localPoolRange: ipnet,
-		ldap:           l,
+		ldap:           ldapPool,
 	}
 	for _, addr := range strings.Split(*httpAddrs, ",") {
 		p.AddHTTPHostMatchRoute(addr, always, t)