/*
 * usocks: SOCKS5 CONNECT proxy on io_uring.
 *
 * One thread runs one ring. recv, send, accept and connect go through the
 * ring; rare calls (socket, setsockopt, shutdown, close, handshake replies)
 * are plain syscalls. Each direction of a relayed connection has one fixed
 * buffer and at most one operation in flight. Domain names are resolved by
 * getaddrinfo() in resolver threads. With -w, worker processes share the
 * listening port through SO_REUSEPORT and nothing else.
 */
#define _GNU_SOURCE
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <linux/io_uring.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <pthread.h>
#include <signal.h>
#include <stdarg.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#define BUFSZ	65536
#define MAXCONN	16384
#define MAXWORK	64
#define BACKLOG	4096	/* clamped to net.core.somaxconn */
#define NRESOLV	16
#define TIMEOUT	10	/* seconds: handshake, lookup wait, connect attempt */

/* In RESOLVE a connection belongs to the resolver threads; nothing drops it. */
enum { FREE, GREET, AUTH, REQUEST, RESOLVE, CONNECT, RELAY, CLOSE };
enum { ACCEPT, RECV, SEND, CONN, RESOLVED, TIMER };

struct conn {
	int fd[2];		/* client, remote */
	int state;
	int inflight;		/* ring operations not yet completed */
	int len[2];		/* bytes in buf[d] */
	int off[2];		/* bytes of buf[d] already sent */
	int eof;		/* bit d: direction d reached EOF */
	int err;		/* errno of the last connect attempt */
	int gai;		/* getaddrinfo() result */
	time_t deadline;	/* for handshake and connect */
	struct sockaddr_storage peer;
	socklen_t peerlen;
	char host[256];
	char port[6];
	struct addrinfo *ai, *cur;
	char buf[2][BUFSZ];	/* direction 0: client to remote, 1: back */
};

static struct {
	int fd;
	unsigned *sqhead, *sqtail, sqmask, sqentries, tail;
	unsigned *cqhead, *cqtail, cqmask;
	struct io_uring_sqe *sqes;
	struct io_uring_cqe *cqes;
} ring;

static struct conn *conns;
static int *freel, nfree, maxconn = 1024;
static int lfd, accepting, workers = 1;
static time_t now;
static const char *user, *pass;

/* Lookups: conn indexes in through queue, out through respipe. */
static pthread_mutex_t qlock = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t qcond = PTHREAD_COND_INITIALIZER;
static int *queue, qhead, qlen;
static int respipe[2], resbuf[256];

__attribute__((noreturn, format(printf, 1, 2)))
static void die(const char *fmt, ...)
{
	va_list ap;

	fputs("usocks: ", stderr);
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	fputc('\n', stderr);
	exit(1);
}

static void setup(unsigned entries, unsigned cqentries)
{
	struct io_uring_params p = { 0 };
	unsigned *array, i;
	size_t sz, cqsz;
	char *m;

	p.flags = IORING_SETUP_SINGLE_ISSUER | IORING_SETUP_DEFER_TASKRUN |
		  IORING_SETUP_SUBMIT_ALL | IORING_SETUP_CQSIZE;
	p.cq_entries = cqentries;
	ring.fd = syscall(SYS_io_uring_setup, entries, &p);
	if (ring.fd < 0) {
		die("io_uring_setup: %s (needs Linux 6.1 or newer with "
		    "io_uring enabled)", strerror(errno));
	}

	sz = p.sq_off.array + p.sq_entries * sizeof(unsigned);
	cqsz = p.cq_off.cqes + p.cq_entries * sizeof(struct io_uring_cqe);
	if (sz < cqsz) {
		sz = cqsz;
	}
	m = mmap(NULL, sz, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_POPULATE,
		 ring.fd, IORING_OFF_SQ_RING);
	if (m == MAP_FAILED) {
		die("mmap: %s", strerror(errno));
	}
	ring.sqes = mmap(NULL, p.sq_entries * sizeof(struct io_uring_sqe),
			 PROT_READ | PROT_WRITE, MAP_SHARED | MAP_POPULATE,
			 ring.fd, IORING_OFF_SQES);
	if (ring.sqes == MAP_FAILED) {
		die("mmap: %s", strerror(errno));
	}

	ring.sqhead = (unsigned *)(m + p.sq_off.head);
	ring.sqtail = (unsigned *)(m + p.sq_off.tail);
	ring.sqmask = *(unsigned *)(m + p.sq_off.ring_mask);
	ring.sqentries = p.sq_entries;
	array = (unsigned *)(m + p.sq_off.array);
	for (i = 0; i < p.sq_entries; i++) {
		array[i] = i;
	}
	ring.cqhead = (unsigned *)(m + p.cq_off.head);
	ring.cqtail = (unsigned *)(m + p.cq_off.tail);
	ring.cqmask = *(unsigned *)(m + p.cq_off.ring_mask);
	ring.cqes = (struct io_uring_cqe *)(m + p.cq_off.cqes);
}

