Refactor for easier testing

cmd/ck-client/ck-client.go

@@ -168,6 +168,10 @@ func main() {
 	if authInfo.Unordered {
 		client.RouteUDP(localConfig, seshMaker)
 	} else {
-		client.RouteTCP(localConfig, seshMaker)
+		listener, err := net.Listen("tcp", localConfig.LocalAddr)
+		if err != nil {
+			log.Fatal(err)
+		}
+		client.RouteTCP(listener, localConfig.Timeout, seshMaker)
 	}
 }

cmd/ck-server/ck-server.go

@@ -155,14 +155,7 @@ func main() {
 		if err != nil {
 			log.Fatal(err)
 		}
-		for {
+		server.Serve(listener, sta)
-			conn, err := listener.Accept()
-			if err != nil {
-				log.Errorf("%v", err)
-				continue
-			}
-			go server.DispatchConnection(conn, sta)
-		}
 	}
 
 	for i, addr := range bindAddr {

internal/client/connector.go

@@ -15,6 +15,7 @@ import (
 
 func MakeSession(connConfig remoteConnConfig, authInfo authInfo, dialer common.Dialer, isAdmin bool) *mux.Session {
 	log.Info("Attempting to start a new session")
+	//TODO: let caller set this
 	if !isAdmin {
 		// sessionID is usergenerated. There shouldn't be a security concern because the scope of
 		// sessionID is limited to its UID.

internal/client/piper.go

@@ -100,14 +100,10 @@ start:
 
 }
 
-func RouteTCP(localConfig localConnConfig, newSeshFunc func() *mux.Session) {
+func RouteTCP(listener net.Listener, streamTimeout time.Duration, newSeshFunc func() *mux.Session) {
-	tcpListener, err := net.Listen("tcp", localConfig.LocalAddr)
-	if err != nil {
-		log.Fatal(err)
-	}
 	var sesh *mux.Session
 	for {
-		localConn, err := tcpListener.Accept()
+		localConn, err := listener.Accept()
 		if err != nil {
 			log.Fatal(err)
 			continue
@@ -142,7 +138,7 @@ func RouteTCP(localConfig localConnConfig, newSeshFunc func() *mux.Session) {
 				}
 			}()
 			//util.Pipe(stream, localConn, localConfig.Timeout)
-			if _, err = common.Copy(stream, localConn, localConfig.Timeout); err != nil {
+			if _, err = common.Copy(stream, localConn, streamTimeout); err != nil {
 				log.Tracef("copying proxy client to stream: %v", err)
 			}
 		}()

internal/server/dispatcher.go

@@ -16,7 +16,18 @@ import (
 
 var b64 = base64.StdEncoding.EncodeToString
 
-func DispatchConnection(conn net.Conn, sta *State) {
+func Serve(l net.Listener, sta *State) {
+	for {
+		conn, err := l.Accept()
+		if err != nil {
+			log.Errorf("%v", err)
+			continue
+		}
+		go dispatchConnection(conn, sta)
+	}
+}
+
+func dispatchConnection(conn net.Conn, sta *State) {
 	remoteAddr := conn.RemoteAddr()
 	var err error
 	buf := make([]byte, 1500)