tcpclient.c 9.26 KB
Newer Older
Michael Graff's avatar
Michael Graff committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
#include <config.h>

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>

#include <isc/assertions.h>
#include <isc/error.h>
#include <isc/mem.h>
#include <isc/task.h>
#include <isc/thread.h>
#include <isc/result.h>
#include <isc/socket.h>
#include <isc/timer.h>

#include <dns/types.h>
#include <dns/result.h>
#include <dns/name.h>
#include <dns/rdata.h>
#include <dns/rdatalist.h>
#include <dns/rdataset.h>
#include <dns/compress.h>

#include <sys/types.h>
#include <sys/socket.h>

#include <netinet/in.h>

#include <arpa/inet.h>

#define LOCK(lp) \
	RUNTIME_CHECK(isc_mutex_lock((lp)) == ISC_R_SUCCESS)
#define UNLOCK(lp) \
	RUNTIME_CHECK(isc_mutex_unlock((lp)) == ISC_R_SUCCESS)

#include "tcpclient.h"

static tcp_cctx_t *tcp_cctx_allocate(isc_mem_t *mctx);
static void tcp_cctx_free(tcp_cctx_t *ctx);

static void tcp_send(isc_task_t *task, isc_event_t *event);
static void tcp_recv_len(isc_task_t *task, isc_event_t *event);
static void tcp_recv_req(isc_task_t *task, isc_event_t *event);
static void tcp_accept(isc_task_t *task, isc_event_t *event);

static tcp_cctx_t *
tcp_cctx_allocate(isc_mem_t *mctx)
{
	tcp_cctx_t *ctx;

	ctx = isc_mem_get(mctx, sizeof(tcp_cctx_t));
	if (ctx == NULL)
		return (NULL);

	ctx->buf = NULL;
	ctx->buflen = 0;
	ctx->slot = 0;
	ctx->mctx = mctx;
	ctx->csock = NULL;

	ctx->count = 0; /* XXX */

	return (ctx);
}

static void
tcp_cctx_free(tcp_cctx_t *ctx)
{
	if (ctx->buf != NULL)
		isc_mem_put(ctx->mctx, ctx->buf, ctx->buflen);
	ctx->buf = NULL;
	isc_mem_put(ctx->mctx, ctx, sizeof(tcp_cctx_t));
}

static void
tcp_restart(isc_task_t *task, tcp_cctx_t *ctx)
{
Michael Graff's avatar
Michael Graff committed
79
#ifdef NOISY
Michael Graff's avatar
Michael Graff committed
80
	printf("Restarting listen on %u\n", ctx->slot);
Michael Graff's avatar
Michael Graff committed
81
#endif
Michael Graff's avatar
Michael Graff committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

	if (ctx->buf != NULL)
		isc_mem_put(ctx->mctx, ctx->buf, ctx->buflen);
	ctx->buf = NULL;
	ctx->buflen = 0;

	if (ctx->csock != NULL)
		isc_socket_detach(&ctx->csock);

	RUNTIME_CHECK(isc_socket_accept(ctx->parent->sock, task,
					tcp_accept, ctx)
		      == ISC_R_SUCCESS);

	isc_mem_stats(ctx->mctx, stdout);
}

static void
tcp_shutdown(isc_task_t *task, isc_event_t *event)
{
	tcp_cctx_t *ctx;
	tcp_listener_t *l;

	ctx = (tcp_cctx_t *)(event->arg);
	l = ctx->parent;

	LOCK(&l->lock);

	if (ctx->csock != NULL)
		isc_socket_detach(&ctx->csock);

	REQUIRE(l->nwactive > 0);

	/*
	 * remove our task from the list of tasks that the listener
	 * maintains by setting things to NULL, then freeing the
	 * pointers we maintain.
	 */
	INSIST(l->tasks[ctx->slot] == task);
	l->tasks[ctx->slot] = NULL;
	l->ctxs[ctx->slot] = NULL;

	l->nwactive--;

	UNLOCK(&l->lock);

Michael Graff's avatar
Michael Graff committed
127
#ifdef NOISY
Michael Graff's avatar
Michael Graff committed
128
	printf("Final shutdown slot %u\n", ctx->slot);
Michael Graff's avatar
Michael Graff committed
129
#endif
Michael Graff's avatar
Michael Graff committed
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
	tcp_cctx_free(ctx);

	isc_event_free(&event);
}