/* Submit queued SQEs; with wait, block until at least one CQE is ready. */
static void enter(unsigned wait)
{
	unsigned n;

	__atomic_store_n(ring.sqtail, ring.tail, __ATOMIC_RELEASE);
	n = ring.tail - __atomic_load_n(ring.sqhead, __ATOMIC_ACQUIRE);
	while (syscall(SYS_io_uring_enter, ring.fd, n, wait,
		       wait ? IORING_ENTER_GETEVENTS : 0, NULL, 0) < 0) {
		if (errno != EINTR) {
			die("io_uring_enter: %s", strerror(errno));
		}
	}
}

static struct io_uring_sqe *sqe(struct conn *c, int op, int d)
{
	struct io_uring_sqe *s;
	unsigned head;

	head = __atomic_load_n(ring.sqhead, __ATOMIC_ACQUIRE);
	if (ring.tail - head == ring.sqentries) {
		enter(0);
		head = __atomic_load_n(ring.sqhead, __ATOMIC_ACQUIRE);
		if (ring.tail - head == ring.sqentries) {
			die("submission queue stuck");
		}
	}
	s = &ring.sqes[ring.tail++ & ring.sqmask];
	memset(s, 0, sizeof(*s));
	s->user_data = (uint64_t)(c - conns) << 8 | d << 4 | op;
	return s;
}

static void queue_accept(void)
{
	struct io_uring_sqe *s;
	struct conn *c;

	if (!nfree) {
		return;
	}
	c = &conns[freel[--nfree]];
	memset(c, 0, offsetof(struct conn, buf));
	c->peerlen = sizeof(c->peer);
	s = sqe(c, ACCEPT, 0);
	s->opcode = IORING_OP_ACCEPT;
	s->fd = lfd;
	s->addr = (uintptr_t)&c->peer;
	s->addr2 = (uintptr_t)&c->peerlen;
	s->accept_flags = SOCK_CLOEXEC;
	accepting = 1;
}

static void queue_recv(struct conn *c, int d)
{
	struct io_uring_sqe *s = sqe(c, RECV, d);

	s->opcode = IORING_OP_RECV;
	s->fd = c->fd[d];
	s->addr = (uintptr_t)(c->buf[d] + c->len[d]);
	s->len = BUFSZ - c->len[d];
	c->inflight++;
}

static void queue_send(struct conn *c, int d)
{
	struct io_uring_sqe *s = sqe(c, SEND, d);

	s->opcode = IORING_OP_SEND;
	s->fd = c->fd[!d];
	s->addr = (uintptr_t)(c->buf[d] + c->off[d]);
	s->len = c->len[d] - c->off[d];
	s->msg_flags = MSG_NOSIGNAL;
	c->inflight++;
}

static void queue_results(void)
{
	struct io_uring_sqe *s = sqe(conns, RESOLVED, 0);

	s->opcode = IORING_OP_READ;
	s->fd = respipe[0];
	s->addr = (uintptr_t)resbuf;
	s->len = sizeof(resbuf);
}

static void queue_timer(void)
{
	static struct __kernel_timespec ts = { .tv_sec = 1 };
	struct io_uring_sqe *s = sqe(conns, TIMER, 0);

	s->opcode = IORING_OP_TIMEOUT;
	s->addr = (uintptr_t)&ts;
	s->len = 1;
}

