diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..78f591a --- /dev/null +++ b/.gitignore @@ -0,0 +1,23 @@ +### Go template +# If you prefer the allow list template instead of the deny list, see community template: +# https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore +# +# Binaries for programs and plugins +*.exe +*.exe~ +*.dll +*.so +*.dylib + +# Test binary, built with `go test -c` +*.test + +# Output of the go coverage tool, specifically when used with LiteIDE +*.out + +# Dependency directories (remove the comment below to include it) +# vendor/ + +# Go workspace file +go.work + diff --git a/cmd/smtprouter/main.go b/cmd/smtprouter/main.go new file mode 100644 index 0000000..77217da --- /dev/null +++ b/cmd/smtprouter/main.go @@ -0,0 +1,24 @@ +package main + +import ( + "context" + "fmt" + + "github.com/patdowney/tcpproxy" +) + +func pretendMatcher(ctx context.Context, hostname string) (tcpproxy.Target, bool) { + fmt.Printf("matched: %v\n", hostname) + return tcpproxy.To("127.0.0.1:4567"), true +} + +func main() { + t := tcpproxy.Proxy{} + + t.AddSMTPSNIRouteFunc(":25", "mx.traffic.lab.dioad.net", pretendMatcher) + + err := t.Run() + if err != nil { + fmt.Printf("%v", err) + } +} diff --git a/cmd/tlsrouter/e2e_test.go b/cmd/tlsrouter/e2e_test.go index 6e54021..0774824 100644 --- a/cmd/tlsrouter/e2e_test.go +++ b/cmd/tlsrouter/e2e_test.go @@ -17,7 +17,7 @@ import ( "testing" "time" - proxyproto "github.com/armon/go-proxyproto" + proxyproto "github.com/pires/go-proxyproto" ) func TestRouting(t *testing.T) { diff --git a/go.mod b/go.mod index af6f76e..323216f 100644 --- a/go.mod +++ b/go.mod @@ -1,5 +1,8 @@ -module github.com/inetaf/tcpproxy +module github.com/patdowney/tcpproxy -go 1.16 +go 1.24 -require github.com/armon/go-proxyproto v0.0.0-20210323213023-7e956b284f0a +require ( + github.com/google/uuid v1.6.0 + github.com/pires/go-proxyproto v0.11.0 +) diff --git a/go.sum b/go.sum index de51fb1..8e6444b 100644 --- a/go.sum +++ b/go.sum @@ -1,2 +1,12 @@ -github.com/armon/go-proxyproto v0.0.0-20210323213023-7e956b284f0a h1:AP/vsCIvJZ129pdm9Ek7bH7yutN3hByqsMoNrWAxRQc= -github.com/armon/go-proxyproto v0.0.0-20210323213023-7e956b284f0a/go.mod h1:QmP9hvJ91BbJmGVGSbutW19IC0Q9phDCLGaomwTJbgU= +github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/pires/go-proxyproto v0.6.0 h1:cLJUPnuQdiNf7P/wbeOKmM1khVdaMgTFDLj8h9ZrVYk= +github.com/pires/go-proxyproto v0.6.0/go.mod h1:Odh9VFOZJCf9G8cLW5o435Xf1J95Jw9Gw5rnCjcwzAY= +github.com/pires/go-proxyproto v0.7.0 h1:IukmRewDQFWC7kfnb66CSomk2q/seBuilHBYFwyq0Hs= +github.com/pires/go-proxyproto v0.7.0/go.mod h1:Vz/1JPY/OACxWGQNIRY2BeyDmpoaWmEP40O9LbuiFR4= +github.com/pires/go-proxyproto v0.8.0 h1:5unRmEAPbHXHuLjDg01CxJWf91cw3lKHc/0xzKpXEe0= +github.com/pires/go-proxyproto v0.8.0/go.mod h1:iknsfgnH8EkjrMeMyvfKByp9TiBZCKZM0jx2xmKqnVY= +github.com/pires/go-proxyproto v0.11.0 h1:gUQpS85X/VJMdUsYyEgyn59uLJvGqPhJV5YvG68wXH4= +github.com/pires/go-proxyproto v0.11.0/go.mod h1:ZKAAyp3cgy5Y5Mo4n9AlScrkCZwUy0g3Jf+slqQVcuU= diff --git a/http.go b/http.go index d28c66f..40f55e5 100644 --- a/http.go +++ b/http.go @@ -19,6 +19,8 @@ import ( "bytes" "context" "net/http" + + "github.com/google/uuid" ) // AddHTTPHostRoute appends a route to the ipPort listener that @@ -27,8 +29,8 @@ import ( // additional routes on ipPort. // // The ipPort is any valid net.Listen TCP address. -func (p *Proxy) AddHTTPHostRoute(ipPort, httpHost string, dest Target) { - p.AddHTTPHostMatchRoute(ipPort, equals(httpHost), dest) +func (p *Proxy) AddHTTPHostRoute(ipPort, httpHost string, dest Target) uuid.UUID { + return p.AddHTTPHostMatchRoute(ipPort, equals(httpHost), dest) } // AddHTTPHostMatchRoute appends a route to the ipPort listener that @@ -37,8 +39,8 @@ func (p *Proxy) AddHTTPHostRoute(ipPort, httpHost string, dest Target) { // 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}) +func (p *Proxy) AddHTTPHostMatchRoute(ipPort string, match Matcher, dest Target) uuid.UUID { + return p.addRoute(ipPort, httpHostMatch{match, dest}) } type httpHostMatch struct { @@ -51,7 +53,7 @@ func (m httpHostMatch) match(br *bufio.Reader) (Target, string) { if m.matcher(context.TODO(), hh) { return m.target, hh } - return nil, "" + return nil, hh } // httpHostHeader returns the HTTP Host header from br without diff --git a/sni.go b/sni.go index c2d37e0..c8a454b 100644 --- a/sni.go +++ b/sni.go @@ -21,6 +21,8 @@ import ( "crypto/tls" "io" "net" + + "github.com/google/uuid" ) // AddSNIRoute appends a route to the ipPort listener that routes to @@ -29,8 +31,8 @@ import ( // ipPort. // // The ipPort is any valid net.Listen TCP address. -func (p *Proxy) AddSNIRoute(ipPort, sni string, dest Target) { - p.AddSNIMatchRoute(ipPort, equals(sni), dest) +func (p *Proxy) AddSNIRoute(ipPort, sni string, dest Target) uuid.UUID { + return p.AddSNIMatchRoute(ipPort, equals(sni), dest) } // AddSNIMatchRoute appends a route to the ipPort listener that routes @@ -39,8 +41,28 @@ func (p *Proxy) AddSNIRoute(ipPort, sni string, dest Target) { // additional routes on ipPort. // // The ipPort is any valid net.Listen TCP address. -func (p *Proxy) AddSNIMatchRoute(ipPort string, matcher Matcher, dest Target) { - p.addRoute(ipPort, sniMatch{matcher: matcher, target: dest}) +func (p *Proxy) AddSNIMatchRoute(ipPort string, matcher Matcher, dest Target) uuid.UUID { + routeId := uuid.New() + + p.addRouteWithId(ipPort, sniMatch{matcher, dest, nil}, routeId) + + return routeId +} + +// AddSMTPSNIRouteFunc +func (p *Proxy) AddSMTPSNIRouteFunc(ipPort string, serverName string, targetLookup SNITargetFunc) uuid.UUID { + return p.AddSNIStartTLSFunc(ipPort, negotiateSMTPStartTLS(serverName), targetLookup) +} + +func (p *Proxy) AddIMAPSNIRouteFunc(ipPort string, targetLookup SNITargetFunc) uuid.UUID { + return p.AddSNIStartTLSFunc(ipPort, negotiateIMAPStartTLS(), targetLookup) +} + +func (p *Proxy) AddSNIStartTLSFunc(ipPort string, negFn NegotiateFunc, targetFn SNITargetFunc) uuid.UUID { + cfg := p.configFor(ipPort) + cfg.negotiateFunc = negFn + + return p.addRoute(ipPort, sniMatch{targetFunc: targetFn}) } // SNITargetFunc is the func callback used by Proxy.AddSNIRouteFunc. @@ -48,8 +70,8 @@ type SNITargetFunc func(ctx context.Context, sniName string) (t Target, ok bool) // AddSNIRouteFunc adds a route to ipPort that matches an SNI request and calls // fn to map its nap to a target. -func (p *Proxy) AddSNIRouteFunc(ipPort string, fn SNITargetFunc) { - p.addRoute(ipPort, sniMatch{targetFunc: fn}) +func (p *Proxy) AddSNIRouteFunc(ipPort string, fn SNITargetFunc) uuid.UUID { + return p.addRoute(ipPort, sniMatch{targetFunc: fn}) } type sniMatch struct { @@ -69,12 +91,12 @@ func (m sniMatch) match(br *bufio.Reader) (Target, string) { if t, ok := m.targetFunc(context.TODO(), sni); ok { return t, sni } - return nil, "" + return nil, sni } if m.matcher(context.TODO(), sni) { return m.target, sni } - return nil, "" + return nil, sni } // clientHelloServerName returns the SNI server name inside the TLS ClientHello, @@ -95,7 +117,7 @@ func clientHelloServerName(br *bufio.Reader) (sni string) { if err != nil { return "" } - tls.Server(sniSniffConn{r: bytes.NewReader(helloBytes)}, &tls.Config{ + _ = tls.Server(sniSniffConn{r: bytes.NewReader(helloBytes)}, &tls.Config{ GetConfigForClient: func(hello *tls.ClientHelloInfo) (*tls.Config, error) { sni = hello.ServerName return nil, nil @@ -112,4 +134,4 @@ type sniSniffConn struct { } func (c sniSniffConn) Read(p []byte) (int, error) { return c.r.Read(p) } -func (sniSniffConn) Write(p []byte) (int, error) { return 0, io.EOF } +func (sniSniffConn) Write(_ []byte) (int, error) { return 0, io.EOF } diff --git a/starttls_imap.go b/starttls_imap.go new file mode 100644 index 0000000..47031ce --- /dev/null +++ b/starttls_imap.go @@ -0,0 +1,45 @@ +package tcpproxy + +import ( + "errors" + "fmt" + "net" + "net/textproto" + "strings" +) + +func negotiateIMAPTLS(c *textproto.Conn) error { + c.Writer.PrintfLine("* OK [CAPABILITY IMAP4rev1 STARTTLS LOGINDISABLED] IMAP4rev1 Service Ready") + + cmdString, err := c.ReadLine() // "STARTTLS" + if err == nil { + tag, cmd, found := strings.Cut(cmdString, " ") + if !found { + return errors.New("") + return err + } + if cmd == "STARTTLS" { + c.Writer.PrintfLine("%s OK Begin TLS negotiation now", tag) + return nil + } + c.Writer.PrintfLine("%s %s Unsupported command", tag, cmd) + return errors.New(fmt.Sprintf("unsupported command %s received", cmd)) + } + + return err +} + +func negotiateIMAPStartTLS() NegotiateFunc { + return func(c net.Conn, cfg *config) bool { + // negotiate STARTTLS + t := textproto.NewConn(c) + + err := negotiateIMAPTLS(t) + if err == nil { + return true + } + + t.Close() + return false + } +} diff --git a/starttls_smtp.go b/starttls_smtp.go new file mode 100644 index 0000000..3690e77 --- /dev/null +++ b/starttls_smtp.go @@ -0,0 +1,55 @@ +package tcpproxy + +import ( + "errors" + "fmt" + "net" + "net/textproto" +) + +func greetSMTP(c *textproto.Conn, serverName string) (string, error) { + c.Writer.PrintfLine("220 %s Service ready", serverName) + l, e := c.ReadLine() // "EHLO + if e != nil { + return "", e + } + + var clientName string + _, e = fmt.Sscanf(l, "EHLO %s", &clientName) + if e != nil { + return "", e + } + + return clientName, nil +} + +func negotiateSMTPTLS(c *textproto.Conn, smtpServerName string) error { + c.Writer.PrintfLine("250-%s G'day!", smtpServerName) + c.Writer.PrintfLine("250 STARTTLS") + cmd, err := c.ReadLine() // "STARTTLS" + if err == nil { + if cmd == "STARTTLS" { + c.Writer.PrintfLine("220 Go ahead") + return nil + } + return errors.New("expecting STARTTLS") + } + + return err +} + +func negotiateSMTPStartTLS(serverName string) NegotiateFunc { + return func(c net.Conn, cfg *config) bool { + // negotiate STARTTLS + t := textproto.NewConn(c) + _, err := greetSMTP(t, serverName) + if err == nil { + err := negotiateSMTPTLS(t, serverName) + if err == nil { + return true + } + } + t.Close() + return false + } +} diff --git a/tcpproxy.go b/tcpproxy.go index d59c434..9f77468 100644 --- a/tcpproxy.go +++ b/tcpproxy.go @@ -56,11 +56,16 @@ import ( "bufio" "context" "errors" - "fmt" + "fmt" "io" "log" + "math/rand" "net" + "sync" "time" + + "github.com/google/uuid" + "github.com/pires/go-proxyproto" ) // Proxy is a proxy. Its zero value is a valid proxy that does @@ -69,6 +74,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 @@ -84,6 +90,15 @@ type Proxy struct { // Matcher reports whether hostname matches the Matcher's criteria. type Matcher func(ctx context.Context, hostname string) bool +// TargetLookup can be used to dynamically lookup a target address +// - returns an empty string and a not nil error if match not found +// - otherwise returns an ip:port string and nil if a match is found +//type TargetLookup func(ctx context.Context, hostname string) (string, error) + +type NegotiateFunc func(gc net.Conn, cfg *config) bool + +type NegotiateBackendFunc func(src *Conn, dst net.Conn) error + // equals is a trivial Matcher that implements string equality. func equals(want string) Matcher { return func(_ context.Context, got string) bool { @@ -93,7 +108,53 @@ func equals(want string) Matcher { // config contains the proxying state for one listener. type config struct { - routes []route + mu sync.Mutex + routes []routeWithId + + enableProxyProtocol bool + negotiateFunc NegotiateFunc + + defaultTarget Target +} + +func (c *config) AddRoute(r route) uuid.UUID { + return c.AddRouteWithId(r, uuid.New()) +} + +func (c *config) AddRouteWithId(r route, id uuid.UUID) uuid.UUID { + c.mu.Lock() + defer c.mu.Unlock() + + c.routes = append(c.routes, routeWithId{id, r}) + + return id +} + +func (c *config) Routes() []routeWithId { + c.mu.Lock() + defer c.mu.Unlock() + + return c.routes +} + +func (c *config) RemoveRouteById(routeId uuid.UUID) { + c.mu.Lock() + defer c.mu.Unlock() + + var newRoutes []routeWithId + + for _, i := range c.routes { + if i.Id != routeId { + newRoutes = append(newRoutes, i) + } + } + + c.routes = newRoutes +} + +type routeWithId struct { + Id uuid.UUID + Route route } // A route matches a connection to a target. @@ -119,6 +180,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) } @@ -128,9 +192,25 @@ func (p *Proxy) configFor(ipPort string) *config { return p.configs[ipPort] } -func (p *Proxy) addRoute(ipPort string, r route) { +func (p *Proxy) configExists(ipPort string) bool { + return p.configs[ipPort] != nil +} + +func (p *Proxy) addRoute(ipPort string, r route) uuid.UUID { cfg := p.configFor(ipPort) - cfg.routes = append(cfg.routes, r) + return cfg.AddRoute(r) +} + +func (p *Proxy) addRouteWithId(ipPort string, r route, id uuid.UUID) { + cfg := p.configFor(ipPort) + cfg.AddRouteWithId(r, id) +} + +func (p *Proxy) removeRouteById(ipPort string, routeId uuid.UUID) { + if p.configExists(ipPort) { + cfg := p.configFor(ipPort) + cfg.RemoveRouteById(routeId) + } } // AddRoute appends an always-matching route to the ipPort listener, @@ -140,8 +220,25 @@ func (p *Proxy) addRoute(ipPort string, r route) { // 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) uuid.UUID { + return p.addRoute(ipPort, fixedTarget{dest}) +} + +// RemoveRouteById removes the specified route from the ipPort listener +// +// This method won't remove an ipPort listener if there are no routes remaining +func (p *Proxy) RemoveRouteById(ipPort string, routeId uuid.UUID) { + p.removeRouteById(ipPort, routeId) +} + +func (p *Proxy) SetDefaultTarget(ipPort string, dest Target) { + cfg := p.configFor(ipPort) + cfg.defaultTarget = dest +} + +func (p *Proxy) EnableProxyProtocol(ipPort string, enableProxyProtocol bool) { + cfg := p.configFor(ipPort) + cfg.enableProxyProtocol = enableProxyProtocol } type fixedTarget struct { @@ -197,8 +294,13 @@ func (p *Proxy) Start() error { p.Close() return err } + + if config.enableProxyProtocol { + ln = &proxyproto.Listener{Listener: ln} + } + 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 @@ -209,37 +311,85 @@ 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) + } +} + +func findRoute(routes []routeWithId, br *bufio.Reader) (Target, string) { + var matchedHostname string + matchedRoutes := make([]Target, 0) + + for _, routeWithId := range routes { + target, hostName := routeWithId.Route.match(br) + matchedHostname = hostName + if target != nil { + matchedRoutes = append(matchedRoutes, target) + } } + + if len(matchedRoutes) == 0 { + return nil, matchedHostname + } + + if len(matchedRoutes) == 1 { + return matchedRoutes[0], matchedHostname + } + + randomIndex := rand.Intn(len(matchedRoutes)) + randomRoute := matchedRoutes[randomIndex] + + return randomRoute, matchedHostname } // 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 { + if cfg.negotiateFunc != nil { + if !cfg.negotiateFunc(c, cfg) { + c.Close() + return false + } + } + br := bufio.NewReader(c) - for _, route := range routes { - if target, hostName := route.match(br); target != nil { - if n := br.Buffered(); n > 0 { - peeked, _ := br.Peek(br.Buffered()) - c = &Conn{ - HostName: hostName, - Peeked: peeked, - Conn: c, - } - } - target.HandleConn(c) - return true + + target, hostName := findRoute(cfg.Routes(), br) + + if n := br.Buffered(); n > 0 { + peeked, _ := br.Peek(br.Buffered()) + c = &Conn{ + HostName: hostName, + Peeked: peeked, + Conn: c, } } + + if target != nil { + target.HandleConn(c) + return true + } + // TODO: hook for this? - log.Printf("tcpproxy: no routes matched conn %v/%v; closing", c.RemoteAddr().String(), c.LocalAddr().String()) + if cfg.defaultTarget != nil { + log.Printf("tcpproxy: no matching routes found for %s. using default target %s", hostName, cfg.defaultTarget) + cfg.defaultTarget.HandleConn(c) + return true + } + + if hostName == "" { + log.Printf("tcpproxy: no routes matched conn %v/%v; closing", c.RemoteAddr().String(), c.LocalAddr().String()) + } else { + log.Printf("tcpproxy: no route matched '%s', conn %v/%v; closing", hostName, c.RemoteAddr().String(), c.LocalAddr().String()) + } + c.Close() return false } @@ -293,11 +443,13 @@ type Target interface { HandleConn(net.Conn) } -// To is shorthand way of writing &tcpproxy.DialProxy{Addr: addr}. +// To is a shorthand way of writing &tcpproxy.DialProxy{Addr: addr}. func To(addr string) *DialProxy { return &DialProxy{Addr: addr} } +type AccessLogger func(net.Conn, net.Conn) + // DialProxy implements Target by dialing a new connection to Addr // and then proxying data back and forth. // @@ -334,6 +486,11 @@ type DialProxy struct { // no graceful downgrade. // If zero, no PROXY header is sent. Currently, version 1 is supported. ProxyProtocolVersion int + + // NegotiateBackendFunc optionally specifies a function to negotiate the backend connection. + NegotiateBackendFunc NegotiateBackendFunc + + AccessLogger AccessLogger } // UnderlyingConn returns c.Conn if c of type *Conn, @@ -395,6 +552,7 @@ func (dp *DialProxy) HandleConn(src net.Conn) { } defer src.Close() + if ka := dp.keepAlivePeriod(); ka > 0 { for _, c := range []net.Conn{src, dst} { if c, ok := tcpConn(c); ok { @@ -404,6 +562,19 @@ func (dp *DialProxy) HandleConn(src net.Conn) { } } + if dp.NegotiateBackendFunc != nil { + srcConn := src.(*Conn) + err = dp.NegotiateBackendFunc(srcConn, dst) + if err != nil { + dp.onDialError()(src, err) + return + } + } + + if dp.AccessLogger != nil { + dp.AccessLogger(src, dst) + } + errc := make(chan error, 2) go proxyCopy(errc, src, dst) go proxyCopy(errc, dst, src) @@ -412,10 +583,15 @@ func (dp *DialProxy) HandleConn(src net.Conn) { } func (dp *DialProxy) sendProxyHeader(w io.Writer, src net.Conn) error { - switch dp.ProxyProtocolVersion { - case 0: + if dp.ProxyProtocolVersion == 0 { return nil - case 1: + } + + var header *proxyproto.Header + if proxyConn, ok := src.(*proxyproto.Conn); ok { + header = proxyConn.ProxyHeader() + } else { + var srcAddr, dstAddr *net.TCPAddr if a, ok := src.RemoteAddr().(*net.TCPAddr); ok { srcAddr = a @@ -424,20 +600,23 @@ func (dp *DialProxy) sendProxyHeader(w io.Writer, src net.Conn) error { dstAddr = a } - if srcAddr == nil || dstAddr == nil { - _, err := io.WriteString(w, "PROXY UNKNOWN\r\n") - return err + transportProtocol := proxyproto.TCPv4 + if srcAddr.IP.To4() == nil { + transportProtocol = proxyproto.TCPv6 } - family := "TCP4" - if srcAddr.IP.To4() == nil { - family = "TCP6" + header = &proxyproto.Header{ + Version: byte(dp.ProxyProtocolVersion), + Command: proxyproto.PROXY, + TransportProtocol: transportProtocol, + SourceAddr: srcAddr, + DestinationAddr: dstAddr, } - _, err := fmt.Fprintf(w, "PROXY %s %s %s %d %d\r\n", family, srcAddr.IP, dstAddr.IP, srcAddr.Port, dstAddr.Port) - return err - default: - return fmt.Errorf("PROXY protocol version %d not supported", dp.ProxyProtocolVersion) } + // After the connection was created write the proxy headers first + _, err := header.WriteTo(w) + + return err } // proxyCopy is the function that copies bytes around. diff --git a/tcpproxy_test.go b/tcpproxy_test.go index 0346a7a..ea41aed 100644 --- a/tcpproxy_test.go +++ b/tcpproxy_test.go @@ -174,6 +174,60 @@ 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{} + + r1Id := cfg.AddRoute(r1) + r2Id := cfg.AddRoute(r2) + + routes := cfg.Routes() + if len(routes) != 2 { + t.Fatalf("got %d routes, want 1", len(routes)) + } + + cfg.RemoveRouteById(r1Id) + + routes = cfg.Routes() + if len(routes) != 1 { + t.Fatalf("got %d routes, want 1", len(routes)) + } + + if routes[0].Id != r2Id { + 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") + r2 := To("bar.com:443") + + p := testProxy(t, front) + r1Id := p.AddRoute(testFrontAddr, r1) + r2Id := 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.RemoveRouteById(testFrontAddr, r2Id) + routes = config.Routes() + if len(routes) != 1 { + t.Fatalf("got %d routes, want 1", len(routes)) + } + + if routes[0].Id != r1Id { + t.Fatalf("got %s, want %s", r1Id, routes[0].Id) + } +} + func TestBufferedClose(t *testing.T) { front := newLocalListener(t) defer front.Close() @@ -460,6 +514,7 @@ func cert(t *testing.T, domain string) tls.Certificate { KeyUsage: x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, BasicConstraintsValid: true, + DNSNames: []string{domain}, } derBytes, err := x509.CreateCertificate(rand.Reader, template, template, &private.PublicKey, private) @@ -495,7 +550,6 @@ func newTLSServer(t *testing.T, domain string) net.Listener { cfg := &tls.Config{ Certificates: []tls.Certificate{cert}, } - cfg.BuildNameToCertificate() conn := tls.Server(rawConn, cfg) if _, err = io.WriteString(conn, domain); err != nil { t.Errorf("writing to tlsconn: %s", err)