gqlgen中initFunc返回错误时WebSocket未关闭的问题求助

问题描述 #

使用gqlgen开发GraphQL API并实现订阅（WebSocket）功能时，遇到异常：当transport.Websocket{}的initFunc返回错误（例如认证失败）时，WebSocket并未关闭，而是持续重复调用initFunc。预期此时WebSocket应被关闭且仅向服务器发起一次请求，请问如何在initFunc执行失败时正确关闭WebSocket？

问题复现代码 #

存在问题的initFunc示例

func webSocketInit(ctx context.Context, initPayload transport.InitPayload) (context.Context, error) {
    // Get the token from payload
    any := initPayload["authToken"]
    token, ok := any.(string)
    if !ok || token == "" {
        fmt.Println("error in init")
        return nil, errors.New("authToken not found in transport payload")
    }

    // Perform token verification and authentication...
    userId := "john.doe" // e.g. userId, err := GetUserFromAuthentication(token)

    // put it in context
    ctxNew := context.WithValue(ctx, "username", userId)

    return ctxNew, nil
}

测试用订阅查询

subscription test {
  subscribe(subscriber:"jj")
}

服务器日志表现

➜  websocket_gqlgen ./server                                  
2022/12/20 09:40:56 connect to http://localhost:8080/ for GraphQL playground
error in init
error in init
error in init
error in init
error in init
...

主函数代码

const defaultPort = "8080"

func main() {
    port := os.Getenv("PORT")
    if port == "" {
        port = defaultPort
    }

    router := chi.NewRouter()

    // CORS setup, allow any for now
    c := cors.New(cors.Options{
        AllowedOrigins:   []string{"*"},
        AllowCredentials: true,
        Debug:            false,
    })

    srv := handler.New(graph.NewExecutableSchema(graph.Config{Resolvers: &graph.Resolver{}}))
    srv.AddTransport(transport.POST{})
    srv.AddTransport(transport.Websocket{
        KeepAlivePingInterval: 10 * time.Second,
        Upgrader: websocket.Upgrader{
            CheckOrigin: func(r *http.Request) bool {
                return true
            },
        },
        InitFunc: func(ctx context.Context, initPayload transport.InitPayload) (context.Context, error) {
            return webSocketInit(ctx, initPayload)
        },
    })
    srv.Use(extension.Introspection{})

    router.Handle("/", playground.Handler("My GraphQL App", "/app"))
    router.Handle("/app", c.Handler(srv))

    log.Printf("connect to http://localhost:%s/ for GraphQL playground", port)
    log.Fatal(http.ListenAndServe(":"+port, router))
}

解决方案 #

当initFunc执行失败时，可通过上下文获取WebSocket连接实例，主动发送关闭帧并关闭连接，避免重复调用。具体修改如下：

修改后的initFunc代码 #

import (
    "fmt"
    "errors"
    "github.com/gorilla/websocket"
    "github.com/99designs/gqlgen/graphql/transport"
)

func webSocketInit(ctx context.Context, initPayload transport.InitPayload) (context.Context, error) {
    // 获取payload中的认证token
    any := initPayload["authToken"]
    token, ok := any.(string)
    if !ok || token == "" {
        fmt.Println("error in init")
        // 从上下文取出WebSocket连接
        if conn, ok := ctx.Value(transport.WebsocketConnKey).(*websocket.Conn); ok {
            // 发送关闭消息，指定关闭码和原因
            closeMsg := websocket.FormatCloseMessage(websocket.ClosePolicyViolation, "authToken is required")
            _ = conn.WriteMessage(websocket.CloseMessage, closeMsg)
            // 关闭连接
            _ = conn.Close()
        }
        return nil, errors.New("authToken not found in transport payload")
    }

    // 执行token验证和逻辑...
    userId := "john.doe" // 示例：userId, err := GetUserFromAuthentication(token)

    // 将用户信息存入上下文
    ctxNew := context.WithValue(ctx, "username", userId)

    return ctxNew, nil
}

说明 #

• transport.WebsocketConnKey是gqlgen提供的上下文键，用于获取当前WebSocket连接实例
• 使用websocket.CloseMessage向客户端发送标准关闭帧，客户端收到后会停止自动重连
• 选择websocket.ClosePolicyViolation作为关闭码，符合认证失败的场景语义

内容的提问来源于stack exchange，提问作者cylon86