From b9670a0fef72e5bfed2ad861af8f187b40298299 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A1=D1=8B=D1=81?= =?UTF-8?q?=D0=BE=D0=B5=D0=B2?= Date: Mon, 23 Oct 2023 21:32:49 +0800 Subject: [PATCH] Improves handling connection errors --- pkg/cli/cli.go | 6 +++++- pkg/ws/ws.go | 12 ++++++++++-- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/pkg/cli/cli.go b/pkg/cli/cli.go index cb9c0fc..a20b1df 100644 --- a/pkg/cli/cli.go +++ b/pkg/cli/cli.go @@ -95,7 +95,11 @@ func (c *CLI) Run(outputFile *os.File) error { continue } - case msg := <-c.wsConn.Messages: + case msg, ok := <-c.wsConn.Messages: + if !ok { + return nil + } + output, err := c.formater.FormatMessage(msg) if err != nil { log.Printf("Fail to format message: %s, %s\n", err, msg.Data) diff --git a/pkg/ws/ws.go b/pkg/ws/ws.go index cf55384..030583b 100644 --- a/pkg/ws/ws.go +++ b/pkg/ws/ws.go @@ -2,8 +2,8 @@ package ws import ( "crypto/tls" - "log" + "github.com/fatih/color" "golang.org/x/net/websocket" ) @@ -56,12 +56,20 @@ func NewWS(url string, opts Options) (*Connection, error) { messages := make(chan Message, WSMessageBufferSize) go func(messages chan Message) { + defer close(messages) + for { var msg string err = websocket.Message.Receive(ws, &msg) if err != nil { - log.Fatal("Fail to read from WS connection:", err) + if err.Error() == "EOF" { + color.New(color.FgRed).Println("Connection closed by the server") + } else { + color.New(color.FgRed).Println("Fail read from connection: ", err) + } + + return } messages <- Message{Type: Response, Data: msg}