1717
1818import java .lang .module .ModuleFinder ;
1919import java .lang .module .ModuleReference ;
20+ import java .util .ArrayList ;
2021import java .util .Collections ;
22+ import java .util .HashMap ;
23+ import java .util .IdentityHashMap ;
24+ import java .util .List ;
2125import java .util .Map ;
2226import java .util .Objects ;
2327import java .util .Optional ;
2428import java .util .Set ;
2529import java .util .function .Function ;
2630import java .util .stream .Collectors ;
31+ import java .util .stream .Stream ;
2732
2833public class PolicyManager {
2934 private static final Logger logger = LogManager .getLogger (ElasticsearchEntitlementChecker .class );
3035
36+ static class ModuleEntitlements {
37+ public static final ModuleEntitlements NONE = new ModuleEntitlements (List .of ());
38+ private final IdentityHashMap <Class <? extends Entitlement >, List <Entitlement >> entitlementsByType ;
39+
40+ ModuleEntitlements (List <Entitlement > entitlements ) {
41+ this .entitlementsByType = entitlements .stream ()
42+ .collect (Collectors .toMap (Entitlement ::getClass , e -> new ArrayList <>(List .of (e )), (a , b ) -> {
43+ a .addAll (b );
44+ return a ;
45+ }, IdentityHashMap ::new ));
46+ }
47+
48+ public boolean hasEntitlement (Class <? extends Entitlement > entitlementClass ) {
49+ return entitlementsByType .containsKey (entitlementClass );
50+ }
51+
52+ public <E extends Entitlement > Stream <E > getEntitlements (Class <E > entitlementClass ) {
53+ return entitlementsByType .get (entitlementClass ).stream ().map (entitlementClass ::cast );
54+ }
55+ }
56+
57+ final Map <Module , ModuleEntitlements > moduleEntitlementsMap = new HashMap <>();
58+
3159 protected final Policy serverPolicy ;
3260 protected final Map <String , Policy > pluginPolicies ;
3361 private final Function <Class <?>, String > pluginResolver ;
@@ -56,27 +84,110 @@ public PolicyManager(Policy defaultPolicy, Map<String, Policy> pluginPolicies, F
5684 this .pluginResolver = pluginResolver ;
5785 }
5886
59- public void checkFlagEntitlement (Class <?> callerClass , FlagEntitlementType type ) {
87+ private static List <Entitlement > lookupEntitlementsForModule (Policy policy , String moduleName ) {
88+ for (int i = 0 ; i < policy .scopes .size (); ++i ) {
89+ var scope = policy .scopes .get (i );
90+ if (scope .name .equals (moduleName )) {
91+ return scope .entitlements ;
92+ }
93+ }
94+ return null ;
95+ }
96+
97+ public void checkExitVM (Class <?> callerClass ) {
98+ checkEntitlementPresent (callerClass , ExitVMEntitlement .class );
99+ }
100+
101+ public void checkCreateClassLoader (Class <?> callerClass ) {
102+ checkEntitlementPresent (callerClass , CreateClassLoaderEntitlement .class );
103+ }
104+
105+ private void checkEntitlementPresent (Class <?> callerClass , Class <? extends Entitlement > entitlementClass ) {
60106 var requestingModule = requestingModule (callerClass );
61107 if (isTriviallyAllowed (requestingModule )) {
62108 return ;
63109 }
64110
65- // TODO: real policy check. For now, we only allow our hardcoded System.exit policy for server.
66- // TODO: this will be checked using policies
67- if (requestingModule .isNamed ()
68- && requestingModule .getName ().equals ("org.elasticsearch.server" )
69- && (type == FlagEntitlementType .SYSTEM_EXIT || type == FlagEntitlementType .CREATE_CLASSLOADER )) {
70- logger .debug ("Allowed: caller [{}] in module [{}] has entitlement [{}]" , callerClass , requestingModule .getName (), type );
111+ ModuleEntitlements entitlements = getEntitlementsOrThrow (callerClass , requestingModule );
112+ if (entitlements .hasEntitlement (entitlementClass )) {
113+ logger .debug (
114+ () -> Strings .format (
115+ "Entitled: caller [%s], module [%s], type [%s]" ,
116+ callerClass ,
117+ requestingModule .getName (),
118+ entitlementClass .getSimpleName ()
119+ )
120+ );
71121 return ;
72122 }
73-
74- // TODO: plugins policy check using pluginResolver and pluginPolicies
75123 throw new NotEntitledException (
76- Strings .format ("Missing entitlement [%s] for caller [%s] in module [%s]" , type , callerClass , requestingModule .getName ())
124+ Strings .format (
125+ "Missing entitlement: caller [%s], module [%s], type [%s]" ,
126+ callerClass ,
127+ requestingModule .getName (),
128+ entitlementClass .getSimpleName ()
129+ )
77130 );
78131 }
79132
133+ ModuleEntitlements getEntitlementsOrThrow (Class <?> callerClass , Module requestingModule ) {
134+ ModuleEntitlements cachedEntitlement = moduleEntitlementsMap .get (requestingModule );
135+ if (cachedEntitlement != null ) {
136+ if (cachedEntitlement == ModuleEntitlements .NONE ) {
137+ throw new NotEntitledException (buildModuleNoPolicyMessage (callerClass , requestingModule ) + "[CACHED]" );
138+ }
139+ return cachedEntitlement ;
140+ }
141+
142+ if (isServerModule (requestingModule )) {
143+ var scopeName = requestingModule .getName ();
144+ return getModuleEntitlementsOrThrow (callerClass , requestingModule , serverPolicy , scopeName );
145+ }
146+
147+ // plugins
148+ var pluginName = pluginResolver .apply (callerClass );
149+ if (pluginName != null ) {
150+ var pluginPolicy = pluginPolicies .get (pluginName );
151+ if (pluginPolicy != null ) {
152+ final String scopeName ;
153+ if (requestingModule .isNamed () == false ) {
154+ scopeName = ALL_UNNAMED ;
155+ } else {
156+ scopeName = requestingModule .getName ();
157+ }
158+ return getModuleEntitlementsOrThrow (callerClass , requestingModule , pluginPolicy , scopeName );
159+ }
160+ }
161+
162+ moduleEntitlementsMap .put (requestingModule , ModuleEntitlements .NONE );
163+ throw new NotEntitledException (buildModuleNoPolicyMessage (callerClass , requestingModule ));
164+ }
165+
166+ private static String buildModuleNoPolicyMessage (Class <?> callerClass , Module requestingModule ) {
167+ return Strings .format ("Missing entitlement policy: caller [%s], module [%s]" , callerClass , requestingModule .getName ());
168+ }
169+
170+ private ModuleEntitlements getModuleEntitlementsOrThrow (Class <?> callerClass , Module module , Policy policy , String moduleName ) {
171+ var entitlements = lookupEntitlementsForModule (policy , moduleName );
172+ if (entitlements == null ) {
173+ // Module without entitlements - remember we don't have any
174+ moduleEntitlementsMap .put (module , ModuleEntitlements .NONE );
175+ throw new NotEntitledException (buildModuleNoPolicyMessage (callerClass , module ));
176+ }
177+ // We have a policy for this module
178+ var classEntitlements = createClassEntitlements (entitlements );
179+ moduleEntitlementsMap .put (module , classEntitlements );
180+ return classEntitlements ;
181+ }
182+
183+ private static boolean isServerModule (Module requestingModule ) {
184+ return requestingModule .isNamed () && requestingModule .getLayer () == ModuleLayer .boot ();
185+ }
186+
187+ private ModuleEntitlements createClassEntitlements (List <Entitlement > entitlements ) {
188+ return new ModuleEntitlements (entitlements );
189+ }
190+
80191 private static Module requestingModule (Class <?> callerClass ) {
81192 if (callerClass != null ) {
82193 Module callerModule = callerClass .getModule ();
@@ -102,10 +213,10 @@ private static Module requestingModule(Class<?> callerClass) {
102213
103214 private static boolean isTriviallyAllowed (Module requestingModule ) {
104215 if (requestingModule == null ) {
105- logger .debug ("Trivially allowed: entire call stack is in composed of classes in system modules" );
216+ logger .debug ("Entitlement trivially allowed: entire call stack is in composed of classes in system modules" );
106217 return true ;
107218 }
108- logger .trace ("Not trivially allowed" );
219+ logger .trace ("Entitlement not trivially allowed" );
109220 return false ;
110221 }
111222
0 commit comments