Merge "Cronet: add TestUploadDataProvider"
diff --git a/Cronet/tests/cts/src/android/net/http/cts/util/TestUploadDataProvider.java b/Cronet/tests/cts/src/android/net/http/cts/util/TestUploadDataProvider.java
new file mode 100644
index 0000000..d047828
--- /dev/null
+++ b/Cronet/tests/cts/src/android/net/http/cts/util/TestUploadDataProvider.java
@@ -0,0 +1,294 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.http.cts.util;
+
+import android.net.http.UploadDataProvider;
+import android.net.http.UploadDataSink;
+import android.os.ConditionVariable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.ClosedChannelException;
+import java.util.ArrayList;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/** An UploadDataProvider implementation used in tests. */
+public class TestUploadDataProvider extends UploadDataProvider {
+    // Indicates whether all success callbacks are synchronous or asynchronous.
+    // Doesn't apply to errors.
+    public enum SuccessCallbackMode {
+        SYNC,
+        ASYNC
+    }
+
+    // Indicates whether failures should throw exceptions, invoke callbacks
+    // synchronously, or invoke callback asynchronously.
+    public enum FailMode {
+        NONE,
+        THROWN,
+        CALLBACK_SYNC,
+        CALLBACK_ASYNC
+    }
+
+    private final ArrayList<byte[]> mReads = new ArrayList<byte[]>();
+    private final SuccessCallbackMode mSuccessCallbackMode;
+    private final Executor mExecutor;
+
+    private boolean mChunked;
+
+    // Index of read to fail on.
+    private int mReadFailIndex = -1;
+    // Indicates how to fail on a read.
+    private FailMode mReadFailMode = FailMode.NONE;
+
+    private FailMode mRewindFailMode = FailMode.NONE;
+
+    private FailMode mLengthFailMode = FailMode.NONE;
+
+    private int mNumReadCalls;
+    private int mNumRewindCalls;
+
+    private int mNextRead;
+    private boolean mStarted;
+    private boolean mReadPending;
+    private boolean mRewindPending;
+    // Used to ensure there are no read/rewind requests after a failure.
+    private boolean mFailed;
+
+    private final AtomicBoolean mClosed = new AtomicBoolean(false);
+    private final ConditionVariable mAwaitingClose = new ConditionVariable(false);
+
+    public TestUploadDataProvider(
+            SuccessCallbackMode successCallbackMode, final Executor executor) {
+        mSuccessCallbackMode = successCallbackMode;
+        mExecutor = executor;
+    }
+
+    // Adds the result to be returned by a successful read request.  The
+    // returned bytes must all fit within the read buffer provided by Cronet.
+    // After a rewind, if there is one, all reads will be repeated.
+    public void addRead(byte[] read) {
+        if (mStarted) {
+            throw new IllegalStateException("Adding bytes after read");
+        }
+        mReads.add(read);
+    }
+
+    public void setReadFailure(int readFailIndex, FailMode readFailMode) {
+        mReadFailIndex = readFailIndex;
+        mReadFailMode = readFailMode;
+    }
+
+    public void setLengthFailure() {
+        mLengthFailMode = FailMode.THROWN;
+    }
+
+    public void setRewindFailure(FailMode rewindFailMode) {
+        mRewindFailMode = rewindFailMode;
+    }
+
+    public void setChunked(boolean chunked) {
+        mChunked = chunked;
+    }
+
+    public int getNumReadCalls() {
+        return mNumReadCalls;
+    }
+
+    public int getNumRewindCalls() {
+        return mNumRewindCalls;
+    }
+
+    /** Returns the cumulative length of all data added by calls to addRead. */
+    @Override
+    public long getLength() throws IOException {
+        if (mClosed.get()) {
+            throw new ClosedChannelException();
+        }
+        if (mLengthFailMode == FailMode.THROWN) {
+            throw new IllegalStateException("Sync length failure");
+        }
+        return getUploadedLength();
+    }
+
+    public long getUploadedLength() {
+        if (mChunked) {
+            return -1;
+        }
+        long length = 0;
+        for (byte[] read : mReads) {
+            length += read.length;
+        }
+        return length;
+    }
+
+    @Override
+    public void read(final UploadDataSink uploadDataSink, final ByteBuffer byteBuffer)
+            throws IOException {
+        int currentReadCall = mNumReadCalls;
+        ++mNumReadCalls;
+        if (mClosed.get()) {
+            throw new ClosedChannelException();
+        }
+        assertIdle();
+
+        if (maybeFailRead(currentReadCall, uploadDataSink)) {
+            mFailed = true;
+            return;
+        }
+
+        mReadPending = true;
+        mStarted = true;
+
+        final boolean finalChunk = (mChunked && mNextRead == mReads.size() - 1);
+        if (mNextRead < mReads.size()) {
+            if ((byteBuffer.limit() - byteBuffer.position()) < mReads.get(mNextRead).length) {
+                throw new IllegalStateException("Read buffer smaller than expected.");
+            }
+            byteBuffer.put(mReads.get(mNextRead));
+            ++mNextRead;
+        } else {
+            throw new IllegalStateException("Too many reads: " + mNextRead);
+        }
+
+        Runnable completeRunnable =
+                new Runnable() {
+                    @Override
+                    public void run() {
+                        mReadPending = false;
+                        uploadDataSink.onReadSucceeded(finalChunk);
+                    }
+                };
+        if (mSuccessCallbackMode == SuccessCallbackMode.SYNC) {
+            completeRunnable.run();
+        } else {
+            mExecutor.execute(completeRunnable);
+        }
+    }
+
+    @Override
+    public void rewind(final UploadDataSink uploadDataSink) throws IOException {
+        ++mNumRewindCalls;
+        if (mClosed.get()) {
+            throw new ClosedChannelException();
+        }
+        assertIdle();
+
+        if (maybeFailRewind(uploadDataSink)) {
+            mFailed = true;
+            return;
+        }
+
+        if (mNextRead == 0) {
+            // Should never try and rewind when rewinding does nothing.
+            throw new IllegalStateException("Unexpected rewind when already at beginning");
+        }
+
+        mRewindPending = true;
+        mNextRead = 0;
+
+        Runnable completeRunnable =
+                new Runnable() {
+                    @Override
+                    public void run() {
+                        mRewindPending = false;
+                        uploadDataSink.onRewindSucceeded();
+                    }
+                };
+        if (mSuccessCallbackMode == SuccessCallbackMode.SYNC) {
+            completeRunnable.run();
+        } else {
+            mExecutor.execute(completeRunnable);
+        }
+    }
+
+    private void assertIdle() {
+        if (mReadPending) {
+            throw new IllegalStateException("Unexpected operation during read");
+        }
+        if (mRewindPending) {
+            throw new IllegalStateException("Unexpected operation during rewind");
+        }
+        if (mFailed) {
+            throw new IllegalStateException("Unexpected operation after failure");
+        }
+    }
+
+    private boolean maybeFailRead(int readIndex, final UploadDataSink uploadDataSink) {
+        if (readIndex != mReadFailIndex) return false;
+
+        switch (mReadFailMode) {
+            case THROWN:
+                throw new IllegalStateException("Thrown read failure");
+            case CALLBACK_SYNC:
+                uploadDataSink.onReadError(new IllegalStateException("Sync read failure"));
+                return true;
+            case CALLBACK_ASYNC:
+                Runnable errorRunnable =
+                        new Runnable() {
+                            @Override
+                            public void run() {
+                                uploadDataSink.onReadError(
+                                        new IllegalStateException("Async read failure"));
+                            }
+                        };
+                mExecutor.execute(errorRunnable);
+                return true;
+            default:
+                return false;
+        }
+    }
+
+    private boolean maybeFailRewind(final UploadDataSink uploadDataSink) {
+        switch (mRewindFailMode) {
+            case THROWN:
+                throw new IllegalStateException("Thrown rewind failure");
+            case CALLBACK_SYNC:
+                uploadDataSink.onRewindError(new IllegalStateException("Sync rewind failure"));
+                return true;
+            case CALLBACK_ASYNC:
+                Runnable errorRunnable =
+                        new Runnable() {
+                            @Override
+                            public void run() {
+                                uploadDataSink.onRewindError(
+                                        new IllegalStateException("Async rewind failure"));
+                            }
+                        };
+                mExecutor.execute(errorRunnable);
+                return true;
+            default:
+                return false;
+        }
+    }
+
+    @Override
+    public void close() throws IOException {
+        if (!mClosed.compareAndSet(false, true)) {
+            throw new AssertionError("Closed twice");
+        }
+        mAwaitingClose.open();
+    }
+
+    public void assertClosed() {
+        mAwaitingClose.block(5000);
+        if (!mClosed.get()) {
+            throw new AssertionError("Was not closed");
+        }
+    }
+}