git: 600745f1e226 - main - pf: bound DIOCGETSTATES memory use

[Kristof Provost]

The branch main has been updated by kp:

URL: https://cgit.FreeBSD.org/src/commit/?id=600745f1e2260e7ed3c2e6183b24388ff38c916c

commit 600745f1e2260e7ed3c2e6183b24388ff38c916c
Author:     Kristof Provost <kp at FreeBSD.org>
AuthorDate: 2021-08-02 07:46:33 +0000
Commit:     Kristof Provost <kp at FreeBSD.org>
CommitDate: 2021-08-02 14:29:23 +0000

    pf: bound DIOCGETSTATES memory use
    
    Similar to what we did earlier for DIOCGETSTATESV2 we only allocate
    enough memory for a handful of states and copy those out, bit by bit,
    rather than allocating memory for all states in one go.
    
    MFC after:      1 week
    Sponsored by:   Rubicon Communications, LLC ("Netgate")
---
 sys/netpfil/pf/pf_ioctl.c | 51 ++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 39 insertions(+), 12 deletions(-)

diff --git a/sys/netpfil/pf/pf_ioctl.c b/sys/netpfil/pf/pf_ioctl.c
index 7c506b79295b..b2d7fc33d8b8 100644
--- a/sys/netpfil/pf/pf_ioctl.c
+++ b/sys/netpfil/pf/pf_ioctl.c
@@ -2916,7 +2916,9 @@ DIOCCHANGERULE_error:
 		struct pfioc_states	*ps = (struct pfioc_states *)addr;
 		struct pf_kstate	*s;
 		struct pfsync_state	*pstore, *p;
-		int i, nr;
+		int			 i, nr;
+		size_t			 slice_count = 16, count;
+		void			*out;
 
 		if (ps->ps_len <= 0) {
 			nr = uma_zone_get_cur(V_pf_state_z);
@@ -2924,34 +2926,59 @@ DIOCCHANGERULE_error:
 			break;
 		}
 
-		p = pstore = malloc(ps->ps_len, M_TEMP, M_WAITOK | M_ZERO);
+		out = ps->ps_states;
+		pstore = mallocarray(slice_count,
+		    sizeof(struct pfsync_state), M_TEMP, M_WAITOK | M_ZERO);
 		nr = 0;
 
 		for (i = 0; i <= pf_hashmask; i++) {
 			struct pf_idhash *ih = &V_pf_idhash[i];
 
+DIOCGETSTATES_retry:
+			p = pstore;
+
+			if (LIST_EMPTY(&ih->states))
+				continue;
+
 			PF_HASHROW_LOCK(ih);
+			count = 0;
+			LIST_FOREACH(s, &ih->states, entry) {
+				if (s->timeout == PFTM_UNLINKED)
+					continue;
+				count++;
+			}
+
+			if (count > slice_count) {
+				PF_HASHROW_UNLOCK(ih);
+				free(pstore, M_TEMP);
+				slice_count = count * 2;
+				pstore = mallocarray(slice_count,
+				    sizeof(struct pfsync_state), M_TEMP,
+				    M_WAITOK | M_ZERO);
+				goto DIOCGETSTATES_retry;
+			}
+
+			if ((nr+count) * sizeof(*p) > ps->ps_len) {
+				PF_HASHROW_UNLOCK(ih);
+				goto DIOCGETSTATES_full;
+			}
+
 			LIST_FOREACH(s, &ih->states, entry) {
 				if (s->timeout == PFTM_UNLINKED)
 					continue;
 
-				if ((nr+1) * sizeof(*p) > ps->ps_len) {
-					PF_HASHROW_UNLOCK(ih);
-					goto DIOCGETSTATES_full;
-				}
 				pfsync_state_export(p, s);
 				p++;
 				nr++;
 			}
 			PF_HASHROW_UNLOCK(ih);
+			error = copyout(pstore, out,
+			    sizeof(struct pfsync_state) * count);
+			if (error)
+				break;
+			out = ps->ps_states + nr;
 		}
 DIOCGETSTATES_full:
-		error = copyout(pstore, ps->ps_states,
-		    sizeof(struct pfsync_state) * nr);
-		if (error) {
-			free(pstore, M_TEMP);
-			break;
-		}
 		ps->ps_len = sizeof(struct pfsync_state) * nr;
 		free(pstore, M_TEMP);