crypto: Hash scatterwalk fixes

index d836e7440ee8..138e2f1bd08f 100644
@@ -671,6 +671,7 @@ sg_page_iter_dma_address(struct sg_dma_page_iter *dma_iter)
 #define SG_MITER_ATOMIC		(1 << 0)	 /* use kmap_atomic */
 #define SG_MITER_TO_SG		(1 << 1)	/* flush back to phys on unmap */
 #define SG_MITER_FROM_SG	(1 << 2)	/* nop */
+#define SG_MITER_LOCAL		(1 << 3)	 /* use kmap_local */
 
 struct sg_mapping_iter {
 	/* the following three fields can be accessed directly */

index 5bb6b8aff232..b58d5ef1a34b 100644
@@ -879,7 +879,7 @@ EXPORT_SYMBOL(sg_miter_skip);
  *   @miter->addr and @miter->length point to the current mapping.
  *
  * Context:
- *   May sleep if !SG_MITER_ATOMIC.
+ *   May sleep if !SG_MITER_ATOMIC && !SG_MITER_LOCAL.
  *
  * Returns:
  *   true if @miter contains the next mapping.  false if end of sg

@@ -901,6 +901,8 @@ bool sg_miter_next(struct sg_mapping_iter *miter)
 
 	if (miter->__flags & SG_MITER_ATOMIC)
 		miter->addr = kmap_atomic(miter->page) + miter->__offset;
+	else if (miter->__flags & SG_MITER_LOCAL)
+		miter->addr = kmap_local_page(miter->page) + miter->__offset;
 	else
 		miter->addr = kmap(miter->page) + miter->__offset;
 

@@ -936,7 +938,9 @@ void sg_miter_stop(struct sg_mapping_iter *miter)
 		if (miter->__flags & SG_MITER_ATOMIC) {
 			WARN_ON_ONCE(!pagefault_disabled());
 			kunmap_atomic(miter->addr);
-		} else
+		} else if (miter->__flags & SG_MITER_LOCAL)
+			kunmap_local(miter->addr);
+		else
 			kunmap(miter->page);
 
 		miter->page = NULL;

@@ -965,7 +969,7 @@ size_t sg_copy_buffer(struct scatterlist *sgl, unsigned int nents, void *buf,
 {
 	unsigned int offset = 0;
 	struct sg_mapping_iter miter;
-	unsigned int sg_flags = SG_MITER_ATOMIC;
+	unsigned int sg_flags = SG_MITER_LOCAL;
 
 	if (to_buffer)
 		sg_flags |= SG_MITER_FROM_SG;

@@ -1080,7 +1084,7 @@ size_t sg_zero_buffer(struct scatterlist *sgl, unsigned int nents,
 {
 	unsigned int offset = 0;
 	struct sg_mapping_iter miter;
-	unsigned int sg_flags = SG_MITER_ATOMIC | SG_MITER_TO_SG;
+	unsigned int sg_flags = SG_MITER_LOCAL | SG_MITER_TO_SG;
 
 	sg_miter_start(&miter, sgl, nents, sg_flags);
 

index c1dcb0dd3a00..79180d28baa9 100644
@@ -67,9 +67,9 @@
 #define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
 
 #include <linux/random.h>
+#include <linux/scatterlist.h>
 #include <linux/skbuff.h>
 #include <linux/slab.h>
-#include <linux/highmem.h>
 #include <linux/lcm.h>
 #include <linux/rtnetlink.h>
 #include <crypto/authenc.h>

@@ -83,26 +83,21 @@
 int crypto_shash_update_sg(struct shash_desc *desc, struct scatterlist *sg,
 			   size_t offset, size_t len)
 {
-	do {
-		int ret;
+	struct sg_mapping_iter miter;
+	size_t i, n;
+	int ret = 0;
 
-		if (offset < sg->length) {
-			struct page *page = sg_page(sg);
-			void *p = kmap_local_page(page);
-			void *q = p + sg->offset + offset;
-			size_t seg = min_t(size_t, len, sg->length - offset);
-
-			ret = crypto_shash_update(desc, q, seg);
-			kunmap_local(p);
-			if (ret < 0)
-				return ret;
-			len -= seg;
-			offset = 0;
-		} else {
-			offset -= sg->length;
-		}
-	} while (len > 0 && (sg = sg_next(sg)));
-	return 0;
+	sg_miter_start(&miter, sg, sg_nents(sg),
+		       SG_MITER_FROM_SG | SG_MITER_LOCAL);
+	for (i = 0; i < len; i += n) {
+		sg_miter_next(&miter);
+		n = min(miter.length, len - i);
+		ret = crypto_shash_update(desc, miter.addr, n);
+		if (ret < 0)
+			break;
+	}
+	sg_miter_stop(&miter);
+	return ret;
 }
 
 static int rfc3961_do_encrypt(struct crypto_sync_skcipher *tfm, void *iv,

index 9c26175c21a8..1fe594880295 100644
@@ -16,6 +16,7 @@
 #include <linux/cryptouser.h>
 #include <linux/err.h>
 #include <linux/kernel.h>
+#include <linux/mm.h>
 #include <linux/module.h>
 #include <linux/sched.h>
 #include <linux/slab.h>

@@ -201,25 +202,36 @@ int shash_ahash_digest(struct ahash_request *req, struct shash_desc *desc)
 	unsigned int nbytes = req->nbytes;
 	struct scatterlist *sg;
 	unsigned int offset;
+	struct page *page;
+	const u8 *data;
 	int err;
 
-	if (ahash_request_isvirt(req))
-		return crypto_shash_digest(desc, req->svirt, nbytes,
-					   req->result);
+	data = req->svirt;
+	if (!nbytes || ahash_request_isvirt(req))
+		return crypto_shash_digest(desc, data, nbytes, req->result);
 
-	if (nbytes &&
-	    (sg = req->src, offset = sg->offset,
-	     nbytes <= min(sg->length, ((unsigned int)(PAGE_SIZE)) - offset))) {
-		void *data;
+	sg = req->src;
+	if (nbytes > sg->length)
+		return crypto_shash_init(desc) ?:
+		       shash_ahash_finup(req, desc);
 
-		data = kmap_local_page(sg_page(sg));
-		err = crypto_shash_digest(desc, data + offset, nbytes,
-					  req->result);
-		kunmap_local(data);
-	} else
-		err = crypto_shash_init(desc) ?:
-		      shash_ahash_finup(req, desc);
+	page = sg_page(sg);
+	offset = sg->offset;
+	data = lowmem_page_address(page) + offset;
+	if (!IS_ENABLED(CONFIG_HIGHMEM))
+		return crypto_shash_digest(desc, data, nbytes, req->result);
 
+	page += offset >> PAGE_SHIFT;
+	offset = offset_in_page(offset);
+
+	if (nbytes > (unsigned int)PAGE_SIZE - offset)
+		return crypto_shash_init(desc) ?:
+		       shash_ahash_finup(req, desc);
+
+	data = kmap_local_page(page);
+	err = crypto_shash_digest(desc, data + offset, nbytes,
+				  req->result);
+	kunmap_local(data);
 	return err;
 }
 EXPORT_SYMBOL_GPL(shash_ahash_digest);