31
|
1 |
|
|
2 |
#include "NetworkTCP.hh"
|
|
3 |
|
|
4 |
#include <cstdlib>
|
|
5 |
#include <cassert>
|
|
6 |
|
|
7 |
NetworkBuffer::NetworkBuffer (NetworkSocket &socket, size_t size_hint) :
|
|
8 |
socket(socket), buf(0), size(0), offset(0) {
|
|
9 |
|
|
10 |
// allocate initial buffer
|
|
11 |
if ((buf = (char *) malloc(size_hint)) == NULL)
|
|
12 |
throw CL_Error("malloc failed");
|
|
13 |
|
|
14 |
// remember size
|
|
15 |
size = size_hint;
|
|
16 |
}
|
32
|
17 |
|
|
18 |
~NetworkBuffer::NetworkBuffer (void) {
|
|
19 |
free(buf);
|
|
20 |
}
|
31
|
21 |
|
32
|
22 |
void NetworkBuffer::resize (size_t suffix_size) {
|
|
23 |
size_t new_size = size;
|
|
24 |
|
|
25 |
// grow new_size until item_size fits
|
|
26 |
while (offset + item_size > new_size)
|
|
27 |
new_size *= 2;
|
31
|
28 |
|
32
|
29 |
// grow if needed
|
|
30 |
if (new_size != size) {
|
|
31 |
// realloc buffer
|
|
32 |
if ((buf = (char *) realloc((void *) buf, new_size)) == NULL)
|
|
33 |
throw CL_Error("realloc failed");
|
|
34 |
|
|
35 |
// update size
|
|
36 |
size = new_size;
|
|
37 |
|
|
38 |
} else if (new_size > (offset + item_size) * 4) {
|
|
39 |
// XXX: shrink?
|
|
40 |
}
|
31
|
41 |
}
|
|
42 |
|
32
|
43 |
void NetworkBuffer::trim (size_t prefix_size) {
|
|
44 |
// update offset
|
|
45 |
offset -= prefix;
|
|
46 |
|
|
47 |
// shift the buffer forwards from (buf + prefix) -> (buf), copying (old_offset - prefix) bytes
|
|
48 |
memmove(buf, buf + ret, offset);
|
|
49 |
}
|
|
50 |
|
|
51 |
bool try_read (size_t item_size) {
|
|
52 |
int ret;
|
|
53 |
size_t to_read = item_size;
|
|
54 |
|
|
55 |
// keept reads at at least NETWORK_CHUNK_SIZE bytes
|
|
56 |
if (to_read < NETWORK_TCP_CHUNK_SIZE)
|
|
57 |
to_read = NETWORK_TCP_CHUNK_SIZE;
|
|
58 |
|
|
59 |
// resize buffer if needed
|
|
60 |
resize(to_read);
|
|
61 |
|
|
62 |
// read once
|
|
63 |
try {
|
|
64 |
ret = socket.recv(buf + offset, to_read);
|
|
65 |
|
|
66 |
} catch (CL_Error &e) {
|
|
67 |
if (errno == EAGAIN)
|
|
68 |
return false;
|
|
69 |
else
|
|
70 |
throw;
|
|
71 |
}
|
|
72 |
|
|
73 |
// update offset
|
|
74 |
offset += ret;
|
|
75 |
|
|
76 |
// did we get enough?
|
|
77 |
if (ret < item_size)
|
|
78 |
return false;
|
|
79 |
else
|
|
80 |
return true;
|
|
81 |
}
|
|
82 |
|
|
83 |
bool peek_prefix (uint16_t &ref) {
|
|
84 |
if (offset < sizeof(uint16_t))
|
|
85 |
return false;
|
|
86 |
|
|
87 |
ret = ntohs(*((uint16_t *) (buf)));
|
|
88 |
|
|
89 |
return true;
|
|
90 |
}
|
|
91 |
|
|
92 |
bool peek_prefix (uint32_t &ref) {
|
|
93 |
if (offset < sizeof(uint32_t))
|
|
94 |
return false;
|
|
95 |
|
|
96 |
ret = ntohl(*((uint32_t *) (buf)));
|
|
97 |
|
|
98 |
return true;
|
|
99 |
}
|
|
100 |
|
|
101 |
template <typename PrefixType> PrefixType read_prefix (char *buf_ptr, size_t buf_max) {
|
|
102 |
PrefixType prefix = 0;
|
|
103 |
size_t missing = 0;
|
|
104 |
|
|
105 |
do {
|
|
106 |
// do we have the prefix?
|
|
107 |
if (peek_prefix(prefix)) {
|
|
108 |
// do we already have the payload?
|
|
109 |
if (offset >= sizeof(PrefixType) + prefix) {
|
|
110 |
break;
|
|
111 |
|
|
112 |
} else {
|
|
113 |
missing = (sizeof(PrefixType) + prefix) - offset;
|
|
114 |
}
|
|
115 |
|
|
116 |
} else {
|
|
117 |
missing = sizeof(PrefixType);
|
|
118 |
}
|
|
119 |
|
|
120 |
// sanity-check
|
|
121 |
assert(missing);
|
|
122 |
|
|
123 |
// try and read the missing data
|
|
124 |
if (try_read(missing) == false) {
|
|
125 |
// if unable to read what we need, return zero.
|
|
126 |
return 0;
|
|
127 |
}
|
|
128 |
|
|
129 |
// assess the situation again
|
|
130 |
} while (true);
|
|
131 |
|
|
132 |
// copy the data over, unless it's too large
|
|
133 |
if (prefix <= buf_max) {
|
|
134 |
memcpy(buf_ptr, buf, prefix);
|
|
135 |
|
|
136 |
// trim the bytes out
|
|
137 |
trim(prefix);
|
|
138 |
|
|
139 |
// return
|
|
140 |
return prefix;
|
|
141 |
|
|
142 |
} else {
|
|
143 |
// trim the bytes out
|
|
144 |
trim(prefix);
|
|
145 |
|
|
146 |
throw CL_Error("recv prefix overflow");
|
|
147 |
}
|
|
148 |
}
|
|
149 |
|
|
150 |
void NetworkBuffer::push_write (char *buf_ptr, size_t buf_size) {
|
31
|
151 |
int ret;
|
|
152 |
|
|
153 |
// try and short-circuit writes unless we have already buffered data
|
|
154 |
if (offset == 0) {
|
|
155 |
try {
|
|
156 |
// attempt to send something
|
|
157 |
ret = socket.send(buf_ptr, buf_size);
|
|
158 |
|
|
159 |
} catch (CL_Error &e) {
|
|
160 |
// ignore EAGAIN, detect this by setting ret to -1
|
|
161 |
if (errno != EAGAIN)
|
|
162 |
throw;
|
|
163 |
|
|
164 |
ret = -1;
|
|
165 |
}
|
|
166 |
|
|
167 |
// if we managed to send something, adjust buf/size and buffer
|
|
168 |
if (ret > 0) {
|
|
169 |
buf_ptr += ret;
|
|
170 |
buf_size -= ret;
|
|
171 |
|
|
172 |
// sanity-check
|
|
173 |
assert(buf_size >= 0);
|
|
174 |
|
|
175 |
// if that was all, we're done
|
|
176 |
if (buf_size == 0)
|
|
177 |
return;
|
|
178 |
}
|
|
179 |
}
|
|
180 |
|
32
|
181 |
// resize to fit buf_size more bytes
|
|
182 |
resize(buf_size);
|
31
|
183 |
|
|
184 |
// copy into our internal buffer
|
|
185 |
memcpy(buf + offset, buf_ptr, buf_size);
|
|
186 |
}
|
|
187 |
|
|
188 |
void NetworkBuffer::flush_write (void) {
|
|
189 |
int ret;
|
|
190 |
|
|
191 |
// ignore if we don't have any data buffered
|
|
192 |
if (offset == 0)
|
|
193 |
return;
|
|
194 |
|
|
195 |
// attempt to write as much as possible
|
|
196 |
try {
|
|
197 |
ret = socket.send(buf, offset)
|
|
198 |
|
|
199 |
} catch (CL_Error &e) {
|
|
200 |
// ignore EAGAIN and just return
|
|
201 |
if (errno == EAGAIN)
|
|
202 |
return;
|
|
203 |
|
|
204 |
else
|
|
205 |
throw;
|
|
206 |
}
|
|
207 |
|
32
|
208 |
// trim the buffer
|
|
209 |
trim(ret);
|
|
210 |
}
|
|
211 |
|
|
212 |
template <typename PrefixType> void NetworkBuffer::write_prefix (char *buf, PrefixType prefix) {
|
|
213 |
push_write(&prefix, sizeof(PrefixType));
|
|
214 |
push_write(buf, prefix);
|
|
215 |
}
|
31
|
216 |
|
|
217 |
|
|
218 |
NetworkTCPTransport::NetworkTCPTransport (CL_Socket socket) :
|
|
219 |
socket(socket), in(NETWORK_TCP_INITIAL_IN_BUF), out(NETWORK_TCP_INITIAL_OUT_BUF) {
|
|
220 |
|
|
221 |
// use nonblocking sockets
|
|
222 |
socket.set_nonblocking(true);
|
|
223 |
|
|
224 |
// connect signals
|
|
225 |
slots.connect(socket.sig_read_triggered(), this, &NetworkTCPTransport::on_read);
|
|
226 |
slots.connect(socket.sig_write_triggered(), this, &NetworkTCPTransport::on_write);
|
|
227 |
slots.connect(socket.sig_disconnected_triggered(), this, &NetworkTCPTransport::on_disconnected);
|
|
228 |
}
|
|
229 |
|
|
230 |
|
|
231 |
void NetworkTCPTransport::on_read (void) {
|
32
|
232 |
uint16_t prefix;
|
31
|
233 |
NetworkPacket packet;
|
32
|
234 |
|
|
235 |
// let the in stream read length-prefixed packets and pass them on to handle_packet
|
|
236 |
while ((prefix = in.read_prefix(packet.get_buf(), packet.get_buf_size)) > 0) {
|
|
237 |
packet.set_data_size(prefix);
|
|
238 |
sig_packet(packet);
|
|
239 |
}
|
31
|
240 |
}
|
|
241 |
|
|
242 |
void NetworkTCPTransport::on_write (void) {
|
|
243 |
// just flush the output buffer
|
|
244 |
out.flush_write();
|
|
245 |
}
|
|
246 |
|
|
247 |
void NetworkTCPTransport::on_disconnected (void) {
|
32
|
248 |
// pass right through
|
|
249 |
sig_disconnect();
|
31
|
250 |
}
|
|
251 |
|
|
252 |
void NetworkTCPTransport::write_packet (const NetworkPacket &packet) {
|
32
|
253 |
uint16_t prefix = packet.get_data_size();
|
|
254 |
|
|
255 |
if (prefix != packet.get_data_size())
|
|
256 |
throw CL_Error("send prefix overflow");
|
|
257 |
|
31
|
258 |
// just write to the output buffer
|
32
|
259 |
out.write_prefix(packet.get_buf(), prefix);
|
31
|
260 |
}
|
|
261 |
|
|
262 |
NetworkTCPServer::NetworkTCPServer (const NetworkAddress &listen_addr) :
|
|
263 |
socket(tcp, ipv4) {
|
|
264 |
|
|
265 |
// bind
|
|
266 |
socket.bind(listen_addr);
|
|
267 |
|
|
268 |
// assign slots
|
|
269 |
slots.connect(socket.sig_read_triggered(), this, &NetworkTCPServer::on_accept);
|
|
270 |
|
|
271 |
// listen
|
|
272 |
socket.listen(NETWORK_LISTEN_BACKLOG);
|
|
273 |
}
|
|
274 |
|
|
275 |
|
|
276 |
void NetworkTCPServer::on_accept (void) {
|
|
277 |
// accept a new socket
|
|
278 |
NetworkSocket client_sock = socket.accept();
|
|
279 |
|
|
280 |
// create a new NetworkTCPTransport
|
|
281 |
NetworkTCPTransport *client = new NetworkTCPTransport(client_sock);
|
|
282 |
|
|
283 |
// let our user handle it
|
32
|
284 |
sig_client(client);
|
31
|
285 |
}
|
|
286 |
|
|
287 |
NetworkTCPClient::NetworkTCPClient (const NetworkAddress &connect_addr) :
|
|
288 |
NetworkTCPTransport(NetworkSocket(tcp, ipv4)) {
|
|
289 |
|
|
290 |
// connect
|
|
291 |
socket.connect(connect_addr);
|
|
292 |
|
|
293 |
}
|