1/*
2 * A very simple controlable traffic generator for TCP testing.
3 *
4 * Copyright 2007, Haiku, Inc. All Rights Reserved.
5 * Distributed under the terms of the MIT License.
6 *
7 * Authors:
8 *      Hugo Santos, hugosantos@gmail.com
9 */
10
11#include <OS.h>
12
13#include <assert.h>
14#include <ctype.h>
15#include <stdio.h>
16#include <stdint.h>
17#include <stdlib.h>
18#include <string.h>
19#include <unistd.h>
20
21#include <arpa/inet.h>
22#include <netinet/in.h>
23#include <sys/socket.h>
24
25
26struct context {
27	int sock;
28	uint8 generator;
29	int index;
30	int8_t buffer[256];
31};
32
33static int process_command(context *ctx);
34
35static int
36number(context *ctx)
37{
38	int result = 0;
39
40	while (isdigit(ctx->buffer[ctx->index])) {
41		result *= 10;
42		result += ctx->buffer[ctx->index] - '0';
43		ctx->index++;
44	}
45
46	return result;
47}
48
49
50static int
51value(context *ctx)
52{
53	if (ctx->buffer[ctx->index] == '[') {
54		ctx->index++;
55		int upper, lower = number(ctx);
56		if (ctx->buffer[ctx->index] == ',') {
57			ctx->index++;
58			upper = number(ctx);
59		} else {
60			upper = lower + 50;
61			lower -= 50;
62		}
63
64		return lower + rand() % (upper - lower + 1);
65	}
66
67	return number(ctx);
68}
69
70
71static int
72repeat(context *ctx)
73{
74	int max, saved, count = number(ctx);
75
76	max = saved = ctx->index;
77	for (int i = 0; i < count; i++) {
78		ctx->index = saved;
79		if (process_command(ctx) < 0)
80			return -1;
81		if (ctx->index > max)
82			max = ctx->index;
83	}
84
85	ctx->index = max;
86	return 0;
87}
88
89
90static void
91send_packet(context *ctx, size_t bytes)
92{
93	uint8_t buffer[1024];
94	uint8_t *ptr = buffer;
95
96	if (bytes > sizeof(buffer))
97		ptr = new uint8_t[bytes];
98
99	for (size_t i = 0; i < bytes; i++) {
100		ptr[i] = ctx->generator + '0';
101		ctx->generator = (ctx->generator + 1) % 10;
102	}
103
104	send(ctx->sock, ptr, bytes, 0);
105
106	if (ptr != buffer)
107		delete [] ptr;
108}
109
110
111static int
112process_command(context *ctx)
113{
114	while (ctx->buffer[ctx->index] != '.') {
115		ctx->index++;
116
117		switch (ctx->buffer[ctx->index - 1]) {
118			case 'r':
119				if (repeat(ctx) < 0)
120					return -1;
121				break;
122
123			case 'b':
124				send_packet(ctx, 1);
125				break;
126
127			case 'p':
128				send_packet(ctx, value(ctx));
129				break;
130
131			case 's':
132				usleep(value(ctx) * 1000);
133				break;
134
135			case 'W':
136			{
137				int value = number(ctx);
138				setsockopt(ctx->sock, SOL_SOCKET, SO_SNDBUF, &value,
139					sizeof(value));
140				break;
141			}
142
143			case 'k':
144				return -1;
145		}
146	}
147
148	return 0;
149}
150
151
152static int
153read_command(context *ctx)
154{
155	int index = 0;
156
157	do {
158		int size = recv(ctx->sock, ctx->buffer + index, 1, 0);
159		if (size == 0)
160			return -1;
161		else if (size < 0)
162			continue;
163
164		index++;
165	} while (ctx->buffer[index - 1] != '.');
166
167	ctx->index = 0;
168	return process_command(ctx);
169}
170
171
172static int32
173handle_client(void *data)
174{
175	context ctx = { *(int *)data, 0 };
176
177	while (read_command(&ctx) == 0);
178
179	fprintf(stderr, "Client %d leaving.\n", ctx.sock);
180
181	close(ctx.sock);
182
183	return 0;
184}
185
186
187int
188main(int argc, char *argv[])
189{
190	int port = 12345;
191
192	for (int i = 1; i < argc; i++) {
193		if (!strcmp(argv[i], "-p")) {
194			i++;
195			assert(i < argc);
196			port = atoi(argv[i]);
197		} else if (!strcmp(argv[i], "-h")) {
198			fprintf(stderr, "tcptester [-p port]\n");
199			return 1;
200		}
201	}
202
203	int sock = socket(AF_INET, SOCK_STREAM, 0);
204
205	if (sock < 0) {
206		perror("socket()");
207		return -1;
208	}
209
210	sockaddr_in sin;
211	memset(&sin, 0, sizeof(sin));
212	sin.sin_family = AF_INET;
213	sin.sin_port = htons(port);
214
215	if (bind(sock, (sockaddr *)&sin, sizeof(sockaddr_in)) < 0) {
216		perror("bind()");
217		return -1;
218	}
219
220	if (listen(sock, 5) < 0) {
221		perror("listen()");
222		return -1;
223	}
224
225	while (1) {
226		sockaddr_in peer;
227		socklen_t peerLen = sizeof(peer);
228
229		int newSock = accept(sock, (sockaddr *)&peer, &peerLen);
230		if (newSock < 0) {
231			perror("accept()");
232			return -1;
233		}
234
235		char buf[64];
236		inet_ntop(AF_INET, &peer.sin_addr, buf, sizeof(buf));
237
238		thread_id newThread = spawn_thread(handle_client, "client",
239			B_NORMAL_PRIORITY, &newSock);
240
241		fprintf(stderr, "New client %d from %s with thread id %ld.\n",
242			newSock, buf, (int32)newThread);
243
244		resume_thread(newThread);
245	}
246
247	return 0;
248}
249
250