1#include <unistd.h>
2#include <memory.h>
3#include <netinet/in.h>
4#include <arpa/inet.h>
5#include <errno.h>
6#include <stdio.h>
7#include <stdlib.h>
8
9const unsigned short TEST_PORT = 40000;
10
11void usage() {
12	printf("server [tcp|udp] [4|6] [local-address]\n");
13	exit(1);
14}
15
16void recvLoop(int fd) {
17	for (;;)	{
18		char buffer[1000];
19		int ret = recv(fd, buffer, sizeof(buffer) - 1, 0);
20		if (ret < 0) {
21			perror("recv");
22			exit(-1);
23		}
24		if (ret == 0) {
25			printf("received EOF!\n");
26			break;
27		} else {
28			buffer[ret] = 0;
29			printf("received %d bytes: \"%s\"\n", ret, buffer);
30		}
31	}
32}
33
34int main(int argc, char *argv[]) {
35	int socketType = SOCK_DGRAM;
36	int socketFamily = AF_INET;
37	if (argc > 1) {
38		if (!strcmp(argv[1], "tcp")) {
39			socketType = SOCK_STREAM;
40		} else if (!strcmp(argv[1], "udp")) {
41			socketType = SOCK_DGRAM;
42		} else {
43			usage();
44		}
45	}
46	if (argc > 2) {
47		switch (atoi(argv[2])) {
48		case 4:
49			socketFamily = AF_INET;
50			break;
51		case 6:
52			socketFamily = AF_INET6;
53			break;
54		default:
55			usage();
56		}
57	}
58
59	sockaddr_storage localAddress;
60	memset(&localAddress, 0, sizeof(localAddress));
61	localAddress.ss_family = socketFamily;
62	((sockaddr_in *) &localAddress)->sin_port = htons(TEST_PORT);
63
64	if (argc > 3) {
65		do {
66			void *dstBuffer = &((sockaddr_in *) &localAddress)->sin_addr;
67			if (inet_pton(AF_INET, argv[3], dstBuffer) == 1) {
68				printf("using IPv4 local address\n");
69				localAddress.ss_family = AF_INET;
70				break;
71			}
72
73			dstBuffer = &((sockaddr_in6 *) &localAddress)->sin6_addr;
74			if (inet_pton(AF_INET6, argv[3], dstBuffer) == 1) {
75				printf("using IPv6 local address\n");
76				localAddress.ss_family = AF_INET6;
77				break;
78			}
79
80			usage();
81		} while (false);
82	}
83
84	int fd = socket(socketFamily, socketType, 0);
85	if (fd < 0) {
86		perror("socket");
87		return -1;
88	}
89
90	if (bind(fd, (sockaddr *)&localAddress, localAddress.ss_family == AF_INET ?
91			 sizeof(sockaddr_in) : sizeof(sockaddr_in6)) < 0) {
92		perror("bind");
93		return -1;
94	}
95
96	switch (socketType) {
97	case SOCK_DGRAM:
98		for (;;) {
99			recvLoop(fd);
100		}
101		break;
102	case SOCK_STREAM:
103		if (listen(fd, 5) < 0) {
104			perror("listen");
105			return 1;
106		}
107		for (;;) {
108			int clientfd = accept(fd, NULL, 0);
109			if (clientfd < 0) {
110				perror("accept");
111				return 1;
112			}
113			printf("TCP server: got some client!\n");
114			if (fork() != 0) {
115				// parent code
116				close(clientfd);
117				continue;
118			}
119			// child code
120			close(fd);
121			recvLoop(clientfd);
122			exit(0);
123		}
124		break;
125	}
126}
127