package com.hdkj.lt.core.ws.interceptor;

import com.hdkj.lt.core.ws.config.WebSocketConfiguration;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.lang.NonNull;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

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

/**
 * @Description:
 * @Author: wenhao
 * @Create: 2024/04/26 14:07
 * @Version: 1.0
 */
@Slf4j
public class ServletWebSocketInterceptor extends HttpSessionHandshakeInterceptor {

    /**
     * 在握手之前执行此方法，继续握手则返回true，中断握手返回false
     *
     * @param request
     * @param response
     * @param wsHandler
     * @param attributes
     * @return
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(@NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response,
                                   @NonNull WebSocketHandler wsHandler, @NonNull Map<String, Object> attributes) throws Exception {
        ServletServerHttpRequest httpRequest = (ServletServerHttpRequest) request;
        HttpServletRequest servletRequest = httpRequest.getServletRequest();
        String sessionKey = servletRequest.getParameter(WebSocketConfiguration.SESSION_KEY);
        if (StringUtils.isBlank(sessionKey)) {
            response.setStatusCode(HttpStatus.BAD_REQUEST);
            return false;
        }
        attributes.put(WebSocketConfiguration.SESSION_KEY, sessionKey);
        return super.beforeHandshake(request, response, wsHandler, attributes);
    }

    /**
     * 在握手之后执行此方法，无论是否握手成功都指明了响应状态码和响应头
     *
     * @param request
     * @param response
     * @param wsHandler
     * @param ex
     */
    @Override
    public void afterHandshake(@NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response, @NonNull WebSocketHandler wsHandler, Exception ex) {
        super.afterHandshake(request, response, wsHandler, ex);
    }

}
