oblivious_http: Release ownership of ObliviousHttpRequest::Context to callers and remove shared_ptr

PiperOrigin-RevId: 481964390
diff --git a/quiche/oblivious_http/buffers/oblivious_http_integration_test.cc b/quiche/oblivious_http/buffers/oblivious_http_integration_test.cc
index 7ed02c1..ac096b7 100644
--- a/quiche/oblivious_http/buffers/oblivious_http_integration_test.cc
+++ b/quiche/oblivious_http/buffers/oblivious_http_integration_test.cc
@@ -84,14 +84,16 @@
   EXPECT_EQ(server_req_decap->GetPlaintextData(), test.request_plaintext);
 
   // Round-trip response flow.
+  auto server_request_context =
+      std::move(server_req_decap.value()).ReleaseContext();
   auto server_resp_encap = ObliviousHttpResponse::CreateServerObliviousResponse(
-      test.response_plaintext,
-      *(server_req_decap->oblivious_http_request_context()));
+      test.response_plaintext, server_request_context);
   EXPECT_TRUE(server_resp_encap.ok());
   ASSERT_FALSE(server_resp_encap->EncapsulateAndSerialize().empty());
+  auto client_request_context =
+      std::move(client_req_encap.value()).ReleaseContext();
   auto client_resp_decap = ObliviousHttpResponse::CreateClientObliviousResponse(
-      server_resp_encap->EncapsulateAndSerialize(),
-      *(client_req_encap->oblivious_http_request_context()));
+      server_resp_encap->EncapsulateAndSerialize(), client_request_context);
   EXPECT_TRUE(client_resp_decap.ok());
   EXPECT_EQ(client_resp_decap->GetPlaintextData(), test.response_plaintext);
 }
diff --git a/quiche/oblivious_http/buffers/oblivious_http_request.cc b/quiche/oblivious_http/buffers/oblivious_http_request.cc
index ca26e4e..4ccd0fa 100644
--- a/quiche/oblivious_http/buffers/oblivious_http_request.cc
+++ b/quiche/oblivious_http/buffers/oblivious_http_request.cc
@@ -7,11 +7,14 @@
 #include <string>
 #include <utility>
 
+#include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
 #include "openssl/hpke.h"
+#include "quiche/common/platform/api/quiche_bug_tracker.h"
 #include "quiche/common/platform/api/quiche_logging.h"
 #include "quiche/common/quiche_crypto_logging.h"
 
@@ -27,8 +30,8 @@
     bssl::UniquePtr<EVP_HPKE_CTX> hpke_context, std::string encapsulated_key,
     const ObliviousHttpHeaderKeyConfig& ohttp_key_config,
     std::string req_ciphertext, std::string req_plaintext)
-    : oblivious_http_request_context_(
-          new Context(std::move(hpke_context), std::move(encapsulated_key))),
+    : oblivious_http_request_context_(absl::make_optional(
+          Context(std::move(hpke_context), std::move(encapsulated_key)))),
       key_config_(ohttp_key_config),
       request_ciphertext_(std::move(req_ciphertext)),
       request_plaintext_(std::move(req_plaintext)) {}
