there should be no functional change with this diff, it just shuffles
the code around.

note that pf_state structs always have valid references to pf_state_keys
until the pf_state itself is destroyed after pf.c r1.1157. this means we
don't have to test if the st->key[PF_SK_STACK] pointers are valid
anymore, we can just use them.

ok?

Index: pf.c
===================================================================
RCS file: /cvs/src/sys/net/pf.c,v
diff -u -p -r1.1209 pf.c
--- pf.c	14 Apr 2025 20:02:34 -0000	1.1209
+++ pf.c	30 Apr 2025 04:38:04 -0000
@@ -7910,15 +7910,20 @@ done:
 
 	if (action == PF_PASS && qid)
 		pd.m->m_pkthdr.pf.qid = qid;
-	if (pd.dir == PF_IN && st && st->key[PF_SK_STACK])
-		pf_mbuf_link_state_key(pd.m, st->key[PF_SK_STACK]);
-	if (pd.dir == PF_OUT && st && st->key[PF_SK_STACK])
-		pf_state_key_link_inpcb(st->key[PF_SK_STACK],
-		    pd.m->m_pkthdr.pf.inp);
+	if (st != NULL) {
+		struct mbuf *m = pd.m;
+		struct inpcb *inp = m->m_pkthdr.pf.inp;
 
-	if (st != NULL && !ISSET(pd.m->m_pkthdr.csum_flags, M_FLOWID)) {
-		pd.m->m_pkthdr.ph_flowid = st->key[PF_SK_WIRE]->hash;
-		SET(pd.m->m_pkthdr.csum_flags, M_FLOWID);
+		if (pd.dir == PF_IN) {
+			KASSERT(inp == NULL);
+			pf_mbuf_link_state_key(m, st->key[PF_SK_STACK]);
+		} else if (pd.dir == PF_OUT)
+			pf_state_key_link_inpcb(st->key[PF_SK_STACK], inp);
+
+		if (!ISSET(m->m_pkthdr.csum_flags, M_FLOWID)) {
+			m->m_pkthdr.ph_flowid = st->key[PF_SK_WIRE]->hash;
+			SET(m->m_pkthdr.csum_flags, M_FLOWID);
+		}
 	}
 
 	/*