[llvm] [Support] Add end/error to decode[US]LEB128AndInc (PR #90006)

[Fangrui Song via llvm-commits]

https://github.com/MaskRay created https://github.com/llvm/llvm-project/pull/90006

Follow-up to #85739 to encourage error checking.


>From 98181cf6a33d34dc8e7b5b689667063abc5f5a73 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i at maskray.me>
Date: Wed, 24 Apr 2024 16:57:26 -0700
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
 =?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.5-bogner
---
 llvm/include/llvm/Support/LEB128.h    | 12 ++++++++----
 llvm/unittests/Support/LEB128Test.cpp | 12 ++++++++++++
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/llvm/include/llvm/Support/LEB128.h b/llvm/include/llvm/Support/LEB128.h
index c4e741549f3ff1..6090bfc3a1aeb8 100644
--- a/llvm/include/llvm/Support/LEB128.h
+++ b/llvm/include/llvm/Support/LEB128.h
@@ -200,16 +200,20 @@ inline int64_t decodeSLEB128(const uint8_t *p, unsigned *n = nullptr,
   return Value;
 }
 
-inline uint64_t decodeULEB128AndInc(const uint8_t *&p) {
+inline uint64_t decodeULEB128AndInc(const uint8_t *&p,
+                                    const uint8_t *end = nullptr,
+                                    const char **error = nullptr) {
   unsigned n;
-  auto ret = decodeULEB128(p, &n);
+  auto ret = decodeULEB128(p, &n, end, error);
   p += n;
   return ret;
 }
 
-inline int64_t decodeSLEB128AndInc(const uint8_t *&p) {
+inline int64_t decodeSLEB128AndInc(const uint8_t *&p,
+                                   const uint8_t *end = nullptr,
+                                   const char **error = nullptr) {
   unsigned n;
-  auto ret = decodeSLEB128(p, &n);
+  auto ret = decodeSLEB128(p, &n, end, error);
   p += n;
   return ret;
 }
diff --git a/llvm/unittests/Support/LEB128Test.cpp b/llvm/unittests/Support/LEB128Test.cpp
index 08b8c5573ce637..c5eea7cfc3731d 100644
--- a/llvm/unittests/Support/LEB128Test.cpp
+++ b/llvm/unittests/Support/LEB128Test.cpp
@@ -155,6 +155,12 @@ TEST(LEB128Test, DecodeInvalidULEB128) {
     EXPECT_NE(Error, nullptr);                                                 \
     EXPECT_EQ(0ul, Actual);                                                    \
     EXPECT_EQ(ERROR_OFFSET, ErrorOffset);                                      \
+    Value = reinterpret_cast<const uint8_t *>(VALUE);                          \
+    Error = nullptr;                                                           \
+    Actual = decodeULEB128AndInc(Value, Value + strlen(VALUE), &Error);        \
+    EXPECT_NE(Error, nullptr);                                                 \
+    EXPECT_EQ(0ul, Actual);                                                    \
+    EXPECT_EQ(ERROR_OFFSET, Value - reinterpret_cast<const uint8_t *>(VALUE)); \
   } while (0)
 
   // Buffer overflow.
@@ -224,6 +230,12 @@ TEST(LEB128Test, DecodeInvalidSLEB128) {
     EXPECT_NE(Error, nullptr);                                                 \
     EXPECT_EQ(0ul, Actual);                                                    \
     EXPECT_EQ(ERROR_OFFSET, ErrorOffset);                                      \
+    Value = reinterpret_cast<const uint8_t *>(VALUE);                          \
+    Error = nullptr;                                                           \
+    Actual = decodeSLEB128AndInc(Value, Value + strlen(VALUE), &Error);        \
+    EXPECT_NE(Error, nullptr);                                                 \
+    EXPECT_EQ(0ul, Actual);                                                    \
+    EXPECT_EQ(ERROR_OFFSET, Value - reinterpret_cast<const uint8_t *>(VALUE)); \
   } while (0)
 
   // Buffer overflow.