/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* 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/. */

/*
 * A service that provides methods for synchronously loading a DOM in various
 * ways.
 */

#include "nsSyncLoadService.h"
#include "nsCOMPtr.h"
#include "nsIChannel.h"
#include "nsIChannelEventSink.h"
#include "nsIAsyncVerifyRedirectCallback.h"
#include "nsIInterfaceRequestor.h"
#include "nsIStreamListener.h"
#include "nsIURI.h"
#include "nsString.h"
#include "nsWeakReference.h"
#include "mozilla/dom/Document.h"
#include "nsIPrincipal.h"
#include "nsContentUtils.h"  // for kLoadAsData
#include "nsThreadUtils.h"
#include "nsNetUtil.h"
#include "nsStreamUtils.h"
#include <algorithm>

using namespace mozilla;

using mozilla::net::ReferrerPolicy;

/**
 * This class manages loading a single XML document
 */

class nsSyncLoader : public nsIStreamListener,
                     public nsIChannelEventSink,
                     public nsIInterfaceRequestor,
                     public nsSupportsWeakReference {
 public:
  nsSyncLoader()
      : mLoading(false), mAsyncLoadStatus(NS_ERROR_NOT_INITIALIZED) {}

  NS_DECL_ISUPPORTS

  nsresult LoadDocument(nsIChannel *aChannel, bool aChannelIsSync,
                        bool aForceToXML, ReferrerPolicy aReferrerPolicy,
                        Document **aResult);

  NS_FORWARD_NSISTREAMLISTENER(mListener->)
  NS_DECL_NSIREQUESTOBSERVER

  NS_DECL_NSICHANNELEVENTSINK

  NS_DECL_NSIINTERFACEREQUESTOR

 private:
  virtual ~nsSyncLoader();

  nsresult PushAsyncStream(nsIStreamListener *aListener);
  nsresult PushSyncStream(nsIStreamListener *aListener);

  nsCOMPtr<nsIChannel> mChannel;
  nsCOMPtr<nsIStreamListener> mListener;
  bool mLoading;
  nsresult mAsyncLoadStatus;
};

class nsForceXMLListener : public nsIStreamListener {
  virtual ~nsForceXMLListener();

 public:
  explicit nsForceXMLListener(nsIStreamListener *aListener);

  NS_DECL_ISUPPORTS
  NS_FORWARD_NSISTREAMLISTENER(mListener->)
  NS_DECL_NSIREQUESTOBSERVER

 private:
  nsCOMPtr<nsIStreamListener> mListener;
};

nsForceXMLListener::nsForceXMLListener(nsIStreamListener *aListener)
    : mListener(aListener) {}

nsForceXMLListener::~nsForceXMLListener() {}

NS_IMPL_ISUPPORTS(nsForceXMLListener, nsIStreamListener, nsIRequestObserver)

NS_IMETHODIMP
nsForceXMLListener::OnStartRequest(nsIRequest *aRequest) {
  nsresult status;
  aRequest->GetStatus(&status);
  nsCOMPtr<nsIChannel> channel = do_QueryInterface(aRequest);
  if (channel && NS_SUCCEEDED(status)) {
    channel->SetContentType(NS_LITERAL_CSTRING("text/xml"));
  }

  return mListener->OnStartRequest(aRequest);
}

NS_IMETHODIMP
nsForceXMLListener::OnStopRequest(nsIRequest *aRequest, nsresult aStatusCode) {
  return mListener->OnStopRequest(aRequest, aStatusCode);
}

nsSyncLoader::~nsSyncLoader() {
  if (mLoading && mChannel) {
    mChannel->Cancel(NS_BINDING_ABORTED);
  }
}

NS_IMPL_ISUPPORTS(nsSyncLoader, nsIStreamListener, nsIRequestObserver,
                  nsIChannelEventSink, nsIInterfaceRequestor,
                  nsISupportsWeakReference)

