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.

1. Is this happening at the gorilla/websocket level?
2. Did I write the function correctly? Maybe reading the response frame also happens at the gorilla level.

CodePudding user response：

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
        }
    }
}