WebSockets with authentication

[colesturza]

Hi everyone,

Over the past few days, I’ve been exploring ways to add authentication to the original Jakarta EE specification implementing extension, but I wasn’t able to find many helpful examples or clear documentation on the topic. I realize that WebSockets Next has become the preferred extension for most use cases now, but for those of us who are still working with the original implementation, I wanted to share my solution with the community. Hopefully, it will be useful for other developers facing similar challenges.

Looking forward to any feedback or additional insights!

package com.acme;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.quarkus.security.identity.SecurityIdentity;
import io.quarkus.websockets.client.runtime.WebSocketPrincipal;
import io.vertx.mutiny.core.Vertx;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.inject.Inject;
import jakarta.websocket.CloseReason;
import jakarta.websocket.OnClose;
import jakarta.websocket.OnError;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.ServerEndpoint;

@ServerEndpoint(value = "/my-web-socket")
@ApplicationScoped
public class DemoAuthWebSocket {

    private static final Logger log = LoggerFactory.getLogger(DemoAuthWebSocket.class);

    private static final String AUTH_REQUIRED = "Authentication required";
    private static final String AUTH_EXPIRED = "Authentication expired";

    @Inject
    Vertx vertx;

    private final Map<String, Session> sessions = new ConcurrentHashMap<>();
    private final Map<Session, String> sessionToUserId = new ConcurrentHashMap<>();
    private final Map<Session, Long> sessionTimers = new ConcurrentHashMap<>();

    @OnOpen
    public void onOpen(Session session) {
        if (!(session.getUserPrincipal() instanceof WebSocketPrincipal principal)) {
            closeSession(session, CloseReason.CloseCodes.VIOLATED_POLICY, AUTH_REQUIRED);
            return;
        }

        SecurityIdentity identity = principal.getSecurityIdentity();
        if (identity == null || identity.isAnonymous()) {
            closeSession(session, CloseReason.CloseCodes.VIOLATED_POLICY, AUTH_REQUIRED);
            return;
        }

        String userId = identity.getPrincipal().getName();
        Long expireAt = identity.getAttribute("quarkus.identity.expire-time");

        if (userId == null || expireAt == null) {
            closeSession(session, CloseReason.CloseCodes.VIOLATED_POLICY, AUTH_REQUIRED);
            return;
        }

        registerSession(userId, session, expireAt);
    }

    @OnClose
    public void onClose(Session session) {
        unregisterSession(session);
    }

    @OnError
    public void onError(Session session, Throwable throwable) {
        log.error("WebSocket error occurred.", throwable);
        unregisterSession(session);
    }

    public void pushToUser(String userId, String message) {
        log.debug("Attempting to send message to user '{}': {}", userId, message);

        Session session = sessions.get(userId);

        if (session == null || !session.isOpen()) {
            log.warn("Message delivery failed: No active session found for user '{}'.", userId);
            return;
        }

        session.getAsyncRemote().sendObject(message, result -> {
            if (result.getException() != null) {
                log.error("Failed to send message to user '{}'.", userId, result.getException());
            } else {
                log.debug("Message successfully sent to user '{}'.", userId);
            }
        });
    }

    private void registerSession(String userId, Session session, Long expiration) {
        log.info("User '{}' connected.", userId);
        sessions.put(userId, session);
        sessionToUserId.put(session, userId);
        scheduleSessionExpiry(userId, session, expiration);
    }

    private void unregisterSession(Session session) {
        cancelExpiryTimer(session);
        String userId = sessionToUserId.remove(session);
        if (userId != null) {
            sessions.remove(userId);
            log.info("User '{}' disconnected.", userId);
        }
    }

    private void scheduleSessionExpiry(String userId, Session session, Long expiration) {
        long expiryMillis = TimeUnit.SECONDS.toMillis(expiration);
        long delay = expiryMillis - System.currentTimeMillis();

        if (delay > 0) {
            long timerId = vertx.setTimer(delay, ignored -> {
                log.warn("Closing session for user '{}' due to expired authentication.", userId);
                closeSession(session, CloseReason.CloseCodes.VIOLATED_POLICY, AUTH_EXPIRED);
            });
            sessionTimers.put(session, timerId);
            log.debug("Scheduled session expiry timer (ID: {}) for user '{}' with delay of {} ms", timerId, userId, delay);
        } else {
            log.warn("Token for user '{}' is already expired. Closing session immediately.", userId);
            closeSession(session, CloseReason.CloseCodes.VIOLATED_POLICY, AUTH_EXPIRED);
        }
    }

    private void cancelExpiryTimer(Session session) {
        Long timerId = sessionTimers.remove(session);
        if (timerId != null) {
            boolean cancelled = vertx.cancelTimer(timerId);
            if (cancelled) {
                log.debug("Cancelled expiry timer (ID: {}) for session", timerId);
            }
        }
    }

    private void closeSession(Session session, CloseReason.CloseCodes code, String reason) {
        try {
            if (session.isOpen()) {
                cancelExpiryTimer(session);
                session.close(new CloseReason(code, reason));
                log.info("Closed session: {} - {}", code.getCode(), reason);
            }
        } catch (IOException e) {
            log.error("Error closing session.", e);
        }
    }
}

[Jaland]

I ran into a similar issue when I was playing around using Quarkus as a backend for my Godot game a few years back.

I just handled it in a similar way to how I think normally rest is handled, by creating a JWT and validating against that

Here is the relevant code

https://github.com/Jaland/godot-lobby-backend/blob/main/src%2Fmain%2Fjava%2Forg%2Flandister%2Flobby%2Fbackend%2Fwebsocket%2FBaseController.java#L93

[sberyozkin]

@colesturza @Jaland Thanks for sharing