static void queue_connect(struct conn *c)
{
	struct io_uring_sqe *s = sqe(c, CONN, 0);

	s->opcode = IORING_OP_CONNECT;
	s->fd = c->fd[1];
	s->addr = (uintptr_t)c->cur->ai_addr;
	s->off = c->cur->ai_addrlen;
	c->inflight++;
}

static void release(struct conn *c)
{
	close(c->fd[0]);
	if (c->fd[1] >= 0) {
		close(c->fd[1]);
	}
	if (c->ai) {
		freeaddrinfo(c->ai);
	}
	c->state = FREE;
	freel[nfree++] = c - conns;
	if (!accepting) {
		queue_accept();
	}
}

/*
 * shutdown() completes pending recvs with 0, fails pending sends and aborts
 * a pending connect. The slot is released once the last completion arrives.
 * Errors are irrelevant here: the socket is closing either way.
 */
static void drop(struct conn *c)
{
	c->state = CLOSE;
	shutdown(c->fd[0], SHUT_RDWR);
	if (c->fd[1] >= 0) {
		shutdown(c->fd[1], SHUT_RDWR);
	}
	if (!c->inflight) {
		release(c);
	}
}

static int fail(struct conn *c, const char *why)
{
	char a[INET6_ADDRSTRLEN] = "?";
	struct sockaddr_in *v4 = (struct sockaddr_in *)&c->peer;
	struct sockaddr_in6 *v6 = (struct sockaddr_in6 *)&c->peer;

	if (c->peer.ss_family == AF_INET) {
		fprintf(stderr, "%s:%u", inet_ntop(AF_INET, &v4->sin_addr,
			a, sizeof(a)), ntohs(v4->sin_port));
	} else if (c->peer.ss_family == AF_INET6) {
		fprintf(stderr, "[%s]:%u", inet_ntop(AF_INET6, &v6->sin6_addr,
			a, sizeof(a)), ntohs(v6->sin6_port));
	} else {
		fputs("unix", stderr);
	}
	if (c->host[0]) {
		fprintf(stderr, strchr(c->host, ':') ? " [%s]:%s" : " %s:%s",
			c->host, c->port);
	}
	fprintf(stderr, ": %s\n", why);
	drop(c);
	return -1;
}

/* Handshake replies are tiny and go to a fresh socket: they cannot block. */
static int reply(struct conn *c, const void *b, int n)
{
	return send(c->fd[0], b, n, MSG_DONTWAIT | MSG_NOSIGNAL) == n ? 0 : -1;
}

static int refuse(struct conn *c, int rep, const char *why)
{
	unsigned char b[10] = { 5, rep, 0, 1 };

	reply(c, b, sizeof(b));
	return fail(c, why);
}

/* Keepalive reaps dead peers; relayed connections have no idle timeout. */
static int sockopts(int fd)
{
	int one = 1, idle = 60, intvl = 10, cnt = 6;
	socklen_t n = sizeof(int);

	if (setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, n) < 0 ||
	    setsockopt(fd, SOL_SOCKET, SO_KEEPALIVE, &one, n) < 0 ||
	    setsockopt(fd, IPPROTO_TCP, TCP_KEEPIDLE, &idle, n) < 0 ||
	    setsockopt(fd, IPPROTO_TCP, TCP_KEEPINTVL, &intvl, n) < 0 ||
	    setsockopt(fd, IPPROTO_TCP, TCP_KEEPCNT, &cnt, n) < 0) {
		return -1;
	}
	return 0;
}

static int rep(int e)
{
	switch (e) {
	case ENETUNREACH:
		return 3;
	case EHOSTUNREACH:
		return 4;
	case ECONNREFUSED:
		return 5;
	case ETIMEDOUT:
		return 6;
	case EAFNOSUPPORT:
		return 8;
	}
	return 1;
}

static void connect_next(struct conn *c)
{
	int fd;

	for (; c->cur; c->cur = c->cur->ai_next) {
		fd = socket(c->cur->ai_family, SOCK_STREAM | SOCK_CLOEXEC, 0);
		if (fd < 0) {
			c->err = errno;
			continue;
		}
		if (sockopts(fd) < 0) {
			c->err = errno;
			close(fd);
			continue;
		}
		c->fd[1] = fd;
		c->deadline = now + TIMEOUT;
		queue_connect(c);
		return;
	}
	refuse(c, rep(c->err), strerror(c->err));
}

