- packet must keep reference to statekey

  this is the second attempt to get it in, the first
  attempt got backed out on Jan 31 2016

  the change also contains fixes contributed by Stefan Kempf
  in earlier iteration.

OK srhen@

1 files changed, 160 insertions, 43 deletions

diff --git a/sys/net/pf.c b/sys/net/pf.c
index 3c7e0f4c807..4e4dca7c3b2 100644
--- a/sys/net/pf.c
+++ b/sys/net/pf.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: pf.c,v 1.966 2016/03/04 22:38:23 sashan Exp $ */
+/*	$OpenBSD: pf.c,v 1.967 2016/03/29 10:34:42 sashan Exp $ */
 
 /*
  * Copyright (c) 2001 Daniel Hartmeier
@@ -231,6 +231,11 @@ int			 pf_step_out_of_anchor(int *, struct pf_ruleset **,
 void			 pf_counters_inc(int, struct pf_pdesc *,
 			    struct pf_state *, struct pf_rule *,
 			    struct pf_rule *);
+void			 pf_state_key_link(struct pf_state_key *,
+			    struct pf_state_key *);
+void			 pf_inpcb_unlink_state_key(struct inpcb *);
+void			 pf_state_key_unlink_reverse(struct pf_state_key *);
+
 #if NPFLOG > 0
 void			 pf_log_matches(struct pf_pdesc *, struct pf_rule *,
 			    struct pf_rule *, struct pf_ruleset *,
@@ -732,6 +737,7 @@ void
 pf_state_key_detach(struct pf_state *s, int idx)
 {
 	struct pf_state_item	*si;
+	struct pf_state_key	*sk;
 
 	if (s->key[idx] == NULL)
 		return;
@@ -745,15 +751,15 @@ pf_state_key_detach(struct pf_state *s, int idx)
 		pool_put(&pf_state_item_pl, si);
 	}
 
-	if (TAILQ_EMPTY(&s->key[idx]->states)) {
-		RB_REMOVE(pf_state_tree, &pf_statetbl, s->key[idx]);
-		if (s->key[idx]->reverse)
-			s->key[idx]->reverse->reverse = NULL;
-		if (s->key[idx]->inp)
-			s->key[idx]->inp->inp_pf_sk = NULL;
-		pool_put(&pf_state_key_pl, s->key[idx]);
-	}
+	sk = s->key[idx];
 	s->key[idx] = NULL;
+	if (TAILQ_EMPTY(&sk->states)) {
+		RB_REMOVE(pf_state_tree, &pf_statetbl, sk);
+		sk->removed = 1;
+		pf_state_key_unlink_reverse(sk);
+		pf_inpcb_unlink_state_key(sk->inp);
+		pf_state_key_unref(sk);
+	}
 }
 
 struct pf_state_key *
@@ -840,6 +846,8 @@ pf_state_key_setup(struct pf_pdesc *pd, struct pf_state_key **skw,
 	sk1->proto = pd->proto;
 	sk1->af = pd->af;
 	sk1->rdomain = pd->rdomain;
+	PF_REF_INIT(sk1->refcnt);
+	sk1->removed = 0;
 	if (rtableid >= 0)
 		wrdom = rtable_l2(rtableid);
 
@@ -871,6 +879,8 @@ pf_state_key_setup(struct pf_pdesc *pd, struct pf_state_key **skw,
 			sk2->proto = pd->proto;
 		sk2->af = pd->naf;
 		sk2->rdomain = wrdom;
+		PF_REF_INIT(sk2->refcnt);
+		sk2->removed = 0;
 	} else
 		sk2 = sk1;
 
@@ -986,7 +996,7 @@ struct pf_state *
 pf_find_state(struct pfi_kif *kif, struct pf_state_key_cmp *key, u_int dir,
     struct mbuf *m)
 {
-	struct pf_state_key	*sk;
+	struct pf_state_key	*sk, *pkt_sk, *inp_sk;
 	struct pf_state_item	*si;
 
 	pf_status.fcounters[FCNT_STATE_SEARCH]++;
@@ -996,31 +1006,47 @@ pf_find_state(struct pfi_kif *kif, struct pf_state_key_cmp *key, u_int dir,
 		addlog("\n");
 	}
 
-	if (dir == PF_OUT && m->m_pkthdr.pf.statekey &&
-	    m->m_pkthdr.pf.statekey->reverse)
-		sk = m->m_pkthdr.pf.statekey->reverse;
-	else if (dir == PF_OUT && m->m_pkthdr.pf.inp &&
-	    m->m_pkthdr.pf.inp->inp_pf_sk)
-		sk = m->m_pkthdr.pf.inp->inp_pf_sk;
-	else {
+	inp_sk = NULL;
+	pkt_sk = NULL;
+	sk = NULL;
+	if (dir == PF_OUT) {
+		/* first if block deals with outbound forwarded packet */
+		pkt_sk = m->m_pkthdr.pf.statekey;
+		if (pf_state_key_isvalid(pkt_sk) &&
+		    pf_state_key_isvalid(pkt_sk->reverse)) {
+			sk = pkt_sk->reverse;
+		} else {
+			pf_pkt_unlink_state_key(m);
+			pkt_sk = NULL;
+		}
+
+		if (pkt_sk == NULL) {
+			/* here we deal with local outbound packet */
+			if (m->m_pkthdr.pf.inp != NULL) {
+				inp_sk = m->m_pkthdr.pf.inp->inp_pf_sk;
+				if (pf_state_key_isvalid(inp_sk))
+					sk = inp_sk;
+				else
+					pf_inpcb_unlink_state_key(
+					    m->m_pkthdr.pf.inp);
+			}
+		}
+	}
+
+	if (sk == NULL) {
 		if ((sk = RB_FIND(pf_state_tree, &pf_statetbl,
 		    (struct pf_state_key *)key)) == NULL)
 			return (NULL);
-		if (dir == PF_OUT && m->m_pkthdr.pf.statekey &&
-		    pf_compare_state_keys(m->m_pkthdr.pf.statekey, sk,
-		    kif, dir) == 0) {
-			m->m_pkthdr.pf.statekey->reverse = sk;
-			sk->reverse = m->m_pkthdr.pf.statekey;
-		} else if (dir == PF_OUT && m->m_pkthdr.pf.inp && !sk->inp) {
-			m->m_pkthdr.pf.inp->inp_pf_sk = sk;
-			sk->inp = m->m_pkthdr.pf.inp;
-		}
+		if (dir == PF_OUT && pkt_sk &&
+		    pf_compare_state_keys(pkt_sk, sk, kif, dir) == 0)
+			pf_state_key_link(sk, pkt_sk);
+		else if (dir == PF_OUT)
+			pf_inp_link(m, m->m_pkthdr.pf.inp);
 	}
 
