Simplify BigInt::get_substring a bit

And forbid 0 length substrings, which did not work correctly anyway.

2 files changed, 40 insertions, 10 deletions

diff --git a/src/lib/math/bigint/bigint.cpp b/src/lib/math/bigint/bigint.cpp
index b996f0fa2..36f1d2883 100644
--- a/src/lib/math/bigint/bigint.cpp
+++ b/src/lib/math/bigint/bigint.cpp
@@ -163,18 +163,19 @@ void BigInt::encode_words(word out[], size_t size) const
 */
 uint32_t BigInt::get_substring(size_t offset, size_t length) const
    {
-   if(length > 32)
-      throw Invalid_Argument("BigInt::get_substring: Substring size " + std::to_string(length) + " too big");
+   if(length == 0 || length > 32)
+      throw Invalid_Argument("BigInt::get_substring invalid substring length");
 
-   uint64_t piece = 0;
-   for(size_t i = 0; i != 8; ++i)
-      {
-      const uint8_t part = byte_at((offset / 8) + (7-i));
-      piece = (piece << 8) | part;
-      }
-
-   const uint64_t mask = (static_cast<uint64_t>(1) << length) - 1;
+   const size_t byte_offset = offset / 8;
    const size_t shift = (offset % 8);
+   const uint32_t mask = 0xFFFFFFFF >> (32 - length);
+
+   const uint8_t b0 = byte_at(byte_offset);
+   const uint8_t b1 = byte_at(byte_offset + 1);
+   const uint8_t b2 = byte_at(byte_offset + 2);
+   const uint8_t b3 = byte_at(byte_offset + 3);
+   const uint8_t b4 = byte_at(byte_offset + 4);
+   const uint64_t piece = make_uint64(0, 0, 0, b4, b3, b2, b1, b0);
 
    return static_cast<uint32_t>((piece >> shift) & mask);
    }
diff --git a/src/tests/test_bigint.cpp b/src/tests/test_bigint.cpp
index ad0af03f4..33c617239 100644
--- a/src/tests/test_bigint.cpp
+++ b/src/tests/test_bigint.cpp
@@ -36,6 +36,7 @@ class BigInt_Unit_Tests final : public Test
          results.push_back(test_random_prime());
          results.push_back(test_encode());
          results.push_back(test_bigint_io());
+         results.push_back(test_get_substring());
 
          return results;
          }
@@ -206,6 +207,34 @@ class BigInt_Unit_Tests final : public Test
          return result;
          }
 
+      Test::Result test_get_substring()
+         {
+         const size_t trials = 1000;
+
+         Test::Result result("BigInt get_substring");
+
+         const Botan::BigInt r(Test::rng(), 250);
+
+         for(size_t s = 1; s <= 32; ++s)
+            {
+            for(size_t trial = 0; trial != trials; ++trial)
+               {
+               const size_t offset = Test::rng().next_byte();
+
+               const uint32_t val = r.get_substring(offset, s);
+
+               Botan::BigInt t = r >> offset;
+               t.mask_bits(s);
+
+               const uint32_t cmp = t.to_u32bit();
+
+               result.test_eq("Same value", size_t(val), size_t(cmp));
+               }
+            }
+
+         return result;
+         }
+
       Test::Result test_bigint_io()
          {
          Test::Result result("BigInt IO operators");