通过HTTP中间件验证WebSocket连接-Golang

Question

Problem Statement:

I am attempting to protect a websocket upgrader http endpoint using basic middleware in Golang, as the WebSocket protocol doesn’t handle authorization or authentication.

Community Suggestions

1. Some have suggested, albeit vaguely "I recommend authenticating the upgrade handshake using the application's code for authenticating HTTP requests."
2. Still others suggest "after connected, client need to send username & password which need to be checked by server. If not match, close the connection", but this seems non-idiomatic.

Strategy:

My failed strategy so far is attempting community strategy 1 above to secure upgrading the connection with a custom header X-Api-Key via middleware, and only upgrade clients who initiate the conversation with a matching key.

The code below results in the client is not using the websocket protocol: 'upgrade' token not found in 'Connection' header on the server side.

The Ask:

I would like to ask for help with understanding:

• If my take on strategy 1 is flawed, how might I improve it? It seems like that by sending the initial auth GET via http, that the subsequent upgrade request via scheme ws is rejected by the server.
• If strategy 2 is feasible, how might this be implemented?

Thoughts and suggestions, examples, gists appreciated, and if I can clarify further or restate please advise.

server.go:

package main

import (
    "flag"
    "log"
    "net/http"

    "github.com/gorilla/websocket"
)

func main() {
    var addr = flag.String("addr", "localhost:8080", "http service address")
    flag.Parse()

    http.Handle("/ws", Middleware(
        http.HandlerFunc(wsHandler),
        authMiddleware,
    ))
    log.Printf("listening on %v", *addr)
    log.Fatal(http.ListenAndServe(*addr, nil))
}

func Middleware(h http.Handler, middleware ...func(http.Handler) http.Handler) http.Handler {
    for _, mw := range middleware {
        h = mw(h)
    }
    return h
}

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
}

func wsHandler(rw http.ResponseWriter, req *http.Request) {
    wsConn, err := upgrader.Upgrade(rw, req, nil)
    if err != nil {
        log.Printf("upgrade err: %v", err)
        return
    }
    defer wsConn.Close()

    for {
        _, message, err := wsConn.ReadMessage()
        if err != nil {
            log.Printf("read err: %v", err)
            break
        }
        log.Printf("recv: %s", message)
    }
}

func authMiddleware(next http.Handler) http.Handler {
    TestApiKey := "test_api_key"
    return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
        var apiKey string
        if apiKey = req.Header.Get("X-Api-Key"); apiKey != TestApiKey {
            log.Printf("bad auth api key: %s", apiKey)
            rw.WriteHeader(http.StatusForbidden)
            return
        }
        next.ServeHTTP(rw, req)
    })
}

client.go:

package main

import (
    "fmt"
    "log"
    "net/http"
    "net/url"

    "github.com/gorilla/websocket"
)

func main() {
    // auth first
    req, err := http.NewRequest("GET", "http://localhost:8080/ws", nil)
    if err != nil {
        log.Fatal(err)
    }
    req.Header.Set("X-Api-Key", "test_api_key")

    resp, err := http.DefaultClient.Do(req)
    if err != nil || resp.StatusCode != http.StatusOK {
        log.Fatalf("auth err: %v", err)
    }
    defer resp.Body.Close()

    // create ws conn
    u := url.URL{Scheme: "ws", Host: "localhost:8080", Path: "/ws"}
    u.RequestURI()
    fmt.Printf("ws url: %s", u.String())
    log.Printf("connecting to %s", u.String())

    conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
    if err != nil {
        log.Fatalf("dial err: %v", err)
    }

    err = conn.WriteMessage(websocket.TextMessage, []byte("hellow websockets"))
    if err != nil {
        log.Fatalf("msg err: %v", err)
    }
}

Answer 1

The client application in the question sends two requests to the websocket endpoint. The first is a an authenticated HTTP request. This request fails to upgrade because it's not a websocket handshake. The second request is an unauthenticated websocket handshake. This request fails because it fails to authenticate.

The fix is to send an authenticated websocket handshake. Pass the auth headers through the last argument to Dial.

func main() {
    u := url.URL{Scheme: "ws", Host: "localhost:8080", Path: "/ws"}
    conn, _, err := websocket.DefaultDialer.Dial(u.String(), http.Header{"X-Api-Key": []string{"test_api_key"}})
    if err != nil {
        log.Fatalf("dial err: %v", err)
    }
    err = conn.WriteMessage(websocket.TextMessage, []byte("hellow websockets"))
    if err != nil {
        log.Fatalf("msg err: %v", err)
    }
}

On the server, authenticate the handshake using the application's code for authenticating HTTP requests.