-	if (dir == PF_OUT) {
-		m->m_pkthdr.pf.statekey = NULL;
-		m->m_pkthdr.pf.inp = NULL;
-	}
+	/* remove firewall data from outbound packet */
+	if (dir == PF_OUT)
+		pf_pkt_addr_changed(m);
 
 	/* list is sorted, if-bound states before floating ones */
 	TAILQ_FOREACH(si, &sk->states, entry)
@@ -6540,12 +6566,14 @@ done:
 		 * IP tunnels.
 		 */
 		KASSERT(pd.m->m_pkthdr.pf.statekey == NULL);
-		pd.m->m_pkthdr.pf.statekey = s->key[PF_SK_STACK];
+		pd.m->m_pkthdr.pf.statekey =
+		    pf_state_key_ref(s->key[PF_SK_STACK]);
 	}
 	if (pd.dir == PF_OUT &&
 	    pd.m->m_pkthdr.pf.inp && !pd.m->m_pkthdr.pf.inp->inp_pf_sk &&
 	    s && s->key[PF_SK_STACK] && !s->key[PF_SK_STACK]->inp) {
-		pd.m->m_pkthdr.pf.inp->inp_pf_sk = s->key[PF_SK_STACK];
+		pd.m->m_pkthdr.pf.inp->inp_pf_sk =
+		    pf_state_key_ref(s->key[PF_SK_STACK]);
 		s->key[PF_SK_STACK]->inp = pd.m->m_pkthdr.pf.inp;
 	}
 
@@ -6716,7 +6744,7 @@ pf_cksum(struct pf_pdesc *pd, struct mbuf *m)
 void
 pf_pkt_addr_changed(struct mbuf *m)
 {
-	m->m_pkthdr.pf.statekey = NULL;
+	pf_pkt_unlink_state_key(m);
 	m->m_pkthdr.pf.inp = NULL;
 }
 
