java-classic-code/websocket/WebSocketConfig.java

package cn.exlive.video.config;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import cn.exlive.video.handler.ProxyWebSocketHandler;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.DependsOn;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

/**
 * 开启WebSocket支持,支持 WebSocket 路由代理配置类
 *
 * @author cxf
 */
@Slf4j
@Setter
@Getter
@EnableWebSocket
@Configuration
public class WebSocketConfig {

    /**
     * 所有 websocket 的特殊处理
     */
    final static Map<String, ApiRoute> WEBSOCKET_ACTION_MAPPING = new ConcurrentHashMap<>();

    /**
     * 添加 WebSocket 地址
     * @param path
     * @param route
     */
    public static void addWebSocketRouter(String path, ApiRoute route) {
        WEBSOCKET_ACTION_MAPPING.put(path, route);
    }

    /**
     * 	注入ServerEndpointExporter，
     * 	这个bean会自动注册使用了@ServerEndpoint注解声明的Websocket endpoint
     */
    // @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }

    @Bean
    @DependsOn("simpleRouteLocator")
    public WebSocketConfigurer webSocketConfigurer() {
        return (WebSocketHandlerRegistry registry) -> {
            for (Map.Entry<String, ApiRoute> entry : WEBSOCKET_ACTION_MAPPING.entrySet()) {
                registry.addHandler(new ProxyWebSocketHandler(entry.getValue().getUrl()), entry.getValue().getApi()) // 设置连接路径和处理
                        .setAllowedOrigins("*");
            }
        };
    }
}