WebSockets

Router.WS registers a WebSocket route. The protocol is implemented in the framework rather than delegated to a library: RFC 6455 framing, masking, UTF-8 validation and the close handshake, with every rule the specification lays down enforced and every violation answered with the status it calls for.

// A WebSocket route is declared like any other: the input is bound from the
// handshake, the dependency resolves before the upgrade, and the handler owns
// the connection until it returns.
r.WS("/items/{item_id}/ws", handlers.ItemSocket,
    muzak.Summary("Talk to an item over a WebSocket"),
    muzak.Needs(core.GetSessionOrToken),
    muzak.WithWebSocket(muzak.WSOptions{
        ReadLimit:    64 << 10,
        PingInterval: 30 * time.Second,
        MessageLimits: []muzak.Quota{
            {Name: "ws-messages", Window: time.Second, Limit: 10},
        },
    }))

// WSItemIn is the input of the item's WebSocket route.
//
// A handshake carries no body, so every field comes from the path, the query
// string, a header or a cookie. The pointer makes the query parameter optional:
// it stays nil when the client did not send one, which is how an absent value is
// told from a zero one.
type WSItemIn struct {
    ItemID string `path:"item_id" doc:"The item being talked about"`
    Q      *int   `query:"q" doc:"An optional number echoed back with each reply"`
}

// ItemSocket answers a WebSocket conversation about one item.
//
// The handshake has already succeeded by the time this runs: the input is
// bound, the guards have passed and the dependency is resolved, so what is left
// is the conversation itself. Returning ends it, and Muzak closes the
// connection; returning nil closes it normally.
func ItemSocket(ctx *muzak.Context, in schemas.WSItemIn, conn *muzak.WSConn) error {
    session := muzak.From[core.SessionOrToken](ctx)

    for {
        message, err := conn.ReadText(ctx.Context())
        if err != nil {
            // The peer closed, or the connection was lost. Either way there is
            // nothing left to say.
            return nil
        }
        if err := conn.WriteText(ctx.Context(), "credential: "+session.Value); err != nil {
            return err
        }
        if in.Q != nil {
            if err := conn.WriteText(ctx.Context(), fmt.Sprintf("q is %d", *in.Q)); err != nil {
                return err
            }
        }
        if err := conn.WriteText(ctx.Context(),
            fmt.Sprintf("you said %q, about item %s", message, in.ItemID)); err != nil {
            return err
        }
    }
}

The handshake is an ordinary GET

Everything that applies to a route applies here. Middleware runs, guards run, dependencies resolve, and the input struct is bound and validated before a single byte is upgraded.

A request that fails any of that is answered with the usual JSON error and never becomes a connection at all, which is the difference between a rejection a client can read and a socket that closes a moment after it opened.

curl -i 'http://localhost:8080/items/plumbus/ws'

{
  "error": {
    "code": "unauthorized",
    "message": "a session cookie or a token query parameter is required",
    "status": 401
  },
  "request_id": "0611f4b2-2f0a-4b57-9c1a-6e6a2e2f9b31"
}

Because a handshake carries no body, an input type with a body field is a registration error rather than a request that mysteriously never arrives.

Where a browser can put a credential

A browser cannot set headers on a WebSocket handshake, so the two places a credential can arrive are a cookie the browser attaches itself and a query parameter the page puts in the URL. A dependency covering both is what makes the route usable from either.

func GetSessionOrToken(ctx *muzak.Context) (SessionOrToken, error) {
    if cookie, err := ctx.Cookie("session"); err == nil && cookie.Value != "" {
        return SessionOrToken{Value: cookie.Value, FromCookie: true}, nil
    }
    if token := ctx.Query("token"); token != "" {
        return SessionOrToken{Value: token}, nil
    }
    return SessionOrToken{}, muzak.Unauthorized("a session cookie or a token query parameter is required")
}

A cookie is ambient authority a browser attaches without being asked, which is exactly why the origin check below cannot be turned off casually.

Reading and writing

Reading and writing are message oriented. A message split across several frames is reassembled and delivered once, whole. Control frames that arrive in between are handled without interrupting it: a ping is answered automatically, and a close is answered and then reported as a *WSCloseError.

type Command struct {
    Action string `json:"action"`
    Target string `json:"target"`
}

func Control(ctx *muzak.Context, in ControlIn, conn *muzak.WSConn) error {
    for {
        var command Command
        if err := conn.ReadJSON(ctx.Context(), &command); err != nil {
            return nil
        }
        if err := conn.WriteJSON(ctx.Context(), handle(command)); err != nil {
            return err
        }
    }
}

ReadJSON decodes with the same rules a request body is decoded with: unknown members, duplicate members and invalid UTF-8 are all rejected. A message that does not decode closes the connection with 1007, because a peer sending malformed JSON on a JSON connection is not going to be understood by carrying on.

The slice Read returns belongs to the caller and is not reused, so it may be kept. The payload passed to Write is not retained, so the caller may reuse it as soon as Write returns.

Concurrency

Writes are serialized, so any number of goroutines may write to one connection and each message goes out intact. A broadcast from many goroutines cannot interleave two messages.

Reads are serialized too, but a second reader is rarely what is wanted: the messages of one connection arrive in order and one loop should consume them.

func Room(ctx *muzak.Context, in RoomIn, conn *muzak.WSConn) error {
    hub := muzak.From[*core.Hub](ctx)

    updates := hub.Subscribe(ctx.Context(), in.Room)
    go func() {
        for update := range updates {
            // Safe from another goroutine: writes are serialized.
            _ = conn.WriteJSON(ctx.Context(), update)
        }
    }()

    for {
        message, err := conn.ReadText(ctx.Context())
        if err != nil {
            return nil
        }
        hub.Publish(in.Room, message)
    }
}

