1 files changed, 81 insertions, 10 deletions

diff --git a/latency.c b/latency.c
index 6181e84..c1782cf 100644
--- a/latency.c
+++ b/latency.c
@@ -2,8 +2,10 @@
 #include <sys/socket.h>
 #include <sys/uio.h>
 #include <errno.h>
+#include <limits.h>
 #include <netdb.h>
 #include <stdio.h>
+#include <stdlib.h>
 #include <unistd.h>
 
 #include "proxy.h"
@@ -53,22 +55,91 @@ static int connect_target(struct addrinfo *target)
 	return fd;
 }
 
-int main(int argc, char **argv)
-{
+struct args {
 	struct addrinfo *source;
 	struct addrinfo *target;
+	unsigned delay;
+};
 
-	if (argc != 3) {
-		fprintf(stderr, "usage: %s targethost targetport\n", argv[0]);
-		return -1;
+static void usage(const char *progname)
+{
+	int width = 20;
+	const char *args[] = {
+		"-d DELAY",     "millisecond delay to add (default: 500ms)",
+		"-p LOCALPORT", "port to listen on (default: 1234)",
+		NULL
+	};
+	fprintf(stderr, "usage: %s [-d DELAY] [-p LOCALPORT] HOST PORT\n\n", progname);
+	for (const char **a = args; *a; a += 2)
+		fprintf(stderr, "\t%-*s %s\n", width, a[0], a[1]);
+}
+
+static int parse_unsigned(const char *arg, unsigned *out)
+{
+	char *end;
+	errno = 0;
+	unsigned long val = strtoul(arg, &end, 10);
+	if (*end) return -1;
+	if (errno && !val) return -1;
+	if (val == ULONG_MAX && errno == ERANGE) return -1;
+	if (val > UINT_MAX) return -1;
+	*out = val;
+	return 0;
+}
+
+static int parse_args(struct args *args, int argc, char **argv)
+{
+	args->delay = 500;
+
+	const char *localport = "1234";
+	const char *host = NULL, *port = NULL;
+
+	extern char *optarg;
+	extern int optind;
+	int opt, error = 0;
+
+	while ((opt = getopt(argc, argv, "+d:p:")) != -1) {
+		switch (opt) {
+		case 'd':
+			if (parse_unsigned(optarg, &args->delay)) {
+				fprintf(stderr, "invalid delay\n");
+				error = 1;
+			}
+			break;
+
+		case 'p':
+			localport = optarg;
+			break;
+
+		default:
+			error = 1;
+		}
 	}
 
-	if (parse_host(NULL, "1234", &source))
+	if (optind <= argc - 2) {
+		host = argv[optind++];
+		port = argv[optind++];
+	} else error = 1;
+
+	if (error) {
+		usage(argv[0]);
+		return -1;
+	}
+	if (parse_host(NULL, localport, &args->source))
+		return -1;
+	if (parse_host(host, port, &args->target))
 		return -1;
-	if (parse_host(argv[1], argv[2], &target))
+
+	return 0;
+}
+
+int main(int argc, char **argv)
+{
+	struct args args;
+	if (parse_args(&args, argc, argv))
 		return -1;
 
-	int listenfd = listen_source(source);
+	int listenfd = listen_source(args.source);
 	if (listenfd < 0) {
 		perror("listen_source");
 		return -1;
@@ -78,14 +149,14 @@ int main(int argc, char **argv)
 		int sourcefd = accept(listenfd, NULL, NULL);
 		if (sourcefd == -1) break;
 
-		int targetfd = connect_target(target);
+		int targetfd = connect_target(args.target);
 		if (targetfd == -1) {
 			perror("connect_target");
 			close(sourcefd);
 			break;
 		}
 
-		proxy(sourcefd, targetfd, 500);
+		proxy(sourcefd, targetfd, args.delay);
 
 		close(targetfd);
 		close(sourcefd);