speedtest-cgi

speedtest.c

/*
 * speedtext fcgi writev v2 2018-08-13
 * (C) [email protected]
 * Licensed under GPL v3.0
 */

#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>
#include <sys/uio.h>
#include <alloca.h>
#include <errno.h>

#define MIB 1048576ULL
// 640k should be enough for anyone
#define MAX_LEN (MIB*128)

#define CHUNK_BITS 12
#define CHUNK_SIZE (1ULL << CHUNK_BITS)
#define CHUNK_DIV(n) ((n) >> CHUNK_BITS)
#define CHUNK_MOD(n) ((n) & ((1 << CHUNK_BITS)-1))

#if 0
# define DPRINT(format, args...) fprintf(stderr, "%s: "format, __FUNCTION__, ##args)
#else
# define DPRINT(format, args...) do { /* nothing */ } while (0)
#endif

void help(const char *err) {
	char url[128];
	snprintf(url, sizeof(url), "%s://%s%s", getenv("REQUEST_SCHEME"), getenv("SERVER_NAME"), getenv("SCRIPT_NAME"));
	printf(
		"Status: 400 %s\n"
		"Content-Type: text/plain\n"
		"\n"
		"error: %s\n\n"
		"usage: %s/<len>[.ext]\n\n"
		"len is a string of one or more positive integers with (optional) unit suffixes\n"
		"  units: k (10^3) K (2^10) m (10^6) M (2^20)\n"
		"ext is an optional content type specifier: bin/binary (default), asc/ascii/txt/text\n"
		"\n"
		"examples: 5k (5*10^3=5000)  13M (13*2^20 = 13631488) 3m2K1 (3*10^6+2*2^10+1=3002049)\n"
		"\n"
		"max. len is %llu (%llu KiB, %llu MiB)\n",
		err, err, url, MAX_LEN, MAX_LEN/1024, MAX_LEN/1024/1024
	);
	exit(0);
}

void err(int code, const char * msg) {
	printf("Status: %1$d %2$s\n\n%1$d %2$s\n", code, msg);
	exit(1);
}

size_t parse_reqlen(char * s, char ** ext) {
	size_t sum = 0;
	// s points either to the leading slash of PATH_INFO or the char after the last num, so we check if there's anything left
	while (*s && *s != '.' && *(++s) && *s != '.') {
		char * endptr;
		size_t num = strtoll(s, &endptr, 10);
		if (num < 0 || num > MAX_LEN || endptr == s)
			help("invalid length");
		int factor;
		switch (*endptr) {
			case '\0':
			case '.' : factor = 1 ; break;
			case 'k' : factor = 1000; break;
			case 'K' : factor = 1024; break;
			case 'm' : factor = 1000*1000; break;
			case 'M' : factor = 1024*1024; break;
			case 'g' : factor = 1000*1000*1000; break;
			case 'G' : factor = 1024*1024*1024; break;
			default:   help("unknown unit");
		}
		sum += num * factor;
		s = endptr;
	}
	if (sum < 0 || sum > MAX_LEN)
		help("length limit");
	if (ext)
		*ext = *s ? s+1 : s;
	return sum;
}

