#include "gtest/gtest.h"

#include "mozilla/SlicedInputStream.h"
#include "mozilla/SpinEventLoopUntil.h"
#include "nsCOMPtr.h"
#include "nsIInputStream.h"
#include "nsIPipe.h"
#include "nsStreamUtils.h"
#include "nsString.h"
#include "nsStringStream.h"
#include "Helpers.h"

using namespace mozilla;

// This helper class is used to call OnInputStreamReady with the right stream
// as argument.
class InputStreamCallback final : public nsIInputStreamCallback {
  nsCOMPtr<nsIAsyncInputStream> mStream;
  nsCOMPtr<nsIInputStreamCallback> mCallback;

 public:
  NS_DECL_THREADSAFE_ISUPPORTS

  InputStreamCallback(nsIAsyncInputStream* aStream,
                      nsIInputStreamCallback* aCallback)
      : mStream(aStream), mCallback(aCallback) {}

  NS_IMETHOD
  OnInputStreamReady(nsIAsyncInputStream* aStream) override {
    return mCallback->OnInputStreamReady(mStream);
  }

 private:
  ~InputStreamCallback() = default;
};

NS_IMPL_ISUPPORTS(InputStreamCallback, nsIInputStreamCallback)

/* We want to ensure that sliced streams work with both seekable and
 * non-seekable input streams.  As our string streams are seekable, we need to
 * provide a string stream that doesn't permit seeking, so we can test the
 * logic that emulates seeking in sliced input streams.
 */
class NonSeekableStringStream final : public nsIAsyncInputStream {
  nsCOMPtr<nsIInputStream> mStream;

 public:
  NS_DECL_THREADSAFE_ISUPPORTS

  explicit NonSeekableStringStream(const nsACString& aBuffer) {
    NS_NewCStringInputStream(getter_AddRefs(mStream), aBuffer);
  }

  explicit NonSeekableStringStream(nsIInputStream* aStream)
      : mStream(aStream) {}

  NS_IMETHOD
  Available(uint64_t* aLength) override { return mStream->Available(aLength); }

  NS_IMETHOD
  Read(char* aBuffer, uint32_t aCount, uint32_t* aReadCount) override {
    return mStream->Read(aBuffer, aCount, aReadCount);
  }

  NS_IMETHOD
  ReadSegments(nsWriteSegmentFun aWriter, void* aClosure, uint32_t aCount,
               uint32_t* aResult) override {
    return mStream->ReadSegments(aWriter, aClosure, aCount, aResult);
  }

  NS_IMETHOD
  Close() override { return mStream->Close(); }

  NS_IMETHOD
  IsNonBlocking(bool* aNonBlocking) override {
    return mStream->IsNonBlocking(aNonBlocking);
  }

  NS_IMETHOD
  CloseWithStatus(nsresult aStatus) override {
    nsCOMPtr<nsIAsyncInputStream> async = do_QueryInterface(mStream);
    if (!async) {
      MOZ_CRASH("This should not happen.");
      return NS_ERROR_FAILURE;
    }

    return async->CloseWithStatus(aStatus);
  }

  NS_IMETHOD
  AsyncWait(nsIInputStreamCallback* aCallback, uint32_t aFlags,
            uint32_t aRequestedCount, nsIEventTarget* aEventTarget) override {
    nsCOMPtr<nsIAsyncInputStream> async = do_QueryInterface(mStream);
    if (!async) {
      MOZ_CRASH("This should not happen.");
      return NS_ERROR_FAILURE;
    }

    RefPtr<InputStreamCallback> callback =
        new InputStreamCallback(this, aCallback);

    return async->AsyncWait(callback, aFlags, aRequestedCount, aEventTarget);
  }

 private:
  ~NonSeekableStringStream() = default;
};

NS_IMPL_ISUPPORTS(NonSeekableStringStream, nsIInputStream, nsIAsyncInputStream)