Ending a connection

The connection is closed when the handler returns, so a handler owns its connection for as long as it runs and never has to arrange the teardown itself.

if !allowed(session, in.Room) {
    return &muzak.WSCloseError{Status: muzak.WSStatusPolicyViolation, Reason: "not a member of this room"}
}

Calling Close yourself is only necessary to choose a status other than normal closure or to end the connection from another goroutine. It is safe to call more than once and from more than one goroutine; only the first call sends anything. The reason is truncated to the 123 bytes a close frame can carry, on a rune boundary.

Telling a goodbye from a violation

Every read and every write returns a *WSCloseError once the connection is finished, which is why the loops above end on any error.

message, err := conn.ReadText(ctx.Context())
if err != nil {
    if status, ok := muzak.WSCloseStatus(err); ok && status == muzak.WSStatusNormalClosure {
        return nil
    }
    return err
}

WSCloseError.Unwrap exposes the transport failure behind an abnormal closure, so errors.Is can test for io.EOF or a network error.

The codes RFC 6455 defines are exported as muzak.WSStatus constants, from WSStatusNormalClosure (1000) to WSStatusTLSHandshake (1015). Three of them describe a local observation rather than something a peer said: WSStatusNoStatusReceived, WSStatusAbnormalClosure and WSStatusTLSHandshake are reported by Muzak but never written to the wire, and passing one to Close closes without a status code. Codes from 4000 to 4999 are free for an application to define.

What a hostile peer cannot do

A WebSocket is the longest-lived thing an unauthenticated stranger can ask a server for, so every direction a peer controls is bounded, and each bound stops something specific.

Nothing a peer sends is echoed into a response header: only a subprotocol the route itself offered can be answered with, and a route offering one that is not a token is refused when the application is built.

Options

app.Include(chat, muzak.WithWebSocket(muzak.WSOptions{
    ReadLimit:      64 << 10,
    PingInterval:   30 * time.Second,
    AllowedOrigins: []string{"https://app.example.com"},
}))

Options can be set application-wide through AppOptions.WebSocket and narrowed for a router or a route with WithWebSocket. Layering works field by field, so a route that raises only the read limit keeps the application's origin policy.

MaxConnections and MaxConnectionsPerIP may only be set on the application: the resource they protect is the process, not a route, so a router or a route that sets either is refused when the application is built.

Keepalive only works while the handler is reading, because a pong is consumed by a read like any other frame. A handler that only ever writes should ping by hand with Ping.

The origin check

A cross-origin handshake is refused by default, and no CORS policy changes that.

A WebSocket handshake is not subject to the same-origin policy and is never preflighted, which is what makes cross-site hijacking possible in the first place: without the check, any page on the internet could open an authenticated connection to your server from a visitor's browser, cookies and all. AppOptions.CORS has no bearing on it, for exactly that reason.

muzak.WSOptions{AllowedOrigins: []string{"https://app.example.com"}}

The server's own origin is always allowed. The single entry "*" allows any origin.

muzak.WSOptions{
    AllowOriginFunc: func(r *http.Request, origin string) bool {
        return strings.HasSuffix(origin, ".example.com")
    },
}

It runs on every handshake an earlier rule did not already allow, so it must be cheap and free of side effects.

InsecureSkipOriginCheck accepts a handshake from anywhere. It is safe only for a connection that carries no ambient authority: one authenticated by a token the client has to present explicitly, never by a cookie, since a browser attaches cookies to a cross-origin handshake without being asked.

Subprotocols

muzak.WSOptions{Subprotocols: []string{"graphql-transport-ws", "graphql-ws"}}

The client's own list is in preference order, so the first of its choices that appears here is the one negotiated. A client asking for something else is answered without the header, which tells it to give up. Read the result with conn.Subprotocol().

Shutdown

Open connections are tracked, so a graceful shutdown tells every peer it is going away with 1001 and waits for the handlers. net/http cannot do that on its own: a hijacked connection is no longer one it knows about.

Documentation

A WebSocket route appears in the OpenAPI document with its parameters and a 101 Switching Protocols response, because that is what OpenAPI can say about a handshake. The conversation itself continues off the document.

Testing

muzak.WSDial is the client half of the same engine, which is what lets a route be tested over a real connection rather than against a second implementation. The test client wraps it:

func TestItemSocket(t *testing.T) {
    client := testclient.New(t, buildApp())

    conn := client.WS("/items/plumbus/ws", testclient.Query("token", "jessica"))

    if err := conn.WriteText(t.Context(), "hello"); err != nil {
        t.Fatalf("WriteText = %v", err)
    }
    reply, err := conn.ReadText(t.Context())
    if err != nil {
        t.Fatalf("ReadText = %v", err)
    }
    if !strings.Contains(reply, "jessica") {
        t.Errorf("reply = %q", reply)
    }
}

func TestItemSocketNeedsCredential(t *testing.T) {
    client := testclient.New(t, buildApp())

    _, response := client.TryWS("/items/plumbus/ws")

    response.AssertStatus(http.StatusUnauthorized)
}

Outside a test, dial directly:

conn, _, err := muzak.WSDial(ctx, "ws://"+app.Addr()+"/items/plumbus/ws", muzak.WSDialOptions{})
if err != nil {
    return err
}
defer conn.Close(muzak.WSStatusNormalClosure, "")

The URL may use ws, wss, http or https. The response is returned alongside the connection so a caller can read the handshake headers, and on failure the status and body the server refused with. WSDial never follows a redirect, because following one would send the handshake's headers to whatever host the answer named.

Where to go next

Server-Sent Events covers the simpler half of what a WebSocket is usually reached for, and Rate Limiting covers the quotas MessageLimits reuses.