[PATCH net-next] tls: Fix improper revert in zerocopy_from_iter

[Doron Roberts-Kedes <doronrk@fb.com>]

The current code is problematic because the iov_iter is reverted and
never advanced in the non-error case. This patch skips the revert in the
non-error case. This patch also fixes the amount by which the iov_iter
is reverted. Currently, iov_iter is reverted by size, which can be
greater than the amount by which the iter was actually advanced.
Instead, mimic the tx path which reverts by the difference before and
after zerocopy_from_iter. 

Fixes: 4718799817c5 ("tls: Fix zerocopy_from_iter iov handling")
Signed-off-by: Doron Roberts-Kedes <doronrk@fb.com>
---
 net/tls/tls_sw.c | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/net/tls/tls_sw.c b/net/tls/tls_sw.c
index 490f2bcc6313..2ea000baebf8 100644
--- a/net/tls/tls_sw.c
+++ b/net/tls/tls_sw.c
@@ -276,7 +276,7 @@ static int zerocopy_from_iter(struct sock *sk, struct iov_iter *from,
 			      int length, int *pages_used,
 			      unsigned int *size_used,
 			      struct scatterlist *to, int to_max_pages,
-			      bool charge, bool revert)
+			      bool charge)
 {
 	struct page *pages[MAX_SKB_FRAGS];
 
@@ -327,8 +327,6 @@ static int zerocopy_from_iter(struct sock *sk, struct iov_iter *from,
 out:
 	*size_used = size;
 	*pages_used = num_elem;
-	if (revert)
-		iov_iter_revert(from, size);
 
 	return rc;
 }
@@ -431,7 +429,7 @@ int tls_sw_sendmsg(struct sock *sk, struct msghdr *msg, size_t size)
 				&ctx->sg_plaintext_size,
 				ctx->sg_plaintext_data,
 				ARRAY_SIZE(ctx->sg_plaintext_data),
-				true, false);
+				true);
 			if (ret)
 				goto fallback_to_reg_send;
 
@@ -811,6 +809,7 @@ int tls_sw_recvmsg(struct sock *sk,
 			    likely(!(flags & MSG_PEEK)))  {
 				struct scatterlist sgin[MAX_SKB_FRAGS + 1];
 				int pages = 0;
+				int orig_chunk = chunk;
 
 				zc = true;
 				sg_init_table(sgin, MAX_SKB_FRAGS + 1);
@@ -820,9 +819,11 @@ int tls_sw_recvmsg(struct sock *sk,
 				err = zerocopy_from_iter(sk, &msg->msg_iter,
 							 to_copy, &pages,
 							 &chunk, &sgin[1],
-							 MAX_SKB_FRAGS,	false, true);
-				if (err < 0)
+							 MAX_SKB_FRAGS,	false);
+				if (err < 0) {
+					iov_iter_revert(&msg->msg_iter, chunk - orig_chunk);
 					goto fallback_to_reg_recv;
+				}
 
 				err = decrypt_skb(sk, skb, sgin);
 				for (; pages > 0; pages--)
-- 
2.17.1