sockstat tcp/udp switches

[Josh Carroll]

> I suggest you use /etc/protocols rather than hard code the protocols.
> This will make the code future-proof.  See getprotoent(3) for the
> correct way to read /etc/protocols.

Thanks Peter, that's a great idea. Below is a new patch that uses
getprotoent(3). For now, to maintain backward compatibility, I've
defined a default list of protocols (tcp, udp, divert), in the event
-P is not specified. This should retain the normal command line
behavior.

Thanks again for the suggestion.

Dmitry - with this change, diver is now the default but can also be
specified as an argument to -P as well.

Thanks,
Josh



--- sockstat.c.orig	Thu Jun  9 23:36:03 2005
+++ sockstat.c	Mon Nov  6 15:12:43 2006
@@ -66,8 +66,18 @@
 static int	 opt_u;		/* Show Unix domain sockets */
 static int	 opt_v;		/* Verbose mode */

+/* maximum length of a given protocol name to accept
+  from user input */
+#define MAX_PROTO_LEN 20
+/* default protocols to use if no -P was defined, currently
+	tcp, udp and divert */
+const char *default_protos[] = {"tcp", "udp", "divert", NULL};
+
+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)))
@@ -104,6 +114,104 @@
 	return (len);
 }

+
+/* this function needs to be updated to reflect additional protocols
+	that are to be supported */
+static int
+get_proto_type(const char *proto)
+{
+	struct protoent *pent;
+
+	if(strlen(proto) == 0)
+		return 0;
+
+	pent = malloc(sizeof(struct protoent));
+	pent = getprotobyname(proto);
+	if(pent == NULL) {
+		printf("Unrecognized protocol: %s\n", proto);
+		return -1;
+	} else {
+		return pent->p_proto;
+	}
+}
+
+
+static void init_protos(int num)
+{
+	int proto_count = 0;
+
+	if(num > 0) {
+		proto_count = num;
+	} else {
+		/* find the maximum number of possible protocols */
+		while(getprotoent() != NULL)
+			proto_count++;
+		endprotoent();
+	}
+
+	protos = malloc(sizeof(int) * proto_count);
+}
+
+
+static int
+parse_protos(const char *protospec)
+{
+	const char *p;
+	char curr_proto[MAX_PROTO_LEN];
+	/* to track the current protocol from the protospec */
+	int pindex = 0;
+	int protos_defined = 0;
+	/* the index into the protos integer array to track found
+		protocols */
+	int proto_index = 0;
+	int proto_num;
+
+	/* 0 argument => get all entries  */
+	init_protos(0);
+
+	p = protospec;
+	while(*p != '\0') {
+		if(*p == ',') {
+			curr_proto[pindex] = '\0';
+			if( (proto_num = get_proto_type(curr_proto)) != -1) {
+				protos[proto_index++] = proto_num;
+				protos_defined++;
+			}
+			pindex = 0;
+		} else {
+			curr_proto[pindex++] = *p;
+		}
+		p++;
+		if(*p == '\0' &&  pindex != 0) {
+			curr_proto[pindex] = '\0';
+			if( (proto_num = get_proto_type(curr_proto)) != -1) {
+				protos[proto_index++] = proto_num;
+				protos_defined++;
+			}
+		}
+		if(pindex == MAX_PROTO_LEN) {
+			printf("Warning: truncating protocol\n");
+			curr_proto[pindex] = '\0';
+			if( (proto_num = get_proto_type(curr_proto)) != -1) {
+				protos[proto_index++] = proto_num;
+				protos_defined++;
+			}
+
+			/* 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++;
+
+			pindex = 0;
+		}
+	}
+	return protos_defined;
+}
+
+
 static void
 parse_ports(const char *portspec)
 {
@@ -573,19 +681,49 @@
 	}
 }

+static int set_default_protos(void)
+{
+	struct protoent *prot;
+	const char **curr_proto;
+	int proto_index = 0;
+	int proto_count = 0;
+
+	/* determine the number of default protocols */
+	for(curr_proto = default_protos; *curr_proto; curr_proto++, proto_count++)
+		;
+
+	init_protos(proto_count);
+
+	for(curr_proto = default_protos; *curr_proto; curr_proto++) {
+		prot = getprotobyname(*curr_proto);
+		if(prot == NULL) {
+			printf("Cannot determine default protocols\n");
+			exit(EXIT_FAILURE);
+		}
+		protos[proto_index++] = prot->p_proto;
+	}
+
+	return proto_index;
+}
+
+
 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[])
 {
-	int o;
+	/* if protos_defined remains -1, no -P was provided, sowe avoid
+		attempting to read from that int array later */
+	int protos_defined = -1;

-	while ((o = getopt(argc, argv, "46clp:uv")) != -1)
+	int o, i;
+
+	while ((o = getopt(argc, argv, "46clp:P:uv")) != -1)
 		switch (o) {
 		case '4':
 			opt_4 = 1;
@@ -602,6 +740,9 @@
 		case 'p':
 			parse_ports(optarg);
 			break;
+		case 'P':
+			protos_defined = parse_protos(optarg);
+			break;
 		case 'u':
 			opt_u = 1;
 			break;
@@ -618,17 +759,25 @@
 	if (argc > 0)
 		usage();

-	if (!opt_4 && !opt_6 && !opt_u)
-		opt_4 = opt_6 = opt_u = 1;
+
+	if (!opt_4 && !opt_6 && !opt_u && protos_defined == -1) {
+		opt_u = 1;
+		/* show tcp, udp and divert by default if no -P was specified */
+		/* restore protos_defined to the number of default protocols */
+		protos_defined = set_default_protos();
+	}
+
+	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);
-		gather_inet(IPPROTO_DIVERT);
+		for(i=0; i < protos_defined; i++)
+			gather_inet(protos[i]);
 	}
-	if (opt_u) {
+
+	if ( opt_u || (protos_defined == -1 && !opt_4 && !opt_6)) {
 		gather_unix(SOCK_STREAM);
 		gather_unix(SOCK_DGRAM);
 	}