WebSockets

Client GET  -> global middleware -> route group middleware -> route handler (checks, then req.webSocket(...) )
                                                                   |
Client <- 101 Switching Protocols (headers from shouldUpgrade) <---+
          ... pipeline replaced ...
onUpgrade(req, ws)  ->  ws.onText / ws.onBinary / ws.onClose ...   (no middleware from here on)

// 1. A route registered directly on a RoutesBuilder (the Application, a group, a controller's RoutesBuilder).
//    It is a GET route for the given path.
@discardableResult
public func webSocket(
    _ path: PathComponent...,                      // also an overload taking [PathComponent]
    maxFrameSize: WebSocketMaxFrameSize = .default,    // 1 << 14 = 16 384 bytes
    shouldUpgrade: @escaping (@Sendable (Request) async throws -> HTTPHeaders?) = { _ in [:] },
    onUpgrade: @Sendable @escaping (Request, WebSocket) async -> ()
) -> Route

// 2. A method on Request, to be returned from your own handler after your own checks.
public func webSocket(
    maxFrameSize: WebSocketMaxFrameSize = .default,
    shouldUpgrade: @escaping (@Sendable (Request) async throws -> HTTPHeaders?) = { _ in [:] },
    onUpgrade: @Sendable @escaping (Request, WebSocket) async -> ()
) -> Response

import Vapor
import WebSocketKit

func echoRoutes(_ app: Application) {
    app.webSocket("echo") { req, ws async in
        req.logger.info("echo socket opened")
        try? await ws.send("connected")
        ws.onText { ws, text async in
            try? await ws.send("echo: \(text)")
        }
        ws.onBinary { ws, buffer async in
            try? await ws.send(Array(buffer.readableBytesView))   // reply with the same bytes
        }
        ws.onClose.whenComplete { _ in req.logger.info("echo socket closed, code \(String(describing: ws.closeCode))") }
    }
}

ECHO -> ["connected", "echo: hi"] [1, 2, 3]
echo socket closed, code Optional(NIOWebSocket.WebSocketErrorCode.goingAway)

NIOCore/NIOLoopBound.swift:172: Precondition failed

import Vapor
import WebSocketKit

/// What the server pushes to browsers. Field names are the JSON keys.
struct ReviewEvent: Content {
    let reviewID: UUID
    let bookID: UUID
    let rating: Int
    let body: String
}

/// What browsers may send.
struct ClientMessage: Decodable, Sendable {
    enum Action: String, Decodable, Sendable { case watch, unwatch }
    let action: Action
    let bookID: UUID?
}

/// Registry of open connections. An actor serialises every mutation of the dictionary.
actor ReviewFeed {
    private struct Connection {
        let socket: WebSocket
        let userID: UUID
        var watching: UUID?   // nil = every book
    }

    private var connections: [UUID: Connection] = [:]
    private let logger: Logger

    init(logger: Logger) { self.logger = logger }

    var count: Int { connections.count }

    func add(_ socket: WebSocket, userID: UUID) -> UUID {
        let id = UUID()
        connections[id] = Connection(socket: socket, userID: userID, watching: nil)
        logger.info("WebSocket \(id) opened for user \(userID); \(connections.count) open")
        return id
    }

    func remove(_ id: UUID) {
        if connections.removeValue(forKey: id) != nil {
            logger.info("WebSocket \(id) removed; \(connections.count) open")
        }
    }

    func watch(_ id: UUID, book: UUID?) {
        connections[id]?.watching = book
    }

    /// Sends the event to every matching connection. A failed send drops that connection only.
    func broadcast(_ event: ReviewEvent) async {
        guard let data = try? JSONEncoder().encode(event) else { return }
        let text = String(decoding: data, as: UTF8.self)
        for (id, connection) in connections where connection.watching == nil || connection.watching == event.bookID {
            do {
                try await connection.socket.send(text)
            } catch {
                logger.warning("Send to \(id) failed: \(error); dropping it")
                connections.removeValue(forKey: id)
            }
        }
    }
}

