Improves handling connection errors

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)

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}