/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim: set ts=8 sts=2 et sw=2 tw=80: */
/* 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 "InputStreamLengthHelper.h"
#include "mozilla/dom/WorkerCommon.h"
#include "nsIAsyncInputStream.h"
#include "nsIInputStream.h"
#include "nsIStreamTransportService.h"

static NS_DEFINE_CID(kStreamTransportServiceCID, NS_STREAMTRANSPORTSERVICE_CID);

namespace mozilla {

namespace {

class AvailableEvent final : public Runnable
{
public:
  AvailableEvent(nsIInputStream* stream,
                 const std::function<void(int64_t aLength)>& aCallback)
    : Runnable("mozilla::AvailableEvent")
    , mStream(stream)
    , mCallback(aCallback)
    , mSize(-1)
  {
    mCallbackTarget = GetCurrentThreadSerialEventTarget();
    MOZ_ASSERT(NS_IsMainThread());
  }

  NS_IMETHOD
  Run() override
  {
    // ping
    if (!NS_IsMainThread()) {
      uint64_t size = 0;
      if (NS_WARN_IF(NS_FAILED(mStream->Available(&size)))) {
        mSize = -1;
      } else {
        mSize = (int64_t)size;
      }

      mStream = nullptr;

      nsCOMPtr<nsIRunnable> self(this); // overly cute
      mCallbackTarget->Dispatch(self.forget(), NS_DISPATCH_NORMAL);
      mCallbackTarget = nullptr;
      return NS_OK;
    }

    // pong
    std::function<void(int64_t aLength)> callback;
    callback.swap(mCallback);
    callback(mSize);
    return NS_OK;
  }

private:
  nsCOMPtr<nsIInputStream> mStream;
  std::function<void(int64_t aLength)> mCallback;
  nsCOMPtr<nsIEventTarget> mCallbackTarget;

  int64_t mSize;
};

} // anonymous

/* static */ bool
InputStreamLengthHelper::GetSyncLength(nsIInputStream* aStream,
                                       int64_t* aLength)
{
  MOZ_ASSERT(aStream);
  MOZ_ASSERT(aLength);

  *aLength = -1;

  // Sync length access.
  nsCOMPtr<nsIInputStreamLength> streamLength = do_QueryInterface(aStream);
  if (streamLength) {
    int64_t length = -1;
    nsresult rv = streamLength->Length(&length);

    // All good!
    if (NS_SUCCEEDED(rv)) {
      *aLength = length;
      return true;
    }

    // Already closed stream or an error occurred.
    if (rv == NS_BASE_STREAM_CLOSED ||
        NS_WARN_IF(rv == NS_ERROR_NOT_AVAILABLE) ||
        NS_WARN_IF(rv != NS_BASE_STREAM_WOULD_BLOCK)) {
      return true;
    }
  }

  nsCOMPtr<nsIAsyncInputStreamLength> asyncStreamLength =
    do_QueryInterface(aStream);
  if (asyncStreamLength) {
    // GetAsyncLength should be used.
    return false;
  }

  // We cannot calculate the length of an async stream.
  nsCOMPtr<nsIAsyncInputStream> asyncStream = do_QueryInterface(aStream);
  if (asyncStream) {
    return false;
  }

  // For main-thread only, we want to avoid calling ::Available() for blocking
  // streams.
  if (NS_IsMainThread()) {
    bool nonBlocking = false;
    if (NS_WARN_IF(NS_FAILED(aStream->IsNonBlocking(&nonBlocking)))) {
      // Let's return -1. There is nothing else we can do here.
      return true;
    }

    if (!nonBlocking) {
      return false;
    }
  }

  // Fallback using available().
  uint64_t available = 0;
  nsresult rv = aStream->Available(&available);
  if (NS_WARN_IF(NS_FAILED(rv))) {
    // Let's return -1. There is nothing else we can do here.
    return true;
  }

  *aLength = (int64_t)available;
  return true;
}

