correct websocket connection closure - go

Question

I wrote a connection closure function. It sends a closing frame and expects the same in response.
func TryCloseNormally(wsConn *websocket.Conn) error {
closeNormalClosure := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
defer wsConn.Close()
if err := wsConn.WriteControl(websocket.CloseMessage, closeNormalClosure, time.Now().Add(time.Second)); err != nil {
return err
}
if err := wsConn.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
return err
}
_, _, err := wsConn.ReadMessage()
if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
return nil
} else {
return errors.New("Websocket doesn't send a close frame in response")
}
}
I wrote a test for this function.
func TestTryCloseNormally(t *testing.T) {
done := make(chan struct{})
exit := make(chan struct{})
ctx := context.Background()
ln, err := net.Listen("tcp", "localhost:")
require.Nil(t, err)
handler := HandlerFunc(func(conn *websocket.Conn) {
for {
_, _, err := conn.ReadMessage()
if err != nil {
require.True(t, websocket.IsCloseError(err, websocket.CloseNormalClosure), err.Error())
return
}
}
})
s, err := makeServer(ctx, handler)
require.Nil(t, err)
go func() {
require.Nil(t, s.Run(ctx, exit, ln))
close(done)
}()
wsConn, _, err := websocket.DefaultDialer.Dial(addr+strconv.Itoa(ln.Addr().(*net.TCPAddr).Port), nil) //nolint:bodyclose
require.Nil(t, err)
require.Nil(t, wsConn.WriteMessage(websocket.BinaryMessage, []byte{'o', 'k'}))
require.Nil(t, TryCloseNormally(wsConn))
close(exit)
<-done
}
To my surprise, it works correctly. Readmessage() reads the closing frame. But in the test, I don't write anything.
Is this happening at the gorilla/websocket level?
Did I write the function correctly? Maybe reading the response frame also happens at the gorilla level.

Answer 1

The function is mostly correct.
Websocket endpoints echo close messages unless the endpoint has already send a close message on its own. See Closing Handshake in the Websocket RFC for more details.
In the normal close scenario, an application should expect to receive a close message after sending a close message.
To handle the case where the peer sent a data message before the sending the close message, read and discard data messages until an error is returned.
func TryCloseNormally(wsConn *websocket.Conn) error {
defer wsConn.Close()
closeNormalClosure := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
if err := wsConn.WriteControl(websocket.CloseMessage, closeNormalClosure, time.Now().Add(time.Second)); err != nil {
return err
}
if err := wsConn.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
return err
}
for {
_, _, err := wsConn.ReadMessage()
if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
return nil
}
if err != nil {
return err
}
}
}