feat: 优化上下文传递

Author: zhou-hao

hsweb-authorization/hsweb-authorization-api/pom.xml

@@ -52,5 +52,11 @@
             <artifactId>jakarta.servlet-api</artifactId>
             <optional>true</optional>
         </dependency>
+
+        <dependency>
+            <groupId>io.micrometer</groupId>
+            <artifactId>context-propagation</artifactId>
+            <optional>true</optional>
+        </dependency>
     </dependencies>
 </project>

hsweb-authorization/hsweb-authorization-api/src/main/java/org/hswebframework/web/authorization/AuthenticationHolder.java

@@ -109,6 +109,18 @@ public static void addSupplier(AuthenticationSupplier supplier) {
         }
     }
 
+    public static void resetCurrent() {
+        CURRENT.remove();
+    }
+
+    public static void makeCurrent(Authentication authentication) {
+        if (authentication == null) {
+            resetCurrent();
+        } else {
+            CURRENT.set(authentication);
+        }
+    }
+
     /**
      * 指定用户权限，执行一个任务。任务执行过程中可通过 {@link Authentication#current()}获取到当前权限.
      *
@@ -119,7 +131,7 @@ public static void addSupplier(AuthenticationSupplier supplier) {
      */
     @SneakyThrows
     public static <T> T executeWith(Authentication current, Callable<T> callable) {
-        Authentication previous = CURRENT.get();
+        Authentication previous = CURRENT.getIfExists();
         try {
             CURRENT.set(current);
             return callable.call();

hsweb-authorization/hsweb-authorization-api/src/main/java/org/hswebframework/web/authorization/context/AuthenticationThreadLocalAccessor.java

@@ -0,0 +1,33 @@
+package org.hswebframework.web.authorization.context;
+
+import io.micrometer.context.ThreadLocalAccessor;
+import org.hswebframework.web.authorization.Authentication;
+import org.hswebframework.web.authorization.AuthenticationHolder;
+
+import javax.annotation.Nonnull;
+
+public class AuthenticationThreadLocalAccessor implements ThreadLocalAccessor<Authentication> {
+
+    static final String KEY = "cp.hs.auth";
+
+    @Override
+    @Nonnull
+    public Object key() {
+        return KEY;
+    }
+
+    @Override
+    public Authentication getValue() {
+        return AuthenticationHolder.get().orElse(null);
+    }
+
+    @Override
+    public void setValue() {
+        AuthenticationHolder.resetCurrent();
+    }
+
+    @Override
+    public void setValue(@Nonnull Authentication value) {
+        AuthenticationHolder.makeCurrent(value);
+    }
+}

hsweb-authorization/hsweb-authorization-api/src/main/resources/META-INF/services/io.micrometer.context.ThreadLocalAccessor

@@ -0,0 +1 @@
+org.hswebframework.web.authorization.context.AuthenticationThreadLocalAccessor

hsweb-authorization/hsweb-authorization-api/src/test/java/org/hswebframework/web/authorization/context/AuthenticationThreadLocalAccessorTest.java

@@ -0,0 +1,38 @@
+package org.hswebframework.web.authorization.context;
+
+import org.hswebframework.web.authorization.Authentication;
+import org.hswebframework.web.authorization.AuthenticationHolder;
+import org.hswebframework.web.authorization.simple.SimpleAuthentication;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Hooks;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class AuthenticationThreadLocalAccessorTest {
+
+
+    @Test
+    void testReadFromReactive() {
+
+        Hooks.enableAutomaticContextPropagation();
+
+        Authentication auth = new SimpleAuthentication();
+
+        Authentication auth2 = AuthenticationHolder.executeWith(
+            auth,
+            () -> Mono
+                .fromCallable(() -> {
+                    // cross context
+                    return Authentication.current().orElse(null);
+                })
+                .subscribeOn(Schedulers.boundedElastic())
+                .contextWrite(c->c)
+                .block());
+
+        assertEquals(auth, auth2);
+
+
+    }
+}

hsweb-authorization/hsweb-authorization-basic/src/main/java/org/hswebframework/web/authorization/basic/web/WebUserTokenInterceptor.java

@@ -39,22 +39,21 @@ public WebUserTokenInterceptor(UserTokenManager userTokenManager,
         this.parser = definitionParser;
 
         enableBasicAuthorization = userTokenParser
-                .stream()
-                .filter(UserTokenForTypeParser.class::isInstance)
-                .anyMatch(parser -> "basic".equalsIgnoreCase(((UserTokenForTypeParser) parser).getTokenType()));
+            .stream()
+            .filter(UserTokenForTypeParser.class::isInstance)
+            .anyMatch(parser -> "basic".equalsIgnoreCase(((UserTokenForTypeParser) parser).getTokenType()));
     }
 
     @Override
     public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
         List<ParsedToken> tokens = userTokenParser
-                .stream()
-                .map(parser -> parser.parseToken(request))
-                .filter(Objects::nonNull)
-                .collect(Collectors.toList());
+            .stream()
+            .map(parser -> parser.parseToken(request))
+            .filter(Objects::nonNull)
+            .toList();
 
         if (tokens.isEmpty()) {
-            if (enableBasicAuthorization && handler instanceof HandlerMethod) {
-                HandlerMethod method = ((HandlerMethod) handler);
+            if (enableBasicAuthorization && handler instanceof HandlerMethod method) {
                 AuthorizeDefinition definition = parser.parse(method.getBeanType(), method.getMethod());
                 if (null != definition) {
                     response.addHeader("WWW-Authenticate", " Basic realm=\"\"");
@@ -69,12 +68,18 @@ public boolean preHandle(HttpServletRequest request, HttpServletResponse respons
                 userToken = userTokenManager.getByToken(token).blockOptional().orElse(null);
             }
             if ((userToken == null || userToken.isExpired()) && parsedToken instanceof AuthorizedToken) {
-                //先踢出旧token
-                userTokenManager.signOutByToken(token).subscribe();
+                userToken =
+                    userTokenManager
+                        .signOutByToken(token)
+                        .then(
+                            userTokenManager
+                                .signIn(parsedToken.getToken(),
+                                        parsedToken.getType(),
+                                        ((AuthorizedToken) parsedToken).getUserId(),
+                                        ((AuthorizedToken) parsedToken)
+                                            .getMaxInactiveInterval())
+                        )
 
-                userToken = userTokenManager
-                        .signIn(parsedToken.getToken(), parsedToken.getType(), ((AuthorizedToken) parsedToken).getUserId(), ((AuthorizedToken) parsedToken)
-                                .getMaxInactiveInterval())
                         .block();
             }
             if (null != userToken) {
@@ -85,4 +90,8 @@ public boolean preHandle(HttpServletRequest request, HttpServletResponse respons
         return true;
     }
 
+    @Override
+    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
+        UserTokenHolder.setCurrent(null);
+    }
 }

hsweb-core/src/main/java/org/hswebframework/web/context/ContextHolder.java

@@ -48,8 +48,11 @@ public static <T> T doInContext(Context context, Callable<T> call) {
         }
     }
 
-    public static <T> Mono<T> currentReactive(Function<ContextView, Mono<T>> handler) {
-        return Mono.deferContextual(ctx -> handler.apply(current().putAll(ctx)));
+    public static <T> Mono<T> wrap(Function<ContextView, Mono<T>> handler) {
+        return Mono.deferContextual(ctx -> {
+            Context context = current().putAll(ctx);
+            return handler.apply(context);
+        });
     }
 
     public static Context current() {

hsweb-core/src/main/java/org/hswebframework/web/i18n/LocaleUtils.java

@@ -7,6 +7,7 @@
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscription;
 import org.springframework.context.MessageSource;
+import org.springframework.context.i18n.LocaleContextHolder;
 import reactor.core.CoreSubscriber;
 import reactor.core.publisher.*;
 import reactor.util.context.Context;
@@ -93,10 +94,8 @@ public static String getMessage(Function<String, String> messageSource,
      */
     public static Locale current() {
         Locale locale = CONTEXT_THREAD_LOCAL.get();
-        if (locale == null) {
-            locale = DEFAULT_LOCALE;
-        }
-        return locale;
+        // fallback to spring
+        return Objects.requireNonNullElseGet(locale, LocaleContextHolder::getLocale);
     }
 
     /**