void handle_get() {
	char * pathinfo = getenv("PATH_INFO");
	if (pathinfo == NULL || pathinfo[0] == '\0' || pathinfo[1] == '\0')
		help("missing length");

	// parse requested data length
	char * ctype;
	size_t reqlen = parse_reqlen(pathinfo, &ctype);

	char header[64];
	ssize_t hlen = 0;
	int want_ascii = 0;
	if (!*ctype)
		ctype = getenv("QUERY_STRING");
	if (!ctype || !*ctype || !strcmp(ctype, "bin") || !strcmp(ctype, "binary")) {
		hlen = snprintf(header, sizeof(header), "Content-Type: application/binary\nContent-Length: %ld\n\n", reqlen);
	} else if (!strcmp(ctype, "asc") || !strcmp(ctype, "ascii") || !strcmp(ctype, "txt") || !strcmp(ctype, "text")) {
		hlen = snprintf(header, sizeof(header), "Content-Type: text/plain\nContent-Length: %ld\n\n", reqlen);
		want_ascii = 1;
	} else {
		help("unknown content type requested");
	}
	DPRINT("CHUNK_SIZE %u reqlen %lu hlen %ld total %ld\n", CHUNK_SIZE, reqlen, hlen, reqlen+hlen);

	unsigned char data[CHUNK_SIZE];
	memset(data, 0xaa, sizeof(data));
	int in = open("/dev/urandom", O_RDONLY);
	if (in < 0) {
		perror("open /dev/urandom");
		err(500, "Internal Server Error");
	}
	size_t want = reqlen > CHUNK_SIZE ? sizeof(data) : reqlen;
	ssize_t got = read(in, data, want);
	if (got < 0) {
		perror("read urandom");
		err(500, "Internal Server Error");
	} else if (got != want) {
		fprintf(stderr, "urandom: wanted %ld bytes but got %ld\n", want, got);
	}
	close(in);
	if (want_ascii)
		for (unsigned char *p=data+want-1; p>=data; --p) {
			if (*p > 126) // don't want DEL (127)
				*p &= 63;
			if (*p < 32)
				*p |= 32;
		}

	// the number of items in iov may be limited, deal with it
	errno = 0;
	long iov_max = sysconf(_SC_IOV_MAX);
	if (iov_max <= 0) {
		if (errno)
			perror("sysconf(_SC_IOV_MAX)");
		err(500, "Internal server error");
	}

	int iov_full = CHUNK_DIV(reqlen) < iov_max ? CHUNK_DIV(reqlen) : iov_max;
	size_t size_full = iov_full << CHUNK_BITS;
	// header + full chunks + (possibly) one partial chunk
	int iov_cnt = 1 + iov_full + (CHUNK_MOD(reqlen) != 0);
	struct iovec * iov = calloc(sizeof(struct iovec), iov_cnt);
	if (!iov) {
		perror("calloc iov");
		err(500, "Internal Server Error");
	}
	DPRINT("iov_full %d  size_full %ld  iov_cnt %d\n", iov_full, size_full, iov_cnt);

	// initialize iov. size of (partial) last iov will be changed in the send loop
	iov[0].iov_base = header;
	iov[0].iov_len = hlen;
	for (int i=1; i < iov_cnt; ++i) {
		iov[i].iov_base = data;
		iov[i].iov_len = CHUNK_SIZE;
	}
	if (CHUNK_MOD(reqlen))
		iov[iov_cnt-1].iov_len = 0;

	// send loop
	ssize_t rem = hlen + reqlen; // remaining bytes to send
	while (rem > 0) {
		int iov_offs, iov_use;
		if (hlen) { // header not sent yet?
			iov_use = iov_cnt; // use all available iovs (may be reduced below)
		} else {
			iov_use = rem >= size_full ?
				iov_full : CHUNK_DIV(rem) + (CHUNK_MOD(rem) != 0);
		}
		iov_offs = iov_cnt - iov_use;
		if (iov_use > iov_max)
			iov_use = iov_max;
		else if (iov_use != iov_full && CHUNK_MOD(rem-hlen))
			iov[iov_cnt-1].iov_len = CHUNK_MOD(rem-hlen);
		ssize_t sent = writev(1, iov+iov_offs, iov_use);
		if (sent < 0) {
			perror("writev");
			break;
		}
		rem -= sent;
		hlen = 0;
		DPRINT("writev(offset %d, cnt %d) = %ld (%ld remaining)\n", iov_offs, iov_use, sent, rem);
	}
	DPRINT("rem: %ld\n", rem);

	close(1);
}

void handle_put() {
	char * lenstr = getenv("CONTENT_LENGTH");
	if (!lenstr) {
		fprintf(stderr, "missing CONTENT_LENGTH\n");
		exit(1);
	}
	size_t len = atoll(lenstr);
	if (len > MAX_LEN)
		help("size limit exceeded");

	struct timeval t0, t1;
	gettimeofday(&t0, NULL);
	char data[CHUNK_SIZE];
	int got;
	do { // we don't read exactly len bytes, but simply until EOF
		got = read(0, data, sizeof(data));
		if (got < 0) {
			perror("body read");
			exit(1);
		}
	} while (got > 0); // 0 means EOF
	gettimeofday(&t1, NULL);

	double tdiff = (double) (t1.tv_sec - t0.tv_sec) + (t1.tv_usec - t0.tv_usec) / 1e6;
	printf("Content-Type: text/plain\n"
		"\n"
		"ok, got %lu octets in %.3lfs (%.1f KiB/s)\n",
		len, tdiff, len/tdiff/1024);
}

int main(int argc, char *argv[], char *envp[]) {
	char * method = getenv("REQUEST_METHOD");
	if (!method) {
		fprintf(stderr, "REQUEST_METHOD not set\n");
		return 1;
	}
	else if (!strcmp(method, "GET"))
		handle_get();
	else if (!strcmp(method, "PUT"))
		handle_put();
	else {
		err(400, "Unsupported Request Method");
	}

	return 0;
}