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