java-classic-code/spring-zuul/WebSocketProxyFilter.java

package com.datasophon.api.configuration;

import com.netflix.zuul.ZuulFilter;
import com.netflix.zuul.context.RequestContext;
import com.netflix.zuul.exception.ZuulException;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 *
 * Zuul 支持 websocket 的代理
 *
 * <pre>
 *
 * Created by zhenqin.
 * User: zhenqin
 * Date: 2023/3/17
 * Time: 上午11:32
 * Vendor: yiidata.com
 *
 * </pre>
 *
 * @author zhenqin
 */
@Component
@Slf4j
public class WebSocketProxyFilter extends ZuulFilter {


    /**
     * 所有 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);
    }

    /**
     * 过滤器类型
     * pre 前置
     * routing
     * post 后置
     * error
     *
     * @return
     */
    @Override
    public String filterType() {
        return "pre";
    }


    /**
     * 执行顺序
     * 数值越小，优先级越高
     *
     * @return
     */
    @Override
    public int filterOrder() {
        return 0;
    }


    /**
     * 执行条件
     * true 开启, 执行 run 的特殊处理
     * false 关闭，否则不需要特殊处理，直接代理
     *
     * @return
     */
    @Override
    public boolean shouldFilter() {
        RequestContext rc = RequestContext.getCurrentContext();
        HttpServletRequest request = rc.getRequest();
        if(request == null) {
            return false;
        }
        final String path = request.getServletPath();
        final String connectionHeader = request.getHeader("Connection");
        final String upgradeHeader = request.getHeader("Upgrade");
        // 检验是不是 websocket 的链接, websocket 连接 Connection: Upgrade, Upgrade: websocket, Websocket, WebSocket
        return StringUtils.equals("Upgrade", connectionHeader) && StringUtils.equalsIgnoreCase("websocket", upgradeHeader) && WEBSOCKET_ACTION_MAPPING.get(path) != null;
    }


    /**
     * 动作（具体操作）
     * 具体逻辑
     *
     * @return
     * @throws ZuulException
     */
    @Override
    public Object run() throws ZuulException {
        RequestContext rc = RequestContext.getCurrentContext();
        HttpServletRequest request = rc.getRequest();
        final String path = request.getServletPath();
        rc.addZuulRequestHeader("Connection", request.getHeader("Connection"));
        rc.addZuulRequestHeader("Upgrade", request.getHeader("Upgrade"));
        //过滤该请求，不进行路由
        rc.setSendZuulResponse(false);
        //返回对应http status状态码
        rc.setResponseStatusCode(401);
        //设置返回的内容
        rc.setResponseBody("{'error':'token is empty'}");
        final ApiRoute apiRoute = WEBSOCKET_ACTION_MAPPING.get(path);
        log.info(path);
        try {
            Thread.sleep(10000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        return apiRoute;
    }
}