[PATCH bpf-next 4/7] dst: Prefetch established socket destinations

[Joe Stringer <joe@xxxxxxxxxxx>]

Enhance the dst_sk_prefetch logic to temporarily store the socket
receive destination, to save the route lookup later on. The dst
reference is kept alive by the caller's socket reference.

Suggested-by: Daniel Borkmann <daniel@xxxxxxxxxxxxx>
Signed-off-by: Joe Stringer <joe@xxxxxxxxxxx>
---
 include/net/dst_metadata.h |  2 +-
 net/core/dst.c             | 20 +++++++++++++++++---
 net/core/filter.c          |  2 +-
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/include/net/dst_metadata.h b/include/net/dst_metadata.h
index 31574c553a07..4f16322b08d5 100644
--- a/include/net/dst_metadata.h
+++ b/include/net/dst_metadata.h
@@ -230,7 +230,7 @@ static inline bool skb_dst_is_sk_prefetch(const struct sk_buff *skb)
 	return dst_is_sk_prefetch(skb_dst(skb));
 }
 
-void dst_sk_prefetch_store(struct sk_buff *skb);
+void dst_sk_prefetch_store(struct sk_buff *skb, struct sock *sk);
 void dst_sk_prefetch_fetch(struct sk_buff *skb);
 
 /**
diff --git a/net/core/dst.c b/net/core/dst.c
index cf1a1d5b6b0a..5068d127d9c2 100644
--- a/net/core/dst.c
+++ b/net/core/dst.c
@@ -346,11 +346,25 @@ EXPORT_SYMBOL(dst_sk_prefetch);
 
 DEFINE_PER_CPU(unsigned long, dst_sk_prefetch_dst);
 
-void dst_sk_prefetch_store(struct sk_buff *skb)
+void dst_sk_prefetch_store(struct sk_buff *skb, struct sock *sk)
 {
-	unsigned long refdst;
+	unsigned long refdst = 0L;
+
+	WARN_ON(!rcu_read_lock_held() &&
+		!rcu_read_lock_bh_held());
+	if (sk_fullsock(sk)) {
+		struct dst_entry *dst = READ_ONCE(sk->sk_rx_dst);
+
+		if (dst)
+			dst = dst_check(dst, 0);
+		if (dst)
+			refdst = (unsigned long)dst | SKB_DST_NOREF;
+	}
+	if (!refdst)
+		refdst = skb->_skb_refdst;
+	if (skb->_skb_refdst != refdst)
+		skb_dst_drop(skb);
 
-	refdst = skb->_skb_refdst;
 	__this_cpu_write(dst_sk_prefetch_dst, refdst);
 	skb_dst_set_noref(skb, (struct dst_entry *)&dst_sk_prefetch.dst);
 }
diff --git a/net/core/filter.c b/net/core/filter.c
index bae0874289d8..db9b7b8b4a04 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c
@@ -5858,7 +5858,7 @@ BPF_CALL_3(bpf_sk_assign, struct sk_buff *, skb, struct sock *, sk, u64, flags)
 	skb_orphan(skb);
 	skb->sk = sk;
 	skb->destructor = sock_edemux;
-	dst_sk_prefetch_store(skb);
+	dst_sk_prefetch_store(skb, sk);
 
 	return 0;
 }
-- 
2.20.1