put the api server on the right listener 🤦

Gopkg.toml

@@ -83,4 +83,7 @@
 
 [[constraint]]
   name = "github.com/gomodule/redigo"
-  version = "2.0.0"
+  version = "2.0.0"
+[[constraint]]
+  name = "github.com/soheilhy/cmux"
+  version = "0.1.4"

api/api.go

@@ -40,6 +40,11 @@ type Handler struct {
 	redisPool *redis.Pool
 }
 
+// NewServer ...
+func NewServer(logger *logrus.Entry, redisPool *redis.Pool) *http.Server {
+	return &http.Server{Handler: NewHandler(logger, redisPool)}
+}
+
 // NewHandler creates a new API handler
 func NewHandler(logger *logrus.Entry, redisPool *redis.Pool) *Handler {
 	r := chi.NewRouter()

local/http2_handler_test.go

@@ -134,7 +134,7 @@ func TestHTTP2Handler(t *testing.T) {
 	assert.NoError(t, err, "Should be no error creating test handler")
 
 	t.Run("Test_dial", func(t *testing.T) {
-		sConnCh := make(chan *net.TCPConn)
+		sConnCh := make(chan net.Conn)
 
 		go func(ln *net.TCPListener) {
 			s, err := ln.AcceptTCP()

net/single_port_listener_factory.go

@@ -4,13 +4,10 @@ import (
 	"crypto/tls"
 	"fmt"
 	"net"
-	"net/http"
 	"strings"
 	"sync"
 
-	"github.com/gomodule/redigo/redis"
 	"github.com/sirupsen/logrus"
-	"github.com/superfly/wormhole/api"
 )
 
 type sharedPortTLSListenerFactory struct {
@@ -20,16 +17,14 @@ type sharedPortTLSListenerFactory struct {
 	fLock   sync.Mutex
 	logger  *logrus.Entry
 
-	stopC      chan struct{}
-	apiHandler *api.Handler
+	stopC chan struct{}
 }
 
 // SharedPortTLSListenerFactoryArgs provides the data needed to create a SharedPortTLSListenerFactory
 type SharedPortTLSListenerFactoryArgs struct {
 	TLSConfig *tls.Config
 	Address   string
 	Logger    *logrus.Logger
-	RedisPool *redis.Pool
 }
 
 // NewSharedPortTLSListenerFactory creates a new listener factory for shared port TLS
@@ -43,11 +38,10 @@ func NewSharedPortTLSListenerFactory(args *SharedPortTLSListenerFactoryArgs) (Li
 	}
 
 	f := &sharedPortTLSListenerFactory{
-		listener:   listener,
-		forward:    make(map[string]*sharedPortTLSListener),
-		stopC:      make(chan struct{}),
-		logger:     args.Logger.WithFields(logrus.Fields{"prefix": "shared_port_tls_listener_factory"}),
-		apiHandler: api.NewHandler(args.Logger.WithFields(logrus.Fields{"prefix": "api_handler"}), args.RedisPool),
+		listener: listener,
+		forward:  make(map[string]*sharedPortTLSListener),
+		stopC:    make(chan struct{}),
+		logger:   args.Logger.WithFields(logrus.Fields{"prefix": "shared_port_tls_listener_factory"}),
 	}
 
 	go func() {
@@ -126,11 +120,6 @@ func (sl *sharedPortTLSListenerFactory) populateCh() error {
 					return
 				}
 
-				if id == "api" {
-					http.Serve(api.NewSingleConnListener(sl.listener, c), sl.apiHandler)
-					return
-				}
-
 				sl.fLock.Lock()
 				ch, ok := sl.forward[id]
 				sl.fLock.Unlock()

net/utils.go

@@ -3,12 +3,13 @@ package net
 import (
 	"crypto/tls"
 	"fmt"
-	"golang.org/x/net/http2"
 	"io"
 	"net"
 	"reflect"
 	"strings"
 	"time"
+
+	"golang.org/x/net/http2"
 )
 
 // CopyDirection describes the direction of data copying in full-duplex link
@@ -52,7 +53,7 @@ type TLSWrapperFunc func(conn net.Conn, cfg *tls.Config) *tls.Conn
 
 // GenericTLSWrap takes a TCP connection, a tls config, and an upgrade function
 // and returns the new connection
-func GenericTLSWrap(conn *net.TCPConn, cfg *tls.Config, tFunc TLSWrapperFunc) (*tls.Conn, error) {
+func GenericTLSWrap(conn net.Conn, cfg *tls.Config, tFunc TLSWrapperFunc) (*tls.Conn, error) {
 	var tConn *tls.Conn
 
 	tCfg := cfg.Clone()
@@ -92,7 +93,7 @@ func GenericTLSWrap(conn *net.TCPConn, cfg *tls.Config, tFunc TLSWrapperFunc) (*
 // While technically the golang implementation will allow us not to perform ALPN,
 // this breaks the http/2 spec. The goal here is to follow the RFC to the letter
 // as documented in http://httpwg.org/specs/rfc7540.html#starting
-func HTTP2ALPNTLSWrap(conn *net.TCPConn, cfg *tls.Config, tFunc TLSWrapperFunc) (*tls.Conn, error) {
+func HTTP2ALPNTLSWrap(conn net.Conn, cfg *tls.Config, tFunc TLSWrapperFunc) (*tls.Conn, error) {
 	protoCfg := cfg.Clone()
 	// TODO: append here
 	protoCfg.NextProtos = []string{http2.NextProtoTLS}

remote.go

@@ -1,6 +1,7 @@
 package wormhole
 
 import (
+	"net"
 	"net/url"
 	"os"
 	"os/signal"
@@ -9,6 +10,8 @@ import (
 
 	"github.com/gomodule/redigo/redis"
 	"github.com/sirupsen/logrus"
+	"github.com/soheilhy/cmux"
+	"github.com/superfly/wormhole/api"
 	"github.com/superfly/wormhole/config"
 	wnet "github.com/superfly/wormhole/net"
 	handler "github.com/superfly/wormhole/remote"
@@ -37,6 +40,15 @@ func StartRemote(cfg *config.ServerConfig) {
 		log.Fatalf("Could not create listener factory: %+v", err)
 	}
 
+	l, err := net.Listen("tcp", ":"+cfg.Port)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	m := cmux.New(l)
+	httpL := m.Match(cmux.HTTP1())
+	tcpL := m.Match(cmux.Any())
+
 	switch cfg.Protocol {
 	case config.SSH:
 		h, err = handler.NewSSHHandler(cfg, registry, redisPool, listenerFactory)
@@ -58,7 +70,12 @@ func StartRemote(cfg *config.ServerConfig) {
 	}
 
 	go handleDeath(h, registry)
-	server.ListenAndServe(":"+cfg.Port, h)
+	go func() {
+		err := api.NewServer(cfg.Logger.WithFields(logrus.Fields{"prefix": "api"}), redisPool).Serve(httpL)
+		log.Println("API Server is done, err?", err)
+	}()
+	go server.Serve(tcpL, h)
+	m.Serve()
 }
 
 func listenerFactoryFromConfig(registry *session.Registry, cfg *config.ServerConfig) (wnet.ListenerFactory, error) {
@@ -82,7 +99,6 @@ func listenerFactoryFromConfig(registry *session.Registry, cfg *config.ServerCon
 			Address:   ":" + cfg.SharedTLSForwardingPort,
 			Logger:    cfg.Logger,
 			TLSConfig: tlsconf.GetDefaultConfig(),
-			RedisPool: redisPool,
 		}
 		sharedL, err := wnet.NewSharedPortTLSListenerFactory(sharedArgs)
 		if err != nil {

remote/handler.go

@@ -6,6 +6,6 @@ import "net"
 // It's the entry point for starting a session, managing a handshake, auth
 // and encryption (e.g. SSH)
 type Handler interface {
-	Serve(*net.TCPConn)
+	Serve(net.Conn)
 	Close()
 }

remote/http2_handler.go

@@ -51,7 +51,7 @@ func NewHTTP2Handler(cfg *config.ServerConfig, registry *session.Registry, pool
 // We are explicit with the *net.TCPConn since we need to be this way - and let the handler and
 // sessions handle wrapping in TLS. Having a TCPConn all the way down will highlight the dangers
 // of sending data over the socket without first wrapping in TLS
-func (h *HTTP2Handler) Serve(conn *net.TCPConn) {
+func (h *HTTP2Handler) Serve(conn net.Conn) {
 	tlsConn, err := h.genericTLSWrap(conn)
 	if err != nil {
 		h.logger.Errorf("error establishing tls session: " + err.Error())
@@ -119,15 +119,15 @@ func (h *HTTP2Handler) Serve(conn *net.TCPConn) {
 	}
 }
 
-func (h *HTTP2Handler) genericTLSWrap(conn *net.TCPConn) (*tls.Conn, error) {
+func (h *HTTP2Handler) genericTLSWrap(conn net.Conn) (*tls.Conn, error) {
 	return wnet.GenericTLSWrap(conn, h.tlsConfig, tls.Server)
 }
 
 // NOTE: The ALPN is a requirement of the spec for HTTP/2 capability discovery
 // While technically the golang implementation will allow us not to perform ALPN,
 // this breaks the http/2 spec. The goal here is to follow the RFC to the letter
 // as documented in http://httpwg.org/specs/rfc7540.html#starting
-func (h *HTTP2Handler) http2ALPNTLSWrap(conn *net.TCPConn) (*tls.Conn, error) {
+func (h *HTTP2Handler) http2ALPNTLSWrap(conn net.Conn) (*tls.Conn, error) {
 	return wnet.HTTP2ALPNTLSWrap(conn, h.tlsConfig, tls.Server)
 }
 

remote/server.go

@@ -1,7 +1,6 @@
 package remote
 
 import (
-	"fmt"
 	"net"
 
 	"github.com/sirupsen/logrus"
@@ -12,35 +11,32 @@ type Server struct {
 	Logger *logrus.Logger
 }
 
-// ListenAndServe accepts incoming wormhole connections and passes them to the handler
-func (s *Server) ListenAndServe(addr string, handler Handler) error {
+// Serve accepts incoming wormhole connections and passes them to the handler
+func (s *Server) Serve(listener net.Listener, handler Handler) error {
 	log := s.Logger.WithFields(logrus.Fields{"prefix": "Server"})
-	listener, err := s.newTCPListener(addr)
-	if err != nil {
-		return fmt.Errorf("Failed to listen on %s (%s)", addr, err.Error())
-	}
 
 	for {
-		tcpConn, err := listener.AcceptTCP()
+		conn, err := listener.Accept()
 		if err != nil {
 			log.Errorf("Failed to accept wormhole connection (%s)", err.Error())
 			break
 		}
-		log.Debugln("Accepted wormhole TCP conn from:", tcpConn.RemoteAddr())
 
-		go handler.Serve(tcpConn)
+		log.Println("Accepted wormhole TCP conn from:", conn.RemoteAddr())
+
+		go handler.Serve(conn)
 	}
 	return nil
 }
 
-func (s *Server) newTCPListener(addr string) (*net.TCPListener, error) {
-	ln, err := net.Listen("tcp", addr)
-	if err != nil {
-		return nil, err
-	}
-	tcpLN, ok := ln.(*net.TCPListener)
-	if !ok {
-		return nil, fmt.Errorf("Could not create tcp listener")
-	}
-	return tcpLN, nil
-}
+// func (s *Server) newTCPListener(addr string) (*net.TCPListener, error) {
+// 	ln, err := net.Listen("tcp", addr)
+// 	if err != nil {
+// 		return nil, err
+// 	}
+// 	tcpLN, ok := ln.(*net.TCPListener)
+// 	if !ok {
+// 		return nil, fmt.Errorf("Could not create tcp listener")
+// 	}
+// 	return tcpLN, nil
+// }

remote/ssh_handler.go

@@ -64,7 +64,7 @@ func NewSSHHandler(cfg *config.ServerConfig, registry *session.Registry, pool *r
 }
 
 // Serve accepts incoming wormhole connections and passes them to the handler
-func (s *SSHHandler) Serve(conn *net.TCPConn) {
+func (s *SSHHandler) Serve(conn net.Conn) {
 	conn.RemoteAddr()
 	ctx, err := s.limiter.Get(ipForConn(conn))
 	if err != nil {

remote/tcp_handler.go

@@ -53,7 +53,7 @@ func NewTCPHandler(cfg *config.ServerConfig, registry *session.Registry, pool *r
 }
 
 // Serve accepts incoming wormhole connections and passes them to the handler
-func (h *TCPHandler) Serve(conn *net.TCPConn) {
+func (h *TCPHandler) Serve(conn net.Conn) {
 	var useConn net.Conn
 	if h.tlsConfig != nil {
 		var err error

session/http2_session_test.go

@@ -56,14 +56,14 @@ func newServerClientTLSConns(alpn bool) (serverTLSConn *tls.Conn, clientTLSConn
 
 	sTLSConnCh := make(chan *tls.Conn)
 
-	var wrapFunc func(*net.TCPConn, *tls.Config, wnet.TLSWrapperFunc) (*tls.Conn, error)
+	var wrapFunc func(net.Conn, *tls.Config, wnet.TLSWrapperFunc) (*tls.Conn, error)
 	if alpn {
 		wrapFunc = wnet.HTTP2ALPNTLSWrap
 	} else {
 		wrapFunc = wnet.GenericTLSWrap
 	}
 
-	go func(sConn *net.TCPConn) {
+	go func(sConn net.Conn) {
 		sTLSConn, err := wrapFunc(sConn, serverTLSConfig, tls.Server)
 		if err != nil {
 			log.Errorf("Error creating tls wrap server")