/* -*- Mode: C; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/* This Source Code Form is subject to the terms of the Mozilla Public

 * License, v. 2.0. If a copy of the MPL was not distributed with this

 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "mozilla/Attributes.h"

#include "mozilla/Base64.h"

#include "nsIScriptableBase64Encoder.h"

#include "nsIInputStream.h"

#include "nsString.h"

#include "gtest/gtest.h"

struct Chunk {

  Chunk(uint32_t l, const char* c)

    : mLength(l), mData(c)

  {}

  uint32_t mLength;

  const char* mData;

};

struct Test {

  Test(Chunk* c, const char* r)

    : mChunks(c), mResult(r)

  {}

  Chunk* mChunks;

  const char* mResult;

};

static Chunk kTest1Chunks[] =

{

   Chunk(9, "Hello sir"),

   Chunk(0, nullptr)

};

static Chunk kTest2Chunks[] =

{

   Chunk(3, "Hel"),

   Chunk(3, "lo "),

   Chunk(3, "sir"),

   Chunk(0, nullptr)

};

static Chunk kTest3Chunks[] =

{

   Chunk(1, "I"),

   Chunk(0, nullptr)

};

static Chunk kTest4Chunks[] =

{

   Chunk(2, "Hi"),

   Chunk(0, nullptr)

};

static Chunk kTest5Chunks[] =

{

   Chunk(1, "B"),

   Chunk(2, "ob"),

   Chunk(0, nullptr)

};

static Chunk kTest6Chunks[] =

{

   Chunk(2, "Bo"),

   Chunk(1, "b"),

   Chunk(0, nullptr)

};

static Chunk kTest7Chunks[] =

{

   Chunk(1, "F"),    // Carry over 1

   Chunk(4, "iref"), // Carry over 2

   Chunk(2, "ox"),   //            1

   Chunk(4, " is "), //            2

   Chunk(2, "aw"),   //            1

   Chunk(4, "esom"), //            2

   Chunk(2, "e!"),

   Chunk(0, nullptr)

};

static Chunk kTest8Chunks[] =

{

   Chunk(5, "ALL T"),

   Chunk(1, "H"),

   Chunk(4, "ESE "),

   Chunk(2, "WO"),

   Chunk(21, "RLDS ARE YOURS EXCEPT"),

   Chunk(9, " EUROPA. "),

   Chunk(25, "ATTEMPT NO LANDING THERE."),

   Chunk(0, nullptr)

};

static Test kTests[] =

  {

    // Test 1, test a simple round string in one chunk

    Test(

      kTest1Chunks,

      "SGVsbG8gc2ly"

    ),

    // Test 2, test a simple round string split into round chunks

    Test(

      kTest2Chunks,

      "SGVsbG8gc2ly"

    ),

    // Test 3, test a single chunk that's 2 short

    Test(

      kTest3Chunks,

      "SQ=="

    ),

    // Test 4, test a single chunk that's 1 short

    Test(

      kTest4Chunks,

      "SGk="

    ),

    // Test 5, test a single chunk that's 2 short, followed by a chunk of 2

    Test(

      kTest5Chunks,

      "Qm9i"

    ),

    // Test 6, test a single chunk that's 1 short, followed by a chunk of 1

    Test(

      kTest6Chunks,

      "Qm9i"

    ),

    // Test 7, test alternating carryovers

    Test(

      kTest7Chunks,

      "RmlyZWZveCBpcyBhd2Vzb21lIQ=="

    ),

    // Test 8, test a longish string

    Test(

      kTest8Chunks,

      "QUxMIFRIRVNFIFdPUkxEUyBBUkUgWU9VUlMgRVhDRVBUIEVVUk9QQS4gQVRURU1QVCBOTyBMQU5ESU5HIFRIRVJFLg=="

    ),

    // Terminator

    Test(

      nullptr,

      nullptr

    )

  };

class FakeInputStream final : public nsIInputStream

{

  ~FakeInputStream() {}

public:

  FakeInputStream()

  : mTestNumber(0),

    mTest(&kTests[0]),

    mChunk(&mTest->mChunks[0]),

    mClosed(false)

  {}

  NS_DECL_ISUPPORTS

  NS_DECL_NSIINPUTSTREAM

  void Reset();

  bool NextTest();

  void CheckTest(nsACString& aResult);

  void CheckTest(nsAString& aResult);

private:

  uint32_t mTestNumber;

  const Test* mTest;

  const Chunk* mChunk;

  bool mClosed;

};

NS_IMPL_ISUPPORTS(FakeInputStream, nsIInputStream)

NS_IMETHODIMP

FakeInputStream::Close()

{

  mClosed = true;

  return NS_OK;

}

NS_IMETHODIMP

FakeInputStream::Available(uint64_t* aAvailable)

{

  *aAvailable = 0;

  if (mClosed)

    return NS_BASE_STREAM_CLOSED;

  const Chunk* chunk = mChunk;

  while (chunk->mLength) {

    *aAvailable += chunk->mLength;

    chunk++;

  }

  return NS_OK;

}

NS_IMETHODIMP

FakeInputStream::Read(char* aBuffer, uint32_t aCount, uint32_t* aOut)

{

  return NS_ERROR_NOT_IMPLEMENTED;

}

NS_IMETHODIMP

FakeInputStream::ReadSegments(nsWriteSegmentFun aWriter,

                              void* aClosure,

                              uint32_t aCount,

                              uint32_t* aRead)

{

  *aRead = 0;

  if (mClosed)

    return NS_BASE_STREAM_CLOSED;

  while (mChunk->mLength) {

    uint32_t written = 0;

    nsresult rv = (*aWriter)(this, aClosure, mChunk->mData,

                             *aRead, mChunk->mLength, &written);

    *aRead += written;

    NS_ENSURE_SUCCESS(rv, rv);

    mChunk++;

  }

  return NS_OK;

}

NS_IMETHODIMP

FakeInputStream::IsNonBlocking(bool* aIsBlocking)

{

  *aIsBlocking = false;

  return NS_OK;

}

void

FakeInputStream::Reset()

{

  mClosed = false;

  mChunk = &mTest->mChunks[0];

}

bool

FakeInputStream::NextTest()

{

  mTestNumber++;

  mTest = &kTests[mTestNumber];

  mChunk = &mTest->mChunks[0];

  mClosed = false;

  return mTest->mChunks ? true : false;

}

void

FakeInputStream::CheckTest(nsACString& aResult)

{

  ASSERT_STREQ(aResult.BeginReading(), mTest->mResult);

}

void

FakeInputStream::CheckTest(nsAString& aResult)

{

  ASSERT_TRUE(aResult.EqualsASCII(mTest->mResult)) <<

    "Actual:   " << aResult.BeginReading() << std::endl <<

    "Expected: " << mTest->mResult;

}

TEST(Base64, StreamEncoder)

{

  nsCOMPtr<nsIScriptableBase64Encoder> encoder =

    do_CreateInstance("@mozilla.org/scriptablebase64encoder;1");

  ASSERT_TRUE(encoder);

  RefPtr<FakeInputStream> stream = new FakeInputStream();

  do {

    nsString wideString;

    nsCString string;

    nsresult rv;

    rv = encoder->EncodeToString(stream, 0, wideString);

    ASSERT_TRUE(NS_SUCCEEDED(rv));

    stream->Reset();

    rv = encoder->EncodeToCString(stream, 0, string);

    ASSERT_TRUE(NS_SUCCEEDED(rv));

    stream->CheckTest(wideString);

    stream->CheckTest(string);

  } while (stream->NextTest());

}

struct EncodeDecodeTestCase

{

    const char* mInput;

    const char* mOutput;

};

static EncodeDecodeTestCase sRFC4648TestCases[] = {

   { "", "" },

   { "f", "Zg==" },

   { "fo", "Zm8=" },

   { "foo", "Zm9v" },

   { "foob", "Zm9vYg==" },

   { "fooba", "Zm9vYmE=" },

   { "foobar", "Zm9vYmFy" },

};

TEST(Base64, RFC4648Encoding)

{

    for (auto& testcase : sRFC4648TestCases) {

        nsDependentCString in(testcase.mInput);

        nsAutoCString out;

        nsresult rv = mozilla::Base64Encode(in, out);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(out.EqualsASCII(testcase.mOutput));

    }

    for (auto& testcase : sRFC4648TestCases) {

        NS_ConvertUTF8toUTF16 in(testcase.mInput);

        nsAutoString out;

        nsresult rv = mozilla::Base64Encode(in, out);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(out.EqualsASCII(testcase.mOutput));

    }

}

TEST(Base64, RFC4648Decoding)

{

    for (auto& testcase : sRFC4648TestCases) {

        nsDependentCString out(testcase.mOutput);

        nsAutoCString in;

        nsresult rv = mozilla::Base64Decode(out, in);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(in.EqualsASCII(testcase.mInput));

    }

    for (auto& testcase : sRFC4648TestCases) {

        NS_ConvertUTF8toUTF16 out(testcase.mOutput);

        nsAutoString in;

        nsresult rv = mozilla::Base64Decode(out, in);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(in.EqualsASCII(testcase.mInput));

    }

}

TEST(Base64, RFC4648DecodingRawPointers)

{

    for (auto& testcase : sRFC4648TestCases) {

        size_t outputLength = strlen(testcase.mOutput);

        size_t inputLength = strlen(testcase.mInput);

        // This will be allocated by Base64Decode.

        char* buffer = nullptr;

        uint32_t binaryLength = 0;

        nsresult rv = mozilla::Base64Decode(testcase.mOutput, outputLength,

                                            &buffer, &binaryLength);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_EQ(binaryLength, inputLength);

        ASSERT_STREQ(testcase.mInput, buffer);

        free(buffer);

    }

}

static EncodeDecodeTestCase sNonASCIITestCases[] = {

    { "\x80", "gA==" },

    { "\xff", "/w==" },

    { "\x80\x80", "gIA=" },

    { "\x80\x81", "gIE=" },

    { "\xff\xff", "//8=" },

    { "\x80\x80\x80", "gICA" },

    { "\xff\xff\xff", "////" },

    { "\x80\x80\x80\x80", "gICAgA==" },

    { "\xff\xff\xff\xff", "/////w==" },

};

TEST(Base64, NonASCIIEncoding)

{

    for (auto& testcase : sNonASCIITestCases) {

        nsDependentCString in(testcase.mInput);

        nsAutoCString out;

        nsresult rv = mozilla::Base64Encode(in, out);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(out.EqualsASCII(testcase.mOutput));

    }

}

TEST(Base64, NonASCIIEncodingWideString)

{

    for (auto& testcase : sNonASCIITestCases) {

        nsAutoString in, out;

        // XXX Handles Latin1 despite the name

        AppendASCIItoUTF16(nsDependentCString(testcase.mInput), in);

        nsresult rv = mozilla::Base64Encode(in, out);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(out.EqualsASCII(testcase.mOutput));

    }

}

TEST(Base64, NonASCIIDecoding)

{

    for (auto& testcase : sNonASCIITestCases) {

        nsDependentCString out(testcase.mOutput);

        nsAutoCString in;

        nsresult rv = mozilla::Base64Decode(out, in);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(in.Equals(testcase.mInput));

    }

}

TEST(Base64, NonASCIIDecodingWideString)

{

    for (auto& testcase : sNonASCIITestCases) {

        nsAutoString in, out;

        // XXX Handles Latin1 despite the name

        AppendASCIItoUTF16(nsDependentCString(testcase.mOutput), out);

        nsresult rv = mozilla::Base64Decode(out, in);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        // Can't use EqualsASCII, because our comparison string isn't ASCII.

        for (size_t i = 0; i < in.Length(); ++i) {

            ASSERT_TRUE(((unsigned int)in[i] & 0xff00) == 0);

            ASSERT_EQ((unsigned char)in[i], (unsigned char)testcase.mInput[i]);

        }

        ASSERT_TRUE(strlen(testcase.mInput) == in.Length());

    }

}

// For historical reasons, our wide string base64 encode routines mask off

// the high bits of non-latin1 wide strings.

TEST(Base64, EncodeNon8BitWideString)

{

    {

        const nsAutoString non8Bit(u"\x1ff");

        nsAutoString out;

        nsresult rv = mozilla::Base64Encode(non8Bit, out);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(out.EqualsLiteral("/w=="));

    }

    {

        const nsAutoString non8Bit(u"\xfff");

        nsAutoString out;

        nsresult rv = mozilla::Base64Encode(non8Bit, out);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(out.EqualsLiteral("/w=="));

    }

}

// For historical reasons, our wide string base64 decode routines mask off

// the high bits of non-latin1 wide strings.

TEST(Base64, DecodeNon8BitWideString)

{

    {

        // This would be "/w==" in a nsCString

        const nsAutoString non8Bit(u"\x12f\x177==");

        const nsAutoString expectedOutput(u"\xff");

        ASSERT_EQ(non8Bit.Length(), 4u);

        nsAutoString out;

        nsresult rv = mozilla::Base64Decode(non8Bit, out);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(out.Equals(expectedOutput));

    }

    {

        const nsAutoString non8Bit(u"\xf2f\xf77==");

        const nsAutoString expectedOutput(u"\xff");

        nsAutoString out;

        nsresult rv = mozilla::Base64Decode(non8Bit, out);

        ASSERT_TRUE(NS_SUCCEEDED(rv));

        ASSERT_TRUE(out.Equals(expectedOutput));

    }

}

TEST(Base64, TruncateOnInvalidDecodeCString)

{

    NS_NAMED_LITERAL_CSTRING(invalid, "@@==");

    nsAutoCString out("I should be truncated!");

    nsresult rv = mozilla::Base64Decode(invalid, out);

    ASSERT_TRUE(NS_FAILED(rv));

    ASSERT_EQ(out.Length(), 0u);

}

TEST(Base64, TruncateOnInvalidDecodeWideString)

{

    NS_NAMED_LITERAL_STRING(invalid, "@@==");

    nsAutoString out(u"I should be truncated!");

    nsresult rv = mozilla::Base64Decode(invalid, out);

    ASSERT_TRUE(NS_FAILED(rv));

    ASSERT_EQ(out.Length(), 0u);

}

// TODO: Add tests for OOM handling.