cosmopolitan/tool/net/echoserver.c

301 lines
10 KiB
C

/*-*- mode:c;indent-tabs-mode:nil;c-basic-offset:2;tab-width:8;coding:utf-8 -*-│
│vi: set net ft=c ts=2 sts=2 sw=2 fenc=utf-8 :vi│
╞══════════════════════════════════════════════════════════════════════════════╡
│ Copyright 2020 Justine Alexandra Roberts Tunney │
│ │
│ This program is free software; you can redistribute it and/or modify │
│ it under the terms of the GNU General Public License as published by │
│ the Free Software Foundation; version 2 of the License. │
│ │
│ This program is distributed in the hope that it will be useful, but │
│ WITHOUT ANY WARRANTY; without even the implied warranty of │
│ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU │
│ General Public License for more details. │
│ │
│ You should have received a copy of the GNU General Public License │
│ along with this program; if not, write to the Free Software │
│ Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA │
│ 02110-1301 USA │
╚─────────────────────────────────────────────────────────────────────────────*/
#include "libc/alg/arraylist.internal.h"
#include "libc/bits/safemacros.internal.h"
#include "libc/calls/calls.h"
#include "libc/calls/struct/iovec.h"
#include "libc/errno.h"
#include "libc/fmt/fmt.h"
#include "libc/log/check.h"
#include "libc/log/log.h"
#include "libc/runtime/gc.h"
#include "libc/runtime/interruptiblecall.h"
#include "libc/runtime/runtime.h"
#include "libc/sock/sock.h"
#include "libc/stdio/stdio.h"
#include "libc/str/str.h"
#include "libc/sysv/consts/af.h"
#include "libc/sysv/consts/exit.h"
#include "libc/sysv/consts/ipproto.h"
#include "libc/sysv/consts/msg.h"
#include "libc/sysv/consts/poll.h"
#include "libc/sysv/consts/so.h"
#include "libc/sysv/consts/sock.h"
#include "libc/sysv/consts/sol.h"
#include "libc/x/x.h"
#include "third_party/getopt/getopt.h"
/**
* @fileoverview Asynchronous TCP/UDP Echo Server.
*
* make -j8 o/default/tool/net/echoserver.com
* o/default/tool/net/echoserver.com udp:0.0.0.0:7 tcp:0.0.0.0:7
*/
enum SocketKind {
kSocketServer,
kSocketClient,
};
struct Message {
struct iovec data;
struct sockaddr_in dest;
uint32_t destsize;
};
struct Messages {
size_t i, n;
struct Message *p;
};
struct Socket {
int64_t fd;
enum SocketKind kind;
int type;
int protocol;
struct sockaddr_in addr;
struct Messages egress; /* LIFO */
};
struct Sockets {
size_t i, n;
struct Socket *p;
};
struct Polls {
size_t i, n;
struct pollfd *p;
};
struct Sockets g_sockets;
struct Polls g_polls;
nodiscard char *DescribeAddress(struct sockaddr_in *addr) {
char ip4buf[16];
return xasprintf("%s:%hu",
inet_ntop(addr->sin_family, &addr->sin_addr.s_addr, ip4buf,
sizeof(ip4buf)),
ntohs(addr->sin_port));
}
nodiscard char *DescribeSocket(struct Socket *s) {
return xasprintf("%s:%s", s->protocol == IPPROTO_UDP ? "udp" : "tcp",
gc(DescribeAddress(&s->addr)));
}
noreturn void ShowUsageAndExit(bool iserror) {
FILE *f = iserror ? stderr : stdout;
int rc = iserror ? EXIT_FAILURE : EXIT_SUCCESS;
fprintf(f, "%s: %s %s\n", "Usage", g_argv[0], "PROTOCOL:ADDR:PORT...");
exit(rc);
}
void GetFlags(int argc, char *argv[]) {
int opt;
while ((opt = getopt(argc, argv, "h")) != -1) {
switch (opt) {
case 'h':
ShowUsageAndExit(false);
default:
ShowUsageAndExit(true);
}
}
if (optind == argc) ShowUsageAndExit(true);
}
void AddSocket(const struct Socket *s) {
struct pollfd pfd;
pfd.fd = s->fd;
pfd.events = POLLIN;
pfd.revents = 0;
CHECK_NE(-1L, append(&g_sockets, s));
CHECK_NE(-1L, append(&g_polls, &pfd));
}
void RemoveSocket(size_t i) {
DCHECK_LT(i, g_sockets.i);
LOGF("removing: %s", gc(DescribeSocket(&g_sockets.p[i])));
CHECK_NE(-1, close(g_sockets.p[i].fd));
while (g_sockets.p[i].egress.i) {
free(g_sockets.p[i].egress.p[g_sockets.p[i].egress.i - 1].data.iov_base);
}
memcpy(&g_sockets.p[i], &g_sockets.p[i + 1],
(intptr_t)&g_sockets.p[g_sockets.i] - (intptr_t)&g_sockets.p[i + 1]);
memcpy(&g_polls.p[i], &g_polls.p[i + 1],
(intptr_t)&g_polls.p[g_polls.i] - (intptr_t)&g_polls.p[i + 1]);
g_sockets.i--;
g_polls.i--;
}
void GetListeningAddressesFromCommandLine(int argc, char *argv[]) {
int i;
for (i = optind; i < argc; ++i) {
struct Socket server;
memset(&server, 0, sizeof(server));
char scheme[4];
unsigned char *ip4 = (unsigned char *)&server.addr.sin_addr.s_addr;
uint16_t port;
if (sscanf(argv[i], "%3s:%hhu.%hhu.%hhu.%hhu:%hu", scheme, &ip4[0], &ip4[1],
&ip4[2], &ip4[3], &port) != 6) {
fprintf(stderr, "error: bad ip4 uri\n");
ShowUsageAndExit(true);
}
server.fd = -1;
server.kind = kSocketServer;
server.addr.sin_family = AF_INET;
server.addr.sin_port = htons(port);
if (strcasecmp(scheme, "tcp") == 0) {
server.type = SOCK_STREAM;
server.protocol = IPPROTO_TCP;
} else if (strcasecmp(scheme, "udp") == 0) {
server.type = SOCK_DGRAM;
server.protocol = IPPROTO_UDP;
} else {
fprintf(stderr, "%s: %s\n", "error", "bad scheme (should be tcp or udp)");
ShowUsageAndExit(true);
}
AddSocket(&server);
}
}
void BeginListeningForIncomingTraffic(void) {
size_t i;
for (i = 0; i < g_sockets.i; ++i) {
int yes = 1;
struct Socket *s = &g_sockets.p[i];
CHECK_NE(-1L,
(g_polls.p[i].fd = s->fd = socket(
s->addr.sin_family, s->type | SOCK_NONBLOCK, s->protocol)));
CHECK_NE(-1L,
setsockopt(s->fd, SOL_SOCKET, SO_REUSEPORT, &yes, sizeof(yes)));
CHECK_NE(-1, bind(s->fd, &s->addr, sizeof(s->addr)));
if (s->protocol == IPPROTO_TCP) {
CHECK_NE(-1, listen(s->fd, 1));
}
uint32_t addrsize = sizeof(s->addr);
CHECK_NE(-1, getsockname(s->fd, &s->addr, &addrsize));
LOGF("listening on %s", gc(DescribeSocket(s)));
}
}
void AcceptConnection(size_t i) {
struct Socket *server = &g_sockets.p[i];
struct Socket client;
memset(&client, 0, sizeof(client));
client.kind = kSocketClient;
client.type = server->type;
client.protocol = server->protocol;
uint32_t addrsize = sizeof(client.addr);
CHECK_NE(-1L, (client.fd = accept4(server->fd, &client.addr, &addrsize,
SOCK_NONBLOCK)));
LOGF("%s accepted %s", gc(DescribeSocket(server)),
gc(DescribeSocket(&client)));
AddSocket(&client);
}
bool ReceiveData(size_t i) {
ssize_t got;
struct Message msg;
bool isudp = g_sockets.p[i].protocol == IPPROTO_UDP;
memset(&msg, 0, sizeof(msg));
msg.destsize = sizeof(msg.dest);
msg.data.iov_len = PAGESIZE;
msg.data.iov_base = xmalloc(msg.data.iov_len);
CHECK_NE(-1L, (got = recvfrom(g_sockets.p[i].fd, msg.data.iov_base,
msg.data.iov_len, 0, isudp ? &msg.dest : NULL,
isudp ? &msg.destsize : NULL)));
if (0 < got && got <= msg.data.iov_len) {
LOGF("%s received %lu bytes from %s", gc(DescribeSocket(&g_sockets.p[i])),
got, gc(DescribeAddress(&msg.dest)));
msg.data.iov_base = xrealloc(msg.data.iov_base, (msg.data.iov_len = got));
append(&g_sockets.p[i].egress, &msg);
g_polls.p[i].events |= POLLOUT;
return true;
} else {
RemoveSocket(i);
free_s(&msg.data.iov_base);
return false;
}
}
void SendData(size_t i) {
ssize_t sent;
struct Socket *s = &g_sockets.p[i];
struct Message *msg = &s->egress.p[s->egress.i - 1];
bool isudp = s->protocol == IPPROTO_UDP;
DCHECK(s->egress.i);
CHECK_NE(-1L, (sent = sendto(s->fd, msg->data.iov_base, msg->data.iov_len, 0,
isudp ? &msg->dest : NULL,
isudp ? msg->destsize : 0)));
LOGF("%s sent %lu bytes to %s", gc(DescribeSocket(s)), msg->data.iov_len,
gc(DescribeAddress(&msg->dest)));
if (!(msg->data.iov_len -= min((size_t)sent, (size_t)msg->data.iov_len))) {
free_s(&msg->data.iov_base);
if (!--s->egress.i) {
g_polls.p[i].events &= ~POLLOUT;
}
}
}
void HandleSomeNetworkTraffic(void) {
size_t i;
int eventcount;
CHECK_GE((eventcount = poll(g_polls.p, g_polls.i, -1)), 0);
for (i = 0; eventcount && i < g_sockets.i; ++i) {
if (!g_polls.p[i].revents) continue;
--eventcount;
if (g_polls.p[i].revents & (POLLERR | POLLHUP | POLLNVAL)) {
CHECK_EQ(kSocketClient, g_sockets.p[i].kind);
RemoveSocket(i);
} else {
if (g_polls.p[i].revents & POLLIN) {
if (g_sockets.p[i].kind == kSocketServer &&
g_sockets.p[i].protocol == IPPROTO_TCP) {
AcceptConnection(i);
} else {
if (!ReceiveData(i)) continue;
}
}
if (g_polls.p[i].revents & POLLOUT) {
SendData(i);
}
}
}
}
void EchoServer(void) {
for (;;) HandleSomeNetworkTraffic();
}
int main(int argc, char *argv[]) {
STATIC_YOINK("isatty");
GetFlags(argc, argv);
GetListeningAddressesFromCommandLine(argc, argv);
BeginListeningForIncomingTraffic();
struct InterruptibleCall icall;
memset(&icall, 0, sizeof(icall));
interruptiblecall(&icall, (void *)EchoServer, 0, 0, 0, 0);
fputc('\r', stderr);
LOGF("%s", "shutting down...");
size_t i;
for (i = g_sockets.i; i; --i) RemoveSocket(i - 1);
return 0;
}