[PATCH bpf-next 2/7] dst: Add socket prefetch metadata destinations

[Joe Stringer <joe@wand.net.nz>]

Metadata destinations were introduced in commit f38a9eb1f77b
("dst: Metadata destinations") to "carry per packet metadata
between forwarding and processing elements via the skb->dst pointer".

The aim of this new METADATA_SK_PREFETCH destination type is to allow
early forwarding elements to store a socket destination for the duration
of receiving into the IP stack, which can be later be identified to
avoid orphaning the skb and losing the prefetched socket in ip_rcv_core().

The destination is stored temporarily in a per-CPU buffer to ensure that
if applications attempt to reach out from loopback address to loopback
address, they may restore the original destination and avoid martian
packet drops.

Signed-off-by: Joe Stringer <joe@wand.net.nz>
---
 include/net/dst_metadata.h | 31 +++++++++++++++++++++++++++++++
 net/core/dst.c             | 30 ++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/include/net/dst_metadata.h b/include/net/dst_metadata.h
index 56cb3c38569a..31574c553a07 100644
--- a/include/net/dst_metadata.h
+++ b/include/net/dst_metadata.h
@@ -9,6 +9,7 @@
 enum metadata_type {
 	METADATA_IP_TUNNEL,
 	METADATA_HW_PORT_MUX,
+	METADATA_SK_PREFETCH,
 };
 
 struct hw_port_info {
@@ -80,6 +81,8 @@ static inline int skb_metadata_dst_cmp(const struct sk_buff *skb_a,
 		return memcmp(&a->u.tun_info, &b->u.tun_info,
 			      sizeof(a->u.tun_info) +
 					 a->u.tun_info.options_len);
+	case METADATA_SK_PREFETCH:
+		return 0;
 	default:
 		return 1;
 	}
@@ -214,4 +217,32 @@ static inline struct metadata_dst *ipv6_tun_rx_dst(struct sk_buff *skb,
 				  0, ip6_flowlabel(ip6h), flags, tunnel_id,
 				  md_size);
 }
+
+extern const struct metadata_dst dst_sk_prefetch;
+
+static inline bool dst_is_sk_prefetch(const struct dst_entry *dst)
+{
+	return dst == &dst_sk_prefetch.dst;
+}
+
+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_fetch(struct sk_buff *skb);
+
+/**
+ * dst_sk_prefetch_reset - reset prefetched socket dst
+ * @skb: buffer
+ *
+ * Reverts the dst back to the originally stored dst if present.
+ */
+static inline void dst_sk_prefetch_reset(struct sk_buff *skb)
+{
+	if (unlikely(skb_dst_is_sk_prefetch(skb)))
+		dst_sk_prefetch_fetch(skb);
+}
+
 #endif /* __NET_DST_METADATA_H */
diff --git a/net/core/dst.c b/net/core/dst.c
index 193af526e908..cf1a1d5b6b0a 100644
--- a/net/core/dst.c
+++ b/net/core/dst.c
@@ -330,3 +330,33 @@ void metadata_dst_free_percpu(struct metadata_dst __percpu *md_dst)
 	free_percpu(md_dst);
 }
 EXPORT_SYMBOL_GPL(metadata_dst_free_percpu);
+
+const struct metadata_dst dst_sk_prefetch = {
+	.dst = {
+		.ops = &md_dst_ops,
+		.input = dst_md_discard,
+		.output = dst_md_discard_out,
+		.flags = DST_NOCOUNT | DST_METADATA,
+		.obsolete = DST_OBSOLETE_NONE,
+		.__refcnt = ATOMIC_INIT(1),
+	},
+	.type = METADATA_SK_PREFETCH,
+};
+EXPORT_SYMBOL(dst_sk_prefetch);
+
+DEFINE_PER_CPU(unsigned long, dst_sk_prefetch_dst);
+
+void dst_sk_prefetch_store(struct sk_buff *skb)
+{
+	unsigned long refdst;
+
+	refdst = skb->_skb_refdst;
+	__this_cpu_write(dst_sk_prefetch_dst, refdst);
+	skb_dst_set_noref(skb, (struct dst_entry *)&dst_sk_prefetch.dst);
+}
+
+void dst_sk_prefetch_fetch(struct sk_buff *skb)
+{
+	skb->_skb_refdst = __this_cpu_read(dst_sk_prefetch_dst);
+}
+EXPORT_SYMBOL(dst_sk_prefetch_fetch);
-- 
2.20.1