package com.gugu.boy.spring.websocket.handler.processor;

import com.gugu.boy.spring.websocket.MessagePayload;
import org.springframework.web.socket.WebSocketSession;

/**
 * 发送socket消息 template.
 *
 * @author zhoumeiqin
 * @date 2019/10/14
 */
public interface PayloadProcessor {
    /**
     * 执行.
     *
     * @param session
     * @param payload
     */
    void process(WebSocketSession session, MessagePayload payload) throws Exception;
}

package com.gugu.boy.spring.websocket.handler;

import com.github.greenyears.core.utils.AssertUtils;
import com.gugu.boy.spring.websocket.handler.processor.PayloadProcessor;
import com.gugu.boy.spring.websocket.handler.processor.ProcessorSupport;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;

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

/**
 * 消息文本处理上下文.
 *
 * @author zhoumeiqin
 * @date 2019/10/14
 */
@Slf4j
@Component
public class PayloadProcessorContext implements CommandLineRunner, ApplicationContextAware {
    private volatile ApplicationContext applicationContext;
    private static final Map<Integer, PayloadProcessor> processorMap = new ConcurrentHashMap<>();

    public PayloadProcessorContext(ApplicationContext applicationContext) {
        this.applicationContext = applicationContext;
    }

    @Override
    public void run(String... args) throws Exception {
        this.initialization();
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    /**
     * 初始化.
     */
    private void initialization() {
        Collection<PayloadProcessor> processors = this.applicationContext.getBeansOfType(PayloadProcessor.class).values();
        processors.forEach(processor -> {
            Class<? extends PayloadProcessor> clazz = processor.getClass();
            ProcessorSupport annotation = clazz.getAnnotation(ProcessorSupport.class);
            if (annotation != null) {
                int[] types = annotation.types();
                for (int type : types) {
                    processorMap.put(type, processor);
                }

            }
        });
    }

    /**
     * 获取实例.
     *
     * @param type
     * @return
     */
    public static PayloadProcessor getInstance(Integer type) {
        PayloadProcessor processor = processorMap.get(type);
        AssertUtils.notNull(processor, "不支持的消息类型");
        return processor;
    }
}

package com.gugu.boy.spring.websocket.interceptor;

import com.gugu.boy.spring.security.GuguUserDetails;
import com.gugu.boy.spring.security.authentication.storage.TokenStorage;
import com.gugu.boy.spring.security.authentication.storage.UserCache;
import com.gugu.boy.spring.security.authentication.token.TokenAuthenticationHandler;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import java.util.ArrayList;
import java.util.Map;

/**
 * websocket握手.
 *
 * @author zmq
 * @date 2019/10/14
 */
@Slf4j
@Component
public class GuguHandshakeInterceptor extends HttpSessionHandshakeInterceptor {

    @Autowired
    private TokenAuthenticationHandler tokenUtils;
    @Autowired
    private TokenStorage tokenStorage;

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request;
        // websocket的token信息不能想
        String token = serverRequest.getServletRequest().getParameter("token");
        boolean authentication = this.authentication(token, attributes);
        if (authentication) {
            log.info("握手成功");
            return super.beforeHandshake(request, response, wsHandler, attributes);
        } else {
            log.warn("握手失败");
        }
        return false;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
    }

    private boolean authentication(String token, Map<String, Object> attributes) {
        if (StringUtils.isBlank(token)) {
            log.warn("缺失参数: token");
            return false;
        }
        String id = tokenUtils.getFromToken(TokenAuthenticationHandler.CLAIM_KEY_ID, token);
        UserCache userCache = tokenStorage.getUserById(id);
        log.info("用户ID【{}】尝试进行握手", id);
        if (id != null && userCache != null) {
            GuguUserDetails userDetails = toUserDetails(userCache);
            if (token.equals(userCache.getToken())) {
                //token中的用户信息和数据库中的用户信息对比成功后将用户信息加入SecurityContextHolder相当于登录
                UsernamePasswordAuthenticationToken authentication = new UsernamePasswordAuthenticationToken(userDetails.getPhone(), null, new ArrayList<>(0));
                authentication.setDetails(userDetails);
                attributes.put("id", userDetails.getId());
                attributes.put("username", userDetails.getPhone());
                attributes.put("userType", userDetails.getUserType());
                SecurityContextHolder.getContext().setAuthentication(authentication);
                return true;
            } else {
                log.warn("用户{}可能在其他设备上被登录了", id);
            }

        }
        log.warn("token信息：{}", token);
        return false;
    }

