From 19d6e29b872232c47190344f3dfded2f73edd8ae Mon Sep 17 00:00:00 2001
From: Mateusz Guzik <mjg@FreeBSD.org>
Date: Thu, 8 Jul 2021 15:11:57 +0200
Subject: [PATCH] pf: add pf_find_state_all_exists

Reviewed by:	kp
Sponsored by:	Rubicon Communications, LLC ("Netgate")
---
 sys/net/pfvar.h        |  2 ++
 sys/netpfil/pf/pf.c    |  9 +++++++++
 sys/netpfil/pf/pf_lb.c | 10 ++++------
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/sys/net/pfvar.h b/sys/net/pfvar.h
index f491e18b03a4..7bbb585106c7 100644
--- a/sys/net/pfvar.h
+++ b/sys/net/pfvar.h
@@ -1591,6 +1591,8 @@ pf_release_staten(struct pf_kstate *s, u_int n)
 extern struct pf_kstate		*pf_find_state_byid(uint64_t, uint32_t);
 extern struct pf_kstate		*pf_find_state_all(struct pf_state_key_cmp *,
 				    u_int, int *);
+extern bool			pf_find_state_all_exists(struct pf_state_key_cmp *,
+				    u_int);
 extern struct pf_ksrc_node	*pf_find_src_node(struct pf_addr *,
 				    struct pf_krule *, sa_family_t, int);
 extern void			 pf_unlink_src_node(struct pf_ksrc_node *);
diff --git a/sys/netpfil/pf/pf.c b/sys/netpfil/pf/pf.c
index 8fae01ce1c23..3ca921ff3287 100644
--- a/sys/netpfil/pf/pf.c
+++ b/sys/netpfil/pf/pf.c
@@ -1453,6 +1453,15 @@ pf_find_state_all(struct pf_state_key_cmp *key, u_int dir, int *more)
 	return (ret);
 }
 
+bool
+pf_find_state_all_exists(struct pf_state_key_cmp *key, u_int dir)
+{
+	struct pf_kstate *s;
+
+	s = pf_find_state_all(key, dir, NULL);
+	return (s != NULL);
+}
+
 /* END state table stuff */
 
 static void
diff --git a/sys/netpfil/pf/pf_lb.c b/sys/netpfil/pf/pf_lb.c
index 5e281eccc144..000ee69d9ae9 100644
--- a/sys/netpfil/pf/pf_lb.c
+++ b/sys/netpfil/pf/pf_lb.c
@@ -244,13 +244,13 @@ pf_get_sport(sa_family_t af, u_int8_t proto, struct pf_krule *r,
 			 * (traceroute -I through nat)
 			 */
 			key.port[1] = sport;
-			if (pf_find_state_all(&key, PF_IN, NULL) == NULL) {
+			if (!pf_find_state_all_exists(&key, PF_IN)) {
 				*nport = sport;
 				return (0);
 			}
 		} else if (low == high) {
 			key.port[1] = htons(low);
-			if (pf_find_state_all(&key, PF_IN, NULL) == NULL) {
+			if (!pf_find_state_all_exists(&key, PF_IN)) {
 				*nport = htons(low);
 				return (0);
 			}
@@ -268,8 +268,7 @@ pf_get_sport(sa_family_t af, u_int8_t proto, struct pf_krule *r,
 			/* low <= cut <= high */
 			for (tmp = cut; tmp <= high && tmp <= 0xffff; ++tmp) {
 				key.port[1] = htons(tmp);
-				if (pf_find_state_all(&key, PF_IN, NULL) ==
-				    NULL) {
+				if (!pf_find_state_all_exists(&key, PF_IN)) {
 					*nport = htons(tmp);
 					return (0);
 				}
@@ -277,8 +276,7 @@ pf_get_sport(sa_family_t af, u_int8_t proto, struct pf_krule *r,
 			tmp = cut;
 			for (tmp -= 1; tmp >= low && tmp <= 0xffff; --tmp) {
 				key.port[1] = htons(tmp);
-				if (pf_find_state_all(&key, PF_IN, NULL) ==
-				    NULL) {
+				if (!pf_find_state_all_exists(&key, PF_IN)) {
 					*nport = htons(tmp);
 					return (0);
 				}