@@ -180,6 +183,11 @@
 // Builds request=[hdr, enc, ct].
 // https://www.ietf.org/archive/id/draft-ietf-ohai-ohttp-03.html#section-4.1-4.5
 std::string ObliviousHttpRequest::EncapsulateAndSerialize() const {
+  if (!oblivious_http_request_context_.has_value()) {
+    QUICHE_BUG(ohttp_encapsulate_after_context_extract)
+        << "EncapsulateAndSerialize cannot be called after ReleaseContext()";
+    return "";
+  }
   return absl::StrCat(key_config_.SerializeOhttpPayloadHeader(),
                       oblivious_http_request_context_->encapsulated_key_,
                       request_ciphertext_);
diff --git a/quiche/oblivious_http/buffers/oblivious_http_request.h b/quiche/oblivious_http/buffers/oblivious_http_request.h
index 0807fc8..75f4780 100644
--- a/quiche/oblivious_http/buffers/oblivious_http_request.h
+++ b/quiche/oblivious_http/buffers/oblivious_http_request.h
@@ -6,6 +6,8 @@
 
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "openssl/hpke.h"
 #include "quiche/oblivious_http/common/oblivious_http_header_key_config.h"
 
 namespace quiche {
@@ -23,6 +25,10 @@
    public:
     ~Context() = default;
 
+    // Movable
+    Context(Context&& other) = default;
+    Context& operator=(Context&& other) = default;
+
    private:
     explicit Context(bssl::UniquePtr<EVP_HPKE_CTX> hpke_context,
                      std::string encapsulated_key);
@@ -41,7 +47,7 @@
     friend class ObliviousHttpResponse_TestEncapsulateWithQuicheRandom_Test;
 
     bssl::UniquePtr<EVP_HPKE_CTX> hpke_context_;
-    const std::string encapsulated_key_;
+    std::string encapsulated_key_;
   };
   // Parse the OHTTP request from the given `encrypted_data`.
   // On success, returns obj that callers will use to `GetPlaintextData`.
@@ -64,6 +70,8 @@
   ~ObliviousHttpRequest() = default;
 
   // Returns serialized OHTTP request bytestring.
+  // @note: This method MUST NOT be called after `ReleaseContext()` has been
+  // called.
   std::string EncapsulateAndSerialize() const;
 
   // Generic Usecase : server-side calls this method after Decapsulation using
@@ -73,8 +81,14 @@
   // Oblivious HTTP request context is created after successful creation of
   // `this` object, and subsequently passed into the `ObliviousHttpResponse` for
   // followup response handling.
-  std::shared_ptr<Context> oblivious_http_request_context() const {
-    return oblivious_http_request_context_;
+  // @returns: This rvalue reference qualified member function transfers the
+  // ownership of `Context` to the caller, and further invokes
+  // ClangTidy:misc-use-after-move warning if callers try to extract `Context`
+  // twice after the fact that the ownership has already been transferred.
+  // @note: Callers shouldn't extract the `Context` until you're done with this
+  // Request and its data.
+  Context ReleaseContext() && {
+    return std::move(oblivious_http_request_context_.value());
   }
 
  private:
@@ -94,7 +108,8 @@
       const ObliviousHttpHeaderKeyConfig& ohttp_key_config,
       absl::string_view seed);
 
-  std::shared_ptr<Context> oblivious_http_request_context_;
+  // This field will be empty after calling `ReleaseContext()`.
+  absl::optional<Context> oblivious_http_request_context_;
   ObliviousHttpHeaderKeyConfig key_config_;
   std::string request_ciphertext_;
   std::string request_plaintext_;
diff --git a/quiche/oblivious_http/buffers/oblivious_http_request_test.cc b/quiche/oblivious_http/buffers/oblivious_http_request_test.cc
index 91ca8e9..bfe7252 100644
--- a/quiche/oblivious_http/buffers/oblivious_http_request_test.cc
+++ b/quiche/oblivious_http/buffers/oblivious_http_request_test.cc
@@ -109,7 +109,8 @@
   // https://www.ietf.org/archive/id/draft-ietf-ohai-ohttp-03.html#appendix-A-10
   constexpr absl::string_view kExpectedEphemeralPublicKey =
       "4b28f881333e7c164ffc499ad9796f877f4e1051ee6d31bad19dec96c208b472";
-  EXPECT_EQ(instance->oblivious_http_request_context()->encapsulated_key_,
+  auto oblivious_request_context = std::move(instance.value()).ReleaseContext();
+  EXPECT_EQ(oblivious_request_context.encapsulated_key_,
             absl::HexStringToBytes(kExpectedEphemeralPublicKey));
 
   // Binary HTTP message.
@@ -145,8 +146,8 @@
   uint16_t aead_id;
   EXPECT_TRUE(reader.ReadUInt16(&aead_id));
   EXPECT_EQ(aead_id, test_aead_id);
-  auto client_encapsulated_key =
-      instance->oblivious_http_request_context()->encapsulated_key_;
+  auto client_request_context = std::move(instance.value()).ReleaseContext();
+  auto client_encapsulated_key = client_request_context.encapsulated_key_;
   EXPECT_EQ(client_encapsulated_key.size(), X25519_PUBLIC_VALUE_LEN);
   auto enc_key_plus_ciphertext = payload_bytes.substr(kHeaderLength);
   auto packed_encapsulated_key =
@@ -163,18 +164,18 @@
   auto encapsulated = CreateClientObliviousRequestWithSeedForTesting(
       "test", GetHpkePublicKey(), ohttp_key_config, GetSeed());
   ASSERT_TRUE(encapsulated.ok());
-  EXPECT_EQ(encapsulated->oblivious_http_request_context()->encapsulated_key_,
+  auto encapsulated_request = encapsulated->EncapsulateAndSerialize();
+  auto ohttp_request_context = std::move(encapsulated.value()).ReleaseContext();
+  EXPECT_EQ(ohttp_request_context.encapsulated_key_,
             GetSeededEncapsulatedKey());
   absl::string_view expected_encrypted_request =
       "9f37cfed07d0111ecd2c34f794671759bcbd922a";
-  EXPECT_NE(encapsulated->oblivious_http_request_context()->hpke_context_,
-            nullptr);
-  size_t encapsulated_key_len = EVP_HPKE_KEM_enc_len(EVP_HPKE_CTX_kem(
-      encapsulated->oblivious_http_request_context()->hpke_context_.get()));
+  EXPECT_NE(ohttp_request_context.hpke_context_, nullptr);
+  size_t encapsulated_key_len = EVP_HPKE_KEM_enc_len(
+      EVP_HPKE_CTX_kem(ohttp_request_context.hpke_context_.get()));
   int encrypted_payload_offset = kHeaderLength + encapsulated_key_len;
-  EXPECT_EQ(
-      encapsulated->EncapsulateAndSerialize().substr(encrypted_payload_offset),
-      absl::HexStringToBytes(expected_encrypted_request));
+  EXPECT_EQ(encapsulated_request.substr(encrypted_payload_offset),
+            absl::HexStringToBytes(expected_encrypted_request));
 }
 
 TEST(ObliviousHttpRequest,
diff --git a/quiche/oblivious_http/buffers/oblivious_http_response_test.cc b/quiche/oblivious_http/buffers/oblivious_http_response_test.cc
index 775ba56..1e8eb34 100644
--- a/quiche/oblivious_http/buffers/oblivious_http_response_test.cc
+++ b/quiche/oblivious_http/buffers/oblivious_http_response_test.cc
@@ -96,9 +96,10 @@
   return bssl_hpke_key;
 }
 
-const ObliviousHttpRequest SetUpObliviousHttpContext(
-    uint8_t key_id, uint16_t kem_id, uint16_t kdf_id, uint16_t aead_id,
-    absl::string_view plaintext) {
+ObliviousHttpRequest SetUpObliviousHttpContext(uint8_t key_id, uint16_t kem_id,
+                                               uint16_t kdf_id,
+                                               uint16_t aead_id,
+                                               absl::string_view plaintext) {
   auto ohttp_key_config = GetOhttpKeyConfig(key_id, kem_id, kdf_id, aead_id);
   auto client_request_encapsulate =
       CreateClientObliviousRequestWithSeedForTesting(
@@ -157,12 +158,13 @@
   absl::string_view encrypted_response =
       "39d5b03c02c97e216df444e4681007105974d4df1585aae05e7b53f3ccdb55d51f711d48"
       "eeefbc1a555d6d928e35df33fd23c23846fa7b083e30692f7b";
+  auto oblivious_context =
+      SetUpObliviousHttpContext(4, EVP_HPKE_DHKEM_X25519_HKDF_SHA256,
+                                EVP_HPKE_HKDF_SHA256, EVP_HPKE_AES_256_GCM,
+                                "test")
+          .ReleaseContext();
   auto decapsulated = ObliviousHttpResponse::CreateClientObliviousResponse(
-      absl::HexStringToBytes(encrypted_response),
-      *(SetUpObliviousHttpContext(4, EVP_HPKE_DHKEM_X25519_HKDF_SHA256,
-                                  EVP_HPKE_HKDF_SHA256, EVP_HPKE_AES_256_GCM,
-                                  "test")
-            .oblivious_http_request_context()));
+      absl::HexStringToBytes(encrypted_response), oblivious_context);
   EXPECT_TRUE(decapsulated.ok());
   auto decrypted = decapsulated->GetPlaintextData();
   EXPECT_EQ(decrypted, "test response");
@@ -192,29 +194,25 @@
   auto server_seeded_request = SetUpObliviousHttpContext(
       6, EVP_HPKE_DHKEM_X25519_HKDF_SHA256, EVP_HPKE_HKDF_SHA256,
       EVP_HPKE_AES_256_GCM, "test");
+  auto server_request_context =
+      std::move(server_seeded_request).ReleaseContext();
   auto server_response_encapsulate =
       ObliviousHttpResponse::CreateServerObliviousResponse(
-          "test response",
-          *(server_seeded_request.oblivious_http_request_context()), &random);
+          "test response", server_request_context, &random);
   EXPECT_TRUE(server_response_encapsulate.ok());
   std::string response_nonce =
       server_response_encapsulate->EncapsulateAndSerialize().substr(
-          0, GetResponseNonceLength(
-                 *(server_seeded_request.oblivious_http_request_context()
-                       ->hpke_context_)));
-  EXPECT_EQ(
-      response_nonce,
-      std::string(GetResponseNonceLength(
-                      *(server_seeded_request.oblivious_http_request_context()
-                            ->hpke_context_)),
-                  'z'));
+          0, GetResponseNonceLength(*(server_request_context.hpke_context_)));
+  EXPECT_EQ(response_nonce,
+            std::string(
+                GetResponseNonceLength(*(server_request_context.hpke_context_)),
+                'z'));
   absl::string_view expected_encrypted_response =
       "2a3271ac4e6a501f51d0264d3dd7d0bc8a06973b58e89c26d6dac06144";
-  EXPECT_EQ(server_response_encapsulate->EncapsulateAndSerialize().substr(
-                GetResponseNonceLength(
-                    *(server_seeded_request.oblivious_http_request_context()
-                          ->hpke_context_))),
-            absl::HexStringToBytes(expected_encrypted_response));
+  EXPECT_EQ(
+      server_response_encapsulate->EncapsulateAndSerialize().substr(
+          GetResponseNonceLength(*(server_request_context.hpke_context_))),
+      absl::HexStringToBytes(expected_encrypted_response));
 }
 
 }  // namespace quiche