// Helper function for creating a seekable nsIInputStream + a SlicedInputStream.
static SlicedInputStream* CreateSeekableStreams(uint32_t aSize, uint64_t aStart,
                                                uint64_t aLength,
                                                nsCString& aBuffer) {
  aBuffer.SetLength(aSize);
  for (uint32_t i = 0; i < aSize; ++i) {
    aBuffer.BeginWriting()[i] = i % 10;
  }

  nsCOMPtr<nsIInputStream> stream;
  NS_NewCStringInputStream(getter_AddRefs(stream), aBuffer);
  return new SlicedInputStream(stream.forget(), aStart, aLength);
}

// Helper function for creating a non-seekable nsIInputStream + a
// SlicedInputStream.
static SlicedInputStream* CreateNonSeekableStreams(uint32_t aSize,
                                                   uint64_t aStart,
                                                   uint64_t aLength,
                                                   nsCString& aBuffer) {
  aBuffer.SetLength(aSize);
  for (uint32_t i = 0; i < aSize; ++i) {
    aBuffer.BeginWriting()[i] = i % 10;
  }

  RefPtr<NonSeekableStringStream> stream = new NonSeekableStringStream(aBuffer);
  return new SlicedInputStream(stream.forget(), aStart, aLength);
}

// Same start, same length.
TEST(TestSlicedInputStream, Simple)
{
  const size_t kBufSize = 4096;

  nsCString buf;
  RefPtr<SlicedInputStream> sis =
      CreateSeekableStreams(kBufSize, 0, kBufSize, buf);

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)kBufSize, length);

  char buf2[kBufSize];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ(count, buf.Length());
  ASSERT_TRUE(nsCString(buf.get()).Equals(nsCString(buf2)));
}

// Simple sliced stream - seekable
TEST(TestSlicedInputStream, Sliced)
{
  const size_t kBufSize = 4096;

  nsCString buf;
  RefPtr<SlicedInputStream> sis = CreateSeekableStreams(kBufSize, 10, 100, buf);

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)100, length);

  char buf2[kBufSize / 2];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)100, count);
  ASSERT_TRUE(nsCString(buf.get() + 10, count).Equals(nsCString(buf2, count)));
}

// Simple sliced stream - non seekable
TEST(TestSlicedInputStream, SlicedNoSeek)
{
  const size_t kBufSize = 4096;

  nsCString buf;
  RefPtr<SlicedInputStream> sis =
      CreateNonSeekableStreams(kBufSize, 10, 100, buf);

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)100, length);

  char buf2[kBufSize / 2];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)100, count);
  ASSERT_TRUE(nsCString(buf.get() + 10, count).Equals(nsCString(buf2, count)));
}

// Big inputStream - seekable
TEST(TestSlicedInputStream, BigSliced)
{
  const size_t kBufSize = 4096 * 40;

  nsCString buf;
  RefPtr<SlicedInputStream> sis =
      CreateSeekableStreams(kBufSize, 4096 * 5, 4096 * 10, buf);

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)4096 * 10, length);

  char buf2[kBufSize / 2];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)4096 * 10, count);
  ASSERT_TRUE(
      nsCString(buf.get() + 4096 * 5, count).Equals(nsCString(buf2, count)));
}

// Big inputStream - non seekable
TEST(TestSlicedInputStream, BigSlicedNoSeek)
{
  const size_t kBufSize = 4096 * 40;

  nsCString buf;
  RefPtr<SlicedInputStream> sis =
      CreateNonSeekableStreams(kBufSize, 4096 * 5, 4096 * 10, buf);

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)4096 * 10, length);

  char buf2[kBufSize / 2];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)4096 * 10, count);
  ASSERT_TRUE(
      nsCString(buf.get() + 4096 * 5, count).Equals(nsCString(buf2, count)));
}

// Available size.
TEST(TestSlicedInputStream, Available)
{
  nsCString buf;
  RefPtr<SlicedInputStream> sis =
      CreateNonSeekableStreams(500000, 4, 400000, buf);

  uint64_t toRead = 400000;
  for (uint32_t i = 0; i < 400; ++i) {
    uint64_t length;
    ASSERT_EQ(NS_OK, sis->Available(&length));
    ASSERT_EQ(toRead, length);

    char buf2[1000];
    uint32_t count;
    ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
    ASSERT_EQ((uint64_t)1000, count);
    ASSERT_TRUE(nsCString(buf.get() + 4 + (1000 * i), count)
                    .Equals(nsCString(buf2, count)));

    toRead -= count;
  }

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)0, length);

  char buf2[4096];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)0, count);
}

