AccessControllerWrapper.java

package security;

import static java.lang.invoke.MethodType.methodType;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.reflect.UndeclaredThrowableException;
import java.security.AccessController;
import java.security.PrivilegedAction;
import java.util.Objects;
import java.util.function.Supplier;

/**
 * {@link AccessController} compatibility wrapper.
 */
public class AccessControllerWrapper {

    /**
     * This method acts equivalently to {@link AccessController#doPrivileged(PrivilegedAction)}
     * but takes additional {@link MethodHandles.Lookup} parameter acting as a {@code caller}
     * for looking up the caller-sensitive {@link AccessController#doPrivileged(PrivilegedAction)}
     * method before dispatching to it.
     * In case there is no {@link AccessController} or {@link PrivilegedAction} classes available,
     * this method simply invokes the provided {@link Supplier}.
     */
    public static <T> T doPrivileged(MethodHandles.Lookup caller, Supplier<T> action) {
        Objects.requireNonNull(action);
        if (!caller.hasFullPrivilegeAccess()) {
            throw new IllegalArgumentException("The caller Lookup has no full privilege access");
        }
        CALLER_LOOKUP.set(caller);
        MethodHandle doPrivileged;
        try {
            doPrivileged = DO_PRIVILEGED_CV.get(caller.lookupClass());
        } finally {
            CALLER_LOOKUP.remove();
        }

        try {
            @SuppressWarnings("unchecked")
            T result = (T) (Object) doPrivileged.invokeExact(action);
            return result;
        } catch (RuntimeException | Error e) {
            throw e;
        } catch (Throwable e) {
            throw new UndeclaredThrowableException(e);
        }
    }

    private static final ThreadLocal<MethodHandles.Lookup> CALLER_LOOKUP = new ThreadLocal<>();

    private static final ClassValue<MethodHandle> DO_PRIVILEGED_CV = new ClassValue<>() {
        @Override
        protected MethodHandle computeValue(Class<?> caller) {
            MethodHandles.Lookup lookup = CALLER_LOOKUP.get();
            MethodHandle doPrivileged;
            try {
                try {
                    MethodHandle acDoPrivileged = lookup
                        .findStatic(AccessController.class, "doPrivileged", methodType(Object.class, PrivilegedAction.class));
                    MethodHandle paWrapperConstructor = MethodHandles
                        .lookup()
                        .findConstructor(PrivilegedActionWrapper.class, methodType(void.class, Supplier.class))
                        .asType(methodType(PrivilegedAction.class, Supplier.class));
                    doPrivileged = MethodHandles
                        .filterArguments(acDoPrivileged, 0, paWrapperConstructor);
                } catch (NoClassDefFoundError e) { // when JVM attempts to resolve AccessController or PrivilegedAction class and there is none
                    doPrivileged = MethodHandles
                        .lookup()
                        .findStatic(AccessControllerWrapper.class, "doPrivilegedNoOp", methodType(Object.class, Supplier.class));
                }
            } catch (IllegalAccessException e) {
                throw (Error) new IllegalAccessError(e.getMessage()).initCause(e);
            } catch (NoSuchMethodException e) {
                throw (Error) new NoSuchMethodError(e.getMessage()).initCause(e);
            }
            return doPrivileged;
        }
    };

    private static final class PrivilegedActionWrapper<T> implements PrivilegedAction<T> {
        private final Supplier<T> action;

        private PrivilegedActionWrapper(Supplier<T> action) {
            this.action = action;
        }

        @Override
        public T run() {
            return action.get();
        }
    }

    private static <T> T doPrivilegedNoOp(Supplier<T> action) {
        return action.get();
    }
}