/* -*- Mode: C++; tab-width: 2; 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/. */

#include "IOActivityMonitor.h"

#include "nsIObserverService.h"

#include "nsPISocketTransportService.h"

#include "nsPrintfCString.h"

#include "nsSocketTransport2.h"

#include "nsSocketTransportService2.h"

#include "nsThreadUtils.h"

#include "mozilla/dom/Promise.h"

#include "mozilla/Services.h"

#include "prerror.h"

#include "prio.h"

#include "prmem.h"

#include <vector>

using namespace mozilla;

using namespace mozilla::net;

mozilla::StaticRefPtr<IOActivityMonitor> gInstance;

static PRDescIdentity sNetActivityMonitorLayerIdentity;

static PRIOMethods sNetActivityMonitorLayerMethods;

static PRIOMethods *sNetActivityMonitorLayerMethodsPtr = nullptr;

// Maximum number of activities entries in the monitoring class

#define MAX_ACTIVITY_ENTRIES 1000

// ActivityMonitorSecret is stored in the activity monitor layer

// and provides a method to get the location.

//

// A location can be :

// - a TCP or UDP socket. The form will be socket://ip:port

// - a File. The form will be file://path

//

// For other cases, the location will be fd://number

class ActivityMonitorSecret final

{

public:

  // constructor used for sockets

  explicit ActivityMonitorSecret(PRFileDesc* aFd) {

    mFd = aFd;

    mLocationSet = false;

  }

  // constructor used for files

  explicit ActivityMonitorSecret(PRFileDesc* aFd, const char* aLocation) {

    mFd = aFd;

    mLocation.AppendPrintf("file://%s", aLocation);

    mLocationSet = true;

  }

  nsCString getLocation() {

    if (!mLocationSet) {

      LazySetLocation();

    }

    return mLocation;

  }

private:

  // Called to set the location using the FD on the first getLocation() usage

  // which is typically when a socket is opened. If done earlier, at

  // construction time, the host won't be bound yet.

  //

  // If the location is a file, it needs to be initialized in the

  // constructor.

  void LazySetLocation() {

    mLocationSet = true;

    PRFileDesc* extract = mFd;

    while (PR_GetDescType(extract) == PR_DESC_LAYERED) {

      if (!extract->lower) {

        break;

      }

      extract = extract->lower;

    }

    PRDescType fdType = PR_GetDescType(extract);

    // we should not use LazySetLocation for files

    MOZ_ASSERT(fdType != PR_DESC_FILE);

    switch (fdType) {

      case PR_DESC_SOCKET_TCP:

      case PR_DESC_SOCKET_UDP: {

        mLocation.AppendPrintf("socket://");

        PRNetAddr addr;

        PRStatus status = PR_GetSockName(mFd, &addr);

        if (NS_WARN_IF(status == PR_FAILURE)) {

          mLocation.AppendPrintf("unknown");

          break;

        }

        // grabbing the host

        char netAddr[mozilla::net::kNetAddrMaxCStrBufSize] = {0};

        status = PR_NetAddrToString(&addr, netAddr, sizeof(netAddr) - 1);

        if (NS_WARN_IF(status == PR_FAILURE) || netAddr[0] == 0) {

          mLocation.AppendPrintf("unknown");

          break;

        }

        mLocation.Append(netAddr);

        // adding the port

        uint16_t port;

        if (addr.raw.family == PR_AF_INET) {

          port = addr.inet.port;

        } else {

          port = addr.ipv6.port;

        }

        mLocation.AppendPrintf(":%d", port);

      } break;

      // for all other cases, we just send back fd://<value>

      default: {

        mLocation.AppendPrintf("fd://%d", PR_FileDesc2NativeHandle(mFd));

      }

    } // end switch

  }

private:

  nsCString mLocation;

  bool mLocationSet;

  PRFileDesc* mFd;

};

// FileDesc2Location converts a PRFileDesc into a "location" by

// grabbing the ActivityMonitorSecret in layer->secret

static nsAutoCString

FileDesc2Location(PRFileDesc *fd)

{

  nsAutoCString location;

  PRFileDesc *monitorLayer = PR_GetIdentitiesLayer(fd, sNetActivityMonitorLayerIdentity);

  if (!monitorLayer) {

    location.AppendPrintf("unknown");

    return location;

  }

  ActivityMonitorSecret* secret = (ActivityMonitorSecret*)monitorLayer->secret;

  location.AppendPrintf("%s", secret->getLocation().get());

  return location;

}

