1717import  java .lang .StackWalker .StackFrame ;
1818import  java .lang .module .ModuleFinder ;
1919import  java .lang .module .ModuleReference ;
20- import  java .util .ArrayList ;
21- import  java .util .HashMap ;
22- import  java .util .IdentityHashMap ;
2320import  java .util .List ;
2421import  java .util .Map ;
25- import  java .util .Objects ;
2622import  java .util .Optional ;
2723import  java .util .Set ;
24+ import  java .util .concurrent .ConcurrentHashMap ;
2825import  java .util .function .Function ;
2926import  java .util .stream .Collectors ;
3027import  java .util .stream .Stream ;
3128
3229import  static  java .lang .StackWalker .Option .RETAIN_CLASS_REFERENCE ;
3330import  static  java .util .Objects .requireNonNull ;
31+ import  static  java .util .stream .Collectors .groupingBy ;
3432
3533public  class  PolicyManager  {
3634    private  static  final  Logger  logger  = LogManager .getLogger (PolicyManager .class );
3735
38-     static  class  ModuleEntitlements  {
39-         public  static  final  ModuleEntitlements  NONE  = new  ModuleEntitlements (List .of ());
40-         private  final  IdentityHashMap <Class <? extends  Entitlement >, List <Entitlement >> entitlementsByType ;
36+     record  ModuleEntitlements (Map <Class <? extends  Entitlement >, List <Entitlement >> entitlementsByType ) {
37+         public  static  final  ModuleEntitlements  NONE  = new  ModuleEntitlements (Map .of ());
4138
42-         ModuleEntitlements ( List < Entitlement >  entitlements )  {
43-             this . entitlementsByType  = entitlements . stream () 
44-                 . collect ( Collectors . toMap ( Entitlement :: getClass ,  e  ->  new   ArrayList <>( List . of ( e )), ( a ,  b ) -> { 
45-                      a . addAll ( b ); 
46-                      return   a ; 
47-                 },  IdentityHashMap :: new ));
39+         ModuleEntitlements  {
40+             entitlementsByType  = Map . copyOf ( entitlementsByType ); 
41+         } 
42+ 
43+         public   static   ModuleEntitlements   from ( List < Entitlement >  entitlements ) { 
44+             return   new   ModuleEntitlements ( entitlements . stream (). collect ( groupingBy ( Entitlement :: getClass ) ));
4845        }
4946
5047        public  boolean  hasEntitlement (Class <? extends  Entitlement > entitlementClass ) {
@@ -56,9 +53,10 @@ public <E extends Entitlement> Stream<E> getEntitlements(Class<E> entitlementCla
5653        }
5754    }
5855
59-     final  Map <Module , ModuleEntitlements > moduleEntitlementsMap  = new  HashMap <>();
56+     final  Map <Module , ModuleEntitlements > moduleEntitlementsMap  = new  ConcurrentHashMap <>();
6057
6158    protected  final  Map <String , List <Entitlement >> serverEntitlements ;
59+     protected  final  List <Entitlement > agentEntitlements ;
6260    protected  final  Map <String , Map <String , List <Entitlement >>> pluginsEntitlements ;
6361    private  final  Function <Class <?>, String > pluginResolver ;
6462
@@ -85,12 +83,14 @@ private static Set<Module> findSystemModules() {
8583    private  final  Module  entitlementsModule ;
8684
8785    public  PolicyManager (
88-         Policy  defaultPolicy ,
86+         Policy  serverPolicy ,
87+         List <Entitlement > agentEntitlements ,
8988        Map <String , Policy > pluginPolicies ,
9089        Function <Class <?>, String > pluginResolver ,
9190        Module  entitlementsModule 
9291    ) {
93-         this .serverEntitlements  = buildScopeEntitlementsMap (requireNonNull (defaultPolicy ));
92+         this .serverEntitlements  = buildScopeEntitlementsMap (requireNonNull (serverPolicy ));
93+         this .agentEntitlements  = agentEntitlements ;
9494        this .pluginsEntitlements  = requireNonNull (pluginPolicies ).entrySet ()
9595            .stream ()
9696            .collect (Collectors .toUnmodifiableMap (Map .Entry ::getKey , e  -> buildScopeEntitlementsMap (e .getValue ())));
@@ -99,15 +99,15 @@ public PolicyManager(
9999    }
100100
101101    private  static  Map <String , List <Entitlement >> buildScopeEntitlementsMap (Policy  policy ) {
102-         return  policy .scopes .stream ().collect (Collectors .toUnmodifiableMap (scope  -> scope .name , scope  -> scope .entitlements ));
102+         return  policy .scopes .stream ().collect (Collectors .toUnmodifiableMap (scope  -> scope .moduleName , scope  -> scope .entitlements ));
103103    }
104104
105105    public  void  checkStartProcess (Class <?> callerClass ) {
106106        neverEntitled (callerClass , "start process" );
107107    }
108108
109109    private  void  neverEntitled (Class <?> callerClass , String  operationDescription ) {
110-         var  requestingModule  = requestingModule (callerClass );
110+         var  requestingModule  = requestingClass (callerClass );
111111        if  (isTriviallyAllowed (requestingModule )) {
112112            return ;
113113        }
@@ -139,49 +139,45 @@ public void checkSetGlobalHttpsConnectionProperties(Class<?> callerClass) {
139139    }
140140
141141    private  void  checkEntitlementPresent (Class <?> callerClass , Class <? extends  Entitlement > entitlementClass ) {
142-         var  requestingModule  = requestingModule (callerClass );
143-         if  (isTriviallyAllowed (requestingModule )) {
142+         var  requestingClass  = requestingClass (callerClass );
143+         if  (isTriviallyAllowed (requestingClass )) {
144144            return ;
145145        }
146146
147-         ModuleEntitlements  entitlements  = getEntitlementsOrThrow ( callerClass ,  requestingModule );
147+         ModuleEntitlements  entitlements  = getEntitlements ( requestingClass );
148148        if  (entitlements .hasEntitlement (entitlementClass )) {
149149            logger .debug (
150150                () -> Strings .format (
151-                     "Entitled: caller  [%s], module [%s], type  [%s]" ,
152-                     callerClass ,
153-                     requestingModule .getName (),
151+                     "Entitled: class  [%s], module [%s], entitlement  [%s]" ,
152+                     requestingClass ,
153+                     requestingClass . getModule () .getName (),
154154                    entitlementClass .getSimpleName ()
155155                )
156156            );
157157            return ;
158158        }
159159        throw  new  NotEntitledException (
160160            Strings .format (
161-                 "Missing entitlement: caller  [%s], module [%s], type  [%s]" ,
162-                 callerClass ,
163-                 requestingModule .getName (),
161+                 "Missing entitlement: class  [%s], module [%s], entitlement  [%s]" ,
162+                 requestingClass ,
163+                 requestingClass . getModule () .getName (),
164164                entitlementClass .getSimpleName ()
165165            )
166166        );
167167    }
168168
169-     ModuleEntitlements  getEntitlementsOrThrow (Class <?> callerClass , Module  requestingModule ) {
170-         ModuleEntitlements  cachedEntitlement  = moduleEntitlementsMap .get (requestingModule );
171-         if  (cachedEntitlement  != null ) {
172-             if  (cachedEntitlement  == ModuleEntitlements .NONE ) {
173-                 throw  new  NotEntitledException (buildModuleNoPolicyMessage (callerClass , requestingModule ) + "[CACHED]" );
174-             }
175-             return  cachedEntitlement ;
176-         }
169+     ModuleEntitlements  getEntitlements (Class <?> requestingClass ) {
170+         return  moduleEntitlementsMap .computeIfAbsent (requestingClass .getModule (), m  -> computeEntitlements (requestingClass ));
171+     }
177172
173+     private  ModuleEntitlements  computeEntitlements (Class <?> requestingClass ) {
174+         Module  requestingModule  = requestingClass .getModule ();
178175        if  (isServerModule (requestingModule )) {
179-             var  scopeName  = requestingModule .getName ();
180-             return  getModuleEntitlementsOrThrow (callerClass , requestingModule , serverEntitlements , scopeName );
176+             return  getModuleScopeEntitlements (requestingClass , serverEntitlements , requestingModule .getName ());
181177        }
182178
183179        // plugins 
184-         var  pluginName  = pluginResolver .apply (callerClass );
180+         var  pluginName  = pluginResolver .apply (requestingClass );
185181        if  (pluginName  != null ) {
186182            var  pluginEntitlements  = pluginsEntitlements .get (pluginName );
187183            if  (pluginEntitlements  != null ) {
@@ -191,60 +187,53 @@ ModuleEntitlements getEntitlementsOrThrow(Class<?> callerClass, Module requestin
191187                } else  {
192188                    scopeName  = requestingModule .getName ();
193189                }
194-                 return  getModuleEntitlementsOrThrow ( callerClass ,  requestingModule , pluginEntitlements , scopeName );
190+                 return  getModuleScopeEntitlements ( requestingClass , pluginEntitlements , scopeName );
195191            }
196192        }
197193
198-         moduleEntitlementsMap .put (requestingModule , ModuleEntitlements .NONE );
199-         throw  new  NotEntitledException (buildModuleNoPolicyMessage (callerClass , requestingModule ));
200-     }
194+         if  (requestingModule .isNamed () == false ) {
195+             // agents are the only thing running non-modular 
196+             return  ModuleEntitlements .from (agentEntitlements );
197+         }
201198
202-     private   static   String   buildModuleNoPolicyMessage ( Class <?>  callerClass ,  Module   requestingModule ) { 
203-         return  Strings . format ( "Missing entitlement policy: caller [%s], module [%s]" ,  callerClass ,  requestingModule . getName ()) ;
199+          logger . warn ( "No applicable entitlement policy for class [{}]" ,  requestingClass . getName ()); 
200+         return  ModuleEntitlements . NONE ;
204201    }
205202
206-     private  ModuleEntitlements  getModuleEntitlementsOrThrow (
203+     private  ModuleEntitlements  getModuleScopeEntitlements (
207204        Class <?> callerClass ,
208-         Module  module ,
209205        Map <String , List <Entitlement >> scopeEntitlements ,
210206        String  moduleName 
211207    ) {
212208        var  entitlements  = scopeEntitlements .get (moduleName );
213209        if  (entitlements  == null ) {
214-             // Module without entitlements - remember we don't have any 
215-             moduleEntitlementsMap .put (module , ModuleEntitlements .NONE );
216-             throw  new  NotEntitledException (buildModuleNoPolicyMessage (callerClass , module ));
210+             logger .warn ("No applicable entitlement policy for module [{}], class [{}]" , moduleName , callerClass );
211+             return  ModuleEntitlements .NONE ;
217212        }
218-         // We have a policy for this module 
219-         var  classEntitlements  = new  ModuleEntitlements (entitlements );
220-         moduleEntitlementsMap .put (module , classEntitlements );
221-         return  classEntitlements ;
213+         return  ModuleEntitlements .from (entitlements );
222214    }
223215
224216    private  static  boolean  isServerModule (Module  requestingModule ) {
225217        return  requestingModule .isNamed () && requestingModule .getLayer () == ModuleLayer .boot ();
226218    }
227219
228220    /** 
229-      * Walks the stack to determine which module's entitlements  should be checked. 
221+      * Walks the stack to determine which class  should be checked for entitlements . 
230222     * 
231-      * @param callerClass when non-null will be used if its module is suitable ; 
223+      * @param callerClass when non-null will be returned ; 
232224     *                    this is a fast-path check that can avoid the stack walk 
233225     *                    in cases where the caller class is available. 
234-      * @return the requesting module , or {@code null} if the entire call stack 
226+      * @return the requesting class , or {@code null} if the entire call stack 
235227     * comes from the entitlement library itself. 
236228     */ 
237-     Module   requestingModule (Class <?> callerClass ) {
229+     Class <?>  requestingClass (Class <?> callerClass ) {
238230        if  (callerClass  != null ) {
239-             var  callerModule  = callerClass .getModule ();
240-             if  (callerModule  != null  && entitlementsModule .equals (callerModule ) == false ) {
241-                 // fast path 
242-                 return  callerModule ;
243-             }
231+             // fast path 
232+             return  callerClass ;
244233        }
245-         Optional <Module >  module  = StackWalker .getInstance (RETAIN_CLASS_REFERENCE )
246-             .walk (frames  -> findRequestingModule (frames .map (StackFrame ::getDeclaringClass )));
247-         return  module .orElse (null );
234+         Optional <Class <?>>  result  = StackWalker .getInstance (RETAIN_CLASS_REFERENCE )
235+             .walk (frames  -> findRequestingClass (frames .map (StackFrame ::getDeclaringClass )));
236+         return  result .orElse (null );
248237    }
249238
250239    /** 
@@ -253,33 +242,25 @@ Module requestingModule(Class<?> callerClass) {
253242     * 
254243     * @throws NullPointerException if the requesting module is {@code null} 
255244     */ 
256-     Optional <Module > findRequestingModule (Stream <Class <?>> classes ) {
257-         return  classes .map (Objects ::requireNonNull )
258-             .map (PolicyManager ::moduleOf )
259-             .filter (m  -> m  != entitlementsModule )  // Ignore the entitlements library itself entirely 
260-             .skip (1 )                            // Skip the sensitive method itself 
245+     Optional <Class <?>> findRequestingClass (Stream <Class <?>> classes ) {
246+         return  classes .filter (c  -> c .getModule () != entitlementsModule )  // Ignore the entitlements library 
247+             .skip (1 )                                           // Skip the sensitive caller method 
261248            .findFirst ();
262249    }
263250
264-     private  static  Module  moduleOf (Class <?> c ) {
265-         var  result  = c .getModule ();
266-         if  (result  == null ) {
267-             throw  new  NullPointerException ("Entitlements system does not support non-modular class ["  + c .getName () + "]" );
268-         } else  {
269-             return  result ;
270-         }
271-     }
272- 
273-     private  static  boolean  isTriviallyAllowed (Module  requestingModule ) {
251+     /** 
252+      * @return true if permission is granted regardless of the entitlement 
253+      */ 
254+     private  static  boolean  isTriviallyAllowed (Class <?> requestingClass ) {
274255        if  (logger .isTraceEnabled ()) {
275256            logger .trace ("Stack trace for upcoming trivially-allowed check" , new  Exception ());
276257        }
277-         if  (requestingModule  == null ) {
258+         if  (requestingClass  == null ) {
278259            logger .debug ("Entitlement trivially allowed: no caller frames outside the entitlement library" );
279260            return  true ;
280261        }
281-         if  (systemModules .contains (requestingModule )) {
282-             logger .debug ("Entitlement trivially allowed from system module [{}]" , requestingModule .getName ());
262+         if  (systemModules .contains (requestingClass . getModule () )) {
263+             logger .debug ("Entitlement trivially allowed from system module [{}]" , requestingClass . getModule () .getName ());
283264            return  true ;
284265        }
285266        logger .trace ("Entitlement not trivially allowed" );
0 commit comments