From 86469161dcb4a47bc014297d0c4fd23cc1d2a04b Mon Sep 17 00:00:00 2001
From: yasuoka <yasuoka@openbsd.org>
Date: Fri, 21 Apr 2017 23:21:02 +0000
Subject: [PATCH] Speed up DIOCKILLSTATES by using the RB tree index if all
 fields used by the tree of given state key are filled.

ok sasha
---
 sys/net/pf_ioctl.c | 60 +++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 57 insertions(+), 3 deletions(-)

diff --git a/sys/net/pf_ioctl.c b/sys/net/pf_ioctl.c
index 887a36481d6..c5ec0026b1f 100644
--- a/sys/net/pf_ioctl.c
+++ b/sys/net/pf_ioctl.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: pf_ioctl.c,v 1.308 2017/03/17 17:19:16 mpi Exp $ */
+/*	$OpenBSD: pf_ioctl.c,v 1.309 2017/04/21 23:21:02 yasuoka Exp $ */
 
 /*
  * Copyright (c) 2001 Daniel Hartmeier
@@ -1444,11 +1444,14 @@ pfioctl(dev_t dev, u_long cmd, caddr_t addr, int flags, struct proc *p)
 
 	case DIOCKILLSTATES: {
 		struct pf_state		*s, *nexts;
-		struct pf_state_key	*sk;
+		struct pf_state_item	*si, *sit;
+		struct pf_state_key	*sk, key;
 		struct pf_addr		*srcaddr, *dstaddr;
 		u_int16_t		 srcport, dstport;
 		struct pfioc_state_kill	*psk = (struct pfioc_state_kill *)addr;
-		u_int			 killed = 0;
+		u_int			 i, killed = 0;
+		const int 		 dirs[] = { PF_IN, PF_OUT };
+		int			 sidx, didx;
 
 		if (psk->psk_pfcmp.id) {
 			if (psk->psk_pfcmp.creatorid == 0)
@@ -1460,6 +1463,57 @@ pfioctl(dev_t dev, u_long cmd, caddr_t addr, int flags, struct proc *p)
 			break;
 		}
 
+		if (psk->psk_af && psk->psk_proto &&
+		    psk->psk_src.port_op == PF_OP_EQ &&
+		    psk->psk_dst.port_op == PF_OP_EQ) {
+
+			key.af = psk->psk_af;
+			key.proto = psk->psk_proto;
+			key.rdomain = psk->psk_rdomain;
+
+			for (i = 0; i < nitems(dirs); i++) {
+				if (dirs[i] == PF_IN) {
+					sidx = 0;
+					didx = 1;
+				} else {
+					sidx = 1;
+					didx = 0;
+				}
+				PF_ACPY(&key.addr[sidx],
+				    &psk->psk_src.addr.v.a.addr, key.af);
+				PF_ACPY(&key.addr[didx],
+				    &psk->psk_dst.addr.v.a.addr, key.af);
+				key.port[sidx] = psk->psk_src.port[0];
+				key.port[didx] = psk->psk_dst.port[0];
+
+				sk = RB_FIND(pf_state_tree, &pf_statetbl, &key);
+				if (sk == NULL)
+					continue;
+
+				TAILQ_FOREACH_SAFE(si, &sk->states, entry, sit)
+					if (((si->s->key[PF_SK_WIRE]->af ==
+					    si->s->key[PF_SK_STACK]->af &&
+					    sk == (dirs[i] == PF_IN ?
+					    si->s->key[PF_SK_WIRE] :
+					    si->s->key[PF_SK_STACK])) ||
+					    (si->s->key[PF_SK_WIRE]->af !=
+					    si->s->key[PF_SK_STACK]->af &&
+					    dirs[i] == PF_IN &&
+					    (sk == si->s->key[PF_SK_STACK] ||
+					    sk == si->s->key[PF_SK_WIRE]))) &&
+					    (!psk->psk_ifname[0] ||
+					    (si->s->kif != pfi_all &&
+					    !strcmp(psk->psk_ifname,
+					    si->s->kif->pfik_name)))) {
+						pf_remove_state(si->s);
+						killed++;
+					}
+			}
+			if (killed)
+				psk->psk_killed = killed;
+			break;
+		}
+
 		for (s = RB_MIN(pf_state_tree_id, &tree_id); s;
 		    s = nexts) {
 			nexts = RB_NEXT(pf_state_tree_id, &tree_id, s);
-- 
2.20.1