Add support to IPFW for classification based on "diverted" status

(that is, input via a divert socket).

3 files changed, 42 insertions, 16 deletions

diff --git a/sys/netinet/ip_divert.c b/sys/netinet/ip_divert.c
index 8f524bb..fa8a6ea 100644
--- a/sys/netinet/ip_divert.c
+++ b/sys/netinet/ip_divert.c
@@ -66,6 +66,7 @@
 #include <netinet/ip.h>
 #include <netinet/ip_divert.h>
 #include <netinet/ip_var.h>
+#include <netinet/ip_fw.h>
 
 /*
  * Divert sockets
@@ -268,6 +269,8 @@ static int
 div_output(struct socket *so, struct mbuf *m,
 	struct sockaddr_in *sin, struct mbuf *control)
 {
+	struct m_tag *mtag;
+	struct divert_tag *dt;
 	int error = 0;
 
 	KASSERT(m->m_pkthdr.rcvif == NULL, ("rcvif not null"));
@@ -275,23 +278,22 @@ div_output(struct socket *so, struct mbuf *m,
 	if (control)
 		m_freem(control);		/* XXX */
 
+	mtag = m_tag_get(PACKET_TAG_DIVERT,
+			sizeof(struct divert_tag), M_NOWAIT);
+	if (mtag == NULL) {
+		error = ENOBUFS;
+		goto cantsend;
+	}
+	dt = (struct divert_tag *)(mtag+1);
+	dt->info = 0;
+	dt->cookie = 0;
+	m_tag_prepend(m, mtag);
+
 	/* Loopback avoidance and state recovery */
 	if (sin) {
-		struct m_tag *mtag;
-		struct divert_tag *dt;
 		int i;
 
-		mtag = m_tag_get(PACKET_TAG_DIVERT,
-				sizeof(struct divert_tag), M_NOWAIT);
-		if (mtag == NULL) {
-			error = ENOBUFS;
-			goto cantsend;
-		}
-		dt = (struct divert_tag *)(mtag+1);
-		dt->info = 0;
 		dt->cookie = sin->sin_port;
-		m_tag_prepend(m, mtag);
-
 		/*
 		 * Find receive interface with the given name, stuffed
 		 * (if it exists) in the sin_zero[] field.
@@ -309,6 +311,7 @@ div_output(struct socket *so, struct mbuf *m,
 		struct ip *const ip = mtod(m, struct ip *);
 		struct inpcb *inp;
 
+		dt->info |= IP_FW_DIVERT_OUTPUT_FLAG;
 		INP_INFO_WLOCK(&divcbinfo);
 		inp = sotoinpcb(so);
 		INP_LOCK(inp);
@@ -341,6 +344,7 @@ div_output(struct socket *so, struct mbuf *m,
 		INP_UNLOCK(inp);
 		INP_INFO_WUNLOCK(&divcbinfo);
 	} else {
+		dt->info |= IP_FW_DIVERT_LOOPBACK_FLAG;
 		if (m->m_pkthdr.rcvif == NULL) {
 			/*
 			 * No luck with the name, check by IP address.
diff --git a/sys/netinet/ip_fw.h b/sys/netinet/ip_fw.h
index 943b505..15ce84b 100644
--- a/sys/netinet/ip_fw.h
+++ b/sys/netinet/ip_fw.h
@@ -135,6 +135,7 @@ enum ipfw_opcodes {		/* arguments (4 byte each)	*/
 	O_ANTISPOOF,		/* none				*/
 	O_JAIL,			/* u32 = id			*/
 	O_ALTQ,			/* u32 = altq classif. qid	*/
+	O_DIVERTED,		/* arg1=bitmap (1:loop, 2:out)	*/
 
 	O_LAST_OPCODE		/* not an opcode!		*/
 };
@@ -415,9 +416,11 @@ typedef struct	_ipfw_table {
  */
 #ifdef _KERNEL
 
-#define	IP_FW_PORT_DYNT_FLAG	0x10000
-#define	IP_FW_PORT_TEE_FLAG	0x20000
-#define	IP_FW_PORT_DENY_FLAG	0x40000
+#define	IP_FW_PORT_DYNT_FLAG		0x00010000
+#define	IP_FW_PORT_TEE_FLAG		0x00020000
+#define	IP_FW_PORT_DENY_FLAG		0x00040000
+#define	IP_FW_DIVERT_LOOPBACK_FLAG	0x00080000
+#define	IP_FW_DIVERT_OUTPUT_FLAG	0x00100000
 
 /*
  * Arguments for calling ipfw_chk() and dummynet_io(). We put them
diff --git a/sys/netinet/ip_fw2.c b/sys/netinet/ip_fw2.c
index 9108d7b..854e2d9 100644
--- a/sys/netinet/ip_fw2.c
+++ b/sys/netinet/ip_fw2.c
@@ -1718,6 +1718,14 @@ ipfw_chk(struct ip_fw_args *args)
 	int ugid_lookup = 0;
 
 	/*
+	 * divinput_flags	If non-zero, set to the IP_FW_DIVERT_*_FLAG
+	 *	associated with a packet input on a divert socket.  This
+	 *	will allow to distinguish traffic and its direction when
+	 *	it originates from a divert socket.
+	 */
+	u_int divinput_flags = 0;
+
+	/*
 	 * oif | args->oif	If NULL, ipfw_chk has been called on the
 	 *	inbound path (ether_input, bdg_forward, ip_input).
 	 *	If non-NULL, ipfw_chk has been called on the outbound path
@@ -1893,8 +1901,11 @@ after_ip_checks:
 		}
 	}
 	/* reset divert rule to avoid confusion later */
-	if (mtag)
+	if (mtag) {
+		divinput_flags = divert_info(mtag) &
+		    (IP_FW_DIVERT_OUTPUT_FLAG | IP_FW_DIVERT_LOOPBACK_FLAG);
 		m_tag_delete(m, mtag);
+	}
 
 	/*
 	 * Now scan the rules, and parse microinstructions for each rule.
@@ -2027,6 +2038,13 @@ check_body:
 				match = (args->eh != NULL);
 				break;
 
+			case O_DIVERTED:
+				match = (cmd->arg1 & 1 && divinput_flags &
+				    IP_FW_DIVERT_LOOPBACK_FLAG) ||
+					(cmd->arg1 & 2 && divinput_flags &
+				    IP_FW_DIVERT_OUTPUT_FLAG);
+				break;
+
 			case O_PROTO:
 				/*
 				 * We do not allow an arg of 0 so the
@@ -2912,6 +2930,7 @@ check_ipfw_struct(struct ip_fw *rule, int size)
 		case O_LAYER2:
 		case O_IN:
 		case O_FRAG:
+		case O_DIVERTED:
 		case O_IPOPT:
 		case O_IPTOS:
 		case O_IPPRECEDENCE: