Using RowWebsocketSession instead of Spring WebsocketSession in class…

…es methods
ep2p · Aug 13, 2020 · ced0499 · ced0499
1 parent f5ee805
commit ced0499
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 11 deletions.
diff --git a/src/main/java/labs/psychogen/row/filter/RowFilter.java b/src/main/java/labs/psychogen/row/filter/RowFilter.java
@@ -1,9 +1,9 @@
 package labs.psychogen.row.filter;
 
+import labs.psychogen.row.domain.RowWebsocketSession;
 import labs.psychogen.row.domain.protocol.RequestDto;
 import labs.psychogen.row.domain.protocol.ResponseDto;
-import org.springframework.web.socket.WebSocketSession;
 
 public interface RowFilter {
-    boolean filter(RequestDto requestDto, ResponseDto responseDto, WebSocketSession webSocketSession) throws Exception;
+    boolean filter(RequestDto requestDto, ResponseDto responseDto, RowWebsocketSession rowWebsocketSession) throws Exception;
 }
diff --git a/src/main/java/labs/psychogen/row/filter/RowFilterChain.java b/src/main/java/labs/psychogen/row/filter/RowFilterChain.java
@@ -1,5 +1,6 @@
 package labs.psychogen.row.filter;
 
+import labs.psychogen.row.domain.RowWebsocketSession;
 import labs.psychogen.row.domain.protocol.RequestDto;
 import labs.psychogen.row.domain.protocol.ResponseDto;
 import lombok.extern.slf4j.Slf4j;
@@ -53,9 +54,9 @@ public void addFilterAfter(RowFilter rowFilter, Class after){
         }
     }
 
-    public void filter(RequestDto requestDto, ResponseDto response, WebSocketSession session) throws Exception {
+    public void filter(RequestDto requestDto, ResponseDto response, RowWebsocketSession rowWebsocketSession) throws Exception {
         for (RowFilter filter : filters) {
-            if (!filter.filter(requestDto, response, session)) {
+            if (!filter.filter(requestDto, response, rowWebsocketSession)) {
                 return;
             }
         }

diff --git a/src/main/java/labs/psychogen/row/filter/RowInvokerFiler.java b/src/main/java/labs/psychogen/row/filter/RowInvokerFiler.java
@@ -1,11 +1,11 @@
 package labs.psychogen.row.filter;
 
 import labs.psychogen.row.domain.RowResponseStatus;
+import labs.psychogen.row.domain.RowWebsocketSession;
 import labs.psychogen.row.domain.protocol.RequestDto;
 import labs.psychogen.row.domain.protocol.ResponseDto;
 import labs.psychogen.row.exception.InvalidPathException;
 import labs.psychogen.row.service.RowInvokerService;
-import org.springframework.web.socket.WebSocketSession;
 
 import java.lang.reflect.InvocationTargetException;
 
@@ -17,7 +17,7 @@ public RowInvokerFiler(RowInvokerService rowInvokerService) {
     }
 
     @Override
-    public boolean filter(RequestDto requestDto, ResponseDto responseDto, WebSocketSession webSocketSession) throws Exception {
+    public boolean filter(RequestDto requestDto, ResponseDto responseDto, RowWebsocketSession rowWebsocketSession) throws Exception {
         try {
             Object invoke = rowInvokerService.invoke(requestDto, responseDto);
             responseDto.setBody(invoke);

diff --git a/src/main/java/labs/psychogen/row/filter/SubscribeFilter.java b/src/main/java/labs/psychogen/row/filter/SubscribeFilter.java
@@ -2,6 +2,7 @@
 
 import labs.psychogen.row.config.Naming;
 import labs.psychogen.row.domain.RowResponseStatus;
+import labs.psychogen.row.domain.RowWebsocketSession;
 import labs.psychogen.row.domain.protocol.RequestDto;
 import labs.psychogen.row.domain.protocol.ResponseDto;
 import labs.psychogen.row.event.Subscription;
@@ -23,7 +24,7 @@ public SubscribeFilter(SubscriberService subscriberService, boolean pre) {
     }
 
     @Override
-    public boolean filter(RequestDto requestDto, ResponseDto responseDto, WebSocketSession webSocketSession) throws Exception {
+    public boolean filter(RequestDto requestDto, ResponseDto responseDto, RowWebsocketSession rowWebsocketSession) throws Exception {
         try {
             if(requestDto.getHeaders() != null && requestDto.getHeaders().containsKey(Naming.UNSUBSCRIBE_HEADER_NAME)){
                 String value = requestDto.getHeaders().get(Naming.UNSUBSCRIBE_HEADER_NAME);

diff --git a/src/main/java/labs/psychogen/row/service/ProtocolService.java b/src/main/java/labs/psychogen/row/service/ProtocolService.java
@@ -8,6 +8,7 @@
 import labs.psychogen.row.context.RowContextHolder;
 import labs.psychogen.row.context.RowUser;
 import labs.psychogen.row.domain.RowResponseStatus;
+import labs.psychogen.row.domain.RowWebsocketSession;
 import labs.psychogen.row.domain.protocol.RequestDto;
 import labs.psychogen.row.domain.protocol.ResponseDto;
 import labs.psychogen.row.filter.RowFilterChain;
@@ -34,10 +35,10 @@ public ProtocolService(RowFilterChain rowFilterChain) {
         validator = Validation.buildDefaultValidatorFactory().getValidator();
     }
 
-    public void handle(WebSocketSession webSocketSession, TextMessage textMessage){
+    public void handle(RowWebsocketSession webSocketSession, TextMessage textMessage){
         log.trace("Received message: " + textMessage.getPayload());
         String payload = textMessage.getPayload();
-        fillContext(webSocketSession);
+        fillContext(webSocketSession.getSession());
         ResponseDto responseDto = ResponseDto.builder().status(RowResponseStatus.OK.getId()).build();
         String requestId = null;
         try {
@@ -67,7 +68,7 @@ public void handle(WebSocketSession webSocketSession, TextMessage textMessage){
         try {
             String responsePayload = objectMapper.writeValueAsString(responseDto);
             log.trace("Sending response: "+ responsePayload);
-            webSocketSession.sendMessage(new TextMessage(responsePayload));
+            webSocketSession.getSession().sendMessage(new TextMessage(responsePayload));
         } catch (IOException e) {
             log.error("Failed to publish response to websocket", e);
         }

diff --git a/src/main/java/labs/psychogen/row/ws/RowWebSocketHandler.java b/src/main/java/labs/psychogen/row/ws/RowWebSocketHandler.java
@@ -49,7 +49,7 @@ public void afterConnectionEstablished(WebSocketSession session) throws Exceptio
 
     @Override
     protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
-        protocolService.handle(session, message);
+        protocolService.handle(rowSessionRegistry.getSession(WebsocketSessionUtil.getUserId(session), session.getId()), message);
         updateHeartbeat(session);
     }