static void
tcp_recv_len(isc_task_t *task, isc_event_t *event)
{
	isc_socket_t *sock;
	isc_socketevent_t *dev;
	tcp_cctx_t *ctx;
	isc_region_t region;

	sock = event->sender;
	dev = (isc_socketevent_t *)event;
	ctx = (tcp_cctx_t *)(event->arg);

Michael Graff's avatar
Michael Graff committed
147
#ifdef NOISY
Michael Graff's avatar
Michael Graff committed
148
149
150
151
152
153
154
	printf("len Task %u (sock %p, base %p, length %d, n %d, result %d)\n",
	       ctx->slot, sock,
	       dev->region.base, dev->region.length,
	       dev->n, dev->result);
	printf("\tFrom: %s port %d\n",
	       inet_ntoa(dev->address.type.sin.sin_addr),
	       ntohs(dev->address.type.sin.sin_port));
Michael Graff's avatar
Michael Graff committed
155
#endif
Michael Graff's avatar
Michael Graff committed
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201

	if (dev->result == ISC_R_CANCELED) {
		isc_task_shutdown(task);

		isc_event_free(&event);

		return;
	}
	if (dev->result != ISC_R_SUCCESS) {
		tcp_restart(task, ctx);

		isc_event_free(&event);

		return;
	}

	/*
	 * Allocate the space needed to complete this request.
	 */
	ctx->buflen = ntohs(ctx->buflen);
	ctx->buf = isc_mem_get(ctx->mctx, ctx->buflen);
	if (ctx->buf == NULL) {
		printf("Out of memory!\n");
		tcp_restart(task, ctx);

		isc_event_free(&event);

		return;
	}

	region.base = ctx->buf;
	region.length = ctx->buflen;

	isc_socket_recv(sock, &region, ISC_FALSE,
			task, tcp_recv_req, event->arg);

	isc_event_free(&event);
}

static void
tcp_recv_req(isc_task_t *task, isc_event_t *event)
{
	isc_socket_t *sock;
	isc_socketevent_t *dev;
	tcp_cctx_t *ctx;
	isc_region_t region;
202
203
204
	unsigned char *cp;
	isc_uint16_t len;
	dns_result_t result;
Michael Graff's avatar
Michael Graff committed
205
206
207
208
209

	sock = event->sender;
	dev = (isc_socketevent_t *)event;
	ctx = (tcp_cctx_t *)(event->arg);

Michael Graff's avatar
Michael Graff committed
210
#ifdef NOISY
Michael Graff's avatar
Michael Graff committed
211
212
213
214
215
216
217
	printf("req Task %u (sock %p, base %p, length %d, n %d, result %d)\n",
	       ctx->slot, sock,
	       dev->region.base, dev->region.length,
	       dev->n, dev->result);
	printf("\tFrom: %s port %d\n",
	       inet_ntoa(dev->address.type.sin.sin_addr),
	       ntohs(dev->address.type.sin.sin_port));
Michael Graff's avatar
Michael Graff committed
218
#endif
Michael Graff's avatar
Michael Graff committed
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235

	if (dev->result == ISC_R_CANCELED) {
		isc_task_shutdown(task);

		isc_event_free(&event);

		return;
	}
	if (dev->result != ISC_R_SUCCESS) {
		tcp_restart(task, ctx);

		isc_event_free(&event);

		return;
	}

	/*
236
237
238
	 * Call the dispatch() function to actually process this packet.
	 * If it returns ISC_R_SUCCESS, we have a packet to transmit.
	 * do so.  If it returns anything else, drop this connection.
Michael Graff's avatar
Michael Graff committed
239
	 */
240
241
242
	region.base = ctx->buf;
	region.length = dev->n;
	result = ctx->parent->dispatch(ctx->mctx, &region, 2);
Michael Graff's avatar
Michael Graff committed
243
244
245
246
247

	if (ctx->buf != region.base) { /* clean up request */
		isc_mem_put(ctx->mctx, ctx->buf, ctx->buflen);
		ctx->buf = NULL;
	}
Michael Graff's avatar
Michael Graff committed
248
249

	/*
250
	 * Failure.  Close TCP client.
Michael Graff's avatar
Michael Graff committed
251
	 */
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
	if (result != DNS_R_SUCCESS) {
		tcp_restart(task, ctx);

		isc_event_free(&event);

		return;
	}

	/*
	 * Success.  Send the packet, after filling in the length at the
	 * front of the packet.
	 */
	len = region.length - 2;
	cp = region.base;
	*cp++ = (len & 0xff00) >> 8;
	*cp++ = (len & 0x00ff);

	isc_socket_send(sock, &region, task, tcp_send, ctx);
Michael Graff's avatar
Michael Graff committed
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327

	isc_event_free(&event);
}

