Vapor 4 Websocket Configuration !

WebSocketVaporServer.swift

import Vapor
import WebSocketKit

extension Application {
    private struct WebsocketClientsKey: StorageKey {
        typealias Value = WebsocketClients
    }

    var wsClients: WebsocketClients {
        get {
            if let existing = self.storage[WebsocketClientsKey.self] {
                return existing
            } else {
                let new = WebsocketClients()
                self.storage[WebsocketClientsKey.self] = new
                return new
            }
        }
        set {
            self.storage[WebsocketClientsKey.self] = newValue
        }
    }
}


import Vapor
import Foundation
import LPGSharedModels

protocol WebSocketMessageHandler {
    func handle(chatOutGoingEvent: ChatOutGoingEvent, using request: Request, on ws: WebSocket) async throws
}

struct ConnectionHandler: WebSocketMessageHandler {
    func handle(chatOutGoingEvent: ChatOutGoingEvent, using request: Request, on ws: WebSocket) async throws {

        let wsClients = request.application.wsClients // Assuming wsClients is accessible via the application
        let user = request.payload.user

        switch chatOutGoingEvent {
        case .connect:
            await wsClients.join(id: user._id, on: ws)
            request.logger.info("WebSocket connected for user \(user.email ?? user.fullName ?? "")")
        case .disconnect:
            await wsClients.leave(id: user._id)
            request.logger.info("WebSocket disconnected for user \(user.email ?? user.fullName ?? "")")
        default:
            break // Other cases are not handled by this handler
        }
    }
}

struct MessageHandler: WebSocketMessageHandler {
    func handle(chatOutGoingEvent: ChatOutGoingEvent, using request: Request, on ws: WebSocket) async throws {

        let wsClients = request.application.wsClients // Assuming wsClients is accessible via the application
        let user = request.payload.user

        switch chatOutGoingEvent {
        case .message(let msg):

            Task {
                try await wsClients.send(msg: msg, req: request)
            }
            request.logger.info("Web socket send msg to conversation id: \(msg.conversationId)")

        default:
            break // Other cases are not handled by this handler
        }
    }
}

struct ConversationsHandler: WebSocketMessageHandler {
    func handle(chatOutGoingEvent: ChatOutGoingEvent, using request: Request, on ws: WebSocket) async throws {

        let wsClients = request.application.wsClients // Assuming wsClients is accessible via the application
        let user = request.payload.user

        switch chatOutGoingEvent {
        case .conversation(let lastMessage):

            Task {
                try await wsClients.send(msg: lastMessage, req: request)
            }

            request.logger.info("conversation conversation: \(lastMessage)")

        default:
            break // Other cases are not handled by this handler
        }
    }
}

actor WebsocketHandle {
    private var connectionHandler: ConnectionHandler
    private var messageHandler: MessageHandler

    init(connectionHandler: ConnectionHandler, messageHandler: MessageHandler) {
        self.connectionHandler = connectionHandler
        self.messageHandler = messageHandler
    }

    func connectionHandler(ws: WebSocket, req: Request) {
        req.eventLoop.execute {
            ws.onText { [self] ws, text in
                // Decode the incoming message to determine its type
                // Assume ChatOutGoingEvent.decode(data:) returns an optional ChatOutGoingEvent
                guard let data = text.data(using: .utf8),
                      let chatOutGoingEvent = ChatOutGoingEvent.decode(data: data) else {
                    req.logger.error("Failed to decode incoming message.")
                    return
                }

                // Based on the event type, delegate to the appropriate handler
                Task {
                    do {

                        switch chatOutGoingEvent {

                        case .connect, .disconnect:
                            try await connectionHandler.handle(chatOutGoingEvent: chatOutGoingEvent, using: req, on: ws)

                        case .message(let msg):
                            try await messageHandler.handle(chatOutGoingEvent: chatOutGoingEvent, using: req, on: ws)

                        case .conversation(let msg):
                            try await connectionHandler.handle(chatOutGoingEvent: chatOutGoingEvent, using: req, on: ws)

                        default:
                            req.logger.error("Unhandled event type.")
                        }

                    } catch {

                        req.logger.error("Error handling message: \(error.localizedDescription)")

                    }
                }
            }
        }
    }
}

struct WebsocketMessage<T: Codable>: Codable {
    var client: UUID = UUID()
    let data: T
}

extension Data {
    func decodeWebsocketMessage<T: Codable>(_ type: T.Type) -> WebsocketMessage<T>? {
        try? JSONDecoder().decode(WebsocketMessage<T>.self, from: self)
    }
}

// configures your application
public func configure(_ app: Application) async throws {
    // MARK: Initialize your WebSocket clients
    app.wsClients = WebsocketClients()
}

import Vapor

func routes(_ app: Application) throws {

    app.get { req in
        try await req.view.render("index")
    }
    
    try app.group("v1") { api in
        let chat = api.grouped("chat")

        // Initialize your handlers
        let connectionHandler = ConnectionHandler()
        let messageHandler = MessageHandler()

        // Initialize the WebsocketHandle with handlers
        let webSocketController = WebsocketHandle(connectionHandler: connectionHandler, messageHandler: messageHandler)

        try chat.register(collection: WebsocketController(wsController: webSocketController))
    }
}