正确的 websocket 连接关闭

Question

我写了一个连接闭包函数。它发送一个关闭帧并期望得到相同的响应。

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

我为这个函数写了一个测试。

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
}

令我惊讶的是，它可以正常工作。 Readmessage() 读取结束帧。但是在测试中，我什么都不写。

1. 这是否发生在 gorilla/websocket 级别？
2. 我是否正确编写了函数？也许读取响应框架也发生在大猩猩级别。

Answer 1

函数基本正确。

Websocket 端点会回显关闭消息，除非端点已经自行发送了关闭消息。有关详细信息，请参阅 Websocket RFC 中的 Closing Handshake。

在正常关闭场景中，应用程序应该在发送关闭消息后收到关闭消息。

为了处理对端在发送关闭消息之前发送数据消息的情况，读取并丢弃数据消息，直到返回错误。

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