Ban sun.misc.Unsafe

README.md

@@ -21,3 +21,4 @@ Supported fields for the `option` list are:
 * `nobigmath` — disables `BigInteger` and `BigDecimal`, raising [appropriate exceptions](https://github.com/DMOJ/java-sandbox-agent/blob/master/src/main/java/ca/dmoj/java/BigIntegerDisallowedException.java) if they are used
 * `unicode` — encodes `System.out` as UTF-8 instead of ASCII, sacrificing performance for Unicode support
 * `nobuf` — sets `System.out` as being line-buffered, for interactive problems
+* `unsafe` — enables `sun.misc.Unsafe`, which is disabled by default

src/main/java/ca/dmoj/java/DisallowedClassRule.java

@@ -0,0 +1,21 @@
+class DisallowedClassRule {
+    protected interface ExceptionFactory {
+        Exception create();
+    }
+
+    protected String disallowedClassName;
+    protected ExceptionFactory disallowedClassException;
+
+    public DisallowedClassRule(String disallowedClassName, ExceptionFactory disallowedClassException) {
+        this.disallowedClassName = disallowedClassName;
+        this.disallowedClassException = disallowedClassException;
+    }
+
+    public String getDisallowedClassName() {
+        return disallowedClassName;
+    }
+
+    public Exception getDisallowedClassException() {
+        return disallowedClassException.create();
+    }
+}

src/main/java/ca/dmoj/java/DisallowedClassesClassFileTransformer.java

@@ -0,0 +1,35 @@
+package ca.dmoj.java;
+
+import java.lang.instrument.ClassFileTransformer;
+import java.lang.instrument.IllegalClassFormatException;
+import java.security.ProtectionDomain;
+
+public class DisallowedClassesClassFileTransformer implements ClassFileTransformer {
+    protected DisallowedClassRule[] disallowedClassRules;
+
+    public DisallowedClassesClassFileTransformer(DisallowedClassRule... disallowedClassRules) {
+        super();
+        this.disallowedClassRules = disallowedClassRules;
+    }
+
+    @Override
+    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
+                            ProtectionDomain protectionDomain, byte[] classfileBuffer)
+                            throws IllegalClassFormatException {
+        if (className == null) return classfileBuffer;
+
+        Exception disallowed = null;
+        for (DisallowedClassRule rule : disallowedClassRules) {
+            // If the class ever loaded it's because a submission used it
+            if (className.startsWith(rule.getDisallowedClassName())) {
+                disallowed = rule.getDisallowedClassException();
+                break;
+            }
+        }
+
+        if (disallowed != null) ExceptionHandler.dumpExceptionAndExit(disallowed);
+
+        // Don't actually retransform anything
+        return classfileBuffer;
+    }
+}

src/main/java/ca/dmoj/java/ExceptionHandler.java

@@ -0,0 +1,9 @@
+package ca.dmoj.java;
+
+class ExceptionHandler {
+    public static void dumpExceptionAndExit(Throwable exception) {
+        System.err.print("7257b50d-e37a-4664-b1a5-b1340b4206c0: ");
+        exception.printStackTrace();
+        System.exit(1);
+    }
+}

src/main/java/ca/dmoj/java/SubmissionAgent.java

@@ -1,19 +1,18 @@
 package ca.dmoj.java;
 
 import java.io.*;
-import java.lang.instrument.ClassFileTransformer;
-import java.lang.instrument.IllegalClassFormatException;
 import java.lang.instrument.Instrumentation;
-import java.security.ProtectionDomain;
 
 public class SubmissionAgent {
     public static void premain(String argv, Instrumentation inst) throws UnsupportedEncodingException {
+        boolean unsafe = false;
         boolean unicode = false;
         boolean noBigMath = false;
         boolean noBuf = false;
 
         if (argv != null) {
             for (String opt : argv.split(",")) {
+                if (opt.equals("unsafe")) unsafe = true;
                 if (opt.equals("unicode")) unicode = true;
                 if (opt.equals("nobigmath")) noBigMath = true;
                 if (opt.equals("nobuf")) noBuf = true;
@@ -22,29 +21,18 @@ public static void premain(String argv, Instrumentation inst) throws Unsupported
 
         final Thread selfThread = Thread.currentThread();
 
+        if (!unsafe) {
+            inst.addTransformer(new DisallowedClassesClassFileTransformer(
+                new DisallowedClassRule("sun/reflect/Unsafe", UnsafeDisallowedException::new),
+                new DisallowedClassRule("sun/misc/Unsafe", UnsafeDisallowedException::new)
+            ), true);
+        }
         if (noBigMath) {
-            inst.addTransformer(new ClassFileTransformer() {
-                @Override
-                public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
-                                        ProtectionDomain protectionDomain, byte[] classfileBuffer)
-                                        throws IllegalClassFormatException {
-                    if (className == null) return classfileBuffer;
-
-                    RuntimeException disallowed = null;
-                    // If the class ever loaded it's because a submission used it
-                    if (className.startsWith("java/math/BigInteger") ||
-                        className.startsWith("java/math/MutableBigInteger")) {
-                        disallowed = new BigIntegerDisallowedException();
-                    } else if (className.startsWith("java/math/BigDecimal")) {
-                        disallowed = new BigDecimalDisallowedException();
-                    }
-
-                    if (disallowed != null) dumpExceptionAndExit(disallowed);
-
-                    // Don't actually retransform anything
-                    return classfileBuffer;
-                }
-            });
+            inst.addTransformer(new DisallowedClassesClassFileTransformer(
+                new DisallowedClassRule("java/math/BigInteger", BigIntegerDisallowedException::new),
+                new DisallowedClassRule("java/math/MutableBigInteger", BigIntegerDisallowedException::new),
+                new DisallowedClassRule("java/math/BigDecimal", BigDecimalDisallowedException::new)
+            ), true);
         }
 
         if (noBuf) {
@@ -61,7 +49,7 @@ public byte[] transform(ClassLoader loader, String className, Class<?> classBein
         selfThread.setUncaughtExceptionHandler(new Thread.UncaughtExceptionHandler() {
             @Override
             public void uncaughtException(Thread thread, Throwable error) {
-                dumpExceptionAndExit(error);
+                ExceptionHandler.dumpExceptionAndExit(error);
             }
         });
 
@@ -75,10 +63,4 @@ public void run() {
             }
         }));
     }
-
-    private static void dumpExceptionAndExit(Throwable exception) {
-        System.err.print("7257b50d-e37a-4664-b1a5-b1340b4206c0: ");
-        exception.printStackTrace();
-        System.exit(1);
-    }
 }

src/main/java/ca/dmoj/java/UnsafeDisallowedException.java

@@ -0,0 +1,4 @@
+package ca.dmoj.java;
+
+public class UnsafeDisallowedException extends RuntimeException {
+}