/// Single-use, short-lived tickets that let a browser open a WebSocket without an Authorization header.
actor TicketStore {
    private var tickets: [String: (userID: UUID, expires: Date)] = [:]
    private let lifetime: TimeInterval

    init(lifetime: TimeInterval = 30) { self.lifetime = lifetime }

    func issue(for userID: UUID) -> String {
        let now = Date()
        tickets = tickets.filter { $0.value.expires > now }          // drop expired ones
        var generator = SystemRandomNumberGenerator()      // cryptographically secure
        let bytes = (0..<32).map { _ in UInt8.random(in: .min ... .max, using: &generator) }
        let ticket = Data(bytes).base64EncodedString()
            .replacingOccurrences(of: "+", with: "-").replacingOccurrences(of: "/", with: "_")
            .replacingOccurrences(of: "=", with: "")
        tickets[ticket] = (userID, now.addingTimeInterval(lifetime))
        return ticket
    }

    /// Returns the user and forgets the ticket, so a second use fails.
    func redeem(_ ticket: String) -> UUID? {
        guard let entry = tickets.removeValue(forKey: ticket), entry.expires > Date() else { return nil }
        return entry.userID
    }
}

struct ReviewFeedKey: StorageKey { typealias Value = ReviewFeed }
struct TicketStoreKey: StorageKey { typealias Value = TicketStore }

extension Application {
    var reviewFeed: ReviewFeed {
        get { storage[ReviewFeedKey.self]! }
        set { storage[ReviewFeedKey.self] = newValue }
    }
    var tickets: TicketStore {
        get { storage[TicketStoreKey.self]! }
        set { storage[TicketStoreKey.self] = newValue }
    }
}

import Fluent
import Vapor
import WebSocketKit
import NIOHTTP1
import NIOWebSocket

let allowedOrigins: Set<String> = ["https://app.example.com", "http://localhost:3000"]

/// Stand-in for real authentication (see xref:backend/vapor/authentication.adoc[Authentication]):
/// the demo trusts an X-Demo-User header holding a user's UUID.
func demoUserID(_ req: Request) throws -> UUID {
    guard let text = req.headers.first(name: "X-Demo-User"), let id = UUID(uuidString: text) else {
        throw Abort(.unauthorized)
    }
    return id
}

func routes(_ app: Application) throws {
    // 1. The browser first asks, over normal authenticated HTTP, for a one-time ticket.
    app.post("ws", "tickets") { req async throws -> [String: String] in
        let userID = try demoUserID(req)
        return ["ticket": await req.application.tickets.issue(for: userID)]
    }

    // 2. Then it opens the socket with ?ticket=...
    app.get("ws", "reviews") { req async throws -> Response in
        // Browsers always send Origin on a WebSocket handshake; check it before anything else.
        guard let origin = req.headers.first(name: .origin), allowedOrigins.contains(origin) else {
            throw Abort(.forbidden, reason: "Origin not allowed")
        }
        guard let ticket = req.query[String.self, at: "ticket"],
              let userID = await req.application.tickets.redeem(ticket) else {
            throw Abort(.unauthorized, reason: "Missing, expired or already used ticket")
        }
        let feed = req.application.reviewFeed
        return req.webSocket(maxFrameSize: .init(integerLiteral: 1 << 12), onUpgrade: { req, ws in
            let id = await feed.add(ws, userID: userID)
            ws.pingInterval = .seconds(25)

            ws.onText { ws, text in
                guard let message = try? JSONDecoder().decode(ClientMessage.self, from: Data(text.utf8)) else {
                    try? await ws.close(code: .unacceptableData)
                    return
                }
                switch message.action {
                case .watch: await feed.watch(id, book: message.bookID)
                case .unwatch: await feed.watch(id, book: nil)
                }
            }
            ws.onPing { _, _ async in req.logger.trace("ping from \(id)") }
            ws.onClose.whenComplete { _ in
                Task { await feed.remove(id) }
            }
        })
    }

    // The same origin check expressed with shouldUpgrade (see "shouldUpgrade" in the page).
    app.get("ws", "strict") { req async throws -> Response in
        req.webSocket(shouldUpgrade: { req in
            guard let origin = req.headers.first(name: .origin), allowedOrigins.contains(origin) else {
                return nil   // refuse the upgrade
            }
            return ["X-Checked-Origin": origin]   // extra headers for the 101 response
        }, onUpgrade: { _, ws in
            try? await ws.send("hello")
            try? await ws.close()
        })
    }

    // 3. A review is saved, then broadcast to the open sockets.
    app.post("books", ":bookID", "reviews") { req async throws -> HTTPStatus in
        struct Input: Content { let rating: Int; let body: String }
        let input = try req.content.decode(Input.self)
        let userID = try demoUserID(req)
        guard let book = try await Book.find(req.parameters.get("bookID"), on: req.db) else { throw Abort(.notFound) }
        let review = Review(rating: input.rating, body: input.body, bookID: try book.requireID(), userID: userID)
        try await review.save(on: req.db)
        await req.application.reviewFeed.broadcast(
            ReviewEvent(reviewID: try review.requireID(), bookID: try book.requireID(), rating: review.rating, body: review.body))
        return .created
    }
}

