Skip to content

Commit

Permalink
Merge pull request #3 from ksysoev/handling_disconnect_properly
Browse files Browse the repository at this point in the history
Improves handling connection errors
  • Loading branch information
ksysoev authored Oct 23, 2023
2 parents a57c6ce + b9670a0 commit 408ece1
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 3 deletions.
6 changes: 5 additions & 1 deletion pkg/cli/cli.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
12 changes: 10 additions & 2 deletions pkg/ws/ws.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@ package ws

import (
"crypto/tls"
"log"

"github.com/fatih/color"
"golang.org/x/net/websocket"
)

Expand Down Expand Up @@ -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}
Expand Down

0 comments on commit 408ece1

Please sign in to comment.