Re: [PATCH v3] bpf, sockmap: Fix cork ownership in tcp_bpf_sendmsg()

diff --git a/net/ipv4/tcp_bpf.c b/net/ipv4/tcp_bpf.c
index cc0bd73f36b6..fcf535b7a754 100644
--- a/net/ipv4/tcp_bpf.c
+++ b/net/ipv4/tcp_bpf.c

@@ -401,7 +401,9 @@ static int tcp_bpf_recvmsg(struct sock *sk, struct msghdr *msg, size_t len,
 }

 static int tcp_bpf_send_verdict(struct sock *sk, struct sk_psock *psock,
-				struct sk_msg *msg, int *copied, int flags)
+				struct sk_msg *msg, int *copied,
+				u32 msg_copied, bool *corked,
+				int flags)
 {
 	bool cork = false, enospc = sk_msg_full(msg), redir_ingress;
 	struct sock *sk_redir;

@@ -428,12 +430,18 @@ static int tcp_bpf_send_verdict(struct sock *sk, struct sk_psock *psock,
 			psock->cork = kzalloc_obj(*psock->cork,
 						  GFP_ATOMIC | __GFP_NOWARN);
 			if (!psock->cork) {
-				sk_msg_free(sk, msg);
-				*copied = 0;
+				int free;
+
+				psock->cork_bytes = 0;
+				free = sk_msg_free(sk, msg);
+				*copied -= min_t(u32, msg_copied, free);
 				return -ENOMEM;
 			}
 		}
-		memcpy(psock->cork, msg, sizeof(*msg));
+		if (psock->cork != msg) {
+			sk_msg_xfer_full(psock->cork, msg);
+			*corked = true;
+		}
 		return 0;
 	}

@@ -480,14 +488,15 @@ static int tcp_bpf_send_verdict(struct sock *sk, struct sk_psock *psock,
 		if (unlikely(ret < 0)) {
 			int free = sk_msg_free(sk, msg);

-			if (!cork)
+			if (cork)
+				*copied -= min_t(u32, msg_copied, free);
+			else
 				*copied -= free;
 		}
 		if (cork) {
 			sk_msg_free(sk, msg);
 			kfree(msg);
 			msg = NULL;
-			ret = 0;
 		}
 		break;
 	case __SK_DROP:

@@ -519,6 +528,7 @@ static int tcp_bpf_sendmsg(struct sock *sk, struct msghdr *msg, size_t size)
 	struct sk_msg tmp, *msg_tx = NULL;
 	int copied = 0, err = 0, ret = 0;
 	struct sk_psock *psock;
+	u32 msg_copied = 0;
 	long timeo;
 	int flags;

@@ -545,9 +555,12 @@ static int tcp_bpf_sendmsg(struct sock *sk, struct msghdr *msg, size_t size)
 		if (!sk_stream_memory_free(sk))
 			goto wait_for_sndbuf;
 		if (psock->cork) {
+			if (msg_tx != psock->cork)
+				msg_copied = 0;
 			msg_tx = psock->cork;
 		} else {
 			msg_tx = &tmp;
+			msg_copied = 0;
 			sk_msg_init(msg_tx);
 		}

@@ -580,21 +595,23 @@ static int tcp_bpf_sendmsg(struct sock *sk, struct msghdr *msg, size_t size)
 			psock->cork_bytes = 0;
 		}

-		err = tcp_bpf_send_verdict(sk, psock, msg_tx, &copied, flags);
+		err = tcp_bpf_send_verdict(sk, psock, msg_tx, &copied,
+					   msg_copied, &corked, flags);
 		if (unlikely(err < 0))
 			goto out_err;
+		if (corked)
+			msg_tx = psock->cork;
 		continue;