Add support for arbitrary matching against HTTP and SNI hostnames.

Add{HTTPHost,SNI}Route remain so that the common case of exact matches
remains trivial to use. Add{HTTPHost,SNI}MatchRoute allow you to specify
your own matching function.

Fixes #9

4 files changed, 52 insertions, 17 deletions

diff --git a/http.go b/http.go
index 601d535..6197da9 100644
--- a/http.go
+++ b/http.go
@@ -17,26 +17,37 @@ package tcpproxy
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"net/http"
 )
 
-// AddHTTPHostRoute appends a route to the ipPort listener that says
-// if the incoming HTTP/1.x Host header name is httpHost, the
-// connection is given to dest. If it doesn't match, rule processing
-// continues for any additional routes on ipPort.
+// AddHTTPHostRoute appends a route to the ipPort listener that
+// routes to dest if the incoming HTTP/1.x Host header name is
+// httpHost. If it doesn't match, rule processing continues for any
+// additional routes on ipPort.
 //
 // The ipPort is any valid net.Listen TCP address.
 func (p *Proxy) AddHTTPHostRoute(ipPort, httpHost string, dest Target) {
-	p.addRoute(ipPort, httpHostMatch{httpHost, dest})
+	p.AddHTTPHostMatchRoute(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
+// for any additional routes on ipPort.
+//
+// The ipPort is any valid net.Listen TCP address.
+func (p *Proxy) AddHTTPHostMatchRoute(ipPort string, match Matcher, dest Target) {
+	p.addRoute(ipPort, httpHostMatch{match, dest})
 }
 
 type httpHostMatch struct {
-	host   string
-	target Target
+	matcher Matcher
+	target  Target
 }
 
 func (m httpHostMatch) match(br *bufio.Reader) Target {
-	if httpHostHeader(br) == m.host {
+	if m.matcher(context.TODO(), httpHostHeader(br)) {
 		return m.target
 	}
 	return nil
diff --git a/sni.go b/sni.go
index 50ab599..44f5796 100644
--- a/sni.go
+++ b/sni.go
@@ -24,10 +24,10 @@ import (
 	"strings"
 )
 
-// AddSNIRoute appends a route to the ipPort listener that says if the
-// incoming TLS SNI server name is sni, the connection is given to
-// dest. If it doesn't match, rule processing continues for any
-// additional routes on ipPort.
+// AddSNIRoute appends a route to the ipPort listener that routes to
+// dest if the incoming TLS SNI server name is sni. If it doesn't
+// match, rule processing continues for any additional routes on
+// ipPort.
 //
 // By default, the proxy will route all ACME tls-sni-01 challenges
 // received on ipPort to all SNI dests. You can disable ACME routing
@@ -35,6 +35,20 @@ import (
 //
 // The ipPort is any valid net.Listen TCP address.
 func (p *Proxy) AddSNIRoute(ipPort, sni string, dest Target) {
+	p.AddSNIMatchRoute(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
+// additional routes on ipPort.
+//
+// By default, the proxy will route all ACME tls-sni-01 challenges
+// received on ipPort to all SNI dests. You can disable ACME routing
+// with AddStopACMESearch.
+//
+// The ipPort is any valid net.Listen TCP address.
+func (p *Proxy) AddSNIMatchRoute(ipPort string, matcher Matcher, dest Target) {
 	cfg := p.configFor(ipPort)
 	if !cfg.stopACME {
 		if len(cfg.acmeTargets) == 0 {
@@ -43,7 +57,7 @@ func (p *Proxy) AddSNIRoute(ipPort, sni string, dest Target) {
 		cfg.acmeTargets = append(cfg.acmeTargets, dest)
 	}
 
-	p.addRoute(ipPort, sniMatch{sni, dest})
+	p.addRoute(ipPort, sniMatch{matcher, dest})
 }
 
 // AddStopACMESearch prevents ACME probing of subsequent SNI routes.
@@ -55,12 +69,12 @@ func (p *Proxy) AddStopACMESearch(ipPort string) {
 }
 
 type sniMatch struct {
-	sni    string
-	target Target
+	matcher Matcher
+	target  Target
 }
 
 func (m sniMatch) match(br *bufio.Reader) Target {
-	if clientHelloServerName(br) == string(m.sni) {
+	if m.matcher(context.TODO(), clientHelloServerName(br)) {
 		return m.target
 	}
 	return nil
diff --git a/tcpproxy.go b/tcpproxy.go
index 02b70f5..8c33604 100644
--- a/tcpproxy.go
+++ b/tcpproxy.go
@@ -81,6 +81,16 @@ type Proxy struct {
 	ListenFunc func(net, laddr string) (net.Listener, error)
 }
 
+// Matcher reports whether hostname matches the Matcher's criteria.
+type Matcher func(ctx context.Context, hostname string) bool
+
+// equals is a trivial Matcher that implements string equality.
+func equals(want string) Matcher {
+	return func(_ context.Context, got string) bool {
+		return want == got
+	}
+}
+
 // config contains the proxying state for one listener.
 type config struct {
 	routes      []route
diff --git a/tcpproxy_test.go b/tcpproxy_test.go
index ac7c917..682214d 100644
--- a/tcpproxy_test.go
+++ b/tcpproxy_test.go
@@ -71,7 +71,7 @@ func TestMatchHTTPHost(t *testing.T) {
 		}
 		t.Run(name, func(t *testing.T) {
 			br := bufio.NewReader(tt.r)
-			r := httpHostMatch{tt.host, noopTarget{}}
+			r := httpHostMatch{equals(tt.host), noopTarget{}}
 			got := r.match(br) != nil
 			if got != tt.want {
 				t.Fatalf("match = %v; want %v", got, tt.want)