$ curl -i -N ... -H "Origin: http://localhost:3000" http://127.0.0.1:8081/ws/strict
HTTP/1.1 101 Switching Protocols
X-Checked-Origin: http://localhost:3000
Upgrade: websocket
...
(then the frames: text "hello", then close 1001)

$ curl -i -N ... -H "Origin: https://evil.example" http://127.0.0.1:8081/ws/strict
curl: (52) Empty reply from server

<script>
  async function openFeed() {
    // 1. normal authenticated HTTP request; cookies or a bearer header are fine here
    const res = await fetch("/ws/tickets", { method: "POST", credentials: "include" });
    const { ticket } = await res.json();
    // 2. the WebSocket handshake carries only the ticket (and the browser's Origin)
    const ws = new WebSocket(`wss://${location.host}/ws/reviews?ticket=${encodeURIComponent(ticket)}`);
    ws.onmessage = (e) => console.log("review", JSON.parse(e.data));
    ws.onclose = (e) => { if (e.code !== 1000) setTimeout(openFeed, 2000); };   // reconnect with a NEW ticket
    ws.onopen = () => ws.send(JSON.stringify({ action: "watch", bookID: "<book id>" }));
  }
  openFeed();
</script>

import Redis
import Vapor

extension RedisChannelName {
    static var reviewCreated: RedisChannelName { "bookshelf:reviews" }
}

/// Publishes to Redis instead of broadcasting locally; every instance (this one included)
/// receives the message from its subscription and broadcasts to its own sockets.
struct ReviewFanout: LifecycleHandler {
    func didBoot(_ application: Application) throws {
        let feed = application.reviewFeed
        let logger = application.logger
        Task {
            do {
                try await application.redis.subscribe(
                    to: [.reviewCreated],
                    messageReceiver: { _, message in
                        guard let text = message.string,
                              let event = try? JSONDecoder().decode(ReviewEvent.self, from: Data(text.utf8)) else {
                            logger.warning("Ignoring malformed review message")
                            return
                        }
                        Task { await feed.broadcast(event) }
                    },
                    onSubscribe: { key, count in logger.info("Subscribed to \(key) (\(count))") },
                    onUnsubscribe: { key, count in logger.warning("Unsubscribed from \(key) (\(count) left)") }
                )
            } catch {
                logger.error("Review subscription failed: \(error)")
            }
        }
    }
}

func publishReview(_ event: ReviewEvent, on req: Request) async throws {
    let text = String(decoding: try JSONEncoder().encode(event), as: UTF8.self)
    _ = try await req.redis.publish(text, to: .reviewCreated).get()
}

func configureRedisFanout(_ app: Application) throws {
    app.redis.configuration = try RedisConfiguration(hostname: Environment.get("REDIS_HOST") ?? "localhost")
    app.lifecycle.use(ReviewFanout())
}