nsresult nsSyncLoader::LoadDocument(nsIChannel *aChannel, bool aChannelIsSync,
                                    bool aForceToXML,
                                    ReferrerPolicy aReferrerPolicy,
                                    Document **aResult) {
  NS_ENSURE_ARG(aChannel);
  NS_ENSURE_ARG_POINTER(aResult);
  *aResult = nullptr;
  nsresult rv = NS_OK;

  mChannel = aChannel;
  nsCOMPtr<nsIHttpChannel> http = do_QueryInterface(mChannel);
  if (http) {
    rv = http->SetRequestHeader(
        NS_LITERAL_CSTRING("Accept"),
        NS_LITERAL_CSTRING(
            "text/xml,application/xml,application/xhtml+xml,*/*;q=0.1"),
        false);
    MOZ_ASSERT(NS_SUCCEEDED(rv));
    nsCOMPtr<nsILoadInfo> loadInfo = aChannel->LoadInfo();
    nsCOMPtr<nsIURI> loaderUri;
    loadInfo->TriggeringPrincipal()->GetURI(getter_AddRefs(loaderUri));
    if (loaderUri) {
      rv = http->SetReferrerWithPolicy(loaderUri, aReferrerPolicy);
      MOZ_ASSERT(NS_SUCCEEDED(rv));
    }
  }

  // Hook us up to listen to redirects and the like.
  // Do this before setting up the cross-site proxy since
  // that installs its own proxies.
  mChannel->SetNotificationCallbacks(this);

  // Get the loadgroup of the channel
  nsCOMPtr<nsILoadGroup> loadGroup;
  rv = aChannel->GetLoadGroup(getter_AddRefs(loadGroup));
  NS_ENSURE_SUCCESS(rv, rv);

  // Create document
  nsCOMPtr<Document> document;
  rv = NS_NewXMLDocument(getter_AddRefs(document));
  NS_ENSURE_SUCCESS(rv, rv);

  // Start the document load. Do this before we attach the load listener
  // since we reset the document which drops all observers.
  nsCOMPtr<nsIStreamListener> listener;
  rv = document->StartDocumentLoad(kLoadAsData, mChannel, loadGroup, nullptr,
                                   getter_AddRefs(listener), true);
  NS_ENSURE_SUCCESS(rv, rv);

  if (aForceToXML) {
    nsCOMPtr<nsIStreamListener> forceListener =
        new nsForceXMLListener(listener);
    listener.swap(forceListener);
  }

  if (aChannelIsSync) {
    rv = PushSyncStream(listener);
  } else {
    rv = PushAsyncStream(listener);
  }

  http = do_QueryInterface(mChannel);
  if (NS_SUCCEEDED(rv) && http) {
    bool succeeded;
    if (NS_FAILED(http->GetRequestSucceeded(&succeeded)) || !succeeded) {
      rv = NS_ERROR_FAILURE;
    }
  }
  mChannel = nullptr;

  // check that the load succeeded
  NS_ENSURE_SUCCESS(rv, rv);

  NS_ENSURE_TRUE(document->GetRootElement(), NS_ERROR_FAILURE);

  document.forget(aResult);

  return NS_OK;
}

nsresult nsSyncLoader::PushAsyncStream(nsIStreamListener *aListener) {
  mListener = aListener;

  mAsyncLoadStatus = NS_OK;

  // Start reading from the channel
  nsresult rv = mChannel->AsyncOpen(this);

  if (NS_SUCCEEDED(rv)) {
    // process events until we're finished.
    mLoading = true;
    nsIThread *thread = NS_GetCurrentThread();
    while (mLoading && NS_SUCCEEDED(rv)) {
      bool processedEvent;
      rv = thread->ProcessNextEvent(true, &processedEvent);
      if (NS_SUCCEEDED(rv) && !processedEvent) rv = NS_ERROR_UNEXPECTED;
    }
  }

  mListener = nullptr;

  NS_ENSURE_SUCCESS(rv, rv);

  // Note that if AsyncOpen failed that's ok -- the only caller of
  // this method nulls out mChannel immediately after we return.

  return mAsyncLoadStatus;
}

nsresult nsSyncLoader::PushSyncStream(nsIStreamListener *aListener) {
  nsCOMPtr<nsIInputStream> in;
  nsresult rv = mChannel->Open(getter_AddRefs(in));
  NS_ENSURE_SUCCESS(rv, rv);

  mLoading = true;
  rv = nsSyncLoadService::PushSyncStreamToListener(in.forget(), aListener,
                                                   mChannel);
  mLoading = false;

  return rv;
}

NS_IMETHODIMP
nsSyncLoader::OnStartRequest(nsIRequest *aRequest) {
  return mListener->OnStartRequest(aRequest);
}

