1 files changed, 177 insertions, 10 deletions
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/runtime/FuzzedDataProviderImpl.java b/agent/src/main/java/com/code_intelligence/jazzer/runtime/FuzzedDataProviderImpl.java
index fe4d8ac7..b7aad33e 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/runtime/FuzzedDataProviderImpl.java
+++ b/agent/src/main/java/com/code_intelligence/jazzer/runtime/FuzzedDataProviderImpl.java
@@ -15,9 +15,119 @@
 package com.code_intelligence.jazzer.runtime;
 
 import com.code_intelligence.jazzer.api.FuzzedDataProvider;
-
-public class FuzzedDataProviderImpl implements FuzzedDataProvider {
-  public FuzzedDataProviderImpl() {}
+import com.github.fmeum.rules_jni.RulesJni;
+import sun.misc.Unsafe;
+
+public class FuzzedDataProviderImpl implements FuzzedDataProvider, AutoCloseable {
+  static {
+    // The replayer loads a standalone version of the FuzzedDataProvider.
+    if (System.getProperty("jazzer.is_replayer") == null) {
+      RulesJni.loadLibrary("jazzer_driver", "/com/code_intelligence/jazzer/driver");
+    }
+    nativeInit();
+  }
+
+  private static native void nativeInit();
+
+  private final boolean ownsNativeData;
+  private long originalDataPtr;
+  private int originalRemainingBytes;
+
+  // Accessed in fuzzed_data_provider.cpp.
+  private long dataPtr;
+  private int remainingBytes;
+
+  private FuzzedDataProviderImpl(long dataPtr, int remainingBytes, boolean ownsNativeData) {
+    this.ownsNativeData = ownsNativeData;
+    this.originalDataPtr = dataPtr;
+    this.dataPtr = dataPtr;
+    this.originalRemainingBytes = remainingBytes;
+    this.remainingBytes = remainingBytes;
+  }
+
+  /**
+   * Creates a {@link FuzzedDataProvider} that consumes bytes from an already existing native array.
+   *
+   * <ul>
+   * <li>{@link #close()} <b>must</b> be called on instances created with this method to free the
+   * native copy of the Java
+   * {@code byte} array.
+   * <li>{@link #setNativeData(long, int)} <b>must not</b> be called on instances created with this
+   * method.
+   *
+   * @param data the raw bytes used as input
+   * @return a {@link FuzzedDataProvider} backed by {@code data}
+   */
+  public static FuzzedDataProviderImpl withJavaData(byte[] data) {
+    return new FuzzedDataProviderImpl(allocateNativeCopy(data), data.length, true);
+  }
+
+  /**
+   * Creates a {@link FuzzedDataProvider} that consumes bytes from an already existing native array.
+   *
+   * <p>The backing array can be set at any time using {@link #setNativeData(long, int)} and is
+   * initially empty.
+   *
+   * @return a {@link FuzzedDataProvider} backed by an empty array.
+   */
+  public static FuzzedDataProviderImpl withNativeData() {
+    return new FuzzedDataProviderImpl(0, 0, false);
+  }
+
+  /**
+   * Replaces the current native backing array.
+   *
+   * <p><b>Must not</b> be called on instances created with {@link #withJavaData(byte[])}.
+   *
+   * @param dataPtr a native pointer to the new backing array
+   * @param dataLength the length of the new backing array
+   */
+  public void setNativeData(long dataPtr, int dataLength) {
+    this.originalDataPtr = dataPtr;
+    this.dataPtr = dataPtr;
+    this.originalRemainingBytes = dataLength;
+    this.remainingBytes = dataLength;
+  }
+
+  /**
+   * Resets the FuzzedDataProvider state to read from the beginning to the end of its current
+   * backing item.
+   */
+  public void reset() {
+    dataPtr = originalDataPtr;
+    remainingBytes = originalRemainingBytes;
+  }
+
+  /**
+   * Releases native memory allocated for this instance (if any).
+   *
+   * <p>While the instance should not be used after this method returns, no usage of {@link
+   * FuzzedDataProvider} methods can result in memory corruption.
+   */
+  @Override
+  public void close() {
+    if (originalDataPtr == 0) {
+      return;
+    }
+    if (ownsNativeData) {
+      UNSAFE.freeMemory(originalDataPtr);
+    }
+    // Prevent double-frees and use-after-frees by effectively making all methods no-ops after
+    // close() has been called.
+    originalDataPtr = 0;
+    originalRemainingBytes = 0;
+    dataPtr = 0;
+    remainingBytes = 0;
+  }
+
+  private static final Unsafe UNSAFE = UnsafeProvider.getUnsafe();
+  private static final long BYTE_ARRAY_OFFSET = UNSAFE.arrayBaseOffset(byte[].class);
+
+  private static long allocateNativeCopy(byte[] data) {
+    long nativeCopy = UNSAFE.allocateMemory(data.length);
+    UNSAFE.copyMemory(data, BYTE_ARRAY_OFFSET, null, nativeCopy, data.length);
+    return nativeCopy;
+  }
 
   @Override public native boolean consumeBoolean();
 
@@ -25,23 +135,51 @@ public class FuzzedDataProviderImpl implements FuzzedDataProvider {
 
   @Override public native byte consumeByte();
 
-  @Override public native byte consumeByte(byte min, byte max);
+  @Override
+  public byte consumeByte(byte min, byte max) {
+    if (min > max) {
+      throw new IllegalArgumentException(
+          String.format("min must be <= max (got min: %d, max: %d)", min, max));
+    }
+    return consumeByteUnchecked(min, max);
+  }
 
   @Override public native short consumeShort();
 
-  @Override public native short consumeShort(short min, short max);
+  @Override
+  public short consumeShort(short min, short max) {
+    if (min > max) {
+      throw new IllegalArgumentException(
+          String.format("min must be <= max (got min: %d, max: %d)", min, max));
+    }
+    return consumeShortUnchecked(min, max);
+  }
 
   @Override public native short[] consumeShorts(int maxLength);
 
   @Override public native int consumeInt();
 
-  @Override public native int consumeInt(int min, int max);
+  @Override
+  public int consumeInt(int min, int max) {
+    if (min > max) {
+      throw new IllegalArgumentException(
+          String.format("min must be <= max (got min: %d, max: %d)", min, max));
+    }
+    return consumeIntUnchecked(min, max);
+  }
 
   @Override public native int[] consumeInts(int maxLength);
 
   @Override public native long consumeLong();
 
-  @Override public native long consumeLong(long min, long max);
+  @Override
+  public long consumeLong(long min, long max) {
+    if (min > max) {
+      throw new IllegalArgumentException(
+          String.format("min must be <= max (got min: %d, max: %d)", min, max));
+    }
+    return consumeLongUnchecked(min, max);
+  }
 
   @Override public native long[] consumeLongs(int maxLength);
 
@@ -49,13 +187,27 @@ public class FuzzedDataProviderImpl implements FuzzedDataProvider {
 
   @Override public native float consumeRegularFloat();
 
-  @Override public native float consumeRegularFloat(float min, float max);
+  @Override
+  public float consumeRegularFloat(float min, float max) {
+    if (min > max) {
+      throw new IllegalArgumentException(
+          String.format("min must be <= max (got min: %f, max: %f)", min, max));
+    }
+    return consumeRegularFloatUnchecked(min, max);
+  }
 
   @Override public native float consumeProbabilityFloat();
 
   @Override public native double consumeDouble();
 
-  @Override public native double consumeRegularDouble(double min, double max);
+  @Override
+  public double consumeRegularDouble(double min, double max) {
+    if (min > max) {
+      throw new IllegalArgumentException(
+          String.format("min must be <= max (got min: %f, max: %f)", min, max));
+    }
+    return consumeRegularDoubleUnchecked(min, max);
+  }
 
   @Override public native double consumeRegularDouble();
 
@@ -63,7 +215,14 @@ public class FuzzedDataProviderImpl implements FuzzedDataProvider {
 
   @Override public native char consumeChar();
 
-  @Override public native char consumeChar(char min, char max);
+  @Override
+  public char consumeChar(char min, char max) {
+    if (min > max) {
+      throw new IllegalArgumentException(
+          String.format("min must be <= max (got min: %c, max: %c)", min, max));
+    }
+    return consumeCharUnchecked(min, max);
+  }
 
   @Override public native char consumeCharNoSurrogates();
 
@@ -80,4 +239,12 @@ public class FuzzedDataProviderImpl implements FuzzedDataProvider {
   @Override public native byte[] consumeRemainingAsBytes();
 
   @Override public native int remainingBytes();
+
+  private native byte consumeByteUnchecked(byte min, byte max);
+  private native short consumeShortUnchecked(short min, short max);
+  private native char consumeCharUnchecked(char min, char max);
+  private native int consumeIntUnchecked(int min, int max);
+  private native long consumeLongUnchecked(long min, long max);
+  private native float consumeRegularFloatUnchecked(float min, float max);
+  private native double consumeRegularDoubleUnchecked(double min, double max);
 }