netwerk/base/NetworkDataCountLayer.cpp

   1 /* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
   2 /* vim:set ts=2 sw=2 sts=2 et cindent: */
   3 /* This Source Code Form is subject to the terms of the Mozilla Public
   4  * License, v. 2.0. If a copy of the MPL was not distributed with this
   5  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
   6
   7 #include "NetworkDataCountLayer.h"
   8 #include "nsSocketTransportService2.h"
   9 #include "prmem.h"
  10 #include "prio.h"
  11
  12 namespace mozilla {
  13 namespace net {
  14
  15 static PRDescIdentity sNetworkDataCountLayerIdentity;
  16 static PRIOMethods sNetworkDataCountLayerMethods;
  17 static PRIOMethods* sNetworkDataCountLayerMethodsPtr = nullptr;
  18
  19 class NetworkDataCountSecret {
  20  public:
  21   NetworkDataCountSecret() = default;
  22
  23   uint64_t mSentBytes = 0;
  24   uint64_t mReceivedBytes = 0;
  25 };
  26
  27 static PRInt32 NetworkDataCountSend(PRFileDesc* fd, const void* buf,
  28                                     PRInt32 amount, PRIntn flags,
  29                                     PRIntervalTime timeout) {
  30   MOZ_RELEASE_ASSERT(fd->identity == sNetworkDataCountLayerIdentity);
  31
  32   NetworkDataCountSecret* secret =
  33       reinterpret_cast<NetworkDataCountSecret*>(fd->secret);
  34
  35   PRInt32 rv =
  36       (fd->lower->methods->send)(fd->lower, buf, amount, flags, timeout);
  37   if (rv > 0) {
  38     secret->mSentBytes += rv;
  39   }
  40   return rv;
  41 }
  42
  43 static PRInt32 NetworkDataCountWrite(PRFileDesc* fd, const void* buf,
  44                                      PRInt32 amount) {
  45   return NetworkDataCountSend(fd, buf, amount, 0, PR_INTERVAL_NO_WAIT);
  46 }
  47
  48 static PRInt32 NetworkDataCountRecv(PRFileDesc* fd, void* buf, PRInt32 amount,
  49                                     PRIntn flags, PRIntervalTime timeout) {
  50   MOZ_RELEASE_ASSERT(fd->identity == sNetworkDataCountLayerIdentity);
  51
  52   NetworkDataCountSecret* secret =
  53       reinterpret_cast<NetworkDataCountSecret*>(fd->secret);
  54
  55   PRInt32 rv =
  56       (fd->lower->methods->recv)(fd->lower, buf, amount, flags, timeout);
  57   if (rv > 0) {
  58     secret->mReceivedBytes += rv;
  59   }
  60   return rv;
  61 }
  62
  63 static PRInt32 NetworkDataCountRead(PRFileDesc* fd, void* buf, PRInt32 amount) {
  64   return NetworkDataCountRecv(fd, buf, amount, 0, PR_INTERVAL_NO_WAIT);
  65 }
  66
  67 static PRStatus NetworkDataCountClose(PRFileDesc* fd) {
  68   if (!fd) {
  69     return PR_FAILURE;
  70   }
  71
  72   PRFileDesc* layer = PR_PopIOLayer(fd, PR_TOP_IO_LAYER);
  73
  74   MOZ_RELEASE_ASSERT(layer && layer->identity == sNetworkDataCountLayerIdentity,
  75                      "NetworkDataCount Layer not on top of stack");
  76
  77   NetworkDataCountSecret* secret =
  78       reinterpret_cast<NetworkDataCountSecret*>(layer->secret);
  79   layer->secret = nullptr;
  80   layer->dtor(layer);
  81   delete secret;
  82   return fd->methods->close(fd);
  83 }
  84
  85 nsresult AttachNetworkDataCountLayer(PRFileDesc* fd) {
  86   if (!sNetworkDataCountLayerMethodsPtr) {
  87     sNetworkDataCountLayerIdentity =
  88         PR_GetUniqueIdentity("NetworkDataCount Layer");
  89     sNetworkDataCountLayerMethods = *PR_GetDefaultIOMethods();
  90     sNetworkDataCountLayerMethods.send = NetworkDataCountSend;
  91     sNetworkDataCountLayerMethods.write = NetworkDataCountWrite;
  92     sNetworkDataCountLayerMethods.recv = NetworkDataCountRecv;
  93     sNetworkDataCountLayerMethods.read = NetworkDataCountRead;
  94     sNetworkDataCountLayerMethods.close = NetworkDataCountClose;
  95     sNetworkDataCountLayerMethodsPtr = &sNetworkDataCountLayerMethods;
  96   }
  97
  98   PRFileDesc* layer = PR_CreateIOLayerStub(sNetworkDataCountLayerIdentity,
  99                                            sNetworkDataCountLayerMethodsPtr);
 100
 101   if (!layer) {
 102     return NS_ERROR_FAILURE;
 103   }
 104
 105   NetworkDataCountSecret* secret = new NetworkDataCountSecret();
 106
 107   layer->secret = reinterpret_cast<PRFilePrivate*>(secret);
 108
 109   PRStatus status = PR_PushIOLayer(fd, PR_NSPR_IO_LAYER, layer);
 110
 111   if (status == PR_FAILURE) {
 112     delete secret;
 113     PR_Free(layer);  // PR_CreateIOLayerStub() uses PR_Malloc().
 114     return NS_ERROR_FAILURE;
 115   }
 116   return NS_OK;
 117 }
 118
 119 void NetworkDataCountSent(PRFileDesc* fd, uint64_t& sentBytes) {
 120   PRFileDesc* ndcFd = PR_GetIdentitiesLayer(fd, sNetworkDataCountLayerIdentity);
 121   MOZ_RELEASE_ASSERT(ndcFd);
 122
 123   NetworkDataCountSecret* secret =
 124       reinterpret_cast<NetworkDataCountSecret*>(ndcFd->secret);
 125   sentBytes = secret->mSentBytes;
 126 }
 127
 128 void NetworkDataCountReceived(PRFileDesc* fd, uint64_t& receivedBytes) {
 129   PRFileDesc* ndcFd = PR_GetIdentitiesLayer(fd, sNetworkDataCountLayerIdentity);
 130   MOZ_RELEASE_ASSERT(ndcFd);
 131
 132   NetworkDataCountSecret* secret =
 133       reinterpret_cast<NetworkDataCountSecret*>(ndcFd->secret);
 134   receivedBytes = secret->mReceivedBytes;
 135 }
 136
 137 }  // namespace net
 138 }  // namespace mozilla