/* 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 "nsIInterfaceRequestor.h"
#include "WebSocketConnection.h"

#include "WebSocketLog.h"
#include "mozilla/net/WebSocketConnectionListener.h"
#include "mozilla/Components.h"
#include "nsIOService.h"
#include "nsITLSSocketControl.h"
#include "nsISocketTransport.h"
#include "nsITransportSecurityInfo.h"
#include "nsSocketTransportService2.h"

namespace mozilla::net {

NS_IMPL_ISUPPORTS(WebSocketConnection, nsIInputStreamCallback,
                  nsIOutputStreamCallback)

WebSocketConnection::WebSocketConnection(nsISocketTransport* aTransport,
                                         nsIAsyncInputStream* aInputStream,
                                         nsIAsyncOutputStream* aOutputStream)
    : mTransport(aTransport),
      mSocketIn(aInputStream),
      mSocketOut(aOutputStream) {
  LOG(("WebSocketConnection ctor %p\n", this));
}

WebSocketConnection::~WebSocketConnection() {
  LOG(("WebSocketConnection dtor %p\n", this));
}

nsresult WebSocketConnection::Init(WebSocketConnectionListener* aListener) {
  NS_ENSURE_ARG_POINTER(aListener);

  mListener = aListener;
  nsresult rv;
  mSocketThread = mozilla::components::SocketTransport::Service(&rv);
  if (NS_FAILED(rv)) {
    return rv;
  }

  if (!mTransport) {
    return NS_ERROR_FAILURE;
  }

  if (XRE_IsParentProcess()) {
    nsCOMPtr<nsIInterfaceRequestor> callbacks = do_QueryInterface(mListener);
    mTransport->SetSecurityCallbacks(callbacks);
  } else {
    // NOTE: we don't use security callbacks in socket process.
    mTransport->SetSecurityCallbacks(nullptr);
  }
  return mTransport->SetEventSink(nullptr, nullptr);
}

void WebSocketConnection::GetIoTarget(nsIEventTarget** aTarget) {
  nsCOMPtr<nsIEventTarget> target = mSocketThread;
  return target.forget(aTarget);
}

void WebSocketConnection::Close() {
  LOG(("WebSocketConnection::Close %p\n", this));
  MOZ_ASSERT(OnSocketThread());

  if (mTransport) {
    mTransport->SetSecurityCallbacks(nullptr);
    mTransport->SetEventSink(nullptr, nullptr);
    mTransport->Close(NS_BASE_STREAM_CLOSED);
    mTransport = nullptr;
  }

  if (mSocketIn) {
    if (mStartReadingCalled) {
      mSocketIn->AsyncWait(nullptr, 0, 0, nullptr);
    }
    mSocketIn = nullptr;
  }

  if (mSocketOut) {
    mSocketOut->AsyncWait(nullptr, 0, 0, nullptr);
    mSocketOut = nullptr;
  }
}

nsresult WebSocketConnection::WriteOutputData(nsTArray<uint8_t>&& aData) {
  MOZ_ASSERT(OnSocketThread());

  if (!mSocketOut) {
    return NS_ERROR_NOT_AVAILABLE;
  }

  mOutputQueue.emplace_back(std::move(aData));
  return OnOutputStreamReady(mSocketOut);
}

nsresult WebSocketConnection::WriteOutputData(const uint8_t* aHdrBuf,
                                              uint32_t aHdrBufLength,
                                              const uint8_t* aPayloadBuf,
                                              uint32_t aPayloadBufLength) {
  MOZ_ASSERT_UNREACHABLE("Should not be called");
  return NS_ERROR_NOT_IMPLEMENTED;
}

nsresult WebSocketConnection::StartReading() {
  MOZ_ASSERT(OnSocketThread());

  if (!mSocketIn) {
    return NS_ERROR_NOT_AVAILABLE;
  }

  MOZ_ASSERT(!mStartReadingCalled, "StartReading twice");
  mStartReadingCalled = true;
  return mSocketIn->AsyncWait(this, 0, 0, mSocketThread);
}

void WebSocketConnection::DrainSocketData() {
  MOZ_ASSERT(OnSocketThread());

  if (!mSocketIn || !mListener) {
    return;
  }

  // If we leave any data unconsumed (including the tcp fin) a RST will be
  // generated The right thing to do here is shutdown(SHUT_WR) and then wait a
  // little while to see if any data comes in.. but there is no reason to delay
  // things for that when the websocket handshake is supposed to guarantee a
  // quiet connection except for that fin.
  char buffer[512];
  uint32_t count = 0;
  uint32_t total = 0;
  nsresult rv;
  do {
    total += count;
    rv = mSocketIn->Read(buffer, 512, &count);
    if (rv != NS_BASE_STREAM_WOULD_BLOCK && (NS_FAILED(rv) || count == 0)) {
      mListener->OnTCPClosed();
    }
  } while (NS_SUCCEEDED(rv) && count > 0 && total < 32000);
}

nsresult WebSocketConnection::GetSecurityInfo(
    nsITransportSecurityInfo** aSecurityInfo) {
  LOG(("WebSocketConnection::GetSecurityInfo() %p\n", this));
  MOZ_ASSERT(OnSocketThread());
  *aSecurityInfo = nullptr;

  if (mTransport) {
    nsCOMPtr<nsITLSSocketControl> tlsSocketControl;
    nsresult rv =
        mTransport->GetTlsSocketControl(getter_AddRefs(tlsSocketControl));
    if (NS_FAILED(rv)) {
      return rv;
    }
    nsCOMPtr<nsITransportSecurityInfo> securityInfo(
        do_QueryInterface(tlsSocketControl));
    if (securityInfo) {
      securityInfo.forget(aSecurityInfo);
    }
  }
  return NS_OK;
}

NS_IMETHODIMP
WebSocketConnection::OnInputStreamReady(nsIAsyncInputStream* aStream) {
  LOG(("WebSocketConnection::OnInputStreamReady() %p\n", this));
  MOZ_ASSERT(OnSocketThread());
  MOZ_ASSERT(mListener);

  // did we we clean up the socket after scheduling InputReady?
  if (!mSocketIn) {
    return NS_OK;
  }

  // this is after the http upgrade - so we are speaking websockets
  uint8_t buffer[2048];
  uint32_t count;
  nsresult rv;

  do {
    rv = mSocketIn->Read((char*)buffer, 2048, &count);
    LOG(("WebSocketConnection::OnInputStreamReady: read %u rv %" PRIx32 "\n",
         count, static_cast<uint32_t>(rv)));

    if (rv == NS_BASE_STREAM_WOULD_BLOCK) {
      mSocketIn->AsyncWait(this, 0, 0, mSocketThread);
      return NS_OK;
    }

    if (NS_FAILED(rv)) {
      mListener->OnError(rv);
      return rv;
    }

    if (count == 0) {
      mListener->OnError(NS_BASE_STREAM_CLOSED);
      return NS_OK;
    }

    rv = mListener->OnDataReceived(buffer, count);
    if (NS_FAILED(rv)) {
      mListener->OnError(rv);
      return rv;
    }
  } while (NS_SUCCEEDED(rv) && mSocketIn && mListener);

  return NS_OK;
}

NS_IMETHODIMP
WebSocketConnection::OnOutputStreamReady(nsIAsyncOutputStream* aStream) {
  LOG(("WebSocketConnection::OnOutputStreamReady() %p\n", this));
  MOZ_ASSERT(OnSocketThread());
  MOZ_ASSERT(mListener);

  if (!mSocketOut) {
    return NS_OK;
  }

  while (!mOutputQueue.empty()) {
    const OutputData& data = mOutputQueue.front();

    char* buffer = reinterpret_cast<char*>(
                       const_cast<uint8_t*>(data.GetData().Elements())) +
                   mWriteOffset;
    uint32_t toWrite = data.GetData().Length() - mWriteOffset;

    uint32_t wrote = 0;
    nsresult rv = mSocketOut->Write(buffer, toWrite, &wrote);
    LOG(("WebSocketConnection::OnOutputStreamReady: write %u rv %" PRIx32,
         wrote, static_cast<uint32_t>(rv)));

    if (rv == NS_BASE_STREAM_WOULD_BLOCK) {
      mSocketOut->AsyncWait(this, 0, 0, mSocketThread);
      return rv;
    }

    if (NS_FAILED(rv)) {
      LOG(("WebSocketConnection::OnOutputStreamReady %p failed %u\n", this,
           static_cast<uint32_t>(rv)));
      mListener->OnError(rv);
      return NS_OK;
    }

    mWriteOffset += wrote;

    if (toWrite == wrote) {
      mWriteOffset = 0;
      mOutputQueue.pop_front();
    } else {
      mSocketOut->AsyncWait(this, 0, 0, mSocketThread);
    }
  }

  return NS_OK;
}

}  // namespace mozilla::net