Security / Wrapped - Propagation

api/src/main/java/io/smallrye/context/spi/WrappingThreadContextSnapshot.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2020 Contributors to the Eclipse Foundation
+ *
+ * See the NOTICE file(s) distributed with this work for additional
+ * information regarding copyright ownership.
+ *
+ * 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 io.smallrye.context.spi;
+
+import java.util.concurrent.Callable;
+
+import org.eclipse.microprofile.context.spi.ThreadContextSnapshot;
+
+/**
+ * An extension of {code ThreadContextSnapshot} which enables the snapshot to
+ * perform propagation by wrapping the task.
+ *
+ * @author <a href="mailto:[email protected]">Darran Lofthouse</a>
+ */
+@FunctionalInterface
+public interface WrappingThreadContextSnapshot extends ThreadContextSnapshot {
+
+    /**
+     * Does this snapshot need to wrap the underlying task instead of directly
+     * manipulating ThreadLocals.
+     *
+     * @return {@code true} if this snapshot needs to wrap the underlying task.
+     */
+    default boolean needsToWrap() {
+        return false;
+    }
+
+    /**
+     * Wrap the provided task to ensure the context being propagated is correctly
+     * established and cleared as the underlying task is called.
+     *
+     * @param task the task to wrap.
+     * @return the wrapped task.
+     */
+    default <T> Callable<T> wrap(Callable<T> callable) {
+        return callable;
+    }
+
+}

core/src/main/java/io/smallrye/context/CleanAutoCloseable.java

@@ -3,8 +3,10 @@
 /**
  * AutoCloseable interface which doesn't throw.
  */
