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
NetworkDataCountSend(PRFileDesc * fd,const void * buf,PRInt32 amount,PRIntn flags,PRIntervalTime timeout)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
NetworkDataCountWrite(PRFileDesc * fd,const void * buf,PRInt32 amount)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
NetworkDataCountRecv(PRFileDesc * fd,void * buf,PRInt32 amount,PRIntn flags,PRIntervalTime timeout)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
NetworkDataCountRead(PRFileDesc * fd,void * buf,PRInt32 amount)63 static PRInt32 NetworkDataCountRead(PRFileDesc* fd, void* buf, PRInt32 amount) {
64 return NetworkDataCountRecv(fd, buf, amount, 0, PR_INTERVAL_NO_WAIT);
65 }
66
NetworkDataCountClose(PRFileDesc * fd)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
AttachNetworkDataCountLayer(PRFileDesc * fd)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
NetworkDataCountSent(PRFileDesc * fd,uint64_t & sentBytes)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
NetworkDataCountReceived(PRFileDesc * fd,uint64_t & receivedBytes)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
139