diff --git a/quiche/oblivious_http/oblivious_http_client_test.cc b/quiche/oblivious_http/oblivious_http_client_test.cc
index 663f4d5..9abfbd3 100644
--- a/quiche/oblivious_http/oblivious_http_client_test.cc
+++ b/quiche/oblivious_http/oblivious_http_client_test.cc
@@ -137,18 +137,21 @@
           *(ConstructHpkeKey(GetHpkePrivateKey(), ohttp_key_config)),
           ohttp_key_config);
   ASSERT_TRUE(decapsulate_req_on_gateway.ok());
+  auto gateway_request_context =
+      std::move(decapsulate_req_on_gateway.value()).ReleaseContext();
   auto encapsulate_resp_on_gateway =
       ObliviousHttpResponse::CreateServerObliviousResponse(
-          "test response",
-          *(decapsulate_req_on_gateway->oblivious_http_request_context()));
+          "test response", gateway_request_context);
   ASSERT_TRUE(encapsulate_resp_on_gateway.ok());
 
   auto client =
       ObliviousHttpClient::Create(GetHpkePublicKey(), ohttp_key_config);
   ASSERT_TRUE(client.ok());
+  auto client_request_context =
+      std::move(encapsulate_req_on_client.value()).ReleaseContext();
   auto decapsulate_resp_on_client = client->DecryptObliviousHttpResponse(
       absl::string_view(encapsulate_resp_on_gateway->EncapsulateAndSerialize()),
-      *(encapsulate_req_on_client->oblivious_http_request_context()));
+      client_request_context);
   ASSERT_TRUE(decapsulate_resp_on_client.ok());
   EXPECT_EQ(decapsulate_resp_on_client->GetPlaintextData(), "test response");
 }