-@FunctionalInterface
-public interface CleanAutoCloseable extends AutoCloseable {
+public interface CleanAutoCloseable<T> extends AutoCloseable {
+
+    T callNoChecked();
+
     /**
      * Close this resource, no exception thrown.
      */

core/src/main/java/io/smallrye/context/SmallRyeThreadContext.java

@@ -68,6 +68,12 @@ public class SmallRyeThreadContext implements ThreadContext {
             .threadLocal(SmallRyeThreadContextStorageDeclaration.class);
 
     private final static CleanAutoCloseable NULL_THREAD_STATE = new CleanAutoCloseable() {
+
+        @Override
+        public Object callNoChecked() {
+            throw new IllegalStateException();
+        }
+
         @Override
         public void close() {
             currentThreadContext.remove();
@@ -84,7 +90,7 @@ public void execute(Runnable command) {
     /**
      * Updates the current @{link SmallRyeThreadContext} in use by the current thread, and returns an
      * object suitable for use in try-with-resource to restore the previous value.
-     * 
+     *
      * @param threadContext the @{link SmallRyeThreadContext} to use
      * @return an object suitable for use in try-with-resource to restore the previous value.
      */
@@ -96,6 +102,12 @@ public static CleanAutoCloseable withThreadContext(SmallRyeThreadContext threadC
             return NULL_THREAD_STATE;
         } else {
             return new CleanAutoCloseable() {
+
+                @Override
+                public Object callNoChecked() {
+                    throw new IllegalStateException();
+                }
+
                 @Override
                 public void close() {
                     currentThreadContext.set(oldValue);
@@ -108,7 +120,7 @@ public void close() {
     /**
      * Invokes the given @{link Runnable} with the current @{link SmallRyeThreadContext} updated to the given value
      * for the current thread.
-     * 
+     *
      * @param threadContext the @{link SmallRyeThreadContext} to use
      * @param f the @{link Runnable} to invoke
      */
@@ -129,7 +141,7 @@ public static void withThreadContext(SmallRyeThreadContext threadContext, Runnab
     /**
      * Returns the current thread's @{link SmallRyeThreadContext} if set, or a @{link SmallRyeThreadContext}
      * which propagates all contexts.
-     * 
+     *
      * @return the current thread's @{link SmallRyeThreadContext} if set, or a @{link SmallRyeThreadContext}
      *         which propagates all contexts.
      */
@@ -140,7 +152,7 @@ public static SmallRyeThreadContext getCurrentThreadContextOrPropagatedContexts(
     /**
      * Returns the current thread's @{link SmallRyeThreadContext} if set, or a @{link SmallRyeThreadContext}
      * which clears all contexts.
-     * 
+     *
      * @return the current thread's @{link SmallRyeThreadContext} if set, or a @{link SmallRyeThreadContext}
      *         which clears all contexts.
      */
@@ -151,7 +163,7 @@ public static SmallRyeThreadContext getCurrentThreadContextOrClearedContexts() {
     /**
      * Returns the current thread's @{link SmallRyeThreadContext} if set, or the given @{link SmallRyeThreadContext}
      * default value.
-     * 
+     *
      * @param defaultValue the default value to use
      * @return the current thread's @{link SmallRyeThreadContext} if set, or the given @{link SmallRyeThreadContext}
      *         default value.
@@ -163,7 +175,7 @@ public static SmallRyeThreadContext getCurrentThreadContext(SmallRyeThreadContex
 
     /**
      * Returns the current thread's @{link SmallRyeThreadContext} if set, or null.
-     * 
+     *
      * @return the current thread's @{link SmallRyeThreadContext} if set, or null.
      */
     public static SmallRyeThreadContext getCurrentThreadContext() {
@@ -200,7 +212,7 @@ public ExecutorService getDefaultExecutor() {
     /**
      * Returns true if this thread context has no context to propagate nor clear, and so
      * will not contextualise anything.
-     * 
+     *
      * @return true if this thread context has no context to propagate nor clear
      */
     public boolean isEmpty() {
@@ -209,7 +221,7 @@ public boolean isEmpty() {
 
     /**
      * Returns true if the given lambda instance is already contextualized
-     * 
+     *
      * @param lambda the lambda to test
      * @return true if the given lambda instance is already contextualized
      */
@@ -236,7 +248,7 @@ public static Builder builder() {
      * or the default executor service as set by
      * {@link SmallRyeContextManager.Builder#withDefaultExecutorService(ExecutorService)},
      * or otherwise have no default executor.
-     * 
+     *
      * If this thread context has no default executor, the new stage and all dependent stages created from it, and so forth,
      * have no default asynchronous execution facility and must raise {@link java.lang.UnsupportedOperationException}
      * for all <code>*Async</code> methods that do not specify an executor. For example,
@@ -284,7 +296,7 @@ public <T> CompletableFuture<T> withContextCapture(CompletableFuture<T> future,
      * or the default executor service as set by
      * {@link SmallRyeContextManager.Builder#withDefaultExecutorService(ExecutorService)},
      * or otherwise have no default executor.
-     * 
+     *
      * If this thread context has no default executor, the new stage and all dependent stages created from it, and so forth,
      * and/or cleared as described in the documentation of the {@link ManagedExecutor} class, except that
      * this ThreadContext instance takes the place of the default asynchronous execution facility in

core/src/main/java/io/smallrye/context/impl/CapturedContextState.java

@@ -1,8 +1,10 @@
 package io.smallrye.context.impl;
 
+import java.util.concurrent.Callable;
+
 import io.smallrye.context.CleanAutoCloseable;
 
 @FunctionalInterface
 public interface CapturedContextState {
-    CleanAutoCloseable begin();
+    <T> CleanAutoCloseable<T> begin(Callable<T> callable);
 }

core/src/main/java/io/smallrye/context/impl/SlowActiveContextState.java

@@ -1,36 +1,55 @@
 package io.smallrye.context.impl;
 
 import java.util.List;
+import java.util.concurrent.Callable;
 
 import org.eclipse.microprofile.context.spi.ThreadContextController;
 import org.eclipse.microprofile.context.spi.ThreadContextSnapshot;
 
 import io.smallrye.context.CleanAutoCloseable;
 import io.smallrye.context.SmallRyeThreadContext;
+import io.smallrye.context.spi.WrappingThreadContextSnapshot;
 
 /**
  * Restores a context and allows you to clean it up (unrestore it).
  */
-public class SlowActiveContextState implements CleanAutoCloseable {
+public class SlowActiveContextState<T> implements CleanAutoCloseable<T> {
 
     private final ThreadContextController[] activeContext;
     private final CleanAutoCloseable activeThreadContext;
+    private final Callable<T> callable;
 
     /**
      * Restores a previously captured context.
-     * 
+     *
      * @param threadContext the thread context
      * @param threadContextSnapshots the captured snapshots
      */
-    public SlowActiveContextState(SmallRyeThreadContext threadContext, List<ThreadContextSnapshot> threadContextSnapshots) {
+    public SlowActiveContextState(SmallRyeThreadContext threadContext, List<ThreadContextSnapshot> threadContextSnapshots,
+            Callable<T> callable) {
         activeContext = new ThreadContextController[threadContextSnapshots.size()];
         int i = 0;
         for (ThreadContextSnapshot threadContextSnapshot : threadContextSnapshots) {
             activeContext[i++] = threadContextSnapshot.begin();
+            if (threadContextSnapshot instanceof WrappingThreadContextSnapshot
+                    && ((WrappingThreadContextSnapshot) threadContextSnapshot).needsToWrap()) {
+                callable = ((WrappingThreadContextSnapshot) threadContextSnapshot).wrap(callable);
+            }
         }
+        this.callable = callable;
         activeThreadContext = SmallRyeThreadContext.withThreadContext(threadContext);
     }
 
+    @Override
+    public T callNoChecked() {
+        try {
+            return callable.call();
+        } catch (Exception e) {
+            Util.rethrow(e);
+            return null;
+        }
+    }
+
     /**
      * Unrestores / clean-up a previously restored context.
      */

core/src/main/java/io/smallrye/context/impl/SlowCapturedContextState.java

@@ -1,6 +1,7 @@
 package io.smallrye.context.impl;
 
 import java.util.List;
+import java.util.concurrent.Callable;
 
 import org.eclipse.microprofile.context.spi.ThreadContextSnapshot;
 
@@ -17,7 +18,7 @@ public class SlowCapturedContextState implements CapturedContextState {
 
     /**
      * Captures the current context according to the given ThreadContext
-     * 
+     *
      * @param threadContext the thread context
      */
     public SlowCapturedContextState(SmallRyeThreadContext threadContext) {
@@ -27,10 +28,10 @@ public SlowCapturedContextState(SmallRyeThreadContext threadContext) {
 
     /**
      * Restores the captured context and returns an instance that can unrestore (cleanup) it.
-     * 
+     *
      * @return the captured context state
      */
-    public SlowActiveContextState begin() {
-        return new SlowActiveContextState(threadContext, threadContextSnapshots);
+    public <T> SlowActiveContextState<T> begin(Callable<T> callable) {
+        return new SlowActiveContextState<T>(threadContext, threadContextSnapshots, callable);
     }
 }

core/src/main/java/io/smallrye/context/impl/Util.java

@@ -0,0 +1,10 @@
+package io.smallrye.context.impl;
+
+/**
+ * @author <a href="mailto:[email protected]">Kabir Khan</a>
+ */
+public class Util {
+    public static <T extends Throwable> void rethrow(Throwable t) throws T {
+        throw (T) t;
+    }
+}

core/src/main/java/io/smallrye/context/impl/wrappers/SlowContextualBiConsumer.java

@@ -17,8 +17,11 @@ public SlowContextualBiConsumer(CapturedContextState state, BiConsumer<T, U> con
 
     @Override
     public void accept(T t, U u) {
-        try (CleanAutoCloseable activeState = state.begin()) {
+        try (CleanAutoCloseable<Void> activeState = state.begin(() -> {
             consumer.accept(t, u);
+            return null;
+        })) {
+            activeState.callNoChecked();
         }
     }
 }

core/src/main/java/io/smallrye/context/impl/wrappers/SlowContextualBiFunction.java

@@ -17,8 +17,8 @@ public SlowContextualBiFunction(CapturedContextState state, BiFunction<T, U, R>
 
     @Override
     public R apply(T t, U u) {
-        try (CleanAutoCloseable activeState = state.begin()) {
-            return function.apply(t, u);
+        try (CleanAutoCloseable<R> activeState = state.begin(() -> function.apply(t, u))) {
+            return activeState.callNoChecked();
         }
     }
 }

core/src/main/java/io/smallrye/context/impl/wrappers/SlowContextualCallable.java

@@ -17,8 +17,8 @@ public SlowContextualCallable(CapturedContextState state, Callable<R> callable)
 
     @Override
     public R call() throws Exception {
-        try (CleanAutoCloseable activeState = state.begin()) {
-            return callable.call();
+        try (CleanAutoCloseable<R> activeState = state.begin(callable)) {
+            return activeState.callNoChecked();
         }
     }
 }

core/src/main/java/io/smallrye/context/impl/wrappers/SlowContextualConsumer.java

@@ -17,8 +17,11 @@ public SlowContextualConsumer(CapturedContextState state, Consumer<T> consumer)
 
     @Override
     public void accept(T t) {
-        try (CleanAutoCloseable activeState = state.begin()) {
+        try (CleanAutoCloseable<Void> activeState = state.begin(() -> {
             consumer.accept(t);
+            return null;
+        })) {
+            activeState.callNoChecked();
         }
     }
 }

core/src/main/java/io/smallrye/context/impl/wrappers/SlowContextualExecutor.java

@@ -15,8 +15,11 @@ public SlowContextualExecutor(CapturedContextState state) {
 
     @Override
     public void execute(Runnable command) {
-        try (CleanAutoCloseable foo = state.begin()) {
+        try (CleanAutoCloseable<Void> foo = state.begin(() -> {
             command.run();
+            return null;
+        })) {
+            foo.callNoChecked();
         }
     }
 

core/src/main/java/io/smallrye/context/impl/wrappers/SlowContextualFunction.java

@@ -17,8 +17,8 @@ public SlowContextualFunction(CapturedContextState state, Function<T, R> functio
 
     @Override
     public R apply(T t) {
-        try (CleanAutoCloseable activeState = state.begin()) {
-            return function.apply(t);
+        try (CleanAutoCloseable<R> activeState = state.begin(() -> function.apply(t))) {
+            return activeState.callNoChecked();
         }
     }
 }

core/src/main/java/io/smallrye/context/impl/wrappers/SlowContextualRunnable.java

@@ -15,8 +15,11 @@ public SlowContextualRunnable(CapturedContextState state, Runnable runnable) {
 
     @Override
     public void run() {
-        try (CleanAutoCloseable activeState = state.begin()) {
+        try (CleanAutoCloseable<Void> activeState = state.begin(() -> {
             runnable.run();
+            return null;
+        })) {
+            activeState.callNoChecked();
         }
     }
 }

core/src/main/java/io/smallrye/context/impl/wrappers/SlowContextualSupplier.java

@@ -17,8 +17,8 @@ public SlowContextualSupplier(CapturedContextState state, Supplier<R> supplier)
 
     @Override
     public R get() {
-        try (CleanAutoCloseable activeState = state.begin()) {
-            return supplier.get();
+        try (CleanAutoCloseable<R> activeState = state.begin(supplier::get)) {
+            return activeState.callNoChecked();
         }
     }
 }