Switch from golang.org/x/net/websocket to nhooyr.io/websocket and handle NotFound errors (#105)

* Switch from golang.org/x/net/websocket to nhooyr.io/websocket

* Do not attach errors that we can handle to the Gin's context

* Add missing newline to "no credentials specified or found, ..." message

* Fix potential NPE in ChooseUsernameAndPassword()

* Fix type in PortForward() error message in "orchard ssh vm"

* Fix potential NPE in Connections()

* Use header.Set() for consistency's sake for Authorization header

internal/command/ssh/vm.go

@@ -56,9 +56,7 @@ func runSSHVM(cmd *cobra.Command, args []string) error {
 
 	wsConn, err := client.VMs().PortForward(cmd.Context(), name, 22, wait)
 	if err != nil {
-		fmt.Printf("failed to forward an SSH port to VM %s: %v\n", name, err)
-
-		return err
+		return fmt.Errorf("%w: failed to setup port-forwarding to the VM %q: %v", ErrFailed, name, err)
 	}
 	defer wsConn.Close()
 
@@ -78,9 +76,6 @@ func runSSHVM(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return fmt.Errorf("%w: failed to establish an SSH connection: %v", ErrFailed, err)
 	}
-	defer func() {
-		_ = sshConn.Close()
-	}()
 
 	sshClient := ssh.NewClient(sshConn, chans, reqs)
 
@@ -188,16 +183,16 @@ func ChooseUsernameAndPassword(
 
 	// Try to get the credentials from the VM's object stored on controller
 	vm, err := client.VMs().Get(ctx, vmName)
-	if err != nil {
-		fmt.Printf("failed to retrieve VM %s's credentials from the API server: %v\n", vmName, err)
-	}
-
-	if vm.Username != "" && vm.Password != "" {
+	if err == nil && vm.Username != "" && vm.Password != "" {
 		return vm.Username, vm.Password
+	} else if err != nil {
+		fmt.Fprintf(os.Stderr, "failed to retrieve VM %s's credentials from the API server: %v\n",
+			vmName, err)
 	}
 
 	// Fall back
-	_, _ = fmt.Fprintf(os.Stderr, "no credentials specified or found, trying default admin:admin credentials...")
+	_, _ = fmt.Fprintf(os.Stderr, "no credentials specified or found, "+
+		"trying default admin:admin credentials...\n")
 
 	return "admin", "admin"
 }

internal/command/vnc/vm.go

@@ -72,6 +72,7 @@ func runVNCVM(cmd *cobra.Command, args []string) (err error) {
 
 					return
 				}
+				defer wsConn.Close()
 
 				if err := proxy.Connections(wsConn, conn); err != nil {
 					fmt.Printf("failed to forward port: %v\n", err)

internal/controller/api_vms_portforward.go

@@ -9,8 +9,8 @@ import (
 	"github.com/cirruslabs/orchard/rpc"
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
-	"golang.org/x/net/websocket"
 	"net/http"
+	"nhooyr.io/websocket"
 	"strconv"
 	"time"
 )
@@ -95,11 +95,26 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
 	// worker will asynchronously start port-forwarding so we wait
 	select {
 	case fromWorkerConnection := <-boomerangConnCh:
-		websocket.Handler(func(wsConn *websocket.Conn) {
-			if err := proxy.Connections(wsConn, fromWorkerConnection); err != nil {
-				controller.logger.Warnf("failed to port-forward: %v", err)
-			}
-		}).ServeHTTP(ctx.Writer, ctx.Request)
+		wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{
+			OriginPatterns: []string{"*"},
+		})
+		if err != nil {
+			return responder.Error(err)
+		}
+
+		expectedMsgType := websocket.MessageBinary
+
+		// Backwards compatibility with older Orchard clients
+		// using "golang.org/x/net/websocket" package
+		if ctx.Request.Header.Get("User-Agent") == "" {
+			expectedMsgType = websocket.MessageText
+		}
+
+		wsConnAsNetConn := websocket.NetConn(ctx, wsConn, expectedMsgType)
+
+		if err := proxy.Connections(wsConnAsNetConn, fromWorkerConnection); err != nil {
+			controller.logger.Warnf("failed to port-forward: %v", err)
+		}
 
 		return responder.Empty()
 	case <-ctx.Done():

internal/proxy/proxy.go

@@ -39,7 +39,7 @@ func Connections(left net.Conn, right net.Conn) (finalErr error) {
 		recordErr(<-rightErrCh)
 	}
 
-	if strings.Contains(finalErr.Error(), "use of closed network connection") {
+	if finalErr != nil && strings.Contains(finalErr.Error(), "use of closed network connection") {
 		finalErr = nil
 	}
 

internal/responder/error.go

@@ -20,9 +20,12 @@ func Error(err error) Responder {
 
 func (responder *ErrorResponder) Respond(c *gin.Context) {
 	var code = http.StatusInternalServerError
+
 	if errors.Is(responder.err, storepkg.ErrNotFound) {
 		code = http.StatusNotFound
+	} else {
+		_ = c.Error(responder.err)
 	}
-	_ = c.Error(responder.err)
+
 	c.Status(code)
 }