    /**
     * 类型转换.
     *
     * @param source
     * @return UserCache
     * @author zhoumeiqin
     * @date 2019-07-30 10:21
     **/
    private GuguUserDetails toUserDetails(UserCache source) {
        GuguUserDetails target = new GuguUserDetails();
        BeanUtils.copyProperties(source, target);
        return target;
    }
}

package com.gugu.boy.spring.websocket.handler;

import com.alibaba.fastjson.JSON;
import com.gugu.boy.spring.security.GuguUserDetails;
import com.gugu.boy.spring.websocket.MessagePayload;
import com.gugu.boy.spring.websocket.PayloadConst;
import com.gugu.boy.spring.websocket.WebSocketContainer;
import com.gugu.boy.spring.websocket.WebSocketTemplate;
import com.gugu.core.exception.LogicException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

/**
 * @author zhoumeiqin
 * @date 2019/10/14
 */
@Slf4j
@Component
public class WebSocketMessageHandler extends TextWebSocketHandler {
    @Autowired
    private WebSocketContainer container;
    @Autowired
    private WebSocketTemplate webSocketTemplate;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        container.connect(session);
        super.afterConnectionEstablished(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        container.close(session, status);
        super.afterConnectionClosed(session, status);
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        GuguUserDetails user = container.getCurrentUser(session);
        MessagePayload payload = JSON.parseObject(message.getPayload(), MessagePayload.class);
        if (payload.getEvent() == PayloadConst.Event.HEARTBEAT) {
            log.info("{} heartbeat", user.getPhone());
            session.sendMessage(message);
        } else {
            try {
                webSocketTemplate.sendMessage(payload);
            } catch (LogicException e) {
                log.warn("触发业务异常: {} - {}", e.getErrorCode(), e.getMessage());
            } catch (Exception e) {
                log.error(",errMsg", e);
            }
        }
    }
}

package com.gugu.boy.spring.websocket;

import com.alibaba.fastjson.JSON;
import com.gugu.boy.spring.security.GuguUserDetails;
import com.gugu.boy.spring.websocket.handler.PayloadProcessorContext;
import com.gugu.boy.spring.websocket.handler.processor.PayloadProcessor;
import com.gugu.core.constants.SystemConstant;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.concurrent.BasicThreadFactory;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.io.IOException;
import java.security.Principal;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * @author zhoumeiqin
 * @date 2019/10/21
 */
@Slf4j
@Component
public class WebSocketContainer {
    private ScheduledExecutorService executor;
    /**
     * 查勘员在线session.
     * <p>key = user id </p>
     */
    private static Map<Long, WebSocketSession> surveySessions = new ConcurrentHashMap<>();

    @PostConstruct
    public void init() {
        this.monitor();
    }

    @PreDestroy
    public void destroy() {
        if (this.executor != null) {
            log.info("销毁定时器");
            this.executor.shutdown();
        }
    }

    /**
     * 监控.
     */
    private void monitor() {
        log.info("定时检测在线人数");
        Runnable command = () -> {
            log.info("在线人数: {}", surveySessions.size());
            log.info("在线人员：{}", surveySessions.keySet());
        };
        executor = new ScheduledThreadPoolExecutor(1,
                new BasicThreadFactory.Builder().namingPattern("socket-mon-%d").daemon(true).build());
        executor.scheduleAtFixedRate(command, 1, 5, TimeUnit.MINUTES);
    }

    /**
     * 查勘员连接socket.
     *
     * @param user
     * @param session
     * @return
     * @throws IOException
     */
    public void connect(GuguUserDetails user, WebSocketSession session) throws IOException {
        WebSocketSession history = surveySessions.get(user.getId());
        if (history != null) {
            log.info("用户[{}]重复连接，关闭已有连接", user.getPhone());
            history.close(WebSocketConst.Close.NORMAL);
        }
        surveySessions.put(user.getId(), session);
        log.info("[{}]登录了, sessionId: {}，当前在线人数：{}", user.getPhone(), session.getId(), surveySessions.size());
    }

