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/DisallowedClassesClassFileTransformer.java

@@ -0,0 +1,44 @@
+package ca.dmoj.java;
+
+import java.lang.instrument.ClassFileTransformer;
+import java.lang.instrument.IllegalClassFormatException;
+import java.security.ProtectionDomain;
+
+public class DisallowedClassesClassFileTransformer implements ClassFileTransformer {
+    protected String[] disallowedClassNames;
+    protected RuntimeException[] disallowedClassExceptions;
+
+    public DisallowedClassesClassFileTransformer(String[] disallowedClassNames, RuntimeException[] disallowedClassExceptions) {
+        super();
+        assert disallowedClassNames.length == disallowedClassExceptions.length;
+        this.disallowedClassNames = disallowedClassNames;
+        this.disallowedClassExceptions = disallowedClassExceptions;
+    }
+
+    @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
+        for (int i = 0; i < this.disallowedClassNames.length; ++i) {
+            if (className.startsWith(this.disallowedClassNames[i])) {
+                disallowed = this.disallowedClassExceptions[i];
+                break;
+            }
+        }
+
+        if (disallowed != null) dumpExceptionAndExit(disallowed);
+
+        // Don't actually retransform anything
+        return classfileBuffer;
+    }
+
+    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/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,27 @@ public static void premain(String argv, Instrumentation inst) throws Unsupported
 
         final Thread selfThread = Thread.currentThread();
 
+        if (!unsafe) {
+            inst.addTransformer(new DisallowedClassesClassFileTransformer(
+                new String[]{
+                    "sun/misc/Unsafe",
+                },
+                new RuntimeException[]{
+                    new UnsafeDisallowedException()
+            }));
+        }
         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 String[]{
+                    "java/math/BigInteger",
+                    "java/math/MutableBigInteger",
+                    "java/math/BigDecimal"
+                },
+                new RuntimeException[]{
+                    new BigIntegerDisallowedException(),
+                    new BigIntegerDisallowedException(),
+                    new BigDecimalDisallowedException(),
+            }));
         }
 
         if (noBuf) {

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

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