Add auth layer handling for API Key

Author: aq-ikhwa-tech

server/api-service/lowcoder-sdk/src/main/java/org/lowcoder/sdk/util/CookieHelper.java

@@ -53,11 +53,6 @@ public String getCookieToken(ServerWebExchange exchange) {
         return getCookieValue(exchange, getCookieName(), "");
     }
 
-    @Nullable
-    public String getJWT(ServerWebExchange exchange) {
-        return getCookieValue(exchange, "JWT", null);
-    }
-
     public String getCookieValue(ServerWebExchange exchange, String cookieName, String defaultValue) {
         MultiValueMap<String, HttpCookie> cookies = exchange.getRequest().getCookies();
         return ofNullable(cookies.getFirst(cookieName))

server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/authentication/util/JWTUtils.java

@@ -5,16 +5,19 @@
 import io.jsonwebtoken.Jwts;
 import io.jsonwebtoken.SignatureAlgorithm;
 import jakarta.annotation.PostConstruct;
-import jakarta.servlet.http.HttpServletRequest;
+import lombok.extern.slf4j.Slf4j;
 import org.lowcoder.domain.user.model.User;
 import org.lowcoder.sdk.config.AuthProperties;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Component;
+import org.springframework.web.server.ServerWebExchange;
+
 import java.util.Random;
 
 import java.util.Date;
 
 @Component
+@Slf4j(topic = "JWTUtils")
 public class JWTUtils {
 
     @Autowired
@@ -43,13 +46,18 @@ public String createToken(User user) {
                 .compact();
     }
 
-    private Claims parseJwtClaims(String token) {
-        return jwtParser.parseClaimsJws(token).getBody();
+    public Claims parseJwtClaims(String token) {
+        try {
+            return jwtParser.parseClaimsJws(token).getBody();
+        } catch (Exception e) {
+            log.warn("Failed to validate token. Exception: ", e);
+            return null;
+        }
     }
 
-    public String resolveToken(HttpServletRequest request) {
+    public String resolveToken(ServerWebExchange exchange) {
 
-        String bearerToken = request.getHeader(TOKEN_HEADER);
+        String bearerToken = exchange.getRequest().getHeaders().getFirst(TOKEN_HEADER);
         if (bearerToken != null && bearerToken.startsWith(TOKEN_PREFIX)) {
             return bearerToken.substring(TOKEN_PREFIX.length());
         }

server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/framework/filter/APIKeyAuthFilter.java

@@ -0,0 +1,54 @@
+package org.lowcoder.api.framework.filter;
+
+import io.jsonwebtoken.Claims;
+import lombok.extern.slf4j.Slf4j;
+import org.lowcoder.api.authentication.util.JWTUtils;
+import org.lowcoder.api.home.SessionUserService;
+import org.lowcoder.sdk.util.CookieHelper;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.WebFilter;
+import org.springframework.web.server.WebFilterChain;
+import reactor.core.publisher.Mono;
+
+import javax.annotation.Nonnull;
+
+import static org.lowcoder.api.authentication.util.AuthenticationUtils.toAuthentication;
+import static org.springframework.security.core.context.ReactiveSecurityContextHolder.withAuthentication;
+
+@Slf4j
+public class APIKeyAuthFilter implements WebFilter {
+
+    private final SessionUserService service;
+
+    private final CookieHelper cookieHelper;
+    private final JWTUtils jwtUtils;
+
+    public APIKeyAuthFilter(SessionUserService service, CookieHelper cookieHelper, JWTUtils jwtUtils) {
+        this.service = service;
+        this.cookieHelper = cookieHelper;
+        this.jwtUtils = jwtUtils;
+    }
+
+    @Nonnull
+    @Override
+    public Mono<Void> filter(@Nonnull ServerWebExchange exchange, WebFilterChain chain) {
+        String cookieToken = cookieHelper.getCookieToken(exchange);
+        if(cookieToken.isEmpty()) {
+            String jwtToken = jwtUtils.resolveToken(exchange);
+            if(jwtToken == null || jwtToken.isEmpty()) {
+                return chain.filter(exchange);
+            } else {
+                Claims claims = jwtUtils.parseJwtClaims(jwtToken);
+                if(claims == null) {
+                    return chain.filter(exchange);
+                } else {
+                    return service.resolveSessionUserForJWT(claims, jwtToken)
+                            .flatMap(user -> chain.filter(exchange).contextWrite(withAuthentication(toAuthentication(user))));
+                }
+
+            }
+        } else {
+            return chain.filter(exchange);
+        }
+    }
+}

server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/framework/security/SecurityConfig.java

@@ -3,6 +3,8 @@
 
 import org.lowcoder.api.authentication.request.AuthRequestFactory;
 import org.lowcoder.api.authentication.service.AuthenticationApiServiceImpl;
+import org.lowcoder.api.authentication.util.JWTUtils;
+import org.lowcoder.api.framework.filter.APIKeyAuthFilter;
 import org.lowcoder.api.framework.filter.UserSessionPersistenceFilter;
 import org.lowcoder.api.home.SessionUserService;
 import org.lowcoder.domain.authentication.AuthenticationService;
@@ -66,6 +68,9 @@ public class SecurityConfig {
     @Autowired
     AuthRequestFactory<AuthRequestContext> authRequestFactory;
 
+    @Autowired
+    JWTUtils jwtUtils;
+
     @Bean
     SecurityWebFilterChain securityWebFilterChain(ServerHttpSecurity http) {
 
@@ -149,6 +154,7 @@ SecurityWebFilterChain securityWebFilterChain(ServerHttpSecurity http) {
         );
 
         http.addFilterBefore(new UserSessionPersistenceFilter(sessionUserService, cookieHelper, authenticationService, authenticationApiService, authRequestFactory), SecurityWebFiltersOrder.AUTHENTICATION);
+        http.addFilterBefore(new APIKeyAuthFilter(sessionUserService, cookieHelper, jwtUtils), SecurityWebFiltersOrder.AUTHENTICATION);
 
         return http.build();
     }

server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/home/SessionUserService.java

@@ -1,5 +1,6 @@
 package org.lowcoder.api.home;
 
+import io.jsonwebtoken.Claims;
 import org.lowcoder.domain.organization.model.OrgMember;
 import org.lowcoder.domain.user.model.User;
 import org.lowcoder.infra.annotation.NonEmptyMono;
@@ -29,5 +30,7 @@ public interface SessionUserService {
 
     Mono<User> resolveSessionUserFromCookie(String token);
 
+    Mono<User> resolveSessionUserForJWT(Claims claims, String token);
+
     Mono<Boolean> tokenExist(String token);
 }

server/api-service/lowcoder-server/src/main/java/org/lowcoder/api/home/SessionUserServiceImpl.java

@@ -8,6 +8,7 @@
 import java.time.Duration;
 import java.util.Objects;
 
+import io.jsonwebtoken.Claims;
 import org.apache.commons.lang3.StringUtils;
 import org.lowcoder.api.usermanagement.UserApiService;
 import org.lowcoder.domain.organization.model.OrgMember;
@@ -139,6 +140,17 @@ public Mono<User> resolveSessionUserFromCookie(String token) {
                 .filter(user -> user.getState() != UserState.DELETED);
     }
 
+    @Override
+    public Mono<User> resolveSessionUserForJWT(Claims claims, String token) {
+        String userId = claims.get("userId").toString();
+        return userService.findById(userId)
+                .filter(user -> user.getState() != UserState.DELETED)
+                .filter(user -> {
+                    long apiKeyFound = user.getApiKeysList().stream().filter(apiKey -> apiKey.getToken().equals(token)).count();
+                    return apiKeyFound > 0;
+                });
+    }
+
     @Override
     public Mono<Boolean> tokenExist(String token) {
         return reactiveTemplate.hasKey(token);