@@ -168,10 +171,11 @@
           *(ConstructHpkeKey(GetHpkePrivateKey(), ohttp_key_config)),
           ohttp_key_config);
   ASSERT_TRUE(decapsulate_req_on_gateway.ok());
+  auto gateway_request_context =
+      std::move(decapsulate_req_on_gateway.value()).ReleaseContext();
   auto encapsulate_resp_on_gateway =
       ObliviousHttpResponse::CreateServerObliviousResponse(
-          "test response",
-          *(decapsulate_req_on_gateway->oblivious_http_request_context()));
+          "test response", gateway_request_context);
   ASSERT_TRUE(encapsulate_resp_on_gateway.ok());
 
   auto client =
@@ -179,7 +183,7 @@
   ASSERT_TRUE(client.ok());
   auto decapsulate_resp_on_client = client->DecryptObliviousHttpResponse(
       absl::string_view(encapsulate_resp_on_gateway->EncapsulateAndSerialize()),
-      *(decapsulate_req_on_gateway->oblivious_http_request_context()));
+      gateway_request_context);
   ASSERT_TRUE(decapsulate_resp_on_client.ok());
   EXPECT_EQ(decapsulate_resp_on_client->GetPlaintextData(), "test response");
 }
@@ -208,17 +212,19 @@
               *(ConstructHpkeKey(GetHpkePrivateKey(), ohttp_key_config_)),
               ohttp_key_config_);
       ASSERT_TRUE(decapsulate_req_on_gateway.ok());
