vailidate tty action

Author: gy2006

core/src/main/java/com/flowci/core/auth/WebAuth.java

@@ -20,18 +20,24 @@
 import com.flowci.core.auth.annotation.Action;
 import com.flowci.core.auth.service.AuthService;
 import com.flowci.core.common.manager.SessionManager;
+import com.flowci.core.user.domain.User;
 import com.flowci.exception.AccessException;
 import com.flowci.exception.AuthenticationException;
+import com.flowci.util.StringHelper;
 import com.google.common.base.Strings;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
 import org.springframework.stereotype.Component;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.method.HandlerMethod;
 import org.springframework.web.servlet.HandlerInterceptor;
 import org.springframework.web.servlet.ModelAndView;
 
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import java.util.Objects;
+import java.util.Optional;
 
 /**
  * @author yang
@@ -51,6 +57,31 @@ public class WebAuth implements HandlerInterceptor {
     @Autowired
     private SessionManager sessionManager;
 
+    /**
+     * Get user object from ws message header
+     * @param headers
+     * @return User object
+     * @exception AuthenticationException if Token header is missing or invalid token
+     */
+    public User validate(MessageHeaders headers) {
+        MultiValueMap<String, String> map = headers.get(StompHeaderAccessor.NATIVE_HEADERS, MultiValueMap.class);
+        if (Objects.isNull(map)) {
+            throw new AuthenticationException("Invalid token");
+        }
+
+        String token = map.getFirst(HeaderToken);
+        if (!StringHelper.hasValue(token)) {
+            throw new AuthenticationException("Invalid token");
+        }
+
+        Optional<User> user = authService.get(token);
+        if (!user.isPresent()) {
+            throw new AuthenticationException("Invalid token");
+        }
+
+        return user.get();
+    }
+
     @Override
     public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
         if (!authService.isEnabled()) {

core/src/main/java/com/flowci/core/auth/service/AuthService.java

@@ -19,6 +19,9 @@
 
 import com.flowci.core.auth.annotation.Action;
 import com.flowci.core.auth.domain.Tokens;
+import com.flowci.core.user.domain.User;
+
+import java.util.Optional;
 
 /**
  * 'login' ->
@@ -52,11 +55,16 @@ public interface AuthService {
     boolean hasPermission(Action action);
 
     /**
-     * Set current user by token
+     * Set current logged in user by token
      * @return Current user object or null if not verified
      */
     boolean set(String token);
 
+    /**
+     * Get current logged in user
+     */
+    Optional<User> get(String token);
+
     /**
      * Set current user from default admin form config properties
      */

core/src/main/java/com/flowci/core/auth/service/AuthServiceImpl.java

@@ -29,6 +29,8 @@
 import com.flowci.exception.AuthenticationException;
 import com.flowci.util.HashingHelper;
 import java.util.Objects;
+import java.util.Optional;
+
 import lombok.extern.log4j.Log4j2;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.cache.Cache;
@@ -128,20 +130,29 @@ public Tokens refresh(Tokens tokens) {
 
     @Override
     public boolean set(String token) {
+        Optional<User> optional = get(token);
+        if (optional.isPresent()) {
+            sessionManager.set(optional.get());
+            return true;
+        }
+        return false;
+    }
+
+    @Override
+    public Optional<User> get(String token) {
         String email = JwtHelper.decode(token);
 
         User user = onlineUsersCache.get(email, User.class);
         if (Objects.isNull(user)) {
-            return false;
+            return Optional.empty();
         }
 
         boolean verify = JwtHelper.verify(token, user, true);
         if (verify) {
-            sessionManager.set(user);
-            return true;
+            return Optional.of(user);
         }
 
-        return false;
+        return Optional.empty();
     }
 
     @Override

core/src/main/java/com/flowci/core/common/config/WebSocketConfig.java

@@ -16,13 +16,22 @@
 
 package com.flowci.core.common.config;
 
+import com.flowci.core.common.helper.ThreadHelper;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.converter.MessageConverter;
+import org.springframework.messaging.simp.config.ChannelRegistration;
 import org.springframework.messaging.simp.config.MessageBrokerRegistry;
+import org.springframework.messaging.support.ChannelInterceptor;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
 import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
 import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
 
+import java.util.List;
+
 /**
  * @author yang
  */
@@ -118,4 +127,16 @@ public void configureMessageBroker(MessageBrokerRegistry registry) {
         registry.enableSimpleBroker("/topic");
         registry.setApplicationDestinationPrefixes("/app");
     }
+
+    @Override
+    public void configureClientInboundChannel(ChannelRegistration registration) {
+        ThreadPoolTaskExecutor executor = ThreadHelper.createTaskExecutor(50, 50, 50, "ws-inbound-");
+        registration.taskExecutor(executor);
+    }
+
+    @Override
+    public void configureClientOutboundChannel(ChannelRegistration registration) {
+        ThreadPoolTaskExecutor executor = ThreadHelper.createTaskExecutor(50, 50, 50, "ws-outbound-");
+        registration.taskExecutor(executor);
+    }
 }

core/src/main/java/com/flowci/core/job/controller/TtyController.java

@@ -1,32 +1,57 @@
 package com.flowci.core.job.controller;
 
+import com.flowci.core.auth.WebAuth;
 import com.flowci.core.job.service.TtyService;
+import com.flowci.core.user.domain.User;
+import com.flowci.exception.AuthenticationException;
 import lombok.extern.log4j.Log4j2;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.messaging.MessageHeaders;
 import org.springframework.messaging.handler.annotation.DestinationVariable;
+import org.springframework.messaging.handler.annotation.MessageExceptionHandler;
 import org.springframework.messaging.handler.annotation.MessageMapping;
 import org.springframework.messaging.handler.annotation.Payload;
+import org.springframework.messaging.simp.annotation.SendToUser;
 import org.springframework.stereotype.Controller;
 
 @Log4j2
 @Controller
 public class TtyController {
 
+    @Autowired
+    private WebAuth webAuth;
+
     @Autowired
     private TtyService ttyService;
 
+    @MessageExceptionHandler(AuthenticationException.class)
+    @SendToUser("/topic/tty/errors")
+    public Exception handleExceptions(AuthenticationException e) {
+        return e;
+    }
+
     @MessageMapping("/tty/{jobId}/open")
-    public void open(@DestinationVariable String jobId) {
+    public void open(@DestinationVariable String jobId, MessageHeaders headers) {
+        validate(headers);
         ttyService.open(jobId);
     }
 
     @MessageMapping("/tty/{jobId}/shell")
-    public void shell(@DestinationVariable String jobId, @Payload String script) {
+    public void shell(@DestinationVariable String jobId, @Payload String script, MessageHeaders headers) {
+        validate(headers);
         ttyService.shell(jobId, script);
     }
 
     @MessageMapping("/tty/{jobId}/close")
-    public void close(@DestinationVariable String jobId) {
+    public void close(@DestinationVariable String jobId, MessageHeaders headers) {
+        validate(headers);
         ttyService.close(jobId);
     }
+
+    private void validate(MessageHeaders headers) {
+        User user = webAuth.validate(headers);
+        if (!user.isAdmin()) {
+            throw new AuthenticationException("Admin permission is required");
+        }
+    }
 }