From c2d6fceaea636a02b87578b294c85f035eafdf17 Mon Sep 17 00:00:00 2001
From: Thomas Holenstein <tholenst@google.com>
Date: Thu, 23 Jan 2025 08:25:48 -0800
Subject: [PATCH] Add SecretBuffer::substr.

PiperOrigin-RevId: 718877508
Change-Id: I1482c12f89d1cfc1b28651251cef6ce94ae45002
---
 tink/internal/BUILD.bazel           |  1 +
 tink/internal/CMakeLists.txt        |  1 +
 tink/internal/secret_buffer.h       | 26 ++++++++++++++++++++++++++
 tink/internal/secret_buffer_test.cc | 26 ++++++++++++++++++++++++++
 4 files changed, 54 insertions(+)

diff --git a/tink/internal/BUILD.bazel b/tink/internal/BUILD.bazel
index 5411549c..2e4c6c76 100644
--- a/tink/internal/BUILD.bazel
+++ b/tink/internal/BUILD.bazel
@@ -1976,6 +1976,7 @@ cc_test(
     deps = [
         ":secret_buffer",
         "//tink/util:test_util",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
     ],
 )
diff --git a/tink/internal/CMakeLists.txt b/tink/internal/CMakeLists.txt
index 2793ce6f..c345d982 100644
--- a/tink/internal/CMakeLists.txt
+++ b/tink/internal/CMakeLists.txt
@@ -1904,6 +1904,7 @@ tink_cc_test(
   DEPS
     tink::internal::secret_buffer
     gmock
+    absl::string_view
     tink::util::test_util
 )
 
diff --git a/tink/internal/secret_buffer.h b/tink/internal/secret_buffer.h
index bbd478ca..41797259 100644
--- a/tink/internal/secret_buffer.h
+++ b/tink/internal/secret_buffer.h
@@ -17,6 +17,7 @@
 #ifndef TINK_INTERNAL_SECRET_BUFFER_H_
 #define TINK_INTERNAL_SECRET_BUFFER_H_
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
@@ -136,6 +137,31 @@ class SecretBuffer {
     swap(capacity_, other.capacity_);
   }
 
+  SecretBuffer substr(
+      size_t pos, size_t count = std::numeric_limits<size_t>::max()) const& {
+    CHECK(pos <= size());
+    count = std::min(count, size() - pos);
+    SecretBuffer result;
+    result.reserve(count);
+    SafeMemCopy(result.data_, data_ + pos, count);
+    result.size_ = count;
+    return result;
+  }
+
+  SecretBuffer substr(size_t pos,
+                      size_t count = std::numeric_limits<size_t>::max()) && {
+    CHECK(pos <= size());
+    count = std::min(count, size() - pos);
+    SecretBuffer result;
+    result.swap(*this);
+    if (pos != 0) {
+      crypto::tink::internal::SafeMemMove(result.data(), result.data() + pos,
+                                          count);
+    }
+    result.resize(count);
+    return result;
+  }
+
   bool operator==(const SecretBuffer& other) const {
     if (size() != other.size()) {
       return false;
diff --git a/tink/internal/secret_buffer_test.cc b/tink/internal/secret_buffer_test.cc
index e40d628a..cabab4a2 100644
--- a/tink/internal/secret_buffer_test.cc
+++ b/tink/internal/secret_buffer_test.cc
@@ -21,6 +21,7 @@
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "absl/strings/string_view.h"
 #include "tink/util/test_util.h"
 
 namespace crypto {
@@ -157,6 +158,31 @@ TEST(SecretBufferTest, SwapWithEmpty) {
   EXPECT_THAT(buffer2.AsStringView(), Eq(""));
 }
 
+TEST(SecretBufferTest, SubstrConstRef) {
+  constexpr absl::string_view kData = "Some arbitrary data";
+  SecretBuffer buffer(kData);
+  for (int start = 0; start <= kData.size(); ++start) {
+    for (int num = 0; num <= kData.size(); ++num) {
+      EXPECT_THAT(buffer.substr(start, num),
+                  Eq(SecretBuffer(kData.substr(start, num))))
+          << "substr(" << start << ", " << num << ")";
+    }
+  }
+}
+
+TEST(SecretBufferTest, SubstrRvalueRef) {
+  constexpr absl::string_view kData = "Some arbitrary data";
+  SecretBuffer buffer(kData);
+  for (int start = 0; start <= kData.size(); ++start) {
+    for (int num = 0; num <= kData.size(); ++num) {
+      SecretBuffer tmp_buffer = buffer;
+      EXPECT_THAT(std::move(tmp_buffer).substr(start, num),
+                  Eq(SecretBuffer(kData.substr(start, num))))
+          << "substr(" << start << ", " << num << ")";
+    }
+  }
+}
+
 }  // namespace
 }  // namespace internal
 }  // namespace tink