// What if start is > then the size of the buffer?
TEST(TestSlicedInputStream, StartBiggerThan)
{
  nsCString buf;
  RefPtr<SlicedInputStream> sis = CreateNonSeekableStreams(500, 4000, 1, buf);

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)0, length);

  char buf2[4096];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)0, count);
}

// What if the length is > than the size of the buffer?
TEST(TestSlicedInputStream, LengthBiggerThan)
{
  nsCString buf;
  RefPtr<SlicedInputStream> sis = CreateNonSeekableStreams(500, 0, 500000, buf);

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)500, length);

  char buf2[4096];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)500, count);
}

// What if the length is 0?
TEST(TestSlicedInputStream, Length0)
{
  nsCString buf;
  RefPtr<SlicedInputStream> sis = CreateNonSeekableStreams(500, 0, 0, buf);

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)0, length);

  char buf2[4096];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)0, count);
}

// Seek test NS_SEEK_SET
TEST(TestSlicedInputStream, Seek_SET)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  RefPtr<SlicedInputStream> sis;
  {
    nsCOMPtr<nsIInputStream> stream;
    NS_NewCStringInputStream(getter_AddRefs(stream), buf);
    sis = new SlicedInputStream(stream.forget(), 1, buf.Length());
  }

  ASSERT_EQ(NS_OK, sis->Seek(nsISeekableStream::NS_SEEK_SET, 1));

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)buf.Length() - 2, length);

  char buf2[4096];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)buf.Length() - 2, count);
  ASSERT_EQ(0, strncmp(buf2, "llo world", count));
}

// Seek test NS_SEEK_CUR
TEST(TestSlicedInputStream, Seek_CUR)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  RefPtr<SlicedInputStream> sis;
  {
    nsCOMPtr<nsIInputStream> stream;
    NS_NewCStringInputStream(getter_AddRefs(stream), buf);

    sis = new SlicedInputStream(stream.forget(), 1, buf.Length());
  }

  ASSERT_EQ(NS_OK, sis->Seek(nsISeekableStream::NS_SEEK_CUR, 1));

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)buf.Length() - 2, length);

  char buf2[3];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)3, count);
  ASSERT_EQ(0, strncmp(buf2, "llo", count));

  ASSERT_EQ(NS_OK, sis->Seek(nsISeekableStream::NS_SEEK_CUR, 1));

  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)3, count);
  ASSERT_EQ(0, strncmp(buf2, "wor", count));
}

// Seek test NS_SEEK_END - length > real one
TEST(TestSlicedInputStream, Seek_END_Bigger)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  RefPtr<SlicedInputStream> sis;
  {
    nsCOMPtr<nsIInputStream> stream;
    NS_NewCStringInputStream(getter_AddRefs(stream), buf);

    sis = new SlicedInputStream(stream.forget(), 2, buf.Length());
  }

  ASSERT_EQ(NS_OK, sis->Seek(nsISeekableStream::NS_SEEK_END, -5));

  nsCOMPtr<nsIInputStream> stream;
  NS_NewCStringInputStream(getter_AddRefs(stream), buf);
  nsCOMPtr<nsISeekableStream> seekStream = do_QueryInterface(stream);
  ASSERT_EQ(NS_OK, seekStream->Seek(nsISeekableStream::NS_SEEK_END, -5));

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)5, length);

  ASSERT_EQ(NS_OK, stream->Available(&length));
  ASSERT_EQ((uint64_t)5, length);

  char buf2[5];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)5, count);
  ASSERT_EQ(0, strncmp(buf2, "world", count));

  ASSERT_EQ(NS_OK, stream->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)5, count);
  ASSERT_EQ(0, strncmp(buf2, "world", count));
}

