Enable add/remove routes after server starts

tcpproxy.go

@@ -60,6 +60,7 @@ import (
 	"io"
 	"log"
 	"net"
+	"sync"
 	"time"
 )
 
@@ -93,11 +94,19 @@ func equals(want string) Matcher {
 
 // config contains the proxying state for one listener.
 type config struct {
-	routes      []route
+	sync.Mutex  // protect r/w of routes
+	nextRouteId int
+	routes      map[int]route
 	acmeTargets []Target // accumulates targets that should be probed for acme.
 	stopACME    bool     // if true, AddSNIRoute doesn't add targets to acmeTargets.
 }
 
+func NewConfig() (cfg *config) {
+	cfg = &config{}
+	cfg.routes = make(map[int]route)
+	return
+}
+
 // A route matches a connection to a target.
 type route interface {
 	// match examines the initial bytes of a connection, looking for a
@@ -122,25 +131,43 @@ func (p *Proxy) configFor(ipPort string) *config {
 		p.configs = make(map[string]*config)
 	}
 	if p.configs[ipPort] == nil {
-		p.configs[ipPort] = &config{}
+		p.configs[ipPort] = NewConfig()
 	}
 	return p.configs[ipPort]
 }
 
-func (p *Proxy) addRoute(ipPort string, r route) {
+func (p *Proxy) addRoute(ipPort string, r route) (routeId int) {
 	cfg := p.configFor(ipPort)
-	cfg.routes = append(cfg.routes, r)
+	cfg.Lock()
+	defer cfg.Unlock()
+	routeId = cfg.nextRouteId
+	cfg.nextRouteId++
+	cfg.routes[routeId] = r
+	return
 }
 
 // AddRoute appends an always-matching route to the ipPort listener,
-// directing any connection to dest.
+// directing any connection to dest. The added route's id is returned
+// for future removal.
 //
 // This is generally used as either the only rule (for simple TCP
 // proxies), or as the final fallback rule for an ipPort.
 //
 // The ipPort is any valid net.Listen TCP address.
-func (p *Proxy) AddRoute(ipPort string, dest Target) {
-	p.addRoute(ipPort, fixedTarget{dest})
+func (p *Proxy) AddRoute(ipPort string, dest Target) (routeId int) {
+	return p.addRoute(ipPort, fixedTarget{dest})
+}
+
+// RemoveRoute removes an existing route for ipPort. If the route is
+// not found, this is an no-op.
+//
+// Both AddRoute and RemoveRoute is go-routine safe.
+func (p *Proxy) RemoveRoute(ipPort string, routeId int) (err error) {
+	cfg := p.configFor(ipPort)
+	cfg.Lock()
+	defer cfg.Unlock()
+	delete(cfg.routes, routeId)
+	return
 }
 
 type fixedTarget struct {
@@ -197,7 +224,7 @@ func (p *Proxy) Start() error {
 			return err
 		}
 		p.lns = append(p.lns, ln)
-		go p.serveListener(errc, ln, config.routes)
+		go p.serveListener(errc, ln, config)
 	}
 	go p.awaitFirstError(errc)
 	return nil
@@ -208,22 +235,24 @@ func (p *Proxy) awaitFirstError(errc <-chan error) {
 	close(p.donec)
 }
 
-func (p *Proxy) serveListener(ret chan<- error, ln net.Listener, routes []route) {
+func (p *Proxy) serveListener(ret chan<- error, ln net.Listener, cfg *config) {
 	for {
 		c, err := ln.Accept()
 		if err != nil {
 			ret <- err
 			return
 		}
-		go p.serveConn(c, routes)
+		go p.serveConn(c, cfg)
 	}
 }
 
 // serveConn runs in its own goroutine and matches c against routes.
 // It returns whether it matched purely for testing.
-func (p *Proxy) serveConn(c net.Conn, routes []route) bool {
+func (p *Proxy) serveConn(c net.Conn, cfg *config) bool {
 	br := bufio.NewReader(c)
-	for _, route := range routes {
+	cfg.Lock()
+	defer cfg.Unlock()
+	for _, route := range cfg.routes {
 		if target := route.match(br); target != nil {
 			if n := br.Buffered(); n > 0 {
 				peeked, _ := br.Peek(br.Buffered())