refactor websocket

Author: Nicolas Chatelain

cmd/proxy/main.go

@@ -22,7 +22,7 @@ var (
 func main() {
 	var allowDomains []string
 	var verboseFlag = flag.Bool("v", false, "enable verbose mode")
-	var listenInterface = flag.String("laddr", "0.0.0.0:11601", "listening address ")
+	var listenInterface = flag.String("laddr", "0.0.0.0:11601", "listening address (prefix with https:// for websocket)")
 	var enableAutocert = flag.Bool("autocert", false, "automatically request letsencrypt certificates, requires port 80 to be accessible")
 	var enableSelfcert = flag.Bool("selfcert", false, "dynamically generate self-signed certificates")
 	var certFile = flag.String("certfile", "certs/cert.pem", "TLS server certificate")

pkg/controller/controller.go

@@ -12,7 +12,6 @@ import (
 	"net/http"
 	"nhooyr.io/websocket"
 	"strings"
-	"time"
 )
 
 type Controller struct {
@@ -34,29 +33,10 @@ type ControllerConfig struct {
 	DomainWhitelist []string
 }
 
-var wsconn net.Conn
-
 func New(config ControllerConfig) Controller {
 	return Controller{Network: "tcp", Connection: make(chan net.Conn, 1024), ControllerConfig: config, startchan: make(chan error), SelfCertCache: "ligolo-selfcerts"}
 }
 
-type ligoloHttpServer struct {
-	// logf controls where logs are sent.
-	logf func(f string, v ...interface{})
-}
-
-func (s ligoloHttpServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-
-	c, err := websocket.Accept(w, r, nil)
-	if err != nil {
-		logrus.Error(err)
-		return
-	}
-	netctx, _ := context.WithTimeout(context.Background(), time.Hour*999999)
-	wsconn = websocket.NetConn(netctx, c, websocket.MessageBinary)
-	return
-}
-
 func (c *Controller) WaitForReady() error {
 	return <-c.startchan
 }
@@ -114,40 +94,31 @@ func (c *Controller) ListenAndServe() {
 		return
 	}
 
-	if strings.Contains(c.Address, "https://") {
+	if strings.HasPrefix(c.Address, "https://") {
 		//SSL websocket protocol
 		listener, err := tls.Listen(c.Network, strings.Replace(c.Address, "https://", "", 1), &tlsConfig)
 		if err != nil {
-			logrus.Fatal(err)
+			c.startchan <- err
+			return
 		}
 		defer listener.Close()
-		close(c.startchan) // Controller is listening.
+
+		c.startchan <- nil
 		logrus.Infof("Listening websocket on %s", c.Address)
 
 		s := &http.Server{
-			Handler: ligoloHttpServer{},
-		}
-		for {
-			//start http handler in go routine
-			go func() {
-				err = s.Serve(listener)
-			}()
-			if err != nil {
-				logrus.Error(err)
-			}
-			//manual waiting until handler got connection and global variable wsconn is set by http handler
-			//this not so gracefully but effective ))
-			for {
-				if wsconn != nil {
-					logrus.Debugf("Got websocket connection %s", wsconn.RemoteAddr())
-					c.Connection <- wsconn
-					wsconn = nil
-					break
+			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				ws, err := websocket.Accept(w, r, nil)
+				if err != nil {
+					logrus.Error(err)
+					return
 				}
-				//add some sleep to reduce CPU usage, because it is in loop
-				time.Sleep(time.Millisecond * 500)
-			}
+				netctx := context.Background()
+
+				c.Connection <- websocket.NetConn(netctx, ws, websocket.MessageBinary)
+			}),
 		}
+		err = s.Serve(listener)
 	} else {
 		//direct listen with legacy ligolo-ng protocol
 		listener, err := tls.Listen(c.Network, c.Address, &tlsConfig)