Implment BytesConsumer::tee

third_party/WebKit/Source/modules/fetch/BytesConsumerTest.cpp

+// Copyright 2016 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "modules/fetch/BytesConsumer.h"
+
+#include "core/testing/DummyPageHolder.h"
+#include "modules/fetch/BytesConsumerTestUtil.h"
+#include "platform/blob/BlobData.h"
+#include "platform/testing/UnitTestHelpers.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "wtf/RefPtr.h"
+
+namespace blink {
+
+namespace {
+
+using Command = BytesConsumerTestUtil::Command;
+using Result = BytesConsumer::Result;
+using ReplayingBytesConsumer = BytesConsumerTestUtil::ReplayingBytesConsumer;
+
+String toString(const Vector<char>& v)
+{
+    return String(v.data(), v.size());
+}
+
+class TestClient final : public GarbageCollectedFinalized<TestClient>, public BytesConsumer::Client {
+    USING_GARBAGE_COLLECTED_MIXIN(TestClient);
+public:
+    void onStateChange() override { ++m_numOnStateChangeCalled; }
+    int numOnStateChangeCalled() const { return m_numOnStateChangeCalled; }
+
+private:
+    int m_numOnStateChangeCalled = 0;
+};
+
+class BytesConsumerTeeTest : public ::testing::Test {
+public:
+    BytesConsumerTeeTest()
+        : m_page(DummyPageHolder::create())
+    {
+    }
+
+    Document* document() { return &m_page->document(); }
+
+private:
+    std::unique_ptr<DummyPageHolder> m_page;
+};
+
+class FakeBlobBytesConsumer : public BytesConsumer {
+public:
+    explicit FakeBlobBytesConsumer(PassRefPtr<BlobDataHandle> handle) : m_blobHandle(handle) {}
+    ~FakeBlobBytesConsumer() override {}
+
+    Result beginRead(const char** buffer, size_t* available) override
+    {
+        if (m_state == PublicState::Closed)
+            return Result::Done;
+        m_blobHandle = nullptr;
+        m_state = PublicState::Errored;
+        return Result::Error;
+    }
+    Result endRead(size_t readSize) override
+    {
+        if (m_state == PublicState::Closed)
+            return Result::Error;
+        m_blobHandle = nullptr;
+        m_state = PublicState::Errored;
+        return Result::Error;
+    }
+    PassRefPtr<BlobDataHandle> drainAsBlobDataHandle(BlobSizePolicy policy)
+    {
+        if (m_state != PublicState::ReadableOrWaiting)
+            return nullptr;
+        DCHECK(m_blobHandle);
+        if (policy == BlobSizePolicy::DisallowBlobWithInvalidSize && m_blobHandle->size() == UINT64_MAX)
+            return nullptr;
+        m_state = PublicState::Closed;
+        return m_blobHandle.release();
+    }
+
+    void setClient(Client*) override {}
+    void clearClient() override {}
+    void cancel() override {}
+    PublicState getPublicState() const override { return m_state; }
+    Error getError() const override { return Error(); }
+    String debugName() const override { return "FakeBlobBytesConsumer"; }
+
+private:
+    PublicState m_state = PublicState::ReadableOrWaiting;
+    RefPtr<BlobDataHandle> m_blobHandle;
+};
+
+TEST_F(BytesConsumerTeeTest, CreateDone)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+    src->add(Command(Command::Done));
+    EXPECT_FALSE(src->isCancelled());
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    auto result1 = (new BytesConsumerTestUtil::Reader(dest1))->run();
+    auto result2 = (new BytesConsumerTestUtil::Reader(dest2))->run();
+
+    EXPECT_EQ(Result::Done, result1.first);
+    EXPECT_TRUE(result1.second.isEmpty());
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(Result::Done, result2.first);
+    EXPECT_TRUE(result2.second.isEmpty());
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest2->getPublicState());
+    EXPECT_FALSE(src->isCancelled());
+
+    // Cancelling does nothing when closed.
+    dest1->cancel();
+    dest2->cancel();
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest2->getPublicState());
+    EXPECT_FALSE(src->isCancelled());
+}
+
+TEST_F(BytesConsumerTeeTest, Read)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Data, "hello, "));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Data, "world"));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Done));
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest2->getPublicState());
+
+    auto result1 = (new BytesConsumerTestUtil::Reader(dest1))->run();
+    auto result2 = (new BytesConsumerTestUtil::Reader(dest2))->run();
+
+    EXPECT_EQ(Result::Done, result1.first);
+    EXPECT_EQ("hello, world", toString(result1.second));
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(Result::Done, result2.first);
+    EXPECT_EQ("hello, world", toString(result2.second));
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest2->getPublicState());
+    EXPECT_FALSE(src->isCancelled());
+}
+
+TEST_F(BytesConsumerTeeTest, TwoPhaseRead)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Data, "hello, "));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Data, "world"));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Done));
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest2->getPublicState());
+
+    auto result1 = (new BytesConsumerTestUtil::TwoPhaseReader(dest1))->run();
+    auto result2 = (new BytesConsumerTestUtil::TwoPhaseReader(dest2))->run();
+
+    EXPECT_EQ(Result::Done, result1.first);
+    EXPECT_EQ("hello, world", toString(result1.second));
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(Result::Done, result2.first);
+    EXPECT_EQ("hello, world", toString(result2.second));
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest2->getPublicState());
+    EXPECT_FALSE(src->isCancelled());
+}
+
+TEST_F(BytesConsumerTeeTest, Error)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+
+    src->add(Command(Command::Data, "hello, "));
+    src->add(Command(Command::Data, "world"));
+    src->add(Command(Command::Error));
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    EXPECT_EQ(BytesConsumer::PublicState::Errored, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::Errored, dest2->getPublicState());
+
+    auto result1 = (new BytesConsumerTestUtil::TwoPhaseReader(dest1))->run();
+    auto result2 = (new BytesConsumerTestUtil::TwoPhaseReader(dest2))->run();
+
+    EXPECT_EQ(Result::Error, result1.first);
+    EXPECT_TRUE(result1.second.isEmpty());
+    EXPECT_EQ(BytesConsumer::PublicState::Errored, dest1->getPublicState());
+    EXPECT_EQ(Result::Error, result2.first);
+    EXPECT_TRUE(result2.second.isEmpty());
+    EXPECT_EQ(BytesConsumer::PublicState::Errored, dest2->getPublicState());
+    EXPECT_FALSE(src->isCancelled());
+
+    // Cancelling does nothing when errored.
+    dest1->cancel();
+    dest2->cancel();
+    EXPECT_EQ(BytesConsumer::PublicState::Errored, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::Errored, dest2->getPublicState());
+    EXPECT_FALSE(src->isCancelled());
+}
+
+TEST_F(BytesConsumerTeeTest, Cancel)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+
+    src->add(Command(Command::Data, "hello, "));
+    src->add(Command(Command::Wait));
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest2->getPublicState());
+
+    EXPECT_FALSE(src->isCancelled());
+    dest1->cancel();
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest2->getPublicState());
+    EXPECT_FALSE(src->isCancelled());
+    dest2->cancel();
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest2->getPublicState());
+    EXPECT_TRUE(src->isCancelled());
+}
+
+TEST_F(BytesConsumerTeeTest, CancelShouldNotAffectTheOtherDestination)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+
+    src->add(Command(Command::Data, "hello, "));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Data, "world"));
+    src->add(Command(Command::Done));
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest2->getPublicState());
+
+    EXPECT_FALSE(src->isCancelled());
+    dest1->cancel();
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest2->getPublicState());
+    EXPECT_FALSE(src->isCancelled());
+
+    auto result2 = (new BytesConsumerTestUtil::TwoPhaseReader(dest2))->run();
+
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest2->getPublicState());
+    EXPECT_EQ(Result::Done, result2.first);
+    EXPECT_EQ("hello, world", toString(result2.second));
+    EXPECT_FALSE(src->isCancelled());
+}
+
+TEST_F(BytesConsumerTeeTest, CancelShouldNotAffectTheOtherDestination2)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+
+    src->add(Command(Command::Data, "hello, "));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Data, "world"));
+    src->add(Command(Command::Error));
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest2->getPublicState());
+
+    EXPECT_FALSE(src->isCancelled());
+    dest1->cancel();
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest2->getPublicState());
+    EXPECT_FALSE(src->isCancelled());
+
+    auto result2 = (new BytesConsumerTestUtil::TwoPhaseReader(dest2))->run();
+
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    EXPECT_EQ(BytesConsumer::PublicState::Errored, dest2->getPublicState());
+    EXPECT_EQ(Result::Error, result2.first);
+    EXPECT_FALSE(src->isCancelled());
+}
+
+TEST_F(BytesConsumerTeeTest, BlobHandle)
+{
+    RefPtr<BlobDataHandle> blobDataHandle = BlobDataHandle::create(BlobData::create(), 12345);
+    BytesConsumer* src = new FakeBlobBytesConsumer(blobDataHandle);
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    RefPtr<BlobDataHandle> destBlobDataHandle1 = dest1->drainAsBlobDataHandle(BytesConsumer::BlobSizePolicy::AllowBlobWithInvalidSize);
+    RefPtr<BlobDataHandle> destBlobDataHandle2 = dest2->drainAsBlobDataHandle(BytesConsumer::BlobSizePolicy::DisallowBlobWithInvalidSize);
+    ASSERT_TRUE(destBlobDataHandle1);
+    ASSERT_TRUE(destBlobDataHandle2);
+    EXPECT_EQ(12345u, destBlobDataHandle1->size());
+    EXPECT_EQ(12345u, destBlobDataHandle2->size());
+}
+
+TEST_F(BytesConsumerTeeTest, BlobHandleWithInvalidSize)
+{
+    RefPtr<BlobDataHandle> blobDataHandle = BlobDataHandle::create(BlobData::create(), -1);
+    BytesConsumer* src = new FakeBlobBytesConsumer(blobDataHandle);
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    RefPtr<BlobDataHandle> destBlobDataHandle1 = dest1->drainAsBlobDataHandle(BytesConsumer::BlobSizePolicy::AllowBlobWithInvalidSize);
+    RefPtr<BlobDataHandle> destBlobDataHandle2 = dest2->drainAsBlobDataHandle(BytesConsumer::BlobSizePolicy::DisallowBlobWithInvalidSize);
+    ASSERT_TRUE(destBlobDataHandle1);
+    ASSERT_FALSE(destBlobDataHandle2);
+    EXPECT_EQ(UINT64_MAX, destBlobDataHandle1->size());
+}
+
+TEST_F(BytesConsumerTeeTest, ConsumerCanBeErroredInTwoPhaseRead)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+    src->add(Command(Command::Data, "a"));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Error));
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+    TestClient* client = new TestClient();
+    dest1->setClient(client);
+
+    const char* buffer = nullptr;
+    size_t available = 0;
+    ASSERT_EQ(Result::Ok, dest1->beginRead(&buffer, &available));
+    ASSERT_EQ(1u, available);
+
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest1->getPublicState());
+    int numOnStateChangeCalled = client->numOnStateChangeCalled();
+    EXPECT_EQ(Result::Error, (new BytesConsumerTestUtil::Reader(dest2))->run().first);
+    EXPECT_EQ(BytesConsumer::PublicState::Errored, dest1->getPublicState());
+    EXPECT_EQ(numOnStateChangeCalled + 1, client->numOnStateChangeCalled());
+    EXPECT_EQ('a', buffer[0]);
+    EXPECT_EQ(Result::Ok, dest1->endRead(available));
+}
+
+TEST_F(BytesConsumerTeeTest, AsyncNotificationShouldBeDispatchedWhenAllDataIsConsumed)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+    src->add(Command(Command::Data, "a"));
+    src->add(Command(Command::Wait));
+    src->add(Command(Command::Done));
+    TestClient* client = new TestClient();
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    dest1->setClient(client);
+
+    const char* buffer = nullptr;
+    size_t available = 0;
+    ASSERT_EQ(Result::Ok, dest1->beginRead(&buffer, &available));
+    ASSERT_EQ(1u, available);
+    EXPECT_EQ('a', buffer[0]);
+
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, src->getPublicState());
+    testing::runPendingTasks();
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, src->getPublicState());
+    // Just for checking UAF.
+    EXPECT_EQ('a', buffer[0]);
+    ASSERT_EQ(Result::Ok, dest1->endRead(1));
+
+    EXPECT_EQ(0, client->numOnStateChangeCalled());
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest1->getPublicState());
+    testing::runPendingTasks();
+    EXPECT_EQ(1, client->numOnStateChangeCalled());
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+}
+
+TEST_F(BytesConsumerTeeTest, AsyncCloseNotificationShouldBeCancelledBySubsequentReadCall)
+{
+    ReplayingBytesConsumer* src = new ReplayingBytesConsumer(document());
+    src->add(Command(Command::Data, "a"));
+    src->add(Command(Command::Done));
+    TestClient* client = new TestClient();
+
+    BytesConsumer* dest1 = nullptr;
+    BytesConsumer* dest2 = nullptr;
+    BytesConsumer::tee(document(), src, &dest1, &dest2);
+
+    dest1->setClient(client);
+
+    const char* buffer = nullptr;
+    size_t available = 0;
+    ASSERT_EQ(Result::Ok, dest1->beginRead(&buffer, &available));
+    ASSERT_EQ(1u, available);
+    EXPECT_EQ('a', buffer[0]);
+
+    testing::runPendingTasks();
+    // Just for checking UAF.
+    EXPECT_EQ('a', buffer[0]);
+    ASSERT_EQ(Result::Ok, dest1->endRead(1));
+    EXPECT_EQ(BytesConsumer::PublicState::ReadableOrWaiting, dest1->getPublicState());
+
+    EXPECT_EQ(Result::Done, dest1->beginRead(&buffer, &available));
+    EXPECT_EQ(0, client->numOnStateChangeCalled());
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+    testing::runPendingTasks();
+    EXPECT_EQ(0, client->numOnStateChangeCalled());
+    EXPECT_EQ(BytesConsumer::PublicState::Closed, dest1->getPublicState());
+}
+
+} // namespace
+
+} // namespace blink