/* static */ void
InputStreamLengthHelper::GetAsyncLength(nsIInputStream* aStream,
                                        const std::function<void(int64_t aLength)>& aCallback)
{
  MOZ_ASSERT(aStream);
  MOZ_ASSERT(aCallback);

  // We don't want to allow this class to be used on workers because we are not
  // using the correct Runnable types.
  MOZ_DIAGNOSTIC_ASSERT(NS_IsMainThread() || !dom::IsCurrentThreadRunningWorker());

  RefPtr<InputStreamLengthHelper> helper =
    new InputStreamLengthHelper(aStream, aCallback);

  // Let's be sure that we don't call ::Available() on main-thread.
  if (NS_IsMainThread()) {
    nsCOMPtr<nsIInputStreamLength> streamLength = do_QueryInterface(aStream);
    nsCOMPtr<nsIAsyncInputStreamLength> asyncStreamLength =
      do_QueryInterface(aStream);
    if (!streamLength && !asyncStreamLength) {
      // We cannot calculate the length of an async stream. We must fix the
      // caller if this happens.
#ifdef DEBUG
      nsCOMPtr<nsIAsyncInputStream> asyncStream = do_QueryInterface(aStream);
      MOZ_DIAGNOSTIC_ASSERT(!asyncStream);
#endif

      bool nonBlocking = false;
      if (NS_SUCCEEDED(aStream->IsNonBlocking(&nonBlocking)) && !nonBlocking) {
        nsCOMPtr<nsIEventTarget> target =
          do_GetService(NS_STREAMTRANSPORTSERVICE_CONTRACTID);
        MOZ_ASSERT(target);

        RefPtr<AvailableEvent> event = new AvailableEvent(aStream, aCallback);
        target->Dispatch(event.forget(), NS_DISPATCH_NORMAL);
        return;
      }
    }
  }

  // Let's go async in order to have similar behaviors for sync and async
  // nsIInputStreamLength implementations.
  GetCurrentThreadSerialEventTarget()->Dispatch(helper, NS_DISPATCH_NORMAL);
}

InputStreamLengthHelper::InputStreamLengthHelper(nsIInputStream* aStream,
                                                 const std::function<void(int64_t aLength)>& aCallback)
  : Runnable("InputStreamLengthHelper")
  , mStream(aStream)
  , mCallback(aCallback)
{
  MOZ_ASSERT(aStream);
  MOZ_ASSERT(aCallback);
}

InputStreamLengthHelper::~InputStreamLengthHelper() = default;

NS_IMETHODIMP
InputStreamLengthHelper::Run()
{
  // Sync length access.
  nsCOMPtr<nsIInputStreamLength> streamLength = do_QueryInterface(mStream);
  if (streamLength) {
    int64_t length = -1;
    nsresult rv = streamLength->Length(&length);

    // All good!
    if (NS_SUCCEEDED(rv)) {
      ExecCallback(length);
      return NS_OK;
    }

    // Already closed stream or an error occurred.
    if (rv == NS_BASE_STREAM_CLOSED ||
        NS_WARN_IF(rv == NS_ERROR_NOT_AVAILABLE) ||
        NS_WARN_IF(rv != NS_BASE_STREAM_WOULD_BLOCK)) {
      ExecCallback(-1);
      return NS_OK;
    }
  }

  // Async length access.
  nsCOMPtr<nsIAsyncInputStreamLength> asyncStreamLength =
    do_QueryInterface(mStream);
  if (asyncStreamLength) {
    nsresult rv =
     asyncStreamLength->AsyncLengthWait(this,
                                        GetCurrentThreadSerialEventTarget());
    if (NS_WARN_IF(NS_FAILED(rv))) {
      ExecCallback(-1);
    }

    return NS_OK;
  }

  // Fallback using available().
  uint64_t available = 0;
  nsresult rv = mStream->Available(&available);
  if (NS_WARN_IF(NS_FAILED(rv))) {
    ExecCallback(-1);
    return NS_OK;
  }

  ExecCallback((int64_t)available);
  return NS_OK;
}

NS_IMETHODIMP
InputStreamLengthHelper::OnInputStreamLengthReady(nsIAsyncInputStreamLength* aStream,
                                                  int64_t aLength)
{
  ExecCallback(aLength);
  return NS_OK;
}

void
InputStreamLengthHelper::ExecCallback(int64_t aLength)
{
  MOZ_ASSERT(mCallback);

  std::function<void(int64_t aLength)> callback;
  callback.swap(mCallback);

  callback(aLength);
}

NS_IMPL_ISUPPORTS_INHERITED(InputStreamLengthHelper, Runnable,
                            nsIInputStreamLengthCallback)

} // mozilla namespace