udpflood.c

/* An adaptation of Eric Dumazet's udpflood tool.  */

#include <stdio.h>
#include <stddef.h>
#include <malloc.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <stdint.h>

#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define _GNU_SOURCE
#include <getopt.h>

static int debug = 0;

typedef union sa_u {
	struct sockaddr_in a4;
	struct sockaddr_in6 a6;
} sa_u;

static int usage(void)
{
	printf("usage: udpflood [ -l count ] [ -m message_size ] [ -c num_ip_addrs ] IP_ADDRESS\n");
	return -1;
}

static uint32_t get_last32h(const sa_u *sa)
{
	if (sa->a4.sin_family == PF_INET)
		return ntohl(sa->a4.sin_addr.s_addr);
	else
		return ntohl(sa->a6.sin6_addr.s6_addr32[3]);
}

static void set_last32h(sa_u *sa, uint32_t last32h)
{
	if (sa->a4.sin_family == PF_INET)
		sa->a4.sin_addr.s_addr = htonl(last32h);
	else
		sa->a6.sin6_addr.s6_addr32[3] = htonl(last32h);
}

static void print_sa(const sa_u *sa, const char *msg)
{
	char buf[sizeof("xxxx:xxxx:xxxx:xxxx:xxxx:xxxx:xxxx:xxxx")];

	if (!debug)
		return;

	switch (sa->a4.sin_family) {
	case PF_INET:
		inet_ntop(PF_INET, &(sa->a4.sin_addr.s_addr), buf,
			  sizeof(buf));
		break;
	case PF_INET6:
		inet_ntop(PF_INET6, sa->a6.sin6_addr.s6_addr, buf, sizeof(buf));
		break;
	}

	printf("%s: %s\n", msg, buf);
}

static long get_diff_ms(const struct timeval *now,
			     const struct timeval *start)
{
	long start_ms, now_ms;
	start_ms = start->tv_sec * 1000 + (start->tv_usec / 1000);
	now_ms = now->tv_sec * 1000 + (now->tv_usec  / 1000);
	return now_ms - start_ms;
}

static int send_packets(const sa_u *start_sa, size_t num_addrs, int count,
			int msg_sz)
{
	char *msg = malloc(msg_sz);
	sa_u cur_sa;
	uint32_t start_addr32h, end_addr32h, cur_addr32h;
	int fd, i, err;
	struct timeval last, now;

	if (!msg)
		return -ENOMEM;

	memset(msg, 0, msg_sz);

	memcpy(&cur_sa, start_sa, sizeof(cur_sa));
	cur_addr32h = start_addr32h = get_last32h(&cur_sa);
	end_addr32h = start_addr32h + num_addrs;

	fd = socket(cur_sa.a4.sin_family, SOCK_DGRAM, IPPROTO_IP);
	if (fd < 0) {
		perror("socket");
		err = fd;
		goto out_nofd;
	}
	err = connect(fd, (struct sockaddr *) &cur_sa, sizeof(cur_sa));
	if (err < 0) {
		perror("connect");
		goto out;
	}

	print_sa(start_sa, "start_addr");
	gettimeofday(&last, NULL);
	for (i = 0; i < count; i++) {
		print_sa(&cur_sa, "sendto");
		err = sendto(fd, msg, msg_sz, 0,
			     (struct sockaddr *) &cur_sa, sizeof(cur_sa));
		if (err < 0) {
			perror("sendto");
			goto out;
		}

		if (++cur_addr32h >= end_addr32h)
			cur_addr32h = start_addr32h;
		set_last32h(&cur_sa, cur_addr32h);

		/*
		 * print timing info for every 65536 fib inserts to
		 * observe the gc effect (mostly for IPv6 fib).
		 */
		if (i && (i & 0xFFFF) == 0) {
		    long diff_ms;
		    gettimeofday(&now, NULL);
		    diff_ms = get_diff_ms(&now, &last);
		    printf("%d %ld.%ld\n", i >> 16,
			   diff_ms / 1000, diff_ms % 1000);
		    memcpy(&last, &now, sizeof(last));
		}
	}

	err = 0;
out:
	close(fd);
out_nofd:
	free(msg);
	return err;
}

int main(int argc, char **argv, char **envp)
{
	int port, msg_sz, count, num_addrs, ret;
	sa_u start_sa;

	port = 6000;
	msg_sz = 32;
	count = 10000000;
	num_addrs = 1;

	while ((ret = getopt(argc, argv, "dl:s:p:c:")) >= 0) {
		switch (ret) {
		case 'l':
			sscanf(optarg, "%d", &count);
			break;
		case 's':
			sscanf(optarg, "%d", &msg_sz);
			break;
		case 'p':
			sscanf(optarg, "%d", &port);
			break;
		case 'c':
			sscanf(optarg, "%d", &num_addrs);
			break;
		case 'd':
			debug = 1;
			break;
		case '?':
			return usage();
		}
	}

	if (num_addrs < 1 || count < 1)
		return usage();

	if (!argv[optind])
		return usage();

	memset(&start_sa, 0, sizeof(start_sa));
	start_sa.a4.sin_port = htons(port);
	if (inet_pton(PF_INET, argv[optind], &start_sa.a4.sin_addr))
		start_sa.a4.sin_family = PF_INET;
	else if (inet_pton(PF_INET6, argv[optind],
			   start_sa.a6.sin6_addr.s6_addr))
		start_sa.a6.sin6_family = PF_INET6;
	else
		return usage();

	return send_packets(&start_sa, num_addrs, count, msg_sz);
}