Add support allow routes to be added after the tcpproxy is running.

Inspired by inetaf#15

Author: Patrick Downey

.gitignore

@@ -1,2 +1,3 @@
+.idea
 tlsrouter
 tlsrouter.test

http.go

@@ -31,6 +31,14 @@ func (p *Proxy) AddHTTPHostRoute(ipPort, httpHost string, dest Target) {
 	p.AddHTTPHostMatchRoute(ipPort, equals(httpHost), dest)
 }
 
+// RemoveHTTPHostRoute removes a route to the ipPort listener
+//
+// The ipPort is any net.Listen TCP address. If it hasn't been registered with
+// tcpproxy this is a noop
+func (p *Proxy) RemoveHTTPHostRoute(ipPort, httpHost string, dest Target) {
+	p.RemoveHTTPHostMatchRoute(ipPort, equals(httpHost), dest)
+}
+
 // AddHTTPHostMatchRoute appends a route to the ipPort listener that
 // routes to dest if the incoming HTTP/1.x Host header name is
 // accepted by matcher. If it doesn't match, rule processing continues
@@ -41,6 +49,14 @@ func (p *Proxy) AddHTTPHostMatchRoute(ipPort string, match Matcher, dest Target)
 	p.addRoute(ipPort, httpHostMatch{match, dest})
 }
 
