Add a special vhost for health checking that uses /etc/nolvs

Also adds another special vhost for debugging/monitoring

Author: quentinmit

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

@@ -28,8 +28,10 @@ func always(context.Context, string) bool {
 }
 
 type ldapTarget struct {
-	localPoolRange *net.IPNet
-	ldap           *ldap.Pool
+	localPoolRange    *net.IPNet
+	ldap              *ldap.Pool
+	statuszServer     *HijackedServer
+	unavailableServer *HijackedServer
 }
 
 // HandleConn is called by tcpproxy after receiving a connection and sniffing the host.
@@ -39,6 +41,17 @@ func (l *ldapTarget) HandleConn(netConn net.Conn) {
 	var pool string
 	var err error
 	if conn, ok := netConn.(*tcpproxy.Conn); ok {
+		switch conn.HostName {
+		case "proxy.scripts.scripts.mit.edu":
+			// Special handling for proxy.scripts.scripts.mit.edu
+			l.statuszServer.HandleConn(netConn)
+			return
+		case "heartbeat.scripts.scripts.mit.edu":
+			if nolvsPresent() {
+				l.unavailableServer.HandleConn(netConn)
+				return
+			}
+		}
 		pool, err = l.ldap.ResolvePool(conn.HostName)
 		if err != nil {
 			log.Printf("resolving %q: %v", conn.HostName, err)
@@ -90,8 +103,10 @@ func main() {
 
 	var p tcpproxy.Proxy
 	t := &ldapTarget{
-		localPoolRange: ipnet,
-		ldap:           ldapPool,
+		localPoolRange:    ipnet,
+		ldap:              ldapPool,
+		statuszServer:     NewHijackedServer(nil),
+		unavailableServer: NewUnavailableServer(),
 	}
 	for _, addr := range strings.Split(*httpAddrs, ",") {
 		p.AddHTTPHostMatchRoute(addr, always, t)

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

@@ -0,0 +1,63 @@
+package main
+
+import (
+	"errors"
+	"net"
+	"net/http"
+	_ "net/http/pprof"
+	"os"
+)
+
+func nolvsPresent() bool {
+	if _, err := os.Stat("/etc/nolvs"); err == nil {
+		return true
+	}
+	return false
+}
+
+// HijackedServer is an HTTP server that serves from hijacked connections instead of a listening socket.
+// Users can call HandleConn to handle any request(s) waiting on that net.Conn.
+type HijackedServer struct {
+	connCh chan net.Conn
+}
+
+// NewHijackedServer constructs a HijackedServer that handles incoming HTTP connections with handler.
+func NewHijackedServer(handler http.Handler) *HijackedServer {
+	s := &HijackedServer{
+		connCh: make(chan net.Conn),
+	}
+	go http.Serve(s, handler)
+	return s
+}
+
+// Accept is called by http.Server to acquire a new connection.
+func (s *HijackedServer) Accept() (net.Conn, error) {
+	c, ok := <-s.connCh
+	if ok {
+		return c, nil
+	}
+	return nil, errors.New("closed")
+}
+
+// Close shuts down the server.
+func (s *HijackedServer) Close() error {
+	close(s.connCh)
+	return nil
+}
+
+// Addr must be present to implement net.Listener
+func (s *HijackedServer) Addr() net.Addr {
+	return nil
+}
+
+// HandleConn instructs the server to take control of c and handle any HTTP request(s) that are waiting.
+func (s *HijackedServer) HandleConn(c net.Conn) {
+	s.connCh <- c
+}
+
+// NewUnavailableServer constructs a HijackedServer that serves 500 Service Unavailable for all requests.
+func NewUnavailableServer() *HijackedServer {
+	return NewHijackedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, "0 proxy nolvs", http.StatusServiceUnavailable)
+	}))
+}