chore: Handle ping msg for shell and port-forward (#241)

Qovery · Jan 8, 2024 · 6ded257 · 6ded257
1 parent aecb34e
commit 6ded257
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 7 deletions.
diff --git a/go.mod b/go.mod
@@ -19,7 +19,7 @@ require (
 	github.com/pkg/errors v0.9.1
 	github.com/posthog/posthog-go v0.0.0-20221221115252-24dfed35d71a
 	github.com/pterm/pterm v0.12.55
-	github.com/qovery/qovery-client-go v0.0.0-20240102155005-3a6753cad4f4
+	github.com/qovery/qovery-client-go v0.0.0-20240104104714-c749f31f31e6
 	github.com/sirupsen/logrus v1.9.0
 	github.com/spf13/cobra v1.6.1
 	github.com/spf13/pflag v1.0.5

diff --git a/go.sum b/go.sum
@@ -198,6 +198,8 @@ github.com/qovery/qovery-client-go v0.0.0-20240102111457-e9d5e9a578cd h1:K3H0JYT
 github.com/qovery/qovery-client-go v0.0.0-20240102111457-e9d5e9a578cd/go.mod h1:5QD7sC1Z6XCCYd31c4XKVwGdEOjvtgG0NDcaVDoWb+o=
 github.com/qovery/qovery-client-go v0.0.0-20240102155005-3a6753cad4f4 h1:jF8XOqAsbbZeEJOjyVDsy8fI9RNCLKNP/xa38mgizpo=
 github.com/qovery/qovery-client-go v0.0.0-20240102155005-3a6753cad4f4/go.mod h1:5QD7sC1Z6XCCYd31c4XKVwGdEOjvtgG0NDcaVDoWb+o=
+github.com/qovery/qovery-client-go v0.0.0-20240104104714-c749f31f31e6 h1:WRC6Gt5bWNax8UT1ZbKrK/ITCsLC/A7cq7bZtBbQSOE=
+github.com/qovery/qovery-client-go v0.0.0-20240104104714-c749f31f31e6/go.mod h1:5QD7sC1Z6XCCYd31c4XKVwGdEOjvtgG0NDcaVDoWb+o=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis=
 github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=

diff --git a/pkg/port-forward.go b/pkg/port-forward.go
@@ -36,12 +36,22 @@ func (w WebsocketPortForward) Write(p []byte) (n int, err error) {
 	return len(p), err
 }
 func (w WebsocketPortForward) Read(p []byte) (n int, err error) {
-	_, msg, err := w.ws.ReadMessage()
-	if err != nil {
-		return 0, err
-	}
+	for {
+		msgType, msg, err := w.ws.ReadMessage()
+		if err != nil {
+			return 0, err
+		}
+
+		if msgType == websocket.CloseMessage {
+			return 0, io.EOF
+		}
 
-	return copy(p, msg), err
+		if msgType != websocket.BinaryMessage {
+			continue
+		}
+
+		return copy(p, msg), err
+	}
 }
 
 func mkWebsocketConn(req *PortForwardRequest) (*WebsocketPortForward, error) {

diff --git a/pkg/shell.go b/pkg/shell.go
@@ -103,7 +103,7 @@ func createWebsocketConn(req *ShellRequest) (*websocket.Conn, error) {
 func readWebsocketConnection(wsConn *websocket.Conn, currentConsole console.Console, done chan struct{}) {
 	defer close(done)
 	for {
-		_, msg, err := wsConn.ReadMessage()
+		msgType, msg, err := wsConn.ReadMessage()
 		if err != nil {
 			var e *websocket.CloseError
 			if errors.As(err, &e) {
@@ -117,6 +117,15 @@ func readWebsocketConnection(wsConn *websocket.Conn, currentConsole console.Cons
 			log.Error("error while reading on websocket:", err)
 			return
 		}
+
+		if msgType == websocket.CloseMessage {
+			return
+		}
+
+		if msgType != websocket.BinaryMessage {
+			continue
+		}
+
 		if _, err = currentConsole.Write(msg); err != nil {
 			log.Error("error while writing in console:", err)
 			return