+      auto gateway_request_context =
+          std::move(decapsulate_req_on_gateway.value()).ReleaseContext();
       auto encapsulate_resp_on_gateway =
           ObliviousHttpResponse::CreateServerObliviousResponse(
-              "test response",
-              *(decapsulate_req_on_gateway->oblivious_http_request_context()));
+              "test response", gateway_request_context);
       ASSERT_TRUE(encapsulate_resp_on_gateway.ok());
       ASSERT_FALSE(
           encapsulate_resp_on_gateway->EncapsulateAndSerialize().empty());
-
+      auto client_request_context =
+          std::move(encrypted_request.value()).ReleaseContext();
       auto decrypted_response = client_.DecryptObliviousHttpResponse(
           encapsulate_resp_on_gateway->EncapsulateAndSerialize(),
-          *(encrypted_request->oblivious_http_request_context()));
+          client_request_context);
       ASSERT_TRUE(decrypted_response.ok());
       ASSERT_FALSE(decrypted_response->GetPlaintextData().empty());
     }
diff --git a/quiche/oblivious_http/oblivious_http_gateway_test.cc b/quiche/oblivious_http/oblivious_http_gateway_test.cc
index 5482c1c..af03b21 100644
--- a/quiche/oblivious_http/oblivious_http_gateway_test.cc
+++ b/quiche/oblivious_http/oblivious_http_gateway_test.cc
@@ -125,9 +125,10 @@
   auto decapsulated_req_on_server = instance->DecryptObliviousHttpRequest(
       encapsualte_request_on_client->EncapsulateAndSerialize());
   ASSERT_TRUE(decapsulated_req_on_server.ok());
+  auto server_request_context =
+      std::move(decapsulated_req_on_server.value()).ReleaseContext();
   auto encapsulate_resp_on_gateway = instance->CreateObliviousHttpResponse(
-      "some response",
-      *(decapsulated_req_on_server->oblivious_http_request_context()));
+      "some response", server_request_context);
   ASSERT_TRUE(encapsulate_resp_on_gateway.ok());
   ASSERT_FALSE(encapsulate_resp_on_gateway->EncapsulateAndSerialize().empty());
 }
@@ -155,18 +156,15 @@
 
   // Extract contexts and handle the response for each corresponding request.
   auto oblivious_request_context_1 =
-      decrypted_request_1->oblivious_http_request_context();
-  EXPECT_NE(oblivious_request_context_1, nullptr);
+      std::move(decrypted_request_1.value()).ReleaseContext();
   auto encrypted_response_1 = instance->CreateObliviousHttpResponse(
-      "test response 1", *(oblivious_request_context_1));
+      "test response 1", oblivious_request_context_1);
   ASSERT_TRUE(encrypted_response_1.ok());
   ASSERT_FALSE(encrypted_response_1->EncapsulateAndSerialize().empty());
-
   auto oblivious_request_context_2 =
-      decrypted_request_1->oblivious_http_request_context();
-  EXPECT_NE(oblivious_request_context_2, nullptr);
+      std::move(decrypted_request_2.value()).ReleaseContext();
   auto encrypted_response_2 = instance->CreateObliviousHttpResponse(
-      "test response 2", *(oblivious_request_context_2));
+      "test response 2", oblivious_request_context_2);
   ASSERT_TRUE(encrypted_response_2.ok());
   ASSERT_FALSE(encrypted_response_2->EncapsulateAndSerialize().empty());
 }
@@ -187,9 +185,10 @@
           gateway_receiver_.DecryptObliviousHttpRequest(request_payload_);
       ASSERT_TRUE(decrypted_request.ok());
       ASSERT_FALSE(decrypted_request->GetPlaintextData().empty());
+      auto gateway_request_context =
+          std::move(decrypted_request.value()).ReleaseContext();
       auto encrypted_response = gateway_receiver_.CreateObliviousHttpResponse(
-          response_payload_,
-          *(decrypted_request->oblivious_http_request_context()));
+          response_payload_, gateway_request_context);
       ASSERT_TRUE(encrypted_response.ok());
       ASSERT_FALSE(encrypted_response->EncapsulateAndSerialize().empty());
     }