NS_IMETHODIMP
nsSyncLoader::OnStopRequest(nsIRequest *aRequest, nsresult aStatusCode) {
  if (NS_SUCCEEDED(mAsyncLoadStatus) && NS_FAILED(aStatusCode)) {
    mAsyncLoadStatus = aStatusCode;
  }
  nsresult rv = mListener->OnStopRequest(aRequest, aStatusCode);
  if (NS_SUCCEEDED(mAsyncLoadStatus) && NS_FAILED(rv)) {
    mAsyncLoadStatus = rv;
  }
  mLoading = false;

  return rv;
}

NS_IMETHODIMP
nsSyncLoader::AsyncOnChannelRedirect(nsIChannel *aOldChannel,
                                     nsIChannel *aNewChannel, uint32_t aFlags,
                                     nsIAsyncVerifyRedirectCallback *callback) {
  MOZ_ASSERT(aNewChannel, "Redirecting to null channel?");

  mChannel = aNewChannel;

  callback->OnRedirectVerifyCallback(NS_OK);
  return NS_OK;
}

NS_IMETHODIMP
nsSyncLoader::GetInterface(const nsIID &aIID, void **aResult) {
  return QueryInterface(aIID, aResult);
}

/* static */
nsresult nsSyncLoadService::LoadDocument(
    nsIURI *aURI, nsContentPolicyType aContentPolicyType,
    nsIPrincipal *aLoaderPrincipal, nsSecurityFlags aSecurityFlags,
    nsILoadGroup *aLoadGroup, nsICookieSettings *aCookieSettings,
    bool aForceToXML, ReferrerPolicy aReferrerPolicy, Document **aResult) {
  nsCOMPtr<nsIChannel> channel;
  nsresult rv =
      NS_NewChannel(getter_AddRefs(channel), aURI, aLoaderPrincipal,
                    aSecurityFlags, aContentPolicyType, aCookieSettings,
                    nullptr,  // PerformanceStorage
                    aLoadGroup);
  NS_ENSURE_SUCCESS(rv, rv);

  if (!aForceToXML) {
    channel->SetContentType(NS_LITERAL_CSTRING("text/xml"));
  }

  bool isChrome = false, isResource = false;
  // if the load needs to enforce CORS, then force the load to be async
  bool isSync =
      !(aSecurityFlags & nsILoadInfo::SEC_REQUIRE_CORS_DATA_INHERITS) &&
      ((NS_SUCCEEDED(aURI->SchemeIs("chrome", &isChrome)) && isChrome) ||
       (NS_SUCCEEDED(aURI->SchemeIs("resource", &isResource)) && isResource));
  RefPtr<nsSyncLoader> loader = new nsSyncLoader();
  return loader->LoadDocument(channel, isSync, aForceToXML, aReferrerPolicy,
                              aResult);
}

/* static */
nsresult nsSyncLoadService::PushSyncStreamToListener(
    already_AddRefed<nsIInputStream> aIn, nsIStreamListener *aListener,
    nsIChannel *aChannel) {
  nsCOMPtr<nsIInputStream> in = std::move(aIn);

  // Set up buffering stream
  nsresult rv;
  nsCOMPtr<nsIInputStream> bufferedStream;
  if (!NS_InputStreamIsBuffered(in)) {
    int64_t chunkSize;
    rv = aChannel->GetContentLength(&chunkSize);
    if (NS_FAILED(rv) || chunkSize < 1) {
      chunkSize = 4096;
    }
    chunkSize = std::min(int64_t(UINT16_MAX), chunkSize);

    rv = NS_NewBufferedInputStream(getter_AddRefs(bufferedStream), in.forget(),
                                   chunkSize);
    NS_ENSURE_SUCCESS(rv, rv);

    in = bufferedStream;
  }

  // Load
  rv = aListener->OnStartRequest(aChannel);
  if (NS_SUCCEEDED(rv)) {
    uint64_t sourceOffset = 0;
    while (1) {
      uint64_t readCount = 0;
      rv = in->Available(&readCount);
      if (NS_FAILED(rv) || !readCount) {
        if (rv == NS_BASE_STREAM_CLOSED) {
          // End of file, but not an error
          rv = NS_OK;
        }
        break;
      }

      if (readCount > UINT32_MAX) readCount = UINT32_MAX;

      rv = aListener->OnDataAvailable(
          aChannel, in, (uint32_t)std::min(sourceOffset, (uint64_t)UINT32_MAX),
          (uint32_t)readCount);
      if (NS_FAILED(rv)) {
        break;
      }
      sourceOffset += readCount;
    }
  }
  if (NS_FAILED(rv)) {
    aChannel->Cancel(rv);
  }
  aListener->OnStopRequest(aChannel, rv);

  return rv;
}