[PATCH v3] riscv: fix strnlen() overflow in Zbb implementation

[gao.rui at zte.com.cn]

riscv: fix strnlen() overflow in Zbb implementation

The RISC-V Zbb optimized strnlen() implementation can return incorrect
results when very large count values are supplied.

The previous implementation calculates an end address based on the
input pointer and count. When count is close to SIZE_MAX, the address
calculation may overflow, resulting in incorrect termination checks and
wrong return values.

This issue was observed while running device-mapper tests:

dmsetup create testname9 --table "0 8 zero"
cat /sys/block/dm-*/dm/name
dmsetup remove testname9

Rework the Zbb implementation to use a word counter instead of an end
address. This removes the dependency on end-address calculations,
avoids overflow entirely, and simplifies the termination condition of
the word scanning loop. The final length is calculated using the saved
original pointer.

Performance was evaluated with string_bench_strnlen:

New Implementation Previous Implementation

len=0 : 71 ns/call 70 ns/call
len=1 : 82 ns/call 81 ns/call
len=7 : 83 ns/call 81 ns/call
len=8 : 83 ns/call 81 ns/call
len=16 : 94 ns/call 90 ns/call
len=31 : 117 ns/call 113 ns/call
len=64 : 175 ns/call 167 ns/call
len=127 : 259 ns/call 258 ns/call
len=512 : 802 ns/call 833 ns/call
len=1024 : 1512 ns/call 1536 ns/call
len=3173 : 4581 ns/call 4640 ns/call
len=4096 : 6256 ns/call 6068 ns/call

Results show comparable performance to the previous implementation
while fixing the overflow issue.

Fixes: 5ba15d419fab ("riscv: lib: add strnlen() implementation")
Suggested-by: DavidLaight <david.laight.linux at gmail.com>
Signed-off-by: Gao Rui <gao.rui at zte.com.cn>

---
v3:
- Keep the original register allocation (t0-t4) and numeric labels.
- Keep the original comments.
- Calculate word count as (remaining - 1) / SZREG + 1, removing the
  branch from the calculation.
v2:
- Rework the Zbb implementation to eliminate end-address overflow
  instead of falling back to the generic path.
- Use a decrementing word counter for word scanning.
- Run KUnit string tests successfully.
- Add string_bench_strnlen benchmark results.
---
 arch/riscv/lib/strnlen.S | 54 +++++++++++++++++--------------
 1 file changed, 29 insertions(+), 25 deletions(-)

diff --git a/arch/riscv/lib/strnlen.S b/arch/riscv/lib/strnlen.S
index a8911605c248..000000000000 100644
--- a/arch/riscv/lib/strnlen.S
+++ b/arch/riscv/lib/strnlen.S
@@ -73,6 +73,9 @@
 	/* If maxlen is 0, return 0. */
 	beqz	a1, 3f

+	/* Save original pointer for final length calculation. */
+	mv	t4, a0
+
 	/* Number of irrelevant bytes in the first word. */
 	andi	t2, a0, SZREG-1

@@ -82,15 +85,6 @@
 	li	t3, SZREG
 	sub	t3, t3, t2
 	slli	t2, t2, 3
-
-	/*
-	 * Aligned boundary.  Use the address of the last valid byte
-	 * (s + count - 1) to avoid loading a word past the count
-	 * boundary in the loop below.  count == 0 is handled above.
-	 */
-	add	t4, a0, a1
-	addi	t4, t4, -1
-	andi	t4, t4, -SZREG

 	/* Get the first word.  */
 	REG_L	t1, 0(t0)
@@ -125,15 +119,26 @@

 	bgtu	t3, a0, 2f

-	/* All remaining bytes are in the first word, no loop needed. */
-	bgeu	t0, t4, 2f
+	/*
+	 * Use a word counter to avoid calculating an end address,
+	 * which might overflow when count is close to SIZE_MAX.
+	 */
+	sub	t2, a1, t3
+	beqz	t2, 3f

-	/* Prepare for the word comparison loop. */
-	addi	t2, t0, SZREG
+	/* Calculate the word count: count = (remaining - 1) / SZREG + 1 */
+	addi	t2, t2, -1
+#if defined(CONFIG_64BIT)
+	srli	t1, t2, 3
+#else
+	srli	t1, t2, 2
+#endif
+	addi	t2, t1, 1
+
 	li	t3, -1

 	/*
-	 * Our critical loop is 4 instructions and processes data in
+	 * Our critical loop is 6 instructions and processes data in
 	 * 4 byte or 8 byte chunks.
 	 */
 	.p2align 3
@@ -141,23 +146,22 @@
 	REG_L	t1, SZREG(t0)
 	addi	t0, t0, SZREG
 	orc.b	t1, t1
-	bgeu	t0, t4, 4f
-	beq	t1, t3, 1b
+	bne	t1, t3, 4f
+	addi	t2, t2, -1
+	bnez	t2, 1b
+
+	/* Scanned all words without finding NUL, return maxlen. */
+	mv	a0, a1
+	ret
+
 4:
 	not	t1, t1
 	CZ	t1, t1
 	srli	t1, t1, 3

-	/* Get number of processed bytes. */
-	sub	t2, t0, t2
-
-	/* Add number of characters in the first word.  */
-	add	a0, a0, t2
-
-	/* Add number of characters in the last word.  */
+	/* Length = (current - original) + offset in final word. */
+	sub	a0, t0, t4
 	add	a0, a0, t1
-
-	/* Ensure the final result does not exceed maxlen. */
 	minu	a0, a0, a1
 2:
 	ret
--
2.27.0