sockstat tcp/udp switches

[Josh Carroll]

Below is an initial attempt at the -P argument implementation.

I'm not sure if the method used to parse out the protocol is "up to
snuff" for the FreeBSD project, but I'm welcome to
suggestions/advice/etc.

I tested the previous behavior and given the exact command line,
sockstat works the same prior to the patch and with the patch. So the
-P option is additive only (so as not to break legacy
scripts/automation that may parse the sockstat output).

I included a limitation on the maximum length of a proto (mostly to
avoid buffer overflows) and 20 is probably way too large, so I can
lower that if need be. The -P argument also will sense and ignore
multiple commas, e.g.:

-P tcp,,udp

Will still just detect tcp and udp. I implemented the protocols with a
sort of lookup table for the protocol types for arguments to -P. So
adding another protocol should be as easy as adding a #define,
incrementing the #defined for NUM_PROTOS and updating the
add_proto_type function to recognize and properly store the new
protocol type. Again, I'm open to any feedback on this particular
approach.

Thanks again for everyone's attetion on this. It's a nice little
project to get me ramped back up into C and I do appreciate all the
feedback!

Regards,
Josh



--- sockstat.c.orig	Thu Nov  2 15:04:39 2006
+++ sockstat.c	Fri Nov  3 15:25:34 2006
@@ -66,8 +66,15 @@
 static int	 opt_u;		/* Show Unix domain sockets */
 static int	 opt_v;		/* Verbose mode */

+#define NUM_PROTOS 2
+#define MAX_PROTO_LEN 20
+#define TCP_PROTO 0
+#define UDP_PROTO 1
+static int *protos;		/* protocols to use */
+
 static int	*ports;

+
 #define INT_BIT (sizeof(int)*CHAR_BIT)
 #define SET_PORT(p) do { ports[p / INT_BIT] |= 1 << (p % INT_BIT); } while (0)
 #define CHK_PORT(p) (ports[p / INT_BIT] & (1 << (p % INT_BIT)))
@@ -105,6 +112,97 @@
 }

 static void
+clear_protos(void)
+{
+	int i;
+
+	for(i=0; i <= NUM_PROTOS; i++)
+		protos[i] = 0;
+
+}
+
+static void
+init_protos(void)
+{
+	int i;
+
+	if((protos = calloc(NUM_PROTOS, sizeof(int))) == NULL)
+		err(1, "calloc()");
+
+	for(i=0; i <= NUM_PROTOS; i++)
+		protos[i] = 0;
+
+}
+
+
+/* this function needs to be updated to reflect additional protocols
+	that are to be supported */
+static int
+add_proto_type(const char *proto)
+{
+	int found = 0;
+
+	if(strlen(proto) == 0)
+		return 0;
+
+	if(strncmp(proto, "tcp", 3) == 0) {
+		protos[TCP_PROTO] = 1;
+		found = 1;
+	}
+	else if(strncmp(proto, "udp", 3) == 0) {
+		protos[UDP_PROTO] = 1;
+		found = 1;
+	} else {
+		printf("Invalid protocol: %s\n", proto);
+	}
+
+	return found;
+}
+
+static int
+parse_protos(const char *protospec)
+{
+	const char *p;
+	char curr_proto[MAX_PROTO_LEN];
+	int proto_index = 0;
+	int protos_defined = 0;
+
+	p = protospec;
+	while(*p != '\0') {
+		if(*p == ',') {
+			curr_proto[proto_index] = '\0';
+			protos_defined += add_proto_type(curr_proto);
+			proto_index = 0;
+		} else {
+			curr_proto[proto_index++] = *p;
+		}
+		p++;
+		if(*p == '\0' &&  proto_index != 0) {
+			curr_proto[proto_index] = '\0';
+			protos_defined += add_proto_type(curr_proto);
+		}
+		if(proto_index == MAX_PROTO_LEN) {
+			printf("Warning: truncating protocol\n");
+			curr_proto[proto_index] = '\0';
+			protos_defined += add_proto_type(curr_proto);
+
+			/* need to seek to the next , or \0 now */
+			while(*p != ',' && *p != '\0')
+				p++;
+
+			/* if we're at a proto boundary (comma), move on */
+			if(*p == ',')
+				p++;
+
+			proto_index = 0;
+		}
+	}
+
+	return protos_defined;
+}
+
+
+static void
 parse_ports(const char *portspec)
 {
 	const char *p, *q;
@@ -576,16 +674,24 @@
 static void
 usage(void)
 {
-	fprintf(stderr, "Usage: sockstat [-46clu] [-p ports]\n");
+	fprintf(stderr, "Usage: sockstat [-46clu] [-p ports] [-P protos]\n");
 	exit(1);
 }

 int
 main(int argc, char *argv[])
 {
+	/* if protos_defined remains -1, no -P was provided, sowe avoid
+		attempting to read from that int array later */
+	int protos_defined = -1;
+
 	int o;

-	while ((o = getopt(argc, argv, "46clp:uv")) != -1)
+	/* initialize all protos to 0, if they're all still unset later
+		then -P wasn't use and the default is to display all of them */
+	init_protos();
+
+	while ((o = getopt(argc, argv, "46clp:P:uv")) != -1)
 		switch (o) {
 		case '4':
 			opt_4 = 1;
@@ -602,6 +708,9 @@
 		case 'p':
 			parse_ports(optarg);
 			break;
+		case 'P':
+			protos_defined = parse_protos(optarg);
+			break;
 		case 'u':
 			opt_u = 1;
 			break;
@@ -618,17 +727,28 @@
 	if (argc > 0)
 		usage();

-	if (!opt_4 && !opt_6 && !opt_u)
-		opt_4 = opt_6 = opt_u = 1;
+	/* if -u was explicityl defined, zero out the protos array */
+	if(opt_u && protos_defined == -1 && !opt_4 && !opt_6) {
+		clear_protos();
+		protos_defined = 0;
+	}
+
+	if (!opt_4 && !opt_6 && !opt_u && protos_defined == -1)
+		opt_u = 1;
+	if (!opt_4 && !opt_6)
+		opt_4 = opt_6 = 1;
 	if (!opt_c && !opt_l)
 		opt_c = opt_l = 1;

 	if (opt_4 || opt_6) {
-		gather_inet(IPPROTO_TCP);
-		gather_inet(IPPROTO_UDP);
+		if(protos[TCP_PROTO] == 1 || protos_defined == -1)
+			gather_inet(IPPROTO_TCP);
+		if(protos[UDP_PROTO] == 1 || protos_defined == -1)
+			gather_inet(IPPROTO_UDP);
 		gather_inet(IPPROTO_DIVERT);
 	}
-	if (opt_u) {
+
+	if ( opt_u || (protos_defined == -1 && !opt_4 && !opt_6)) {
 		gather_unix(SOCK_STREAM);
 		gather_unix(SOCK_DGRAM);
 	}