@@ -6724,25 +6752,40 @@ struct inpcb *
 pf_inp_lookup(struct mbuf *m)
 {
 	struct inpcb *inp = NULL;
+	struct pf_state_key *sk = m->m_pkthdr.pf.statekey;
 
-	if (m->m_pkthdr.pf.statekey) {
+	if (!pf_state_key_isvalid(sk))
+		pf_pkt_unlink_state_key(m);
+	else
 		inp = m->m_pkthdr.pf.statekey->inp;
-		if (inp && inp->inp_pf_sk)
-			KASSERT(m->m_pkthdr.pf.statekey == inp->inp_pf_sk);
-	}
+
+	if (inp && inp->inp_pf_sk)
+		KASSERT(m->m_pkthdr.pf.statekey == inp->inp_pf_sk);
+
 	return (inp);
 }
 
 void
 pf_inp_link(struct mbuf *m, struct inpcb *inp)
 {
-	if (m->m_pkthdr.pf.statekey && inp &&
-	    !m->m_pkthdr.pf.statekey->inp && !inp->inp_pf_sk) {
-		m->m_pkthdr.pf.statekey->inp = inp;
-		inp->inp_pf_sk = m->m_pkthdr.pf.statekey;
+	struct pf_state_key *sk = m->m_pkthdr.pf.statekey;
+
+	if (!pf_state_key_isvalid(sk)) {
+		pf_pkt_unlink_state_key(m);
+		return;
+	}
+
+	/*
+	 * we don't need to grab PF-lock here. At worst case we link inp to
+	 * state, which might be just being marked as deleted by another
+	 * thread.
+	 */
+	if (inp && !sk->inp && !inp->inp_pf_sk) {
+		sk->inp = inp;
+		inp->inp_pf_sk = pf_state_key_ref(sk);
 	}
 	/* The statekey has finished finding the inp, it is no longer needed. */
-	m->m_pkthdr.pf.statekey = NULL;
+	pf_pkt_unlink_state_key(m);
 }
 
 void
@@ -6750,10 +6793,21 @@ pf_inp_unlink(struct inpcb *inp)
 {
 	if (inp->inp_pf_sk) {
 		inp->inp_pf_sk->inp = NULL;
-		inp->inp_pf_sk = NULL;
+		pf_inpcb_unlink_state_key(inp);
 	}
 }
 
+void
+pf_state_key_link(struct pf_state_key *sk, struct pf_state_key *pkt_sk)
+{
+	/*
+	 * Assert will not wire as long as we are called by pf_find_state()
+	 */
+	KASSERT((pkt_sk->reverse == NULL) && (sk->reverse == NULL));
+	pkt_sk->reverse = pf_state_key_ref(sk);
+	sk->reverse = pf_state_key_ref(pkt_sk);
+}
+
 #if NPFLOG > 0
 void
 pf_log_matches(struct pf_pdesc *pd, struct pf_rule *rm, struct pf_rule *am,
@@ -6770,3 +6824,66 @@ pf_log_matches(struct pf_pdesc *pd, struct pf_rule *rm, struct pf_rule *am,
 			PFLOG_PACKET(pd, PFRES_MATCH, rm, am, ruleset, ri->r);
 }
 #endif	/* NPFLOG > 0 */
+
+struct pf_state_key *
+pf_state_key_ref(struct pf_state_key *sk)
+{
+	if (sk != NULL)
+		PF_REF_TAKE(sk->refcnt);
+
+	return (sk);
+}
+
+void
+pf_state_key_unref(struct pf_state_key *sk)
+{
+	if ((sk != NULL) && PF_REF_RELE(sk->refcnt)) {
+		/* state key must be removed from tree */
+		KASSERT(!pf_state_key_isvalid(sk));
+		/* state key must be unlinked from reverse key */
+		KASSERT(sk->reverse == NULL);
+		/* state key must be unlinked from socket */
+		KASSERT((sk->inp == NULL) || (sk->inp->inp_pf_sk == NULL));
+		sk->inp = NULL;
+		pool_put(&pf_state_key_pl, sk);
+	}
+}
+
+int
+pf_state_key_isvalid(struct pf_state_key *sk)
+{
+	return ((sk != NULL) && (sk->removed == 0));
+}
+
+void
+pf_pkt_unlink_state_key(struct mbuf *m)
+{
+	pf_state_key_unref(m->m_pkthdr.pf.statekey);
+	m->m_pkthdr.pf.statekey = NULL;
+}
+
+void
+pf_pkt_state_key_ref(struct mbuf *m)
+{
+	pf_state_key_ref(m->m_pkthdr.pf.statekey);
+}
+
+void
+pf_inpcb_unlink_state_key(struct inpcb *inp)
+{
+	if (inp != NULL) {
+		pf_state_key_unref(inp->inp_pf_sk);
+		inp->inp_pf_sk = NULL;
+	}
+}
+
+void
+pf_state_key_unlink_reverse(struct pf_state_key *sk)
+{
+	if ((sk != NULL) && (sk->reverse != NULL)) {
+		pf_state_key_unref(sk->reverse->reverse);
+		sk->reverse->reverse = NULL;
+		pf_state_key_unref(sk->reverse);
+		sk->reverse = NULL;
+	}
+}