static void lookup(struct conn *c)
{
	static const struct addrinfo hints = {
		.ai_flags = AI_NUMERICHOST | AI_NUMERICSERV,
		.ai_socktype = SOCK_STREAM,
	};
	int r;

	r = getaddrinfo(c->host, c->port, &hints, &c->ai);
	if (r == EAI_NONAME) {
		c->deadline = now + TIMEOUT;
		pthread_mutex_lock(&qlock);
		queue[(qhead + qlen++) % maxconn] = c - conns;
		pthread_cond_signal(&qcond);
		pthread_mutex_unlock(&qlock);
		return;
	}
	if (r) {
		refuse(c, 1, gai_strerror(r));
		return;
	}
	c->state = CONNECT;
	c->cur = c->ai;
	connect_next(c);
}

static void *resolver(void *arg)
{
	static const struct addrinfo hints = {
		.ai_flags = AI_NUMERICSERV,
		.ai_socktype = SOCK_STREAM,
	};
	struct timespec ts;
	struct conn *c;
	int i;

	(void)arg;
	for (;;) {
		pthread_mutex_lock(&qlock);
		while (!qlen) {
			pthread_cond_wait(&qcond, &qlock);
		}
		i = queue[qhead];
		qhead = (qhead + 1) % maxconn;
		qlen--;
		pthread_mutex_unlock(&qlock);

		c = &conns[i];
		/* A request that waited past its deadline is not looked up. */
		clock_gettime(CLOCK_MONOTONIC, &ts);
		if (ts.tv_sec < c->deadline) {
			c->gai = getaddrinfo(c->host, c->port, &hints, &c->ai);
		} else {
			c->gai = EAI_AGAIN;
		}
		if (c->gai) {
			c->ai = NULL;
		}
		if (write(respipe[1], &i, sizeof(i)) != sizeof(i)) {
			die("write: %s", strerror(errno));
		}
	}
	return NULL;
}

static void resolved(int res)
{
	struct conn *c;
	int i;

	/* Writes of one int are atomic, so reads return whole ints. */
	if (res <= 0 || res % (int)sizeof(int)) {
		die("read: %s", res < 0 ? strerror(-res) : "short read");
	}
	for (i = 0; i < res / (int)sizeof(int); i++) {
		c = &conns[resbuf[i]];
		if (c->gai) {
			refuse(c, 4, now >= c->deadline ? "lookup timeout" :
			       gai_strerror(c->gai));
			continue;
		}
		c->state = CONNECT;
		c->cur = c->ai;
		connect_next(c);
	}
	queue_results();
}

/* Constant time in the contents of s. */
static int equal(const char *s, const unsigned char *p, int n)
{
	int i, m = strlen(s);
	unsigned char d = m != n;

	for (i = 0; i < n; i++) {
		d |= p[i] ^ (unsigned char)s[i < m ? i : 0];
	}
	return !d;
}