//

// Wrappers around the socket APIS

//

static PRStatus

nsNetMon_Connect(PRFileDesc *fd, const PRNetAddr *addr, PRIntervalTime timeout)

{

  return fd->lower->methods->connect(fd->lower, addr, timeout);

}

static PRStatus

nsNetMon_Close(PRFileDesc *fd)

{

  if (!fd) {

    return PR_FAILURE;

  }

  PRFileDesc* layer = PR_PopIOLayer(fd, PR_TOP_IO_LAYER);

  MOZ_RELEASE_ASSERT(layer &&

                     layer->identity == sNetActivityMonitorLayerIdentity,

                     "NetActivityMonitor Layer not on top of stack");

  if (layer->secret) {

    delete (ActivityMonitorSecret *)layer->secret;

    layer->secret = nullptr;

  }

  layer->dtor(layer);

  return fd->methods->close(fd);

}

static int32_t

nsNetMon_Read(PRFileDesc *fd, void *buf, int32_t len)

{

  int32_t ret = fd->lower->methods->read(fd->lower, buf, len);

  if (ret >= 0) {

    IOActivityMonitor::Read(fd, len);

  }

  return ret;

}

static int32_t

nsNetMon_Write(PRFileDesc *fd, const void *buf, int32_t len)

{

  int32_t ret = fd->lower->methods->write(fd->lower, buf, len);

  if (ret > 0) {

    IOActivityMonitor::Write(fd, len);

  }

  return ret;

}

static int32_t

nsNetMon_Writev(PRFileDesc *fd,

                const PRIOVec *iov,

                int32_t size,

                PRIntervalTime timeout)

{

  int32_t ret = fd->lower->methods->writev(fd->lower, iov, size, timeout);

  if (ret > 0) {

    IOActivityMonitor::Write(fd, size);

  }

  return ret;

}

static int32_t

nsNetMon_Recv(PRFileDesc *fd,

              void *buf,

              int32_t amount,

              int flags,

              PRIntervalTime timeout)

{

  int32_t ret = fd->lower->methods->recv(fd->lower, buf, amount, flags, timeout);

  if (ret > 0) {

    IOActivityMonitor::Read(fd, amount);

  }

  return ret;

}

static int32_t

nsNetMon_Send(PRFileDesc *fd,

              const void *buf,

              int32_t amount,

              int flags,

              PRIntervalTime timeout)

{

  int32_t ret = fd->lower->methods->send(fd->lower, buf, amount, flags, timeout);

  if (ret > 0) {

    IOActivityMonitor::Write(fd, amount);

  }

  return ret;

}

static int32_t

nsNetMon_RecvFrom(PRFileDesc *fd,

                  void *buf,

                  int32_t amount,

                  int flags,

                  PRNetAddr *addr,

                  PRIntervalTime timeout)

{

  int32_t ret = fd->lower->methods->recvfrom(fd->lower, buf, amount, flags,

                                             addr, timeout);

  if (ret > 0) {

    IOActivityMonitor::Read(fd, amount);

  }

  return ret;

}

static int32_t

nsNetMon_SendTo(PRFileDesc *fd,

                const void *buf,

                int32_t amount,

                int flags,

                const PRNetAddr *addr,

                PRIntervalTime timeout)

{

  int32_t ret = fd->lower->methods->sendto(fd->lower, buf, amount, flags,

                                           addr, timeout);

  if (ret > 0) {

    IOActivityMonitor::Write(fd, amount);

  }

  return ret;

}

static int32_t

nsNetMon_AcceptRead(PRFileDesc *listenSock,

                    PRFileDesc **acceptedSock,

                    PRNetAddr **peerAddr,

                    void *buf,

                    int32_t amount,

                    PRIntervalTime timeout)

{

  int32_t ret = listenSock->lower->methods->acceptread(listenSock->lower, acceptedSock,

                                                       peerAddr, buf, amount, timeout);

  if (ret > 0) {

    IOActivityMonitor::Read(listenSock, amount);

  }

  return ret;

}

//

// Class IOActivityMonitor

//

NS_IMPL_ISUPPORTS(IOActivityMonitor, nsINamed)

IOActivityMonitor::IOActivityMonitor()

  : mLock("IOActivityMonitor::mLock")

