....

src/main/java/cn/com/qmth/examcloud/ws/core/Message.java

@@ -0,0 +1,22 @@
+package cn.com.qmth.examcloud.ws.core;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * 接口ID 注解
+ *
+ * @author WANGWEI
+ * @date 2019年3月15日
+ * @Copyright (c) 2018-? http://qmth.com.cn All Rights Reserved.
+ */
+@Target({ElementType.PARAMETER})
+@Retention(RetentionPolicy.RUNTIME)
+@Documented
+public @interface Message {
+
+	String value() default "";
+}

src/main/java/cn/com/qmth/examcloud/ws/core/SessionHolder.java

@@ -21,60 +21,34 @@ public class SessionHolder {
 
 	private static final Map<Session, String> SESESION_ID_MAP = Maps.newConcurrentMap();
 
-	/**
-	 * 获取 Session
-	 *
-	 * @author WANGWEI
-	 * @param sessionId
-	 * @return
-	 */
-	public static Session getSession(String sessionId) {
-		Session session = ID_SESESION_MAP.get(sessionId);
-		return session;
-	}
+	private static final Map<Session, String> SESSION_KEY_MAP = Maps.newConcurrentMap();
 
-	/**
-	 * 获取 sessionId
-	 *
-	 * @author WANGWEI
-	 * @param session
-	 * @return
-	 */
-	public static String getSessionId(Session session) {
-		return SESESION_ID_MAP.get(session);
-	}
+	private static final Map<Session, String> SESSION_TOKEN_MAP = Maps.newConcurrentMap();
 
 	/**
 	 * 重新设置session
 	 *
 	 * @author WANGWEI
-	 * @param sessionId
+	 * @param path
+	 * @param key
+	 * @param token
 	 * @param session
 	 */
-	public static void setSession(String sessionId, Session session) {
+	public static void setSession(String path, String key, String token, Session session) {
+		String sessionId = key + ":" + path;
 		Session oldSession = ID_SESESION_MAP.get(sessionId);
 		if (null != oldSession && !oldSession.equals(session)) {
-			IOUtils.closeQuietly(oldSession);
 			SESESION_ID_MAP.remove(oldSession);
+			SESSION_KEY_MAP.remove(oldSession);
+			SESSION_TOKEN_MAP.remove(oldSession);
+			IOUtils.closeQuietly(oldSession);
 		}
 
 		ID_SESESION_MAP.put(sessionId, session);
 		SESESION_ID_MAP.put(session, sessionId);
-	}
 
-	/**
-	 * 删除 session
-	 *
-	 * @author WANGWEI
-	 * @param sessionId
-	 */
-	public static void delSession(String sessionId) {
-		Session session = ID_SESESION_MAP.get(sessionId);
-		ID_SESESION_MAP.remove(sessionId);
-		if (null != session) {
-			IOUtils.closeQuietly(session);
-			SESESION_ID_MAP.remove(session);
-		}
+		SESSION_KEY_MAP.put(session, key);
+		SESSION_TOKEN_MAP.put(session, token);
 	}
 
 	/**
@@ -88,7 +62,45 @@ public class SessionHolder {
 		ID_SESESION_MAP.remove(sessionId);
 
 		SESESION_ID_MAP.remove(session);
+		SESSION_KEY_MAP.remove(session);
+		SESSION_TOKEN_MAP.remove(session);
+
 		IOUtils.closeQuietly(session);
 	}
 
+	/**
+	 * 获取 session
+	 *
+	 * @author WANGWEI
+	 * @param path
+	 * @param key
+	 * @return
+	 */
+	public Session getSession(String path, String key) {
+		String sessionId = key + ":" + path;
+		Session session = ID_SESESION_MAP.get(sessionId);
+		return session;
+	}
+
+	/**
+	 * 获取 key
+	 *
+	 * @author WANGWEI
+	 * @param session
+	 * @return
+	 */
+	public static String getKey(Session session) {
+		return SESSION_KEY_MAP.get(session);
+	}
+
+	/**
+	 * 获取 token
+	 *
+	 * @author WANGWEI
+	 * @param session
+	 * @return
+	 */
+	public static String getToken(Session session) {
+		return SESSION_TOKEN_MAP.get(session);
+	}
 }

src/main/java/cn/com/qmth/examcloud/ws/core/WebSocketServerEndpoint.java

@@ -1,6 +1,8 @@
 package cn.com.qmth.examcloud.ws.core;
 
 import java.io.IOException;
+import java.lang.annotation.Annotation;
+import java.lang.reflect.Method;
 import java.util.List;
 import java.util.Map;
 
@@ -21,6 +23,7 @@ import org.springframework.stereotype.Component;
 import cn.com.qmth.examcloud.api.commons.security.bean.User;
 import cn.com.qmth.examcloud.commons.logging.ExamCloudLog;
 import cn.com.qmth.examcloud.commons.logging.ExamCloudLogFactory;
+import cn.com.qmth.examcloud.commons.util.JsonUtil;
 import cn.com.qmth.examcloud.web.redis.RedisClient;
 
 /**
@@ -80,22 +83,7 @@ public class WebSocketServerEndpoint {
 			return;
 		}
 
-		String sessionId = buildSessionId(user, path);
-
-		SessionHolder.setSession(sessionId, session);
-
-	}
-
-	/**
-	 * 构建 sessionId
-	 *
-	 * @author WANGWEI
-	 * @param user
-	 * @param path
-	 * @return
-	 */
-	private String buildSessionId(User user, String path) {
-		return user.getKey() + ":" + path;
+		SessionHolder.setSession(path, user.getKey(), user.getToken(), session);
 	}
 
 	/**
@@ -126,6 +114,70 @@ public class WebSocketServerEndpoint {
 	@OnMessage
 	public void onMessage(Session session, @PathParam("path") String path, String message) {
 		WS_LOG.debug("[onMessage]. path=" + path + ". message=" + message);
+		String key = SessionHolder.getKey(session);
+		String token = SessionHolder.getToken(session);
+
+		User user = redisClient.get(key, User.class);
+
+		if (null == user) {
+			WS_LOG.error("[onMessage-FAIL]. path=" + path + ". no login.");
+			IOUtils.closeQuietly(session);
+			return;
+		} else if (!token.equals(user.getToken())) {
+			WS_LOG.error("[onMessage-FAIL]. path=" + path + ". token is wrong.");
+			IOUtils.closeQuietly(session);
+			return;
+		}
+
+		Method method = MessageHandlerHolder.getMethod(path);
+		Object bean = MessageHandlerHolder.getBean(path);
+
+		Annotation[][] an2 = method.getParameterAnnotations();
+
+		boolean hasMessageParam = false;
+		int messageParamIndex = 0;
+		OUTER : for (int i = 0; i < an2.length; i++) {
+			Annotation[] an1 = an2[i];
+			for (Annotation an : an1) {
+				if (an.annotationType().equals(Message.class)) {
+					hasMessageParam = true;
+					messageParamIndex = i;
+					break OUTER;
+				}
+			}
+		}
+
+		Class<?>[] parameterTypes = method.getParameterTypes();
+		Object[] args = new Object[parameterTypes.length];
+		for (int i = 0; i < parameterTypes.length; i++) {
+			Class<?> curType = parameterTypes[i];
+			if (hasMessageParam && i == messageParamIndex) {
+				args[i] = JsonUtil.fromJson(message, curType);
+				continue;
+			}
+			if (curType.equals(User.class)) {
+				args[i] = user;
+				continue;
+			}
+		}
+
+		Object result = null;
+		try {
+			result = method.invoke(bean, args);
+		} catch (Exception e) {
+			WS_LOG.error("[onMessage-FAIL]. path=" + path + "", e);
+		}
+
+		if (null != result) {
+			try {
+				session.getBasicRemote().sendText(JsonUtil.toJson(result));
+			} catch (IOException e) {
+				WS_LOG.error("[onMessage-FAIL]. path=" + path + "", e);
+				IOUtils.closeQuietly(session);
+				return;
+			}
+		}
+
 	}
 
 	@OnError

src/main/java/cn/com/qmth/examcloud/ws/handler/TestHandler.java

@@ -2,16 +2,17 @@ package cn.com.qmth.examcloud.ws.handler;
 
 import org.springframework.stereotype.Component;
 
+import cn.com.qmth.examcloud.api.commons.security.bean.User;
+import cn.com.qmth.examcloud.ws.core.Message;
 import cn.com.qmth.examcloud.ws.core.MessageHandler;
-import cn.com.qmth.examcloud.ws.core.MessageIn;
 
 @Component
 @MessageHandler("test")
 public class TestHandler {
 
 	@MessageHandler
-	public MessageIn test(MessageIn messageIn) {
-		return messageIn;
+	public String test(User user, @Message String message) {
+		return "";
 	}
 
 }