// Seek test NS_SEEK_END - length < real one
TEST(TestSlicedInputStream, Seek_END_Lower)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  RefPtr<SlicedInputStream> sis;
  {
    nsCOMPtr<nsIInputStream> stream;
    NS_NewCStringInputStream(getter_AddRefs(stream), buf);

    sis = new SlicedInputStream(stream.forget(), 2, 6);
  }

  ASSERT_EQ(NS_OK, sis->Seek(nsISeekableStream::NS_SEEK_END, -3));

  uint64_t length;
  ASSERT_EQ(NS_OK, sis->Available(&length));
  ASSERT_EQ((uint64_t)3, length);

  char buf2[5];
  uint32_t count;
  ASSERT_EQ(NS_OK, sis->Read(buf2, sizeof(buf2), &count));
  ASSERT_EQ((uint64_t)3, count);
  ASSERT_EQ(0, strncmp(buf2, " wo", count));
}

// Check the nsIAsyncInputStream interface
TEST(TestSlicedInputStream, NoAsyncInputStream)
{
  const size_t kBufSize = 4096;

  nsCString buf;
  nsCOMPtr<nsIInputStream> sis =
      CreateSeekableStreams(kBufSize, 0, kBufSize, buf);

  // If the stream is not asyncInputStream, also SIS is not.
  nsCOMPtr<nsIAsyncInputStream> async = do_QueryInterface(sis);
  ASSERT_TRUE(!async);
}

TEST(TestSlicedInputStream, AsyncInputStream)
{
  nsCOMPtr<nsIAsyncInputStream> reader;
  nsCOMPtr<nsIAsyncOutputStream> writer;

  const uint32_t segmentSize = 1024;
  const uint32_t numSegments = 1;

  nsresult rv = NS_NewPipe2(getter_AddRefs(reader), getter_AddRefs(writer),
                            true, true,  // non-blocking - reader, writer
                            segmentSize, numSegments);
  ASSERT_TRUE(NS_SUCCEEDED(rv));

  nsTArray<char> inputData;
  testing::CreateData(segmentSize, inputData);

  // We have to wrap the reader because it implements only a partial
  // nsISeekableStream interface. When ::Seek() is called, it does a MOZ_CRASH.
  nsCOMPtr<nsIInputStream> sis;
  {
    RefPtr<NonSeekableStringStream> wrapper =
        new NonSeekableStringStream(reader);

    sis = new SlicedInputStream(wrapper.forget(), 500, 500);
  }

  nsCOMPtr<nsIAsyncInputStream> async = do_QueryInterface(sis);
  ASSERT_TRUE(!!async);

  RefPtr<testing::InputStreamCallback> cb = new testing::InputStreamCallback();

  rv = async->AsyncWait(cb, 0, 0, nullptr);
  ASSERT_TRUE(NS_SUCCEEDED(rv));

  ASSERT_FALSE(cb->Called());

  uint32_t numWritten = 0;
  rv = writer->Write(inputData.Elements(), inputData.Length(), &numWritten);
  ASSERT_TRUE(NS_SUCCEEDED(rv));

  ASSERT_TRUE(cb->Called());

  inputData.RemoveElementsAt(0, 500);
  inputData.RemoveElementsAt(500, 24);

  testing::ConsumeAndValidateStream(async, inputData);
}

TEST(TestSlicedInputStream, QIInputStreamLength)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  for (int i = 0; i < 4; i++) {
    nsCOMPtr<nsIInputStream> sis;
    {
      RefPtr<testing::LengthInputStream> stream =
          new testing::LengthInputStream(buf, i % 2, i > 1);

      sis = new SlicedInputStream(stream.forget(), 0, 5);
    }

    {
      nsCOMPtr<nsIInputStreamLength> qi = do_QueryInterface(sis);
      ASSERT_EQ(!!(i % 2), !!qi);
    }

    {
      nsCOMPtr<nsIAsyncInputStreamLength> qi = do_QueryInterface(sis);
      ASSERT_EQ(i > 1, !!qi);
    }
  }
}

