WebSocketConfig.java

import lombok.RequiredArgsConstructor;
import lombok.val;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.env.Environment;
import org.springframework.lang.NonNull;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.config.StompBrokerRelayRegistration;
import org.springframework.session.Session;
import org.springframework.session.web.socket.config.annotation.AbstractSessionWebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.StompWebSocketEndpointRegistration;

import javax.annotation.PostConstruct;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Optional;

@RequiredArgsConstructor
@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig<S extends Session> extends AbstractSessionWebSocketMessageBrokerConfigurer<S> {
    private final Environment environment;

    private RelayConfigurer relayConfigurer;

    @PostConstruct
    protected void init() throws URISyntaxException {
        val relayUrl = environment.getProperty("relay.url", String.class);
        if (relayUrl != null) {
            relayConfigurer = new RelayConfigurer.UriAdapter(new URI(relayUrl));
        }
    }

    @Override
    public void configureMessageBroker(@NonNull MessageBrokerRegistry registry) {
        val destinationPrefixes = new String[]{"/queue/", "/topic/"};

        if (relayConfigurer == null) {
            // use the simple broker when a relay url is not configured
            registry.enableSimpleBroker(destinationPrefixes);
        } else {
            // use the stomp broker when a relay url is configured
            var relay = registry.enableStompBrokerRelay(destinationPrefixes)
                .setUserDestinationBroadcast("/topic/unresolved-user-dest")
                .setUserRegistryBroadcast("/topic/simp-user-registry");
            relayConfigurer.configure(relay);
        }

        registry.setApplicationDestinationPrefixes("/app");
    }

    @Override
    protected void configureStompEndpoints(StompEndpointRegistry registry) {
        applyCustomizations(registry.addEndpoint("/websocket"));
        applyCustomizations(registry.addEndpoint("/websocket/")).withSockJS();
    }

    private StompWebSocketEndpointRegistration applyCustomizations(StompWebSocketEndpointRegistration endpointRegistration) {
        // add handshake interceptors, set handshake handler, set allowed origins etc.
        return endpointRegistration.setAllowedOriginPatterns("*");
    }
}

interface RelayConfigurer {
    Optional<String> getUsername();

    Optional<String> getPassword();

    Optional<String> getRelayHost();

    Optional<Integer> getRelayPort();

    Optional<String> getVirtualHost();

    default void configure(StompBrokerRelayRegistration registration) {
        val log = LoggerFactory.getLogger(RelayConfigurer.class);

        getUsername().ifPresentOrElse(username -> {
            log.debug("relay.username resolved to {}", username);
            registration.setClientLogin(username);
            registration.setSystemLogin(username);
        }, () -> log.debug("relay.username could not be resolved)"));

        getPassword().ifPresentOrElse(password -> {
            log.debug("relay.password resolved to {}", password);
            registration.setClientPasscode(password);
            registration.setSystemPasscode(password);
        }, () -> log.debug("relay.password could not be resolved)"));

        getRelayHost().ifPresentOrElse(relayHost -> {
            log.debug("relay.host resolved to {}", relayHost);
            registration.setRelayHost(relayHost);
        }, () -> log.debug("relay.host could not be resolved"));

        getRelayPort().ifPresentOrElse(relayPort -> {
            log.debug("relay.port resolved to {}", relayPort);
            registration.setRelayPort(relayPort);
        }, () -> log.debug("relay.port could not be resolved"));

        getVirtualHost().ifPresentOrElse(vHost -> {
            log.debug("relay.vHost resolved to {}", vHost);
            registration.setVirtualHost(vHost);
        }, () -> log.debug("relay.vHost could not be resolved"));
    }

    @RequiredArgsConstructor
    class UriAdapter implements RelayConfigurer {
        private final URI uri;

        public Optional<String> getUsername() {
            return Optional.ofNullable(uri.getUserInfo()).map(userInfo -> {
                int idx = userInfo.indexOf(":");
                return idx < 0 ? userInfo : userInfo.substring(0, idx);
            });
        }

        public Optional<String> getPassword() {
            return Optional.ofNullable(uri.getUserInfo()).flatMap(userInfo -> {
                int idx = userInfo.indexOf(":");
                return idx < 0 ? Optional.empty() : Optional.of(userInfo.substring(idx + 1));
            });
        }

        public Optional<String> getRelayHost() {
            return Optional.ofNullable(uri.getHost());
        }

        public Optional<Integer> getRelayPort() {
            return uri.getPort() < 0 ? Optional.empty() : Optional.of(uri.getPort());
        }

        public Optional<String> getVirtualHost() {
            return Optional.ofNullable(uri.getPath())
                .map(vHost -> vHost.startsWith("/") ? vHost.substring(1) : vHost)
                .filter(vHost -> vHost.length() > 0);
        }
    }
}