+// RemoveHTTPHostMatchRoute removes a route to the ipPort listener
+//
+// The ipPort is any net.Listen TCP address. If it hasn't been registered with
+// tcpproxy this is a noop
+func (p *Proxy) RemoveHTTPHostMatchRoute(ipPort string, match Matcher, dest Target) {
+	p.removeRoute(ipPort, httpHostMatch{match, dest})
+}
+
 type httpHostMatch struct {
 	matcher Matcher
 	target  Target

sni.go

@@ -38,6 +38,14 @@ func (p *Proxy) AddSNIRoute(ipPort, sni string, dest Target) {
 	p.AddSNIMatchRoute(ipPort, equals(sni), dest)
 }
 
+// RemoveSNIRoute removes a route from the ipPort listener
+//
+// The ipPort is any net.Listen TCP address, if it hasn't been registered with
+// tcpproxy this is a noop.
+func (p *Proxy) RemoveSNIRoute(ipPort, sni string, dest Target) {
+	p.RemoveSNIMatchRoute(ipPort, equals(sni), dest)
+}
+
 // AddSNIMatchRoute appends a route to the ipPort listener that routes
 // to dest if the incoming TLS SNI server name is accepted by
 // matcher. If it doesn't match, rule processing continues for any
@@ -60,6 +68,21 @@ func (p *Proxy) AddSNIMatchRoute(ipPort string, matcher Matcher, dest Target) {
 	p.addRoute(ipPort, sniMatch{matcher, dest})
 }
 
+// RemoveSNIMatchRoute removes a route from the ipPort listener
+//
+// The ipPort is any net.Listen TCP address, if it hasn't been registered with
+// tcpproxy this is a noop.
+func (p *Proxy) RemoveSNIMatchRoute(ipPort string, matcher Matcher, dest Target) {
+
+	if p.configExists(ipPort) {
+		cfg := p.configFor(ipPort)
+
+		p.removeRoute(ipPort, &acmeMatch{cfg})
+
+		p.removeRoute(ipPort, sniMatch{matcher, dest})
+	}
+}
+
 // AddStopACMESearch prevents ACME probing of subsequent SNI routes.
 // Any ACME challenges on ipPort for SNI routes previously added
 // before this call will still be proxied to all possible SNI

tcpproxy.go

@@ -60,6 +60,7 @@ import (
 	"io"
 	"log"
 	"net"
+	"sync"
 	"time"
 )
 
@@ -69,6 +70,7 @@ import (
 // The order that routes are added in matters; each is matched in the order
 // registered.
 type Proxy struct {
+	mu sync.Mutex
 	configs map[string]*config // ip:port => config
 
 	lns   []net.Listener
@@ -81,6 +83,14 @@ type Proxy struct {
 	ListenFunc func(net, laddr string) (net.Listener, error)
 }
 
+//
+//func (p *Proxy) Configs() map[string]*config {
+//	p.mu.Lock()
+//	defer p.mu.Unlock()
+//
+//	return p.configs
+//}
+
 // Matcher reports whether hostname matches the Matcher's criteria.
 type Matcher func(ctx context.Context, hostname string) bool
 
@@ -93,11 +103,42 @@ func equals(want string) Matcher {
 
 // config contains the proxying state for one listener.
 type config struct {
+	mu     sync.Mutex
+
 	routes      []route
 	acmeTargets []Target // accumulates targets that should be probed for acme.
 	stopACME    bool     // if true, AddSNIRoute doesn't add targets to acmeTargets.
 }
 
+func (c *config) AddRoute(r route) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.routes = append(c.routes, r)
+}
+
+func (c *config) Routes() []route {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	return c.routes
+}
+
+func (c *config) RemoveRoute(r route) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	var newRoutes []route
+
+	for _, i := range c.routes {
+		if i != r {
+			newRoutes = append(newRoutes, i)
+		}
+	}
+
+	c.routes = newRoutes
+}
+
 // A route matches a connection to a target.
 type route interface {
 	// match examines the initial bytes of a connection, looking for a
@@ -121,6 +162,9 @@ func (p *Proxy) netListen() func(net, laddr string) (net.Listener, error) {
 }
 
 func (p *Proxy) configFor(ipPort string) *config {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
 	if p.configs == nil {
 		p.configs = make(map[string]*config)
 	}
@@ -130,9 +174,20 @@ func (p *Proxy) configFor(ipPort string) *config {
 	return p.configs[ipPort]
 }
 
+func (p *Proxy) configExists(ipPort string) bool {
+	return p.configs[ipPort] != nil
+}
+
 func (p *Proxy) addRoute(ipPort string, r route) {
 	cfg := p.configFor(ipPort)
-	cfg.routes = append(cfg.routes, r)
+	cfg.AddRoute(r)
+}
+
+func (p *Proxy) removeRoute(ipPort string, r route) {
+	if p.configExists(ipPort) {
+		cfg := p.configFor(ipPort)
+		cfg.RemoveRoute(r)
+	}
 }
 
 // AddRoute appends an always-matching route to the ipPort listener,
@@ -146,6 +201,13 @@ func (p *Proxy) AddRoute(ipPort string, dest Target) {
 	p.addRoute(ipPort, fixedTarget{dest})
 }
 
+// RemoveRoute removes the specified target from the ipPort listener
+//
+// This method won't remove an ipPort listener if there are no routes remaining
+func (p *Proxy) RemoveRoute(ipPort string, dest Target) {
+	p.removeRoute(ipPort, fixedTarget{dest})
+}
+
 type fixedTarget struct {
 	t Target
 }
@@ -200,7 +262,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
@@ -211,22 +273,22 @@ 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 {
+	for _, route := range cfg.Routes() {
 		if target, hostName := route.match(br); target != nil {
 			if n := br.Buffered(); n > 0 {
 				peeked, _ := br.Peek(br.Buffered())

tcpproxy_test.go

@@ -175,6 +175,61 @@ func testProxy(t *testing.T, front net.Listener) *Proxy {
 	}
 }
 
+func TestConfigAddRemoveRoute(t *testing.T) {
+	r1 := fixedTarget{To("foo.com:443")}
+	r2 := fixedTarget{To("bar.com:443")}
+
+	cfg := &config{}
+
+	cfg.AddRoute(r1)
+	cfg.AddRoute(r2)
+
+	routes := cfg.Routes()
+	if len(routes) != 2 {
+		t.Fatalf("got %d routes, want 1", len(routes))
+	}
+
+	cfg.RemoveRoute(r1)
+
+	routes = cfg.Routes()
+	if len(routes) != 1 {
+		t.Fatalf("got %d routes, want 1", len(routes))
+	}
+
+	if routes[0] != r2 {
+		t.Fatalf("got %s, want %s", r2, routes[0])
+	}
+}
+
+func TestProxyAddRemoveRoute(t *testing.T) {
+	front := newLocalListener(t)
+	defer front.Close()
+
+	r1 := To("foo.com:443")
+	r1Fixed := fixedTarget{r1}
+	r2 := To("bar.com:443")
+
+	p := testProxy(t, front)
+	p.AddRoute(testFrontAddr, r1)
+	p.AddRoute(testFrontAddr, r2)
+
+	config := p.configFor(testFrontAddr)
+	routes := config.Routes()
+	if len(routes) != 2 {
+		t.Fatalf("got %d routes, want 2", len(routes))
+	}
+
+	p.RemoveRoute(testFrontAddr, r2)
+	routes = config.Routes()
+	if len(routes) != 1 {
+		t.Fatalf("got %d routes, want 1", len(routes))
+	}
+
+	if routes[0] != r1Fixed {
+		t.Fatalf("got %s, want %s", r1Fixed, routes[0])
+	}
+}
+
 func TestProxyAlwaysMatch(t *testing.T) {
 	front := newLocalListener(t)
 	defer front.Close()