static void
tcp_accept(isc_task_t *task, isc_event_t *event)
{
	isc_region_t region;
	isc_socket_newconnev_t *dev;
	isc_socket_t *sock;
	tcp_cctx_t *ctx;

	sock = event->sender;
	dev = (isc_socket_newconnev_t *)event;
	ctx = (tcp_cctx_t *)(event->arg);

	/*
	 * If we get an error, close the socket.  This routine will actually
	 * close the socket and restart a listen on the parent socket for
	 * this task.  If, however, the result is that the I/O was canceled,
	 * we are being asked to shut down.  Do so.
	 */
	if (dev->result == ISC_R_CANCELED) {
		isc_task_shutdown(task);

		isc_event_free(&event);

		return;
	}
	if (dev->result != ISC_R_SUCCESS) {
		tcp_restart(task, ctx);

		isc_event_free(&event);

		return;
	}

	ctx->csock = dev->newsocket;

	/*
	 * New connection.  Start the read.  In this case, the first read
	 * goes into the length field.
	 */
	region.length = 2;
	region.base = (unsigned char *)&ctx->buflen;

	RUNTIME_CHECK(isc_socket_recv(ctx->csock, &region, ISC_FALSE, task,
				      tcp_recv_len, ctx)
		      == ISC_R_SUCCESS);

	isc_event_free(&event);
}

static void
tcp_send(isc_task_t *task, isc_event_t *event)
{
	isc_socket_t *sock;
	isc_socketevent_t *dev;
328
329
	tcp_cctx_t *ctx;
	isc_region_t region;
Michael Graff's avatar
Michael Graff committed
330
331
332

	sock = event->sender;
	dev = (isc_socketevent_t *)event;
333
	ctx = (tcp_cctx_t *)(event->arg);
Michael Graff's avatar
Michael Graff committed
334

Michael Graff's avatar
Michael Graff committed
335
#ifdef NOISY
336
337
	printf("tcp_send: task %u\n\t(base %p, length %d, n %d, result %d)\n",
	       ctx->slot, dev->region.base, dev->region.length,
Michael Graff's avatar
Michael Graff committed
338
	       dev->n, dev->result);
Michael Graff's avatar
Michael Graff committed
339
#endif
Michael Graff's avatar
Michael Graff committed
340

341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
	/*
	 * release memory regardless of outcome.
	 */
	isc_mem_put(ctx->mctx, dev->region.base, dev->region.length);

	if (dev->result == ISC_R_CANCELED) {
		isc_task_shutdown(task);

		isc_event_free(&event);

		return;
	}
	if (dev->result != ISC_R_SUCCESS) {
		tcp_restart(task, ctx);

		isc_event_free(&event);

		return;
	}

	/*
	 * Queue up another receive.
	 */
	region.base = (unsigned char *)&ctx->buflen;
	region.length = 2;
	isc_socket_recv(sock, &region, ISC_FALSE, task, tcp_recv_len, ctx);
Michael Graff's avatar
Michael Graff committed
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393

	isc_event_free(&event);
}

