package com.njzscloud.common.mvc.ws.support;

import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;

/**
 * websocket Token 拦截器<br/>
 * Sec-WebSocket-Protocol: [TOKEN]<br/>
 * TOKEN 已在 Security 中校验过，此处不校验，仅配合 Websocket 完成子协议回写
 */
@Slf4j
public class TokenHandshakeInterceptor implements HandshakeInterceptor {
    private static final String SEC_WEB_SOCKET_PROTOCOL = "Sec-WebSocket-Protocol";

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse serverHttpResponse, WebSocketHandler wsHandler, Map<String, Object> attributes) {
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        try {
            HttpServletRequest httpRequest = ((ServletServerHttpRequest) request).getServletRequest();
            HttpServletResponse httpResponse = ((ServletServerHttpResponse) response).getServletResponse();
            String token = httpRequest.getHeader(SEC_WEB_SOCKET_PROTOCOL);
            if (StrUtil.isNotEmpty(token)) {
                httpResponse.addHeader(SEC_WEB_SOCKET_PROTOCOL, token);
            }
        } catch (ClassCastException e) {
            log.error("[Websocket 拦截器] 类型转换失败: {}、{}", request.getClass(), response.getClass(), e);
        }
    }
}
