Add JNDI helper API

README.md

@@ -35,12 +35,12 @@ In Maven:
 <dependency>
   <groupId>io.github.pixee</groupId>
   <artifactId>java-security-toolkit</artifactId>
-  <version>1.1.3</version>
+  <version>1.2.0</version>
 </dependency>
 ```
 In Gradle:
 ```kotlin
-implementation("io.github.pixee:java-security-toolkit:1.1.3")
+implementation("io.github.pixee:java-security-toolkit:1.2.0")
 ```
 
 ## Contributing 

build.gradle.kts

@@ -91,7 +91,7 @@ tasks.named(java11SourceSet.jarTaskName) {
 }
 
 group = "io.github.pixee"
-version = "1.1.3"
+version = "1.2.0"
 description = "java-security-toolkit"
 
 

src/main/java/io/github/pixee/security/JNDI.java

@@ -0,0 +1,89 @@
+package io.github.pixee.security;
+
+import javax.naming.Context;
+import javax.naming.NamingException;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/** Offers utilities to defend against JNDI attacks by controlling allowed resources. */
+public final class JNDI {
+
+    private JNDI() {}
+
+    /**
+     * Looks up a resource in the context, only allowing resources non-URL-based resources and "java:" resources.
+     */
+    public static LimitedContext limitedContext(final Context context) {
+        return new ProtocolLimitedContext(context, J8ApiBridge.setOf(UrlProtocol.JAVA));
+    }
+
+    /**
+     * Looks up a resource in the context, only allowing resources from the specified protocols.
+     */
+    public static LimitedContext limitedContextByProtocol(final Context context, final Set<UrlProtocol> allowedProtocols) {
+        return new ProtocolLimitedContext(context, allowedProtocols);
+    }
+
+    /**
+     * Looks up a resource in the context, only allowing resources with the given names.
+     */
+    public static LimitedContext limitedContextByResourceName(final Context context, final Set<String> allowedResourceNames) {
+        return new NameLimitedContext(context, allowedResourceNames);
+    }
+
+    /** A lookalike method for {@link Context} that allows sandboxing resolution. */
+    public interface LimitedContext {
+        /**
+         * Looks up a resource in the context, but only allows resources that are in the allowed set.
+         *
+         * @param resource the resource to look up
+         * @return the object bound to the resource
+         * @throws NamingException if the resource is not allowed or if the lookup fails as per {@link Context#lookup(String)}
+         */
+        Object lookup(final String resource) throws NamingException;
+    }
+
+    /** A context which limits protocols. */
+    private static class ProtocolLimitedContext implements LimitedContext {
+        private final Context context;
+        private final Set<UrlProtocol> allowedProtocols;
+
+        private ProtocolLimitedContext(final Context context, final Set<UrlProtocol> allowedProtocols) {
+            this.context = Objects.requireNonNull(context);
+            this.allowedProtocols = Objects.requireNonNull(allowedProtocols);
+        }
+
+        @Override
+        public Object lookup(final String resource) throws NamingException {
+            Set<String> allowedProtocolPrefixes = allowedProtocols.stream().map(UrlProtocol::getKey).map(p -> p + ":").collect(Collectors.toSet());
+            String canonicalResource = resource.toLowerCase().trim();
+            if(canonicalResource.contains(":")) {
+                if (allowedProtocolPrefixes.stream().anyMatch(canonicalResource::startsWith)) {
+                    return context.lookup(resource);
+                } else {
+                    throw new SecurityException("Unexpected JNDI resource protocol: " + resource);
+                }
+            }
+            return context.lookup(resource);
+        }
+    }
+
+    /** A context which only allows pre-defined resource names. */
+    private static class NameLimitedContext implements LimitedContext {
+        private final Context context;
+        private final Set<String> allowedResourceNames;
+
+        private NameLimitedContext(final Context context, final Set<String> allowedResourceNames) {
+            this.context = Objects.requireNonNull(context);
+            this.allowedResourceNames = Objects.requireNonNull(allowedResourceNames);
+        }
+        @Override
+        public Object lookup(final String resource) throws NamingException {
+            if(allowedResourceNames.contains(resource)) {
+                return context.lookup(resource);
+            }
+            throw new SecurityException("Unexpected JNDI resource name: " + resource);
+        }
+    }
+}

src/main/java/io/github/pixee/security/UrlProtocol.java

@@ -49,7 +49,16 @@ public enum UrlProtocol {
   TELNET("telnet"),
 
   /** Classpath */
-  CLASSPATH("classpath");
+  CLASSPATH("classpath"),
+
+  /** LDAP */
+  LDAP("ldap"),
+
+  /** Java */
+  JAVA("java"),
+
+  /** RMI */
+  RMI("rmi");
 
   private final String key;
 

src/main/java/module-info.java

@@ -4,6 +4,7 @@
 
     requires org.apache.commons.io;
     requires java.xml;
+    requires java.naming;
     requires java.desktop;
     requires java.base;
 }

src/test/java/io/github/pixee/security/JNDITest.java

@@ -0,0 +1,71 @@
+package io.github.pixee.security;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import javax.naming.Context;
+import javax.naming.NamingException;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.*;
+
+final class JNDITest {
+
+    private Context context;
+    private final Object NAMED_OBJECT = new Object();
+    private final Object JAVA_OBJECT = new Object();
+    private final Object LDAP_OBJECT = new Object();
+    private final Object RMI_OBJECT = new Object();
+
+    @BeforeEach
+    void setup() throws NamingException {
+        context = mock(Context.class);
+        when(context.lookup("simple_name")).thenReturn(NAMED_OBJECT);
+        when(context.lookup("java:comp/env")).thenReturn(JAVA_OBJECT);
+        when(context.lookup("ldap://localhost:1389/ou=system")).thenReturn(LDAP_OBJECT);
+        when(context.lookup("rmi://localhost:1099/evil")).thenReturn(RMI_OBJECT);
+    }
+
+    @Test
+    void it_limits_resources_by_name() throws NamingException {
+        JNDI.LimitedContext limitedContext = JNDI.limitedContextByResourceName(context, J8ApiBridge.setOf("simple_name"));
+        assertThat(limitedContext.lookup("simple_name"), is(NAMED_OBJECT));
+        assertThrows(SecurityException.class, () -> limitedContext.lookup("anything_else"));
+        verify(context, times(1)).lookup(anyString());
+    }
+
+    @Test
+    void it_limits_resources_by_protocol() throws NamingException {
+        JNDI.LimitedContext onlyJavaContext = JNDI.limitedContextByProtocol(context, J8ApiBridge.setOf(UrlProtocol.JAVA));
+        assertThat(onlyJavaContext.lookup("java:comp/env"), is(JAVA_OBJECT));
+
+        // confirm protocols protections dont restrict simple name lookups
+        assertThat(onlyJavaContext.lookup("simple_name"), is(NAMED_OBJECT));
+        assertThrows(SecurityException.class, () -> onlyJavaContext.lookup("ldap://localhost:1389/ou=system"));
+        assertThrows(SecurityException.class, () -> onlyJavaContext.lookup("rmi://localhost:1099/evil"));
+
+        JNDI.LimitedContext onlyLdapContext = JNDI.limitedContextByProtocol(context, J8ApiBridge.setOf(UrlProtocol.LDAP));
+        assertThat(onlyLdapContext.lookup("ldap://localhost:1389/ou=system"), is(LDAP_OBJECT));
+        assertThrows(SecurityException.class, () -> onlyLdapContext.lookup("java:comp/env"));
+        assertThrows(SecurityException.class, () -> onlyLdapContext.lookup("rmi://localhost:1099/evil"));
+
+        JNDI.LimitedContext onlyLdapAndJavaContext = JNDI.limitedContextByProtocol(context, J8ApiBridge.setOf(UrlProtocol.JAVA, UrlProtocol.LDAP));
+        assertThat(onlyLdapAndJavaContext.lookup("ldap://localhost:1389/ou=system"), is(LDAP_OBJECT));
+        assertThat(onlyLdapAndJavaContext.lookup("java:comp/env"), is(JAVA_OBJECT));
+        assertThrows(SecurityException.class, () -> onlyLdapAndJavaContext.lookup("rmi://localhost:1099/evil"));
+    }
+
+    @Test
+    void default_limits_rmi_and_ldap() throws NamingException {
+        JNDI.LimitedContext defaultLimitedContext = JNDI.limitedContext(context);
+        assertThat(defaultLimitedContext.lookup("java:comp/env"), is(JAVA_OBJECT));
+
+        // confirm simple name lookups still work
+        assertThat(defaultLimitedContext.lookup("simple_name"), is(NAMED_OBJECT));
+        assertThrows(SecurityException.class, () -> defaultLimitedContext.lookup("rmi://localhost:1099/evil"));
+        assertThrows(SecurityException.class, () -> defaultLimitedContext.lookup("ldap://localhost:1389/ou=system"));
+    }
+
+}