Vault.java

package org.openjdk.meta;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.lang.invoke.MethodHandles.Lookup;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;

public class Vault {
  private Vault() {  /* empty */ }
  private static final Vault INSTANCE = new Vault();
  public static Vault getInstance() { return INSTANCE; }
  
  private static final int MODES = Lookup.PRIVATE | Lookup.PACKAGE | Lookup.PROTECTED | Lookup.PUBLIC;
  private static void checkMode(Lookup lookup) {
    if ((lookup.lookupModes() & MODES) != MODES) {
      throw new IllegalArgumentException("invalid lookup");
    }
  }
  
  public interface AccessControl {
    public boolean grantAccess(Class<?> type);
  }
  
  private final ClassValue<AtomicReference<Function<Class<?>, Optional<Lookup>>>> cache = new ClassValue<AtomicReference<Function<Class<?>, Optional<Lookup>>>>() {
    @Override
    protected AtomicReference<Function<Class<?>, Optional<Lookup>>> computeValue(Class<?> type) {
      return new AtomicReference<>(__ -> Optional.empty());
    }
  };
  
  public void grantAccess(Lookup lookup, AccessControl accessControl) {
    Objects.requireNonNull(lookup);
    Objects.requireNonNull(accessControl);
    checkMode(lookup);
    cache.get(lookup.lookupClass())
         .updateAndGet(fun -> clientClass -> fun.apply(clientClass)
                                   .or(() -> Optional.of(lookup).filter(__ -> accessControl.grantAccess(clientClass))));
  }
  
  public Optional<Lookup> getAccess(Class<?> declaringClass, Lookup clientLookup) {
    Objects.requireNonNull(declaringClass);
    checkMode(clientLookup);
    try {
      Class.forName(declaringClass.getName(), true, declaringClass.getClassLoader());  // force init, maybe use unsafe here ?
    } catch (ClassNotFoundException e) {
      throw (NoClassDefFoundError)new NoClassDefFoundError().initCause(e);
    } 
    return cache.get(declaringClass).get().apply(clientLookup.lookupClass());
  }
  
  //////////////////////////////////////////////////////////////////////////////////////////////
  
  
  static class A {
    private void foo() { System.out.println("hello"); }  
    
    static {
      // grant access to B
      Vault.getInstance().grantAccess(MethodHandles.lookup(), type -> type == B.class);
    }
  }
  static class B {
    public static void bar() {
      // access A members from B 
      Lookup lookup = Vault.getInstance().getAccess(A.class, MethodHandles.lookup()).get();
      MethodHandle mh;
      try {
        mh = lookup.findVirtual(A.class, "foo", MethodType.methodType(void.class));
      } catch (NoSuchMethodException | IllegalAccessException e) {
        throw new AssertionError(e);
      }
      try {
        mh.invokeExact(new A());
      } catch (Throwable e) {
        throw new AssertionError(e);
      }
    }
  }
  
  public static void main(String[] args) {
    B.bar();
  }
}