tcp_listener_t *
tcp_listener_allocate(isc_mem_t *mctx, u_int nwmax)
{
	tcp_listener_t *l;

	l = isc_mem_get(mctx, sizeof(tcp_listener_t));
	if (l == NULL)
		return (NULL);

	if (isc_mutex_init(&l->lock) != ISC_R_SUCCESS) {
		isc_mem_put(mctx, l, sizeof(tcp_listener_t));

		UNEXPECTED_ERROR(__FILE__, __LINE__,
				 "isc_mutex_init() failed");

		return (NULL);
	}

	l->tasks = isc_mem_get(mctx, sizeof(isc_task_t *) * nwmax);
	RUNTIME_CHECK(l->tasks != NULL); /* XXX should be non-fatal? */
	l->ctxs = isc_mem_get(mctx, sizeof(tcp_cctx_t *) * nwmax);
	RUNTIME_CHECK(l->ctxs != NULL);  /* XXX should be non-fatal? */

394
395
396
397
	l->sock = NULL;
	l->nwstart = 0;
	l->nwkeep = 0;
	l->nwmax = nwmax;
Michael Graff's avatar
Michael Graff committed
398
	l->mctx = mctx;
399
400
	l->dispatch = NULL;
	l->nwactive = 0;
Michael Graff's avatar
Michael Graff committed
401
402
403
404
405
406
407

	return (l);
}

isc_result_t
tcp_listener_start(tcp_listener_t *l,
		   isc_socket_t *sock, isc_taskmgr_t *tmgr,
408
409
410
		   u_int nwstart, u_int nwkeep, u_int nwtimeout,
		   dns_result_t (*dispatch)(isc_mem_t *, isc_region_t *,
					    unsigned int))
Michael Graff's avatar
Michael Graff committed
411
412
413
{
	u_int i;

414
415
416
	(void)nwkeep;		/* Make compiler happy. */
	(void)nwtimeout;	/* Make compiler happy. */

Michael Graff's avatar
Michael Graff committed
417
418
	LOCK(&l->lock);
	INSIST(l->nwactive == 0);
419
	INSIST(dispatch != NULL);
Michael Graff's avatar
Michael Graff committed
420

421
	l->dispatch = dispatch;
Michael Graff's avatar
Michael Graff committed
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
	l->sock = sock;
	RUNTIME_CHECK(isc_socket_listen(sock, 0) == ISC_R_SUCCESS);

	for (i = 0 ; i < nwstart ; i++) {
		l->tasks[i] = NULL;
		RUNTIME_CHECK(isc_task_create(tmgr, NULL, 0, &l->tasks[i])
			      == ISC_R_SUCCESS);

		l->ctxs[i] = tcp_cctx_allocate(l->mctx);
		RUNTIME_CHECK(l->ctxs[i] != NULL);

		l->ctxs[i]->parent = l;
		l->ctxs[i]->slot = i;

		RUNTIME_CHECK(isc_task_onshutdown(l->tasks[i], tcp_shutdown,
						  l->ctxs[i])
			      == ISC_R_SUCCESS);

		RUNTIME_CHECK(isc_socket_accept(sock, l->tasks[i],
						tcp_accept, l->ctxs[i])
			      == ISC_R_SUCCESS);

		l->nwactive++;
	}

	UNLOCK(&l->lock);

	return (ISC_R_SUCCESS);
}