    /**
     * 连接.
     *
     * @param session
     * @throws IOException
     */
    public void connect(WebSocketSession session) throws Exception {
        GuguUserDetails user = this.getCurrentUser(session);
        if (user != null) {
            this.connect(user, session);
            this.afterSocketConnect(user.getId());
        } else {
            log.warn("当前用户未知");
        }
    }

    /**
     * socket连接错误.
     *
     * @param userId
     */
    public void afterSocketConnect(Long userId) throws Exception {
        MessagePayload payload = new MessagePayload();
        payload.setEvent(PayloadConst.Event.TO_SINGLE);
        payload.setType(PayloadConst.Type.SOCKET_CONNECT);
        payload.setToUser(userId);
        payload.setToUserType(SystemConstant.SURVEYOR);
        this.sendMessage(userId, payload);
    }

    /**
     * close.
     *
     * @param userId
     */
    public void close(Long userId) throws IOException {
        WebSocketSession session = surveySessions.get(userId);
        if (session == null) {
            log.warn("用户[{}]不在线", userId);
            return;
        }
        session.close(WebSocketConst.Close.NORMAL);
        surveySessions.remove(userId);
    }

    /**
     * 关闭. 默认4000，主动关闭，app不会进行连接重试.
     *
     * @param user
     */
    public void close(GuguUserDetails user) throws IOException {
        this.close(user, WebSocketConst.Close.NORMAL);
    }

    /**
     * 关闭.
     *
     * @param user
     * @param status
     */
    public void close(GuguUserDetails user, CloseStatus status) throws IOException {
        if (user != null) {
            log.info("连接关闭：{}, 关闭状态：{}", user.getPhone(), status);
            this.close(user.getId());
        } else {
            log.warn("当前用户未知");
        }
    }

    /**
     * 关闭连接.
     *
     * @param session
     * @param status
     */
    public void close(WebSocketSession session, CloseStatus status) throws IOException {
        GuguUserDetails user = this.getCurrentUser(session);
        this.close(user, status);
    }

    /**
     * 获取session.
     *
     * @param userId
     * @return
     */
    public WebSocketSession getSession(Long userId) {
        return surveySessions.get(userId);
    }

    /**
     * 发送消息给单独用户.
     *
     * @param userId
     * @param payload
     * @throws Exception
     */
    public synchronized void sendMessage(Long userId, MessagePayload payload) throws Exception {
        log.info("发送消息给用户：{}, 数据：{}", userId, JSON.toJSONString(payload));
        WebSocketSession session = this.getSession(userId);
        if (session == null) {
            log.warn("当前用户:{}不存在或者不在线", userId);
            return;
        }
        PayloadProcessor processor = PayloadProcessorContext.getInstance(payload.getType());
        processor.process(session, payload);
    }

    /**
     * 发送消息给所有用户.
     *
     * @param payload
     */
    public synchronized void sendMessage(MessagePayload payload) {
        log.info("发送消息给所有人.: {}", JSON.toJSONString(payload));
        PayloadProcessor processor = PayloadProcessorContext.getInstance(payload.getType());
        surveySessions.forEach((key, value) -> {
            try {
                processor.process(value, payload);
            } catch (Exception e) {
                log.error("发送给用户消息失败: {}", JSON.toJSONString(value.getPrincipal()), e);
            }
        });
    }

    /**
     * 获取当前用户.
     *
     * @param session
     * @return
     */
    public GuguUserDetails getCurrentUser(WebSocketSession session) {
        Principal principal = session.getPrincipal();
        if (principal instanceof UsernamePasswordAuthenticationToken) {
            UsernamePasswordAuthenticationToken authenticationToken = (UsernamePasswordAuthenticationToken) principal;
            return (GuguUserDetails) authenticationToken.getDetails();
        } else {
            log.warn("未知的连接：{}", JSON.toJSONString(session.getPrincipal()));
        }
        return null;
    }

    /**
     * 是否在线.
     *
     * @param userId
     * @return
     */
    public boolean isOnline(Long userId) {
        return surveySessions.containsKey(userId);
    }
}