sockstat tcp/udp switches

[Josh Carroll]

Below is the patch using strsep instead. I also updated the
gather_inet function to print a message that a protocol is not
supported if it exists in /etc/protocols, but does not have a case in
the gather_inet function.

patch is here if you'd rather just fetch it:

http://pflog.net/~floyd/sockstat.patch

Thanks,
Josh


--- sockstat.c.orig	Thu Jun  9 23:36:03 2005
+++ sockstat.c	Tue Nov  7 10:49:28 2006
@@ -66,8 +66,15 @@
 static int	 opt_u;		/* Show Unix domain sockets */
 static int	 opt_v;		/* Verbose mode */

+/* 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 +111,68 @@
 	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 = 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();
+	}
+
+	if ((protos = malloc(sizeof(int) * proto_count)) == NULL)
+		err(1, "malloc()");
+}
+
+
+static int
+parse_protos(char *protospec)
+{
+	char **prot;
+	char *tmp = protospec;
+	int protos_defined = 0, proto_type, proto_index = 0;
+
+	init_protos(0);
+
+	while( (*prot = strsep(&tmp, ",")) != NULL) {
+		/* handle ,, */
+		if(strlen(*prot) > 0) {
+			proto_type = get_proto_type(*prot);
+			if(proto_type != -1) {
+				protos[proto_index++] = proto_type;
+				protos_defined++;
+			}
+		}
+	}
+	return protos_defined;
+}
+
+
 static void
 parse_ports(const char *portspec)
 {
@@ -188,6 +257,7 @@
 	size_t len, bufsize;
 	void *buf;
 	int hash, retry, vflag;
+	struct protoent *pent;

 	vflag = 0;
 	if (opt_4)
@@ -209,7 +279,9 @@
 		protoname = "div";
 		break;
 	default:
-		abort();
+		pent = getprotobynumber(proto);
+		printf("protocol `%s' not supported.\n", pent->p_name);
+		exit(EXIT_FAILURE);
 	}

 	buf = NULL;
@@ -264,7 +336,9 @@
 			so = &xip->xi_socket;
 			break;
 		default:
-			abort();
+			pent = getprotobynumber(proto);
+			printf("protocol `%s' not supported.\n", pent->p_name);
+			exit(EXIT_FAILURE);
 		}
 		if ((inp->inp_vflag & vflag) == 0)
 			continue;
@@ -573,19 +647,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;
+
+	int o, i;

-	while ((o = getopt(argc, argv, "46clp:uv")) != -1)
+	while ((o = getopt(argc, argv, "46clp:P:uv")) != -1)
 		switch (o) {
 		case '4':
 			opt_4 = 1;
@@ -602,6 +706,9 @@
 		case 'p':
 			parse_ports(optarg);
 			break;
+		case 'P':
+			protos_defined = parse_protos(optarg);
+			break;
 		case 'u':
 			opt_u = 1;
 			break;
@@ -618,17 +725,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);
 	}