/* Bytes consumed from buf[0]; 0: more input needed; -1: dropped. */
static int parse(struct conn *c)
{
	unsigned char *p = (unsigned char *)c->buf[0];
	int n = c->len[0], m, ok, i;

	if (n < 1) {
		return 0;
	}
	switch (c->state) {
	case GREET:
		if (p[0] != 5) {
			return fail(c, "bad version");
		}
		if (n < 2 || n < 2 + p[1]) {
			return 0;
		}
		m = user ? 2 : 0;
		if (!memchr(p + 2, m, p[1])) {
			reply(c, "\5\377", 2);
			return fail(c, "no acceptable auth method");
		}
		if (reply(c, (unsigned char[]){ 5, m }, 2) < 0) {
			return fail(c, "reply failed");
		}
		c->state = user ? AUTH : REQUEST;
		return 2 + p[1];
	case AUTH:
		/* RFC 1929: 1, ulen, user, plen, pass */
		if (p[0] != 1) {
			return fail(c, "bad auth version");
		}
		if (n < 2 || n < 3 + p[1] || n < 3 + p[1] + p[2 + p[1]]) {
			return 0;
		}
		ok = equal(user, p + 2, p[1]);
		ok &= equal(pass, p + 3 + p[1], p[2 + p[1]]);
		if (!ok) {
			reply(c, "\1\1", 2);
			return fail(c, "auth failed");
		}
		if (reply(c, "\1\0", 2) < 0) {
			return fail(c, "reply failed");
		}
		c->state = REQUEST;
		return 3 + p[1] + p[2 + p[1]];
	case REQUEST:
		if (p[0] != 5) {
			return fail(c, "bad version");
		}
		if (n < 4) {
			return 0;
		}
		switch (p[3]) {
		case 1:
			m = 4 + 4 + 2;
			break;
		case 3:
			if (n < 5) {
				return 0;
			}
			m = 4 + 1 + p[4] + 2;
			break;
		case 4:
			m = 4 + 16 + 2;
			break;
		default:
			return refuse(c, 8, "bad address type");
		}
		if (n < m) {
			return 0;
		}
		if (p[1] != 1) {
			return refuse(c, 7, "command not supported");
		}
		if (p[3] == 1) {
			inet_ntop(AF_INET, p + 4, c->host, sizeof(c->host));
		} else if (p[3] == 4) {
			inet_ntop(AF_INET6, p + 4, c->host, sizeof(c->host));
		} else {
			/* Printable ASCII only: names reach logs. */
			for (i = 0; i < p[4]; i++) {
				if (p[5 + i] <= ' ' || p[5 + i] > '~') {
					break;
				}
			}
			if (!p[4] || i < p[4]) {
				return refuse(c, 1, "bad domain name");
			}
			memcpy(c->host, p + 5, p[4]);
		}
		snprintf(c->port, sizeof(c->port), "%u",
			 p[m - 2] << 8 | p[m - 1]);
		c->state = RESOLVE;
		return m;
	}
	return 0;
}

static void handshake(struct conn *c, int res)
{
	int n;

	if (res <= 0) {
		fail(c, res ? strerror(-res) : "closed during handshake");
		return;
	}
	c->len[0] += res;
	while ((n = parse(c)) > 0) {
		c->len[0] -= n;
		memmove(c->buf[0], c->buf[0] + n, c->len[0]);
	}
	if (n < 0) {
		return;
	}
	/* Bytes left after the request are early data for the remote. */
	if (c->state == RESOLVE) {
		lookup(c);
	} else {
		queue_recv(c, 0);
	}
}

static void connected(struct conn *c, int res)
{
	unsigned char b[22] = { 5, 0, 0 };
	struct sockaddr_storage ss;
	struct sockaddr_in *v4 = (struct sockaddr_in *)&ss;
	struct sockaddr_in6 *v6 = (struct sockaddr_in6 *)&ss;
	socklen_t len = sizeof(ss);
	int n;

	if (res < 0) {
		c->err = now >= c->deadline ? ETIMEDOUT : -res;
		close(c->fd[1]);
		c->fd[1] = -1;
		c->cur = c->cur->ai_next;
		connect_next(c);
		return;
	}
	if (getsockname(c->fd[1], (struct sockaddr *)&ss, &len) < 0) {
		refuse(c, 1, strerror(errno));
		return;
	}
	if (ss.ss_family == AF_INET) {
		b[3] = 1;
		memcpy(b + 4, &v4->sin_addr, 4);
		memcpy(b + 8, &v4->sin_port, 2);
		n = 10;
	} else {
		b[3] = 4;
		memcpy(b + 4, &v6->sin6_addr, 16);
		memcpy(b + 20, &v6->sin6_port, 2);
		n = 22;
	}
	if (reply(c, b, n) < 0) {
		fail(c, "reply failed");
		return;
	}
	c->state = RELAY;
	if (c->len[0]) {
		queue_send(c, 0);
	} else {
		queue_recv(c, 0);
	}
	queue_recv(c, 1);
}

