Unsubscribe subscriptions when user role is updated

Author: Goldmensch

backend/src/main/java/club/devcord/devmarkt/DevelopmentStartup.java

@@ -16,14 +16,11 @@
 
 package club.devcord.devmarkt;
 
-import club.devcord.devmarkt.repositories.TemplateRepo;
-import io.micronaut.context.BeanContext;
 import io.micronaut.context.annotation.Requires;
 import io.micronaut.context.event.ApplicationEventListener;
 import io.micronaut.context.event.StartupEvent;
 import io.micronaut.json.JsonMapper;
 import io.micronaut.security.token.jwt.generator.JwtTokenGenerator;
-import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
 import java.util.Map;
 import org.slf4j.Logger;

backend/src/main/java/club/devcord/devmarkt/auth/InvocationDataCustomizer.java

@@ -16,6 +16,7 @@
 
 package club.devcord.devmarkt.auth;
 
+import club.devcord.devmarkt.ws.SessionMetaData;
 import graphql.ExecutionInput;
 import io.micronaut.configuration.graphql.GraphQLExecutionInputCustomizer;
 import io.micronaut.http.HttpRequest;
@@ -28,17 +29,26 @@
 public class InvocationDataCustomizer implements GraphQLExecutionInputCustomizer {
 
   private final UserProvider provider;
+  private final SessionMetaData sessionMetaData;
 
-  public InvocationDataCustomizer(UserProvider provider) {
+  public InvocationDataCustomizer(UserProvider provider, SessionMetaData sessionMetaData) {
     this.provider = provider;
+    this.sessionMetaData = sessionMetaData;
   }
 
   @Override
   public Publisher<ExecutionInput> customize(ExecutionInput executionInput, HttpRequest httpRequest,
       MutableHttpResponse<String> httpResponse) {
     httpRequest.getHeaders().getAuthorization()
         .flatMap(provider::validate)
-        .ifPresent(user -> executionInput.getGraphQLContext().put("user", user));
+        .ifPresent(user -> {
+          executionInput.getGraphQLContext().put("user", user);
+          var session = sessionMetaData.getHttpRequestWebSocketSessionMap().get(httpRequest);
+          if (session != null) {
+            sessionMetaData.getUserSessions().put(user.id(), session);
+          }
+        });
+    sessionMetaData.getHttpRequestWebSocketSessionMap().remove(httpRequest);
     return Mono.just(executionInput);
   }
 }

backend/src/main/java/club/devcord/devmarkt/services/UserService.java

@@ -25,16 +25,19 @@
 import club.devcord.devmarkt.responses.Success;
 import club.devcord.devmarkt.responses.failure.user.ErrorCode;
 import club.devcord.devmarkt.util.Admins;
+import club.devcord.devmarkt.ws.ReflectiveUnsubscriber;
 import jakarta.inject.Singleton;
 import java.util.Optional;
 
 @Singleton
 public class UserService {
 
   private final UserRepo repo;
+  private final ReflectiveUnsubscriber reflectiveUnsubscriber;
 
-  public UserService(UserRepo repo) {
+  public UserService(UserRepo repo, ReflectiveUnsubscriber reflectiveUnsubscriber) {
     this.repo = repo;
+    this.reflectiveUnsubscriber = reflectiveUnsubscriber;
   }
 
   public Optional<User> findDirect(UserId userId) {
@@ -51,6 +54,7 @@ public boolean delete(UserId userId) {
     if (Admins.isAdminUserId(userId)) {
       return false;
     }
+    reflectiveUnsubscriber.unsubscribeSubscriptions(userId);
     return repo.deleteOneById(userId) >= 1;
   }
 
@@ -72,7 +76,7 @@ public Response<User> updateRole(UserId userId, Role role) {
     if (Admins.isAdminUserId(userId)) {
       return new Failure<>(ErrorCode.ADMIN_USER_CANT_BE_MODIFIED);
     }
-
+    reflectiveUnsubscriber.unsubscribeSubscriptions(userId);
     var updated = repo.updateById(userId, role);
     return updated != 0
         ? new Success<>(new User(-1, userId, role))

backend/src/main/java/club/devcord/devmarkt/ws/CustomGraphQLWsController.java

@@ -39,10 +39,13 @@ public class CustomGraphQLWsController {
 
   private final GraphQLWsController controller;
   private final GraphQLJsonSerializer serializer;
+  private final SessionMetaData sessionMetaData;
 
-  public CustomGraphQLWsController(GraphQLWsController controller, GraphQLJsonSerializer serializer) {
+  public CustomGraphQLWsController(GraphQLWsController controller, GraphQLJsonSerializer serializer,
+      SessionMetaData sessionMetaData) {
     this.controller = controller;
     this.serializer = serializer;
+    this.sessionMetaData = sessionMetaData;
   }
 
   @OnOpen
@@ -53,11 +56,13 @@ public void onOpen(WebSocketSession session, HttpRequest<?> request) {
   @OnMessage
   public Publisher<GraphQLWsResponse> onMessage(String message, WebSocketSession session) {
     var type = serializer.deserialize(message, MessageType.class);
+    var request = session.get("httpRequest", HttpRequest.class).get();
     if (ClientType.GQL_CONNECTION_INIT.getType().equals(type.type())) {
-      var request = session.get("httpRequest", HttpRequest.class).get();
+      var auth = (String) type.payload().get("Authorization");
       request.mutate()
-          .header("Authorization", (String) type.payload().get("Authorization"));
+          .header("Authorization", auth);
     }
+    sessionMetaData.getHttpRequestWebSocketSessionMap().put(request, session);
     return controller.onMessage(message, session);
   }
 

backend/src/main/java/club/devcord/devmarkt/ws/ReflectiveUnsubscriber.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2022 Contributors to the Devmarkt-Backend project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package club.devcord.devmarkt.ws;
+
+import club.devcord.devmarkt.entities.auth.UserId;
+import io.micronaut.configuration.graphql.ws.GraphQLWsRequest;
+import io.micronaut.configuration.graphql.ws.GraphQLWsResponse;
+import io.micronaut.context.BeanContext;
+import io.micronaut.websocket.WebSocketSession;
+import jakarta.inject.Singleton;
+import java.util.Map;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+
+// Dirty solution to unsubscribe a user's subscriptions (uses reflections)
+@Singleton
+public class ReflectiveUnsubscriber{
+
+  private final Object wsState;
+  private final SessionMetaData metaData;
+
+  public ReflectiveUnsubscriber(BeanContext beanContext, SessionMetaData metaData) {
+    this.metaData = metaData;
+    try {
+      var stateClazz = Class.forName("io.micronaut.configuration.graphql.ws.GraphQLWsState");
+      this.wsState = beanContext.getBean(stateClazz);
+    } catch (ClassNotFoundException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  public void unsubscribeSubscriptions(UserId userId) {
+    var session = metaData.getUserSessions().get(userId);
+    if (session == null) return;
+    try {
+      var activeOperationsField = wsState.getClass().getDeclaredField("activeOperations");
+      activeOperationsField.setAccessible(true);
+      var activeOperations = ((Map<String, ?>) activeOperationsField.get(wsState)).get(session.getId());
+      var operationIdsField = activeOperations.getClass().getDeclaredField("activeOperations");
+      operationIdsField.setAccessible(true);
+      var operationIds = ((Map<String, ?>) operationIdsField.get(activeOperations)).keySet();
+      var method = wsState.getClass().getDeclaredMethod("stopOperation", GraphQLWsRequest.class,
+          WebSocketSession.class);
+      method.setAccessible(true);
+      for (var id : operationIds) {
+        var request = new GraphQLWsRequest();
+        request.setId(id);
+        Flux.from((Publisher<GraphQLWsResponse>) method.invoke(wsState, request, session))
+            .subscribe(session::sendSync);
+      }
+    } catch (Throwable e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+}

backend/src/main/java/club/devcord/devmarkt/ws/SessionMetaData.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2022 Contributors to the Devmarkt-Backend project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package club.devcord.devmarkt.ws;
+
+import club.devcord.devmarkt.entities.auth.UserId;
+import io.micronaut.http.HttpRequest;
+import io.micronaut.websocket.WebSocketSession;
+import jakarta.inject.Singleton;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+@Singleton
+public class SessionMetaData {
+  private final Map<UserId, WebSocketSession> userSessions = new ConcurrentHashMap<>();
+  private final Map<HttpRequest<?>, WebSocketSession> httpRequestWebSocketSessionMap
+      = new ConcurrentHashMap<>();
+
+  public Map<UserId, WebSocketSession> getUserSessions() {
+    return userSessions;
+  }
+
+  public Map<HttpRequest<?>, WebSocketSession> getHttpRequestWebSocketSessionMap() {
+    return httpRequestWebSocketSessionMap;
+  }
+}

backend/src/main/resources/db/migrations/V3__Create_Application.sql

@@ -21,7 +21,7 @@ CREATE TABLE applications
     id           SERIAL PRIMARY KEY,
     process_time VARCHAR,
     status       application_status NOT NULL,
-    user_id      INT                NOT NULL REFERENCES users (id),
+    user_id      INT                NOT NULL REFERENCES users (id) ON DELETE CASCADE,
     template_id  INT                NOT NULL REFERENCES templates (id)
 );