TEST(TestSlicedInputStream, InputStreamLength)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  nsCOMPtr<nsIInputStream> sis;
  {
    RefPtr<testing::LengthInputStream> stream =
        new testing::LengthInputStream(buf, true, false);

    sis = new SlicedInputStream(stream.forget(), 0, 5);
  }

  nsCOMPtr<nsIInputStreamLength> qi = do_QueryInterface(sis);
  ASSERT_TRUE(!!qi);

  int64_t size;
  nsresult rv = qi->Length(&size);
  ASSERT_EQ(NS_OK, rv);
  ASSERT_EQ(5, size);
}

TEST(TestSlicedInputStream, NegativeInputStreamLength)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  nsCOMPtr<nsIInputStream> sis;
  {
    RefPtr<testing::LengthInputStream> stream =
        new testing::LengthInputStream(buf, true, false, NS_OK, true);

    sis = new SlicedInputStream(stream.forget(), 0, 5);
  }

  nsCOMPtr<nsIInputStreamLength> qi = do_QueryInterface(sis);
  ASSERT_TRUE(!!qi);

  int64_t size;
  nsresult rv = qi->Length(&size);
  ASSERT_EQ(NS_OK, rv);
  ASSERT_EQ(-1, size);
}

TEST(TestSlicedInputStream, AsyncInputStreamLength)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  nsCOMPtr<nsIInputStream> sis;
  {
    RefPtr<testing::LengthInputStream> stream =
        new testing::LengthInputStream(buf, false, true);

    sis = new SlicedInputStream(stream.forget(), 0, 5);
  }

  nsCOMPtr<nsIAsyncInputStreamLength> qi = do_QueryInterface(sis);
  ASSERT_TRUE(!!qi);

  RefPtr<testing::LengthCallback> callback = new testing::LengthCallback();

  nsresult rv = qi->AsyncLengthWait(callback, GetCurrentSerialEventTarget());
  ASSERT_EQ(NS_OK, rv);

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil([&]() { return callback->Called(); }));
  ASSERT_EQ(5, callback->Size());
}

TEST(TestSlicedInputStream, NegativeAsyncInputStreamLength)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  nsCOMPtr<nsIInputStream> sis;
  {
    RefPtr<testing::LengthInputStream> stream =
        new testing::LengthInputStream(buf, false, true, NS_OK, true);

    sis = new SlicedInputStream(stream.forget(), 0, 5);
  }

  nsCOMPtr<nsIAsyncInputStreamLength> qi = do_QueryInterface(sis);
  ASSERT_TRUE(!!qi);

  RefPtr<testing::LengthCallback> callback = new testing::LengthCallback();

  nsresult rv = qi->AsyncLengthWait(callback, GetCurrentSerialEventTarget());
  ASSERT_EQ(NS_OK, rv);

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil([&]() { return callback->Called(); }));
  ASSERT_EQ(-1, callback->Size());
}

TEST(TestSlicedInputStream, AbortLengthCallback)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  nsCOMPtr<nsIInputStream> sis;
  {
    RefPtr<testing::LengthInputStream> stream =
        new testing::LengthInputStream(buf, false, true, NS_OK, true);

    sis = new SlicedInputStream(stream.forget(), 0, 5);
  }

  nsCOMPtr<nsIAsyncInputStreamLength> qi = do_QueryInterface(sis);
  ASSERT_TRUE(!!qi);

  RefPtr<testing::LengthCallback> callback1 = new testing::LengthCallback();
  nsresult rv = qi->AsyncLengthWait(callback1, GetCurrentSerialEventTarget());
  ASSERT_EQ(NS_OK, rv);

  RefPtr<testing::LengthCallback> callback2 = new testing::LengthCallback();
  rv = qi->AsyncLengthWait(callback2, GetCurrentSerialEventTarget());
  ASSERT_EQ(NS_OK, rv);

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil([&]() { return callback2->Called(); }));
  ASSERT_TRUE(!callback1->Called());
  ASSERT_EQ(-1, callback2->Size());
}