static void received(struct conn *c, int d, int res)
{
	if (res < 0) {
		drop(c);
		return;
	}
	if (res == 0) {
		c->eof |= 1 << d;
		if (c->eof == 3 || shutdown(c->fd[!d], SHUT_WR) < 0) {
			drop(c);
		}
		return;
	}
	c->len[d] = res;
	c->off[d] = 0;
	queue_send(c, d);
}

static void sent(struct conn *c, int d, int res)
{
	if (res < 0) {
		drop(c);
		return;
	}
	c->off[d] += res;
	if (c->off[d] < c->len[d]) {
		queue_send(c, d);
		return;
	}
	c->len[d] = 0;
	queue_recv(c, d);
}

static void accepted(struct conn *c, int res)
{
	accepting = 0;
	if (res < 0) {
		/* Errors of the pending connection; accept(2) says retry. */
		switch (-res) {
		case ECONNABORTED:
		case EPERM:
		case EPROTO:
		case ENOPROTOOPT:
		case ENETDOWN:
		case ENETUNREACH:
		case EHOSTDOWN:
		case EHOSTUNREACH:
		case ENONET:
		case EOPNOTSUPP:
			break;
		default:
			die("accept: %s", strerror(-res));
		}
		freel[nfree++] = c - conns;
		queue_accept();
		return;
	}
	c->fd[0] = res;
	c->fd[1] = -1;
	c->state = GREET;
	c->deadline = now + TIMEOUT;
	queue_accept();
	if (c->peer.ss_family != AF_UNIX && sockopts(c->fd[0]) < 0) {
		fail(c, strerror(errno));
		return;
	}
	queue_recv(c, 0);
}

/* A connect past its deadline is shut down, which fails it with ETIMEDOUT. */
static void sweep(void)
{
	struct conn *c;
	int i;

	for (i = 0; i < maxconn; i++) {
		c = &conns[i];
		if (now < c->deadline) {
			continue;
		}
		switch (c->state) {
		case GREET:
		case AUTH:
		case REQUEST:
			fail(c, "handshake timeout");
			break;
		case CONNECT:
			shutdown(c->fd[1], SHUT_RDWR);
			break;
		}
	}
}

static void complete(uint64_t ud, int res)
{
	struct conn *c = &conns[ud >> 8];
	int d = ud >> 4 & 1, op = ud & 15;

	if (op == ACCEPT) {
		accepted(c, res);
		return;
	}
	if (op == RESOLVED) {
		resolved(res);
		return;
	}
	if (op == TIMER) {
		if (res != -ETIME) {
			die("timer: %s", strerror(-res));
		}
		sweep();
		queue_timer();
		return;
	}
	c->inflight--;
	if (c->state == CLOSE) {
		if (!c->inflight) {
			release(c);
		}
		return;
	}
	switch (op) {
	case RECV:
		if (c->state == RELAY) {
			received(c, d, res);
		} else {
			handshake(c, res);
		}
		break;
	case SEND:
		sent(c, d, res);
		break;
	case CONN:
		connected(c, res);
		break;
	}
}

static int number(const char *s, int min, int max, const char *what)
{
	char *end;
	long n;

	errno = 0;
	n = strtol(s, &end, 10);
	if (errno || end == s || *end || n < min || n > max) {
		die("bad %s: %s", what, s);
	}
	return n;
}