{

  RefPtr<IOActivityMonitor> mon(gInstance);

  MOZ_ASSERT(!mon, "multiple IOActivityMonitor instances!");

}

// static

void

IOActivityMonitor::RequestActivities(dom::Promise* aPromise)

{

  MOZ_ASSERT(aPromise);

  RefPtr<IOActivityMonitor> mon(gInstance);

  if (!IsActive()) {

    aPromise->MaybeReject(NS_ERROR_FAILURE);

    return;

  }

  mon->RequestActivitiesInternal(aPromise);

}

void

IOActivityMonitor::RequestActivitiesInternal(dom::Promise* aPromise)

{

  nsresult result = NS_OK;

  FallibleTArray<dom::IOActivityDataDictionary> activities;

  {

    mozilla::MutexAutoLock lock(mLock);

    // Remove inactive activities

    for (auto iter = mActivities.Iter(); !iter.Done(); iter.Next()) {

      dom::IOActivityDataDictionary* activity = &iter.Data();

      if (activity->mRx == 0 && activity->mTx == 0) {

        iter.Remove();

      } else {

        if (NS_WARN_IF(!activities.AppendElement(iter.Data(), fallible))) {

          result = NS_ERROR_OUT_OF_MEMORY;

          break;

        }

      }

    }

  }

  if (NS_WARN_IF(NS_FAILED(result))) {

    aPromise->MaybeReject(result);

    return;

  }

  aPromise->MaybeResolve(activities);

}

// static

NS_IMETHODIMP

IOActivityMonitor::GetName(nsACString& aName)

{

  aName.AssignLiteral("IOActivityMonitor");

  return NS_OK;

}

bool

IOActivityMonitor::IsActive()

{

  return gInstance != nullptr;

}

nsresult

IOActivityMonitor::Init()

{

  if (IsActive()) {

    return NS_ERROR_ALREADY_INITIALIZED;

  }

  RefPtr<IOActivityMonitor> mon = new IOActivityMonitor();

  nsresult rv = mon->InitInternal();

  if (NS_SUCCEEDED(rv)) {

    gInstance = mon;

  }

  return rv;

}

nsresult

IOActivityMonitor::InitInternal()

{

  // wraps the socket APIs

  if (!sNetActivityMonitorLayerMethodsPtr) {

    sNetActivityMonitorLayerIdentity =

      PR_GetUniqueIdentity("network activity monitor layer");

    sNetActivityMonitorLayerMethods  = *PR_GetDefaultIOMethods();

    sNetActivityMonitorLayerMethods.connect    = nsNetMon_Connect;

    sNetActivityMonitorLayerMethods.read       = nsNetMon_Read;

    sNetActivityMonitorLayerMethods.write      = nsNetMon_Write;

    sNetActivityMonitorLayerMethods.writev     = nsNetMon_Writev;

    sNetActivityMonitorLayerMethods.recv       = nsNetMon_Recv;

    sNetActivityMonitorLayerMethods.send       = nsNetMon_Send;

    sNetActivityMonitorLayerMethods.recvfrom   = nsNetMon_RecvFrom;

    sNetActivityMonitorLayerMethods.sendto     = nsNetMon_SendTo;

    sNetActivityMonitorLayerMethods.acceptread = nsNetMon_AcceptRead;

    sNetActivityMonitorLayerMethods.close      = nsNetMon_Close;

    sNetActivityMonitorLayerMethodsPtr = &sNetActivityMonitorLayerMethods;

  }

  return NS_OK;

}

nsresult

IOActivityMonitor::Shutdown()

{

  RefPtr<IOActivityMonitor> mon(gInstance);

  if (!mon) {

    return NS_ERROR_NOT_INITIALIZED;

  }

  return mon->ShutdownInternal();

}

nsresult

IOActivityMonitor::ShutdownInternal()

{

  mozilla::MutexAutoLock lock(mLock);

  mActivities.Clear();

  gInstance = nullptr;

  return NS_OK;

}

nsresult

IOActivityMonitor::MonitorSocket(PRFileDesc *aFd)