import Vapor
import WebSocketKit

/// Opens the review feed as a client, prints every event and returns the socket.
func followReviews(host: String, port: Int, ticket: String, origin: String,
                   on group: any EventLoopGroup, logger: Logger) async throws {
    try await WebSocket.connect(
        scheme: "ws", host: host, port: port, path: "/ws/reviews", query: "ticket=\(ticket)",
        headers: ["Origin": origin], on: group
    ) { ws async in
        ws.onText { _, text async in logger.info("event: \(text)") }
        ws.onClose.whenComplete { _ in logger.info("feed closed") }
        try? await ws.send(#"{"action":"unwatch"}"#)
    }
}

# plain connection, type lines to send text frames, see replies on stdout
websocat "ws://127.0.0.1:8080/echo"

# with an Origin header, as a browser would send it, and a ticket from the HTTP endpoint
TICKET=$(curl -s -X POST -H "X-Demo-User: 11111111-1111-1111-1111-111111111111" http://127.0.0.1:8080/ws/tickets \
  | sed 's/.*"ticket":"\([^"]*\)".*/\1/')
websocat -H "Origin: http://localhost:3000" "ws://127.0.0.1:8080/ws/reviews?ticket=$TICKET"

import Fluent
import NIOConcurrencyHelpers
import Testing
import Vapor
import VaporTesting
import WebSocketKit
@testable import App

struct NewReview: Content { let rating: Int; let body: String }

@Suite("Review feed", .serialized)
struct FeedTests {
    func withRunningApp(_ body: (Application, Book, User) async throws -> Void) async throws {
        try await withApp(configure: { app in
            app.environment.commandInput.arguments = ["serve"]   // do not parse the test runner's arguments
            app.http.server.configuration.port = 8765
            try await configure(app)
        }) { app in
            let author = Author(name: "Ada"); try await author.save(on: app.db)
            let book = Book(title: "T", isbn: "1", publishedYear: 2026, priceCents: 1, status: .published, authorID: try author.requireID())
            try await book.save(on: app.db)
            let user = User(email: "a@example.com", passwordHash: "x"); try await user.save(on: app.db)
            try await app.startup()
            try await body(app, book, user)
        }
    }

    func connect(_ app: Application, user: User, origin: String = "http://localhost:3000",
                 onText: @escaping @Sendable (String) -> Void = { _ in }) async throws -> WebSocket {
        let ticket = await app.tickets.issue(for: try user.requireID())
        let holder = NIOLockedValueBox<WebSocket?>(nil)
        try await WebSocket.connect(to: "ws://127.0.0.1:8765/ws/reviews?ticket=\(ticket)",
                                    headers: ["Origin": origin], on: app.eventLoopGroup) { ws async in
            ws.onText { _, text async in onText(text) }
            holder.withLockedValue { $0 = ws }
        }
        try await Task.sleep(for: .milliseconds(200))
        return try #require(holder.withLockedValue { $0 })
    }

    @Test("broadcast, watch filter and close")
    func broadcastAndFilter() async throws {
        try await withRunningApp { app, book, user in
            let all = NIOLockedValueBox<[String]>([])
            let filtered = NIOLockedValueBox<[String]>([])
            let a = try await connect(app, user: user) { text in all.withLockedValue { $0.append(text) } }
            let b = try await connect(app, user: user) { text in filtered.withLockedValue { $0.append(text) } }
            #expect(await app.reviewFeed.count == 2)
            try await b.send(#"{"action":"watch","bookID":"\#(UUID().uuidString)"}"#)   // some other book
            try await Task.sleep(for: .milliseconds(200))
            await app.reviewFeed.broadcast(ReviewEvent(reviewID: UUID(), bookID: try book.requireID(), rating: 5, body: "Great"))
            try await Task.sleep(for: .milliseconds(200))
            print("ALL:", all.withLockedValue { $0 })
            print("FILTERED:", filtered.withLockedValue { $0 })
            #expect(all.withLockedValue { $0.count } == 1)
            #expect(filtered.withLockedValue { $0.count } == 0)
            try await a.close()
            try await b.close()
            try await Task.sleep(for: .milliseconds(300))
            #expect(await app.reviewFeed.count == 0)
        }
    }

    @Test("invalid JSON closes with 1007")
    func badMessage() async throws {
        try await withRunningApp { app, _, user in
            let ws = try await connect(app, user: user)
            try await ws.send("not json")
            try await Task.sleep(for: .milliseconds(300))
            print("CLOSECODE bad json:", String(describing: ws.closeCode))
            #expect(ws.isClosed)
        }
    }

    @Test("oversized frame closes the socket")
    func bigFrame() async throws {
        try await withRunningApp { app, _, user in
            let ws = try await connect(app, user: user)
            try await ws.send(String(repeating: "x", count: 5000))
            try await Task.sleep(for: .milliseconds(300))
            print("CLOSECODE big frame:", String(describing: ws.closeCode), "closed:", ws.isClosed)
            #expect(await app.reviewFeed.count == 0)
        }
    }

    @Test("wrong origin is refused with 403")
    func wrongOrigin() async throws {
        try await withRunningApp { app, _, user in
            do {
                _ = try await connect(app, user: user, origin: "https://evil.example")
                Issue.record("should have been refused")
            } catch {
                print("WRONG ORIGIN error:", error)
            }
            #expect(await app.reviewFeed.count == 0)
        }
    }

    @Test("POST a review reaches the socket")
    func postReview() async throws {
        try await withRunningApp { app, book, user in
            let got = NIOLockedValueBox<[String]>([])
            let ws = try await connect(app, user: user) { text in got.withLockedValue { $0.append(text) } }
            try await app.testing().test(.POST, "books/\(try book.requireID())/reviews",
                                         headers: ["X-Demo-User": try user.requireID().uuidString],
                                         beforeRequest: { req in try req.content.encode(NewReview(rating: 4, body: "Solid")) },
                                         afterResponse: { res async in #expect(res.status == .created) })
            try await Task.sleep(for: .milliseconds(200))
            print("POSTED ->", got.withLockedValue { $0 })
            #expect(got.withLockedValue { $0.count } == 1)
            try await ws.close()
        }
    }

    @Test("echo")
    func echo() async throws {
        try await withRunningApp { app, _, _ in
            let got = NIOLockedValueBox<[String]>([])
            let bin = NIOLockedValueBox<[UInt8]>([])
            let holder = NIOLockedValueBox<WebSocket?>(nil)
            try await WebSocket.connect(to: "ws://127.0.0.1:8765/echo", on: app.eventLoopGroup) { ws async in
                ws.onText { _, text async in got.withLockedValue { $0.append(text) } }
                ws.onBinary { _, buf async in bin.withLockedValue { $0 = Array(buf.readableBytesView) } }
                holder.withLockedValue { $0 = ws }
            }
            let ws = try #require(holder.withLockedValue { $0 })
            try await ws.send("hi")
            try await ws.send([1, 2, 3])
            try await Task.sleep(for: .milliseconds(300))
            print("ECHO ->", got.withLockedValue { $0 }, bin.withLockedValue { $0 })
            try await ws.close()
            try await Task.sleep(for: .milliseconds(200))
        }
    }

    @Test("a ticket works once")
    func ticketOnce() async throws {
        try await withRunningApp { app, _, user in
            let ticket = await app.tickets.issue(for: try user.requireID())
            #expect(await app.tickets.redeem(ticket) != nil)
            #expect(await app.tickets.redeem(ticket) == nil)
        }
    }
}

map $http_upgrade $connection_upgrade {
    default upgrade;
    ''      close;
}

server {
    listen 443 ssl;
    server_name app.example.com;

    location /ws/ {
        proxy_pass http://127.0.0.1:8080;
        proxy_http_version 1.1;
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection $connection_upgrade;
        proxy_set_header Host $host;
        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
        proxy_read_timeout 3600s;      # default 60 s closes quiet sockets; pings every 25 s also keep it alive
    }
}