/* s is addr:port, [addr6]:port or an absolute unix socket path. */
static int listener(const char *s)
{
	struct addrinfo hints = {
		.ai_flags = AI_NUMERICHOST | AI_NUMERICSERV | AI_PASSIVE,
		.ai_socktype = SOCK_STREAM,
	};
	struct sockaddr_un un = { .sun_family = AF_UNIX };
	struct addrinfo *ai = NULL;
	struct sockaddr *sa;
	struct stat st;
	char host[INET6_ADDRSTRLEN + 2];
	const char *port = strrchr(s, ':');
	socklen_t len;
	size_t n;
	int fd, one = 1;

	if (s[0] == '/') {
		n = strlen(s);
		if (n >= sizeof(un.sun_path)) {
			die("listen path too long: %s", s);
		}
		memcpy(un.sun_path, s, n + 1);
		/*
		 * A socket left by an earlier run would make bind fail. Only
		 * one that refuses connections is stale; a full backlog
		 * (EAGAIN) means a live listener.
		 */
		if (lstat(s, &st) == 0 && S_ISSOCK(st.st_mode)) {
			fd = socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK |
				    SOCK_CLOEXEC, 0);
			if (fd < 0) {
				die("socket: %s", strerror(errno));
			}
			if (!connect(fd, (struct sockaddr *)&un, sizeof(un)) ||
			    errno == EAGAIN) {
				die("listen path in use: %s", s);
			}
			if (errno != ECONNREFUSED) {
				die("connect %s: %s", s, strerror(errno));
			}
			close(fd);
			if (unlink(s) < 0) {
				die("unlink %s: %s", s, strerror(errno));
			}
		}
		sa = (struct sockaddr *)&un;
		len = sizeof(un);
	} else {
		if (!port || port - s >= (ptrdiff_t)sizeof(host)) {
			die("bad listen address: %s", s);
		}
		n = port - s;
		memcpy(host, s, n);
		host[n] = 0;
		if (n >= 2 && host[0] == '[' && host[n - 1] == ']') {
			host[n - 1] = 0;
			memmove(host, host + 1, n - 1);
		}
		number(port + 1, 1, 65535, "port");
		if (getaddrinfo(host, port + 1, &hints, &ai)) {
			die("bad listen address: %s", s);
		}
		sa = ai->ai_addr;
		len = ai->ai_addrlen;
	}
	fd = socket(sa->sa_family, SOCK_STREAM | SOCK_CLOEXEC, 0);
	if (fd < 0) {
		die("socket: %s", strerror(errno));
	}
	if (sa->sa_family != AF_UNIX &&
	    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) < 0) {
		die("setsockopt: %s", strerror(errno));
	}
	/* Each worker binds the port; the kernel spreads new connections. */
	if (workers > 1 &&
	    setsockopt(fd, SOL_SOCKET, SO_REUSEPORT, &one, sizeof(one)) < 0) {
		die("setsockopt: %s", strerror(errno));
	}
	if (bind(fd, sa, len) < 0) {
		die("bind %s: %s", s, strerror(errno));
	}
	if (listen(fd, BACKLOG) < 0) {
		die("listen: %s", strerror(errno));
	}
	if (ai) {
		freeaddrinfo(ai);
	}
	return fd;
}

static void usage(void)
{
	fputs("usage: usocks [-l addr:port|path] [-m maxconn] [-w workers] "
	      "[-u user -p pass]\n", stderr);
	exit(2);
}

