Re: [PATCH bpf v3 1/2] bpf, sockmap: settle sk_forward_alloc for strparser SK_PASS

diff --git a/net/core/skmsg.c b/net/core/skmsg.c
index 2521b643fa05..347cb168f749 100644
--- a/net/core/skmsg.c
+++ b/net/core/skmsg.c

@@ -586,21 +586,24 @@ static int sk_psock_skb_ingress_enqueue(struct sk_buff *skb,
 }

 static int sk_psock_skb_ingress_self(struct sk_psock *psock, struct sk_buff *skb,
-				     u32 off, u32 len, bool take_ref);
+				     u32 off, u32 len, bool take_ref,
+				     bool settle_fwd_alloc);
+static int sk_psock_skb_ingress_self_backlog(struct sk_psock *psock,
+					     struct sk_buff *skb,
+					     u32 off, u32 len, bool take_ref,
+					     bool settle_fwd_alloc);

@@ -618,34 +621,100 @@ static int sk_psock_skb_ingress(struct sk_psock *psock, struct sk_buff *skb,
 	return err;
 }

-/* Puts an skb on the ingress queue of the socket already assigned to the
- * skb. In this case we do not need to check memory limits or skb_set_owner_r
- * because the skb is already accounted for here.
+static int sk_psock_skb_ingress_self_assign(struct sock *sk,
+					    struct sk_buff *skb,
+					    bool settle_fwd_alloc)
+{
+	/* Leave skbs already receive-accounted to sk untouched. */
+	if (skb->sk == sk && skb->destructor == sock_rfree)
+		return 0;
+
+	if (settle_fwd_alloc) {
+		sock_owned_by_me(sk);
+
+		if (!sk_rmem_schedule(sk, skb, 0))
+			return -EAGAIN;
+	}
+
+	skb_set_owner_r(skb, sk);
+	return 0;
+}
+
+/* Puts an skb on the ingress queue for psock->sk.
+ *
+ * If the skb already has receive ownership for this socket, leave socket
+ * memory accounting untouched. Otherwise, before assigning receive ownership
+ * to an unowned strparser SK_PASS skb, settle any existing sk_forward_alloc
+ * deficit from earlier clone charges.
  */
 static int sk_psock_skb_ingress_self(struct sk_psock *psock, struct sk_buff *skb,
-				     u32 off, u32 len, bool take_ref)
+				     u32 off, u32 len, bool take_ref,
+				     bool settle_fwd_alloc)
 {
+	struct sock *sk = psock->sk;
 	struct sk_msg *msg = alloc_sk_msg(GFP_ATOMIC);
+	int err;
+
+	if (unlikely(!msg))
+		return -EAGAIN;
+
+	err = sk_psock_skb_ingress_self_assign(sk, skb, settle_fwd_alloc);
+	if (err)
+		goto free;
+
+	/* This is used in tcp_bpf_recvmsg_parser() to determine whether the
+	 * data originates from the socket's own protocol stack. No need to
+	 * refcount sk because msg's lifetime is bound to sk via the ingress_msg.
+	 */
+	msg->sk = sk;
+	err = sk_psock_skb_ingress_enqueue(skb, off, len, psock, sk, msg,
+					   take_ref);
+	if (err < 0)
+		goto free;
+
+	return err;
+free:
+	kfree(msg);
+	return err;
+}
+
+static int sk_psock_skb_ingress_self_backlog(struct sk_psock *psock,
+					     struct sk_buff *skb,
+					     u32 off, u32 len, bool take_ref,
+					     bool settle_fwd_alloc)
+{
 	struct sock *sk = psock->sk;
+	struct sk_msg *msg = alloc_sk_msg(GFP_ATOMIC);
 	int err;

 	if (unlikely(!msg))
 		return -EAGAIN;
-	skb_set_owner_r(skb, sk);
+
+	lock_sock(sk);
+	err = sk_psock_skb_ingress_self_assign(sk, skb, settle_fwd_alloc);
+	release_sock(sk);
+	if (err)
+		goto free;

 	/* This is used in tcp_bpf_recvmsg_parser() to determine whether the
 	 * data originates from the socket's own protocol stack. No need to
 	 * refcount sk because msg's lifetime is bound to sk via the ingress_msg.
 	 */
 	msg->sk = sk;
-	err = sk_psock_skb_ingress_enqueue(skb, off, len, psock, sk, msg, take_ref);
+	err = sk_psock_skb_ingress_enqueue(skb, off, len, psock, sk, msg,
+					   take_ref);
 	if (err < 0)
-		kfree(msg);
+		goto free;
+
+	return err;
+free:
+	kfree(msg);
 	return err;
 }