Authentication and authorisation in Vapor 3 is very easy with Vapor's auth library. You simply add the library into your project and use one of the middlewares to authorize users. In case of the TokenAuthenticationMiddleware, protecting your route is as easy as defining it as follows.

import Authentication
import Vapor

enum RESTRoutes {
    static func create(router: Router, chatSessionManager: ChatSessionManager) throws {
        let authenticatedRouter = router
            .grouped(User.tokenAuthMiddleware())
            .grouped(User.guardAuthMiddleware())
		
        // Chats

        let chatRouter = authenticatedRouter.grouped("chats")

        let chatController = ChatController(chatSessionManager: chatSessionManager)
		
        chatRouter.get("", Chat.parameter, use: chatController.detail)
    }
}

Unfortunately this is not as easy to do when your route allows WebSocket connections.

The problem

WebSocket connection is initiated by a special HTTP request that asks for the connection to be upgraded. In practice (using a great CLI tool wsta) this looks as follows.

 ✘ wsta ws://localhost:8080/chats/1 --head --verbose --header "Authorization: Bearer token"
WebSocket upgrade request
---
Host: localhost:8080
Connection: Upgrade
Upgrade: websocket
Sec-WebSocket-Version: 13
Sec-WebSocket-Key: nc17bUjVtTrmKMEVQxHqyA==
Origin: http://localhost
Authorization:  Bearer token


WebSocket upgrade response
---
101 Switching Protocols
Upgrade: websocket
Sec-WebSocket-Accept: c6JyJMP6GtBzjJSnU4zkUrooxJU=
Connection: upgrade


Connected to ws://localhost:8080/chats/1

The problem is that if the upgrade is successful and allowed by your application, middleware on the route is not executed. In other words even if your route allows both RESTful and WebSocket communication and the REST route has middleware enabled, the middleware execution will be skipped for WebSocket connection.

Since no middleware is executed your route is not protected and even code like try req.requireAuthenticated(User.self) that you would normally execute in your Controller will obviously not work.

The solution

Vapor documentation contains the following example for WebSockets.

// Create a new NIO websocket server
let wss = NIOWebSocketServer.default()

// Add WebSocket upgrade support to GET /echo
wss.get("echo") { ws, req in
    // Add a new on text callback
    ws.onText { ws, text in
        // Simply echo any received text
        ws.send(text)
    }
}

// Register our server
services.register(wss, as: WebSocketServer.self)

This does not help us much of course but it gives us some insight into how routing works with WebSockets.

The get function used above on NIOWebSocketServer is defined as follows.

func get(_ path: PathComponentsRepresentable..., use closure: @escaping (WebSocket, Request) throws -> ()) -> Route<WebSocketResponder>

Obviously we cannot do much with that but fortunately there is a register function that we can use.

public func register(route: Route<WebSocketResponder>) {
    routes.append(route)
    router.register(route: route)
}

Using this function we can register a Route with a WebSocketResponder attached to it.

A WebSocketResponder is a struct containing two functions shouldUpgrade and onUpgrade that are used to handle the incoming HTTP Upgrade request.

Typically what you would do is to use the shouldUpgrade function to do your authentication check but there is a problem. This function is synchronous and requires a synchronous return of headers if the upgrade was successful and nil if it was unsuccessful.

In Vapor pretty much everything is asynchronous though - including database access - and since we need to look into the database to check whether a token is valid, this is a no-go. I hope this will be improved in Vapor 4.

In onUpgrade function we can do anything we want though and that includes the authorisation functionality.

Let's define our route.

import Vapor

enum WebSockets {
    static func create(webSocketServer: NIOWebSocketServer, chatSessionManager: ChatSessionManager) throws {
        let chatController = ChatController(chatSessionManager: chatSessionManager)
        
        let responder = WebSocketResponder(
            shouldUpgrade: { _ in return [:] },
            onUpgrade: { ws, req in
                WebSocketAuthenticationMiddleware.handle(
                    webSocket: ws,
                    request: req,
                    handler: chatController.webSocketDetail
                )
            }
        )
        
        let route: Route<WebSocketResponder> = .init(path: ["chats", Chat.parameter], output: responder)
        webSocketServer.register(route: route)
    }
}

And let's create our WebSocket "Middleware".

import Vapor

struct WebSocketAuthenticationMiddleware {
    static func handle(
        webSocket: WebSocket,
        request: Request,
        handler: @escaping WebSocketAuthenticationResponder.Handler
    ) {
        let authenticationRequiredResponder = WebSocketAuthenticationResponder(
            webSocket: webSocket,
            handler: handler
        )
        
        let middleware = User.tokenAuthMiddleware()
        
        do {
            let _ = try middleware.respond(to: request, chainingTo: authenticationRequiredResponder)
        } catch {
            webSocket.close()
        }
    }
}

Notice that in the code above we're using the tokenAuthMiddleware that we would normally use in our RESTful routes. If the middleware respond call fails it means that the user is not authenticated and we close the WebSocket.

The last missing piece is the WebSocketAuthenticationResponder the code above uses.

import Vapor

struct WebSocketAuthenticationResponder: Responder {
    typealias Handler = (WebSocket, Request) throws -> ()
    
    let webSocket: WebSocket
    let handler: Handler
    
    func respond(to req: Request) throws -> EventLoopFuture<Response> {
        do {
            let _: User = try req.requireAuthenticated()
        } catch {
            Self.onWebSocketError(webSocket: webSocket, error: Abort(.unauthorized))
        }
        
        do {
            try handler(webSocket, req)
        } catch {
            Self.onWebSocketError(webSocket: webSocket, error: error)
        }
        
        let response = Response(http: HTTPResponse(status: .accepted), using: req)
        
        let promise = req.eventLoop.newPromise(Response.self)
        promise.succeed(result: response)
        return promise.futureResult
    }
}

// MARK: - Utility conformances

extension WebSocketAuthenticationResponder: WebSocketHandling {}

This is a bit hacky but it works. We basically create our own Responder that returns an empty HTTP Accepted response so that we have something to chain our middleware to.

If you're wondering that the WebSocketHandling protocol does you can take a look at it below. It basically provides a simple static function that serializes the error to JSON, sends it to the WebSocket and closes the connection.

import Vapor

protocol WebSocketHandling {}

extension WebSocketHandling {
    static func onWebSocketError(webSocket: WebSocket, error: Error) {
        defer {
            webSocket.close()
        }
        
        let error = WebsocketError(error: error.localizedDescription)
        
        guard let json = try? JSONEncoder().encode(error) else {
            return
        }
        
        webSocket.send(json)
    }
}

And that's basically it. With a bit of effort you can use your Authentication Middleware with WebSockets!