{

  RefPtr<IOActivityMonitor> mon(gInstance);

  if (!IsActive()) {

    return NS_OK;

  }

  PRFileDesc* layer;

  PRStatus status;

  layer = PR_CreateIOLayerStub(sNetActivityMonitorLayerIdentity,

                               sNetActivityMonitorLayerMethodsPtr);

  if (!layer) {

    return NS_ERROR_FAILURE;

  }

  ActivityMonitorSecret* secret = new ActivityMonitorSecret(aFd);

  layer->secret = reinterpret_cast<PRFilePrivate *>(secret);

  status = PR_PushIOLayer(aFd, PR_NSPR_IO_LAYER, layer);

  if (status == PR_FAILURE) {

    delete secret;

    PR_Free(layer); // PR_CreateIOLayerStub() uses PR_Malloc().

    return NS_ERROR_FAILURE;

  }

  return NS_OK;

}

nsresult

IOActivityMonitor::MonitorFile(PRFileDesc *aFd, const char* aPath)

{

  RefPtr<IOActivityMonitor> mon(gInstance);

  if (!IsActive()) {

    return NS_OK;

  }

  PRFileDesc* layer;

  PRStatus status;

  layer = PR_CreateIOLayerStub(sNetActivityMonitorLayerIdentity,

                               sNetActivityMonitorLayerMethodsPtr);

  if (!layer) {

    return NS_ERROR_FAILURE;

  }

  ActivityMonitorSecret* secret = new ActivityMonitorSecret(aFd, aPath);

  layer->secret = reinterpret_cast<PRFilePrivate *>(secret);

  status = PR_PushIOLayer(aFd, PR_NSPR_IO_LAYER, layer);

  if (status == PR_FAILURE) {

    delete secret;

    PR_Free(layer); // PR_CreateIOLayerStub() uses PR_Malloc().

    return NS_ERROR_FAILURE;

  }

  return NS_OK;

}

bool

IOActivityMonitor::IncrementActivity(const nsACString& aLocation, uint32_t aRx, uint32_t aTx)

{

  mLock.AssertCurrentThreadOwns();

  if (auto entry = mActivities.Lookup(aLocation)) {

    // already registered

    entry.Data().mTx += aTx;

    entry.Data().mRx += aRx;

    return true;

  }

  // Creating a new IOActivity. Notice that mActivities will

  // grow indefinitely, which is OK since we won't have

  // but a few hundreds entries at the most, but we

  // want to assert we have at the most 1000 entries

  MOZ_ASSERT(mActivities.Count() < MAX_ACTIVITY_ENTRIES);

  dom::IOActivityDataDictionary activity;

  activity.mLocation.Assign(aLocation);

  activity.mTx = aTx;

  activity.mRx = aRx;

  if (NS_WARN_IF(!mActivities.Put(aLocation, activity, fallible))) {

    return false;

  }

  return true;

}

nsresult

IOActivityMonitor::Write(const nsACString& aLocation, uint32_t aAmount)

{

  RefPtr<IOActivityMonitor> mon(gInstance);

  if (!mon) {

    return NS_ERROR_FAILURE;

  }

  return mon->WriteInternal(aLocation, aAmount);

}

nsresult

IOActivityMonitor::Write(PRFileDesc *fd, uint32_t aAmount)

{

  RefPtr<IOActivityMonitor> mon(gInstance);

  if (!mon) {

    return NS_ERROR_FAILURE;

  }

  return mon->Write(FileDesc2Location(fd), aAmount);

}

nsresult

IOActivityMonitor::WriteInternal(const nsACString& aLocation, uint32_t aAmount)

{

  mozilla::MutexAutoLock lock(mLock);

  if (!IncrementActivity(aLocation, aAmount, 0)) {

    return NS_ERROR_FAILURE;

  }

  return NS_OK;

}

nsresult

IOActivityMonitor::Read(PRFileDesc *fd, uint32_t aAmount)

{

  RefPtr<IOActivityMonitor> mon(gInstance);

  if (!mon) {

    return NS_ERROR_FAILURE;

  }

  return mon->Read(FileDesc2Location(fd), aAmount);

}

nsresult

IOActivityMonitor::Read(const nsACString& aLocation, uint32_t aAmount)

{

  RefPtr<IOActivityMonitor> mon(gInstance);

  if (!mon) {

    return NS_ERROR_FAILURE;

  }

  return mon->ReadInternal(aLocation, aAmount);

}

nsresult

IOActivityMonitor::ReadInternal(const nsACString& aLocation, uint32_t aAmount)

{

  mozilla::MutexAutoLock lock(mLock);

  if (!IncrementActivity(aLocation, 0, aAmount)) {

    return NS_ERROR_FAILURE;

  }

  return NS_OK;

}