Route connections based on LDAP

Author: quentinmit

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

@@ -8,20 +8,24 @@ import (
 	"net"
 	"strings"
 
+	ldap "gopkg.in/ldap.v3"
 	"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")
 	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")
 )
 
 func always(context.Context, string) bool {
 	return true
 }
 
 type ldapTarget struct {
+	ldap *ldap.Conn
 }
 
 func (l *ldapTarget) HandleConn(netConn net.Conn) {
@@ -52,14 +56,36 @@ func (l *ldapTarget) HandleConn(netConn net.Conn) {
 }
 
 func (l *ldapTarget) resolvePool(hostname string) (string, error) {
-	return "18.4.86.22", nil
+	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()
+
 	var p tcpproxy.Proxy
-	t := &ldapTarget{}
+	t := &ldapTarget{ldap: l}
 	for _, addr := range strings.Split(*httpAddrs, ",") {
 		p.AddHTTPHostMatchRoute(addr, always, t)
 	}