int main(int argc, char **argv)
{
	static char errbuf[1024];
	const char *addr = getenv("USOCKS_LISTEN");
	const char *max = getenv("USOCKS_MAXCONN");
	const char *work = getenv("USOCKS_WORKERS");
	struct io_uring_cqe *cqe;
	struct timespec ts;
	uint64_t ud;
	struct rlimit rl;
	rlim_t nofile;
	unsigned head, tail, cqentries;
	pthread_t t;
	pid_t pid = 0, ppid;
	int i, o, res, status;

	/* One write per log line, so lines from workers do not mix. */
	setvbuf(stderr, errbuf, _IOLBF, sizeof(errbuf));
	user = getenv("USOCKS_USER");
	pass = getenv("USOCKS_PASS");
	while ((o = getopt(argc, argv, "l:m:w:u:p:")) != -1) {
		switch (o) {
		case 'l':
			addr = optarg;
			break;
		case 'm':
			max = optarg;
			break;
		case 'w':
			work = optarg;
			break;
		case 'u':
			user = optarg;
			break;
		case 'p':
			/* Blank argv so later ps and /proc reads miss it. */
			if (!(pass = strdup(optarg))) {
				die("out of memory");
			}
			memset(optarg, 0, strlen(optarg));
			break;
		default:
			usage();
		}
	}
	if (optind != argc) {
		usage();
	}
	if (!addr) {
		addr = "127.0.0.1:1080";
	}
	if (max) {
		maxconn = number(max, 1, MAXCONN, "maxconn");
	}
	if (work) {
		workers = number(work, 1, MAXWORK, "workers");
	}
	if (workers > 1 && addr[0] == '/') {
		die("workers need addr:port, not a unix socket path");
	}
	if (!user != !pass) {
		die("user and password must be set together");
	}
	if (user && (!*user || !*pass || strlen(user) > 255 ||
		     strlen(pass) > 255)) {
		die("user and password must be 1 to 255 bytes");
	}

	/* Two per conn, a few per lookup in flight, the rest fixed. */
	nofile = (rlim_t)maxconn * 2 + NRESOLV * 4 + 16;
	if (getrlimit(RLIMIT_NOFILE, &rl) < 0) {
		die("getrlimit: %s", strerror(errno));
	}
	if (rl.rlim_cur < nofile) {
		rl.rlim_cur = nofile;
		if (setrlimit(RLIMIT_NOFILE, &rl) < 0) {
			die("maxconn %d needs RLIMIT_NOFILE %llu, hard "
			    "limit is %llu", maxconn,
			    (unsigned long long)nofile,
			    (unsigned long long)rl.rlim_max);
		}
	}

	/* With THP one touched slot pins a whole 2 MB page. */
	if (prctl(PR_SET_THP_DISABLE, 1, 0, 0, 0) < 0) {
		die("prctl: %s", strerror(errno));
	}

	/*
	 * Fork before any ring or thread exists. When one worker exits, the
	 * parent exits and PDEATHSIG ends the others. Inherited signal state
	 * must not interfere: an ignored SIGCHLD makes wait() wait for all
	 * children, and SIGKILL cannot be ignored or blocked.
	 */
	if (workers > 1) {
		signal(SIGCHLD, SIG_DFL);
		ppid = getpid();
		for (i = 0; i < workers; i++) {
			if ((pid = fork()) < 0) {
				die("fork: %s", strerror(errno));
			}
			if (!pid) {
				break;
			}
		}
		if (pid) {
			if ((pid = wait(&status)) < 0) {
				die("wait: %s", strerror(errno));
			}
			if (WIFSIGNALED(status)) {
				die("worker %d killed by signal %d", (int)pid,
				    WTERMSIG(status));
			}
			die("worker %d exited with status %d", (int)pid,
			    WEXITSTATUS(status));
		}
		if (prctl(PR_SET_PDEATHSIG, SIGKILL) < 0) {
			die("prctl: %s", strerror(errno));
		}
		if (getppid() != ppid) {
			exit(1);
		}
	}
	lfd = listener(addr);
	conns = calloc(maxconn, sizeof(*conns));
	freel = calloc(maxconn, sizeof(*freel));
	queue = calloc(maxconn, sizeof(*queue));
	if (!conns || !freel || !queue) {
		die("out of memory");
	}
	for (i = maxconn - 1; i >= 0; i--) {
		freel[nfree++] = i;
	}

	if (pipe2(respipe, O_CLOEXEC) < 0) {
		die("pipe: %s", strerror(errno));
	}
	for (i = 0; i < NRESOLV; i++) {
		if ((errno = pthread_create(&t, NULL, resolver, NULL))) {
			die("pthread_create: %s", strerror(errno));
		}
	}

	/* Room for all ops in flight: 2 per conn, accept, results, timer. */
	cqentries = 256;
	while (cqentries < (unsigned)maxconn * 2 + 3) {
		cqentries *= 2;
	}
	setup(256, cqentries);
	queue_accept();
	queue_results();
	queue_timer();
	fprintf(stderr, "usocks: listening on %s\n", addr);

	for (;;) {
		enter(1);
		clock_gettime(CLOCK_MONOTONIC, &ts);
		now = ts.tv_sec;
		head = *ring.cqhead;
		tail = __atomic_load_n(ring.cqtail, __ATOMIC_ACQUIRE);
		while (head != tail) {
			cqe = &ring.cqes[head & ring.cqmask];
			ud = cqe->user_data;
			res = cqe->res;
			/*
			 * Free the slot before handling: a handler may submit,
			 * and submissions can complete at once.
			 */
			head++;
			__atomic_store_n(ring.cqhead, head, __ATOMIC_RELEASE);
			complete(ud, res);
		}
	}
}