From ae99fd0e07d2bdfee11230e2981722ab9bf35c4c Mon Sep 17 00:00:00 2001
From: Luigi Rizzo <luigi@FreeBSD.org>
Date: Fri, 12 Nov 2010 13:05:17 +0000
Subject: [PATCH] The first customer of the SO_USER_COOKIE option: the
 "sockarg" ipfw option matches packets associated to a local socket and with a
 non-zero so_user_cookie value. The value is made available as tablearg, so it
 can be used as a skipto target or pipe number in ipfw/dummynet rules.

Code by Paul Joe, manpage by me.

Submitted by:	Paul Joe
MFC after:	1 week
---
 sbin/ipfw/ipfw.8                 | 11 +++++++++++
 sbin/ipfw/ipfw2.c                |  7 +++++++
 sbin/ipfw/ipfw2.h                |  1 +
 sys/netinet/ip_fw.h              |  3 +++
 sys/netinet/ipfw/ip_fw2.c        | 33 ++++++++++++++++++++++++++++++++
 sys/netinet/ipfw/ip_fw_sockopt.c |  1 +
 6 files changed, 56 insertions(+)

diff --git a/sbin/ipfw/ipfw.8 b/sbin/ipfw/ipfw.8
index a954c1d5415f..a984f70cb9e7 100644
--- a/sbin/ipfw/ipfw.8
+++ b/sbin/ipfw/ipfw.8
@@ -1510,6 +1510,17 @@ interface.
 Matches TCP packets that have the SYN bit set but no ACK bit.
 This is the short form of
 .Dq Li tcpflags\ syn,!ack .
+.It Cm sockarg
+Matches packets that are associated to a local socket and
+for which the SO_USER_COOKIE socket option has been set
+to a non-zero value. As a side effect, the value of the
+option is made available as
+.Cm tablearg
+value, which in turn can be used as
+.Cm skipto
+or
+.Cm pipe
+number.
 .It Cm src-ip Ar ip-address
 Matches IPv4 packets whose source IP is one of the address(es)
 specified as an argument.
diff --git a/sbin/ipfw/ipfw2.c b/sbin/ipfw/ipfw2.c
index f313b51b968f..9f2fe6961400 100644
--- a/sbin/ipfw/ipfw2.c
+++ b/sbin/ipfw/ipfw2.c
@@ -266,6 +266,7 @@ static struct _s_x rule_options[] = {
 	{ "estab",		TOK_ESTAB },
 	{ "established",	TOK_ESTAB },
 	{ "setup",		TOK_SETUP },
+	{ "sockarg",		TOK_SOCKARG },
 	{ "tcpdatalen",		TOK_TCPDATALEN },
 	{ "tcpflags",		TOK_TCPFLAGS },
 	{ "tcpflgs",		TOK_TCPFLAGS },
@@ -1338,6 +1339,9 @@ show_ipfw(struct ip_fw *rule, int pcwidth, int bcwidth)
 			case O_FIB:
 				printf(" fib %u", cmd->arg1 );
 				break;
+			case O_SOCKARG:
+				printf(" sockarg");
+				break;
 
 			case O_IN:
 				printf(cmd->len & F_NOT ? " out" : " in");
@@ -3531,6 +3535,9 @@ ipfw_add(char *av[])
 			fill_cmd(cmd, O_FIB, 0, strtoul(*av, NULL, 0));
 			av++;
 			break;
+		case TOK_SOCKARG:
+			fill_cmd(cmd, O_SOCKARG, 0, 0);
+			break;
 
 		case TOK_LOOKUP: {
 			ipfw_insn_u32 *c = (ipfw_insn_u32 *)cmd;
diff --git a/sbin/ipfw/ipfw2.h b/sbin/ipfw/ipfw2.h
index 8566cdef0218..2ba091fd5634 100644
--- a/sbin/ipfw/ipfw2.h
+++ b/sbin/ipfw/ipfw2.h
@@ -199,6 +199,7 @@ enum tokens {
 	TOK_FIB,
 	TOK_SETFIB,
 	TOK_LOOKUP,
+	TOK_SOCKARG,
 };
 /*
  * the following macro returns an error message if we run out of
diff --git a/sys/netinet/ip_fw.h b/sys/netinet/ip_fw.h
index cf5d8d03a90e..fdcc5fd4c1df 100644
--- a/sys/netinet/ip_fw.h
+++ b/sys/netinet/ip_fw.h
@@ -192,10 +192,13 @@ enum ipfw_opcodes {		/* arguments (4 byte each)	*/
 
 	O_SETFIB,		/* arg1=FIB number */
 	O_FIB,			/* arg1=FIB desired fib number */
+	
+	O_SOCKARG,		/* socket argument */
 
 	O_LAST_OPCODE		/* not an opcode!		*/
 };
 
+
 /*
  * The extension header are filtered only for presence using a bit
  * vector with a flag for each header.
diff --git a/sys/netinet/ipfw/ip_fw2.c b/sys/netinet/ipfw/ip_fw2.c
index c291089c6f85..43b2d1114295 100644
--- a/sys/netinet/ipfw/ip_fw2.c
+++ b/sys/netinet/ipfw/ip_fw2.c
@@ -1801,6 +1801,39 @@ do {								\
 					match = 1;
 				break;
 
+			case O_SOCKARG:	{
+				struct inpcb *inp = args->inp;
+				struct inpcbinfo *pi;
+				
+				if (is_ipv6) /* XXX can we remove this ? */
+					break;
+
+				if (proto == IPPROTO_TCP)
+					pi = &V_tcbinfo;
+				else if (proto == IPPROTO_UDP)
+					pi = &V_udbinfo;
+				else
+					break;
+
+				/* For incomming packet, lookup up the 
+				inpcb using the src/dest ip/port tuple */
+				if (inp == NULL) {
+					INP_INFO_RLOCK(pi);
+					inp = in_pcblookup_hash(pi, 
+						src_ip, htons(src_port),
+						dst_ip, htons(dst_port),
+						0, NULL);
+					INP_INFO_RUNLOCK(pi);
+				}
+				
+				if (inp && inp->inp_socket) {
+					tablearg = inp->inp_socket->so_user_cookie;
+					if (tablearg)
+						match = 1;
+				}
+				break;
+			}
+
 			case O_TAGGED: {
 				struct m_tag *mtag;
 				uint32_t tag = (cmd->arg1 == IP_FW_TABLEARG) ?
diff --git a/sys/netinet/ipfw/ip_fw_sockopt.c b/sys/netinet/ipfw/ip_fw_sockopt.c
index c50572873a16..0c903eedf5c7 100644
--- a/sys/netinet/ipfw/ip_fw_sockopt.c
+++ b/sys/netinet/ipfw/ip_fw_sockopt.c
@@ -572,6 +572,7 @@ check_ipfw_struct(struct ip_fw *rule, int size)
 		case O_IPTOS:
 		case O_IPPRECEDENCE:
 		case O_IPVER:
+		case O_SOCKARG:
 		case O_TCPWIN:
 		case O_TCPFLAGS:
 		case O_TCPOPTS: