Use QUICHE_ASSIGN_OR_RETURN in oblivious_http

PiperOrigin-RevId: 842776051
diff --git a/quiche/oblivious_http/buffers/oblivious_http_request.cc b/quiche/oblivious_http/buffers/oblivious_http_request.cc
index f6d3803..88fff38 100644
--- a/quiche/oblivious_http/buffers/oblivious_http_request.cc
+++ b/quiche/oblivious_http/buffers/oblivious_http_request.cc
@@ -3,22 +3,19 @@
 #include <stddef.h>
 #include <stdint.h>
 
-#include <memory>
 #include <optional>
 #include <string>
 #include <utility>
 
-#include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/strings/escaping.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.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"
 #include "quiche/common/quiche_data_reader.h"
+#include "quiche/common/quiche_status_utils.h"
 #include "quiche/oblivious_http/common/oblivious_http_definitions.h"
 #include "quiche/oblivious_http/common/oblivious_http_header_key_config.h"
 
@@ -47,19 +44,17 @@
     const ObliviousHttpHeaderKeyConfig& ohttp_key_config,
     absl::string_view request_label) {
   QuicheDataReader reader(encrypted_data);
-  absl::StatusOr<ObliviousHttpRequest::Context> gateway_context =
+  QUICHE_ASSIGN_OR_RETURN(
+      ObliviousHttpRequest::Context gateway_context,
       ObliviousHttpRequest::DecodeEncapsulatedRequestHeader(
-          reader, gateway_key, ohttp_key_config, request_label);
-  if (!gateway_context.ok()) {
-    return gateway_context.status();
-  }
+          reader, gateway_key, ohttp_key_config, request_label));
 
   absl::string_view ciphertext_received = reader.ReadRemainingPayload();
   // Decrypt the message.
   std::string decrypted(ciphertext_received.size(), '\0');
   size_t decrypted_len;
   if (!EVP_HPKE_CTX_open(
-          gateway_context->hpke_context_.get(),
+          gateway_context.hpke_context_.get(),
           reinterpret_cast<uint8_t*>(decrypted.data()), &decrypted_len,
           decrypted.size(),
           reinterpret_cast<const uint8_t*>(ciphertext_received.data()),
@@ -69,8 +64,8 @@
   }
   decrypted.resize(decrypted_len);
   return ObliviousHttpRequest(
-      std::move(gateway_context->hpke_context_),
-      std::string(gateway_context->encapsulated_key_), ohttp_key_config,
+      std::move(gateway_context.hpke_context_),
+      std::string(gateway_context.encapsulated_key_), ohttp_key_config,
       std::string(ciphertext_received), std::move(decrypted));
 }
 
@@ -101,22 +96,19 @@
   if (plaintext_payload.empty() || hpke_public_key.empty()) {
     return absl::InvalidArgumentError("Invalid input.");
   }
-  absl::StatusOr<Context> context = CreateHpkeSenderContext(
-      hpke_public_key, ohttp_key_config, seed, request_label);
-  if (!context.ok()) {
-    return context.status();
-  }
-  std::string encapsulated_key = context->encapsulated_key_;
+  QUICHE_ASSIGN_OR_RETURN(
+      Context context,
+      CreateHpkeSenderContext(hpke_public_key, ohttp_key_config, seed,
+                              request_label));
+  std::string encapsulated_key = context.encapsulated_key_;
   // EncryptChunk with `is_final_chunk` set to false is the same implementation
   // as encrypting the full request.
-  absl::StatusOr<std::string> ciphertext =
-      EncryptChunk(plaintext_payload, *context, /*is_final_chunk=*/false);
-  if (!ciphertext.ok()) {
-    return ciphertext.status();
-  }
+  QUICHE_ASSIGN_OR_RETURN(
+      std::string ciphertext,
+      EncryptChunk(plaintext_payload, context, /*is_final_chunk=*/false));
   return ObliviousHttpRequest(
-      std::move(context->hpke_context_), std::move(encapsulated_key),
-      ohttp_key_config, std::move(*ciphertext), std::move(plaintext_payload));
+      std::move(context.hpke_context_), std::move(encapsulated_key),
+      ohttp_key_config, std::move(ciphertext), std::move(plaintext_payload));
 }
 
 absl::StatusOr<ObliviousHttpRequest::Context>
@@ -240,10 +232,7 @@
     return SslErrorAsStatus("Failed to initialize Gateway/Server's Context.");
   }
 
-  auto is_hdr_ok = ohttp_key_config.ParseOhttpPayloadHeader(reader);
-  if (!is_hdr_ok.ok()) {
-    return is_hdr_ok;
-  }
+  QUICHE_RETURN_IF_ERROR(ohttp_key_config.ParseOhttpPayloadHeader(reader));
 
   size_t enc_key_len = EVP_HPKE_KEM_enc_len(EVP_HPKE_KEY_kem(&gateway_key));
 
diff --git a/quiche/oblivious_http/buffers/oblivious_http_response.cc b/quiche/oblivious_http/buffers/oblivious_http_response.cc
index 3755db3..cc0e86d 100644
--- a/quiche/oblivious_http/buffers/oblivious_http_response.cc
+++ b/quiche/oblivious_http/buffers/oblivious_http_response.cc
@@ -4,7 +4,6 @@
 #include <stdint.h>
 
 #include <algorithm>
-#include <memory>
 #include <string>
 #include <utility>
 
@@ -13,11 +12,13 @@
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "openssl/aead.h"
+#include "openssl/digest.h"
 #include "openssl/hkdf.h"
 #include "openssl/hpke.h"
 #include "quiche/common/platform/api/quiche_bug_tracker.h"
 #include "quiche/common/quiche_crypto_logging.h"
 #include "quiche/common/quiche_random.h"
+#include "quiche/common/quiche_status_utils.h"
 #include "quiche/oblivious_http/buffers/oblivious_http_request.h"
 #include "quiche/oblivious_http/common/oblivious_http_definitions.h"
 #include "quiche/oblivious_http/common/oblivious_http_header_key_config.h"
@@ -55,16 +56,13 @@
     return absl::InvalidArgumentError("Empty encrypted_data input param.");
   }
 
-  absl::StatusOr<CommonAeadParamsResult> aead_params_st =
-      GetCommonAeadParams(oblivious_http_request_context);
-  if (!aead_params_st.ok()) {
-    return aead_params_st.status();
-  }
+  QUICHE_ASSIGN_OR_RETURN(CommonAeadParamsResult aead_params_st,
+                          GetCommonAeadParams(oblivious_http_request_context));
 
   // secret_len = [max(Nn, Nk)] where Nk and Nn are the length of AEAD
   // key and nonce associated with HPKE context.
   // https://www.rfc-editor.org/rfc/rfc9458.html#section-4.4-2.1.1
-  size_t secret_len = aead_params_st.value().secret_len;
+  size_t secret_len = aead_params_st.secret_len;
   if (encrypted_data.size() < secret_len) {
     return absl::InvalidArgumentError(
         absl::StrCat("Invalid input response. Failed to parse required minimum "
@@ -78,26 +76,22 @@
   absl::string_view encrypted_response =
       absl::string_view(encrypted_data).substr(secret_len);
 
-  absl::StatusOr<AeadContextData> aead_context_data =
-      GetAeadContextData(oblivious_http_request_context, *aead_params_st,
-                         resp_label, response_nonce);
-  if (!aead_context_data.ok()) {
-    return aead_context_data.status();
-  }
+  QUICHE_ASSIGN_OR_RETURN(
+      AeadContextData aead_context_data,
+      GetAeadContextData(oblivious_http_request_context, aead_params_st,
+                         resp_label, response_nonce));
 
   // Decrypt with initialized AEAD context.
   // response, error = Open(aead_key, aead_nonce, "", ct)
   // https://www.rfc-editor.org/rfc/rfc9458.html#section-4.4-6
   // DecryptChunk with `is_final_chunk` as false is the same implementation as
   // decrypting the full encrypted response.
-  absl::StatusOr<std::string> decrypted =
-      DecryptChunk(encrypted_response, *aead_context_data,
-                   aead_context_data->aead_nonce, /*is_final_chunk=*/false);
-  if (!decrypted.ok()) {
-    return decrypted.status();
-  }
+  QUICHE_ASSIGN_OR_RETURN(
+      std::string decrypted,
+      DecryptChunk(encrypted_response, aead_context_data,
+                   aead_context_data.aead_nonce, /*is_final_chunk=*/false));
   ObliviousHttpResponse oblivious_response(std::move(encrypted_data),
-                                           std::move(*decrypted));
+                                           std::move(decrypted));
   return oblivious_response;
 }
 
@@ -113,12 +107,9 @@
   if (plaintext_payload.empty()) {
     return absl::InvalidArgumentError("Empty plaintext_payload input param.");
   }
-  absl::StatusOr<CommonAeadParamsResult> aead_params_st =
-      GetCommonAeadParams(oblivious_http_request_context);
-  if (!aead_params_st.ok()) {
-    return aead_params_st.status();
-  }
-  const size_t nonce_size = aead_params_st->secret_len;
+  QUICHE_ASSIGN_OR_RETURN(CommonAeadParamsResult aead_params_st,
+                          GetCommonAeadParams(oblivious_http_request_context));
+  const size_t nonce_size = aead_params_st.secret_len;
   const size_t max_encrypted_data_size =
       nonce_size + plaintext_payload.size() +
       EVP_AEAD_max_overhead(EVP_HPKE_AEAD_aead(EVP_HPKE_CTX_aead(
@@ -132,24 +123,22 @@
 
   // Steps (1, 3 to 5) + AEAD context SetUp before 6th step is performed in
   // CommonOperations.
-  auto common_ops_st = CommonOperationsToEncapDecap(
-      response_nonce, oblivious_http_request_context, response_label,
-      aead_params_st.value().aead_key_len,
-      aead_params_st.value().aead_nonce_len, aead_params_st.value().secret_len);
-  if (!common_ops_st.ok()) {
-    return common_ops_st.status();
-  }
+  QUICHE_ASSIGN_OR_RETURN(
+      CommonOperationsResult common_ops_st,
+      CommonOperationsToEncapDecap(
+          response_nonce, oblivious_http_request_context, response_label,
+          aead_params_st.aead_key_len, aead_params_st.aead_nonce_len,
+          aead_params_st.secret_len));
 
   // ct = Seal(aead_key, aead_nonce, "", response)
   // https://www.rfc-editor.org/rfc/rfc9458.html#section-4.4-2.2.1
   size_t ciphertext_len;
   if (!EVP_AEAD_CTX_seal(
-          common_ops_st.value().aead_ctx.get(),
+          common_ops_st.aead_ctx.get(),
           reinterpret_cast<uint8_t*>(encrypted_data.data() + nonce_size),
           &ciphertext_len, encrypted_data.size() - nonce_size,
-          reinterpret_cast<const uint8_t*>(
-              common_ops_st.value().aead_nonce.data()),
-          aead_params_st.value().aead_nonce_len,
+          reinterpret_cast<const uint8_t*>(common_ops_st.aead_nonce.data()),
+          aead_params_st.aead_nonce_len,
           reinterpret_cast<const uint8_t*>(plaintext_payload.data()),
           plaintext_payload.size(), nullptr, 0)) {
     return SslErrorAsStatus(
@@ -334,15 +323,14 @@
   // Steps (1, 3 to 5) + AEAD context SetUp before 6th step is performed in
   // CommonOperations.
   // https://www.rfc-editor.org/rfc/rfc9458.html#section-4.4-3
-  auto common_ops_st = CommonOperationsToEncapDecap(
-      response_nonce, oblivious_http_request_context, response_label,
-      aead_params.aead_key_len, aead_params.aead_nonce_len,
-      aead_params.secret_len);
-  if (!common_ops_st.ok()) {
-    return common_ops_st.status();
-  }
-  return AeadContextData{std::move(common_ops_st.value().aead_ctx),
-                         std::move(common_ops_st.value().aead_nonce)};
+  QUICHE_ASSIGN_OR_RETURN(
+      CommonOperationsResult common_ops_st,
+      CommonOperationsToEncapDecap(
+          response_nonce, oblivious_http_request_context, response_label,
+          aead_params.aead_key_len, aead_params.aead_nonce_len,
+          aead_params.secret_len));
+  return AeadContextData{std::move(common_ops_st.aead_ctx),
+                         std::move(common_ops_st.aead_nonce)};
 }
 
 // Encrypts the chunk following
diff --git a/quiche/oblivious_http/buffers/oblivious_http_response.h b/quiche/oblivious_http/buffers/oblivious_http_response.h
index 508f2c0..57ae7be 100644
--- a/quiche/oblivious_http/buffers/oblivious_http_response.h
+++ b/quiche/oblivious_http/buffers/oblivious_http_response.h
@@ -79,7 +79,7 @@
   static absl::StatusOr<ObliviousHttpResponse> CreateServerObliviousResponse(
       std::string plaintext_payload,
       ObliviousHttpRequest::Context& oblivious_http_request_context,
-      absl::string_view resp_label =
+      absl::string_view response_label =
           ObliviousHttpHeaderKeyConfig::kOhttpResponseLabel,
       QuicheRandom* quiche_random = nullptr);
 
diff --git a/quiche/oblivious_http/common/oblivious_http_header_key_config.cc b/quiche/oblivious_http/common/oblivious_http_header_key_config.cc
index 26b3bf9..6393fd7 100644
--- a/quiche/oblivious_http/common/oblivious_http_header_key_config.cc
+++ b/quiche/oblivious_http/common/oblivious_http_header_key_config.cc
@@ -23,6 +23,7 @@
 #include "quiche/common/platform/api/quiche_logging.h"
 #include "quiche/common/quiche_data_reader.h"
 #include "quiche/common/quiche_data_writer.h"
+#include "quiche/common/quiche_status_utils.h"
 
 namespace quiche {
 namespace {
@@ -80,10 +81,7 @@
 ObliviousHttpHeaderKeyConfig::Create(uint8_t key_id, uint16_t kem_id,
                                      uint16_t kdf_id, uint16_t aead_id) {
   ObliviousHttpHeaderKeyConfig instance(key_id, kem_id, kdf_id, aead_id);
-  auto is_config_ok = instance.ValidateKeyConfig();
-  if (!is_config_ok.ok()) {
-    return is_config_ok;
-  }
+  QUICHE_RETURN_IF_ERROR(instance.ValidateKeyConfig());
   return instance;
 }
 
@@ -434,11 +432,7 @@
   }
 
   // Public key length depends on the kem_id.
-  auto maybe_key_length = KeyLength(kem_id);
-  if (!maybe_key_length.ok()) {
-    return maybe_key_length.status();
-  }
-  const int key_length = maybe_key_length.value();
+  QUICHE_ASSIGN_OR_RETURN(uint16_t key_length, KeyLength(kem_id));
   std::string key_str(key_length, '\0');
   if (!reader.ReadBytes(key_str.data(), key_length)) {
     return absl::InvalidArgumentError("Invalid key_config!");
@@ -461,14 +455,12 @@
       return absl::InvalidArgumentError("Invalid key_config!");
     }
 
-    absl::StatusOr<ObliviousHttpHeaderKeyConfig> maybe_cfg =
-        ObliviousHttpHeaderKeyConfig::Create(key_id, kem_id, kdf_id, aead_id);
-    if (!maybe_cfg.ok()) {
-      // TODO(kmg): Add support to ignore key types in the server response that
-      // aren't supported by the client.
-      return maybe_cfg.status();
-    }
-    configs[key_id].emplace_back(std::move(maybe_cfg.value()));
+    // TODO(kmg): Add support to ignore key types in the server response that
+    // aren't supported by the client.
+    QUICHE_ASSIGN_OR_RETURN(
+        ObliviousHttpHeaderKeyConfig cfg,
+        ObliviousHttpHeaderKeyConfig::Create(key_id, kem_id, kdf_id, aead_id));
+    configs[key_id].emplace_back(std::move(cfg));
   }
   return absl::OkStatus();
 }
diff --git a/quiche/oblivious_http/oblivious_http_client.cc b/quiche/oblivious_http/oblivious_http_client.cc
index 07fbe78..3bd9bba 100644
--- a/quiche/oblivious_http/oblivious_http_client.cc
+++ b/quiche/oblivious_http/oblivious_http_client.cc
@@ -17,6 +17,7 @@
 #include "quiche/common/quiche_crypto_logging.h"
 #include "quiche/common/quiche_data_reader.h"
 #include "quiche/common/quiche_data_writer.h"
+#include "quiche/common/quiche_status_utils.h"
 #include "quiche/oblivious_http/buffers/oblivious_http_request.h"
 #include "quiche/oblivious_http/buffers/oblivious_http_response.h"
 #include "quiche/oblivious_http/common/oblivious_http_chunk_handler.h"
@@ -124,22 +125,18 @@
         absl::StrCat("Invalid input received in method parameters. ",
                      is_valid_input.message()));
   }
-  absl::StatusOr<ObliviousHttpRequest::Context> hpke_sender_context =
+  QUICHE_ASSIGN_OR_RETURN(
+      ObliviousHttpRequest::Context hpke_sender_context,
       ObliviousHttpRequest::CreateHpkeSenderContext(
           hpke_public_key, ohttp_key_config, seed,
-          ObliviousHttpHeaderKeyConfig::kChunkedOhttpRequestLabel);
-  if (!hpke_sender_context.ok()) {
-    return hpke_sender_context.status();
-  }
-  absl::StatusOr<ObliviousHttpResponse::CommonAeadParamsResult> aead_params =
-      ObliviousHttpResponse::GetCommonAeadParams(*hpke_sender_context);
-  if (!aead_params.ok()) {
-    return aead_params.status();
-  }
+          ObliviousHttpHeaderKeyConfig::kChunkedOhttpRequestLabel));
+  QUICHE_ASSIGN_OR_RETURN(
+      ObliviousHttpResponse::CommonAeadParamsResult aead_params,
+      ObliviousHttpResponse::GetCommonAeadParams(hpke_sender_context));
 
   return ChunkedObliviousHttpClient(ohttp_key_config,
-                                    *std::move(hpke_sender_context),
-                                    *std::move(aead_params), chunk_handler);
+                                    std::move(hpke_sender_context),
+                                    std::move(aead_params), chunk_handler);
 }
 
 absl::StatusOr<std::string> ChunkedObliviousHttpClient::EncryptRequestChunk(
@@ -157,12 +154,10 @@
 
 absl::StatusOr<std::string> ChunkedObliviousHttpClient::EncryptRequestChunkImpl(
     absl::string_view plaintext_payload, bool is_final_chunk) {
-  absl::StatusOr<std::string> encrypted_data =
+  QUICHE_ASSIGN_OR_RETURN(
+      std::string encrypted_data,
       ObliviousHttpRequest::EncryptChunk(plaintext_payload,
-                                         hpke_sender_context_, is_final_chunk);
-  if (!encrypted_data.ok()) {
-    return encrypted_data.status();
-  }
+                                         hpke_sender_context_, is_final_chunk));
 
   std::string maybe_key_header_data = "";
   if (request_current_section_ == RequestMessageSection::kHeader) {
@@ -173,12 +168,12 @@
   }
 
   uint8_t chunk_var_int_length =
-      QuicheDataWriter::GetVarInt62Len(encrypted_data->size());
+      QuicheDataWriter::GetVarInt62Len(encrypted_data.size());
   if (chunk_var_int_length == 0) {
     return absl::InvalidArgumentError(
         "Encrypted data is too large to be represented as a varint.");
   }
-  uint64_t chunk_var_int = encrypted_data->size();
+  uint64_t chunk_var_int = encrypted_data.size();
   if (is_final_chunk) {
     request_current_section_ = RequestMessageSection::kEnd;
     chunk_var_int_length =
@@ -191,7 +186,7 @@
   }
 
   std::string request_buffer(maybe_key_header_data.size() +
-                                 chunk_var_int_length + encrypted_data->size(),
+                                 chunk_var_int_length + encrypted_data.size(),
                              '\0');
 
   QuicheDataWriter writer(request_buffer.size(), request_buffer.data());
@@ -204,7 +199,7 @@
     return absl::InternalError(
         "Failed to write encrypted chunk length to buffer.");
   }
-  if (!writer.WriteStringPiece(*encrypted_data)) {
+  if (!writer.WriteStringPiece(encrypted_data)) {
     return absl::InternalError(
         "Failed to write encrypted chunk to request buffer.");
   }
@@ -239,23 +234,20 @@
         return absl::OutOfRangeError("Not enough data to read response nonce.");
       }
 
-      absl::StatusOr<ObliviousHttpResponse::AeadContextData> aead_context_data =
+      QUICHE_ASSIGN_OR_RETURN(
+          ObliviousHttpResponse::AeadContextData aead_context_data,
           ObliviousHttpResponse::GetAeadContextData(
               hpke_sender_context_, aead_params_,
               ObliviousHttpHeaderKeyConfig::kChunkedOhttpResponseLabel,
-              response_nonce);
-      if (!aead_context_data.ok()) {
-        return aead_context_data.status();
-      }
-      aead_context_data_.emplace(*std::move(aead_context_data));
+              response_nonce));
+      aead_context_data_.emplace(std::move(aead_context_data));
 
-      absl::StatusOr<ObliviousHttpResponse::ChunkCounter>
-          response_chunk_counter = ObliviousHttpResponse::ChunkCounter::Create(
-              aead_context_data_->aead_nonce);
-      if (!response_chunk_counter.ok()) {
-        return response_chunk_counter.status();
-      }
-      response_chunk_counter_.emplace(*std::move(response_chunk_counter));
+      QUICHE_ASSIGN_OR_RETURN(
+          ObliviousHttpResponse::ChunkCounter response_chunk_counter,
+          ObliviousHttpResponse::ChunkCounter::Create(
+              aead_context_data_->aead_nonce));
+
+      response_chunk_counter_.emplace(std::move(response_chunk_counter));
       UpdateCheckpoint(reader, response_checkpoint);
       response_current_section_ = ResponseMessageSection::kChunk;
     }
@@ -294,20 +286,18 @@
             return absl::OutOfRangeError("Not enough data to read chunk.");
           }
 
-          absl::StatusOr<std::string> decrypted_chunk =
+          QUICHE_ASSIGN_OR_RETURN(
+              std::string decrypted_chunk,
               ObliviousHttpResponse::DecryptChunk(
                   chunk, aead_context_data,
-                  response_chunk_counter.GetChunkNonce(), is_final_chunk);
-          if (!decrypted_chunk.ok()) {
-            return decrypted_chunk.status();
-          }
+                  response_chunk_counter.GetChunkNonce(), is_final_chunk));
 
           response_chunk_counter.Increment();
           if (chunk_handler_ == nullptr) {
             return absl::InternalError("Chunk handler is null.");
           }
           absl::Status handler_status =
-              chunk_handler_->OnDecryptedChunk(*decrypted_chunk);
+              chunk_handler_->OnDecryptedChunk(decrypted_chunk);
           if (!handler_status.ok()) {
             return handler_status;
           }
@@ -334,20 +324,18 @@
       }
       ObliviousHttpResponse::ChunkCounter& response_chunk_counter =
           *response_chunk_counter_;
-      absl::StatusOr<std::string> decrypted_chunk =
+      QUICHE_ASSIGN_OR_RETURN(
+          std::string decrypted_chunk,
           ObliviousHttpResponse::DecryptChunk(
               reader.PeekRemainingPayload(), aead_context_data,
               response_chunk_counter.GetChunkNonce(),
-              /*is_final_chunk=*/true);
-      if (!decrypted_chunk.ok()) {
-        return decrypted_chunk.status();
-      }
+              /*is_final_chunk=*/true));
 
       if (chunk_handler_ == nullptr) {
         return absl::InternalError("Chunk handler is null.");
       }
       absl::Status handler_status =
-          chunk_handler_->OnDecryptedChunk(*decrypted_chunk);
+          chunk_handler_->OnDecryptedChunk(decrypted_chunk);
       if (!handler_status.ok()) {
         return handler_status;
       }
diff --git a/quiche/oblivious_http/oblivious_http_gateway.cc b/quiche/oblivious_http/oblivious_http_gateway.cc
index 5cb7b20..1d3b335 100644
--- a/quiche/oblivious_http/oblivious_http_gateway.cc
+++ b/quiche/oblivious_http/oblivious_http_gateway.cc
@@ -7,7 +7,6 @@
 #include <utility>
 
 #include "absl/base/attributes.h"
-#include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
@@ -18,8 +17,8 @@
 #include "quiche/common/quiche_crypto_logging.h"
 #include "quiche/common/quiche_data_reader.h"
 #include "quiche/common/quiche_data_writer.h"
-#include "quiche/common/quiche_endian.h"
 #include "quiche/common/quiche_random.h"
+#include "quiche/common/quiche_status_utils.h"
 #include "quiche/oblivious_http/buffers/oblivious_http_request.h"
 #include "quiche/oblivious_http/common/oblivious_http_chunk_handler.h"
 #include "quiche/oblivious_http/common/oblivious_http_header_key_config.h"
@@ -64,13 +63,11 @@
     absl::string_view hpke_private_key,
     const ObliviousHttpHeaderKeyConfig& ohttp_key_config,
     QuicheRandom* quiche_random) {
-  absl::StatusOr<bssl::UniquePtr<EVP_HPKE_KEY>> recipient_key =
-      CreateServerRecipientKey(hpke_private_key, ohttp_key_config);
-  if (!recipient_key.ok()) {
-    return recipient_key.status();
-  }
+  QUICHE_ASSIGN_OR_RETURN(
+      bssl::UniquePtr<EVP_HPKE_KEY> recipient_key,
+      CreateServerRecipientKey(hpke_private_key, ohttp_key_config));
   if (quiche_random == nullptr) quiche_random = QuicheRandom::GetInstance();
-  return ObliviousHttpGateway(std::move(*recipient_key), ohttp_key_config,
+  return ObliviousHttpGateway(std::move(recipient_key), ohttp_key_config,
                               quiche_random);
 }
 
@@ -105,17 +102,14 @@
     absl::string_view hpke_private_key,
     const ObliviousHttpHeaderKeyConfig& ohttp_key_config,
     ObliviousHttpChunkHandler& chunk_handler, QuicheRandom* quiche_random) {
-  absl::StatusOr<bssl::UniquePtr<EVP_HPKE_KEY>> recipient_key =
-      CreateServerRecipientKey(hpke_private_key, ohttp_key_config);
-  if (!recipient_key.ok()) {
-    return recipient_key.status();
-  }
+  QUICHE_ASSIGN_OR_RETURN(
+      bssl::UniquePtr<EVP_HPKE_KEY> recipient_key,
+      CreateServerRecipientKey(hpke_private_key, ohttp_key_config));
   if (quiche_random == nullptr) {
     quiche_random = QuicheRandom::GetInstance();
   }
-  return ChunkedObliviousHttpGateway(std::move(*recipient_key),
-                                     ohttp_key_config, chunk_handler,
-                                     quiche_random);
+  return ChunkedObliviousHttpGateway(std::move(recipient_key), ohttp_key_config,
+                                     chunk_handler, quiche_random);
 }
 
 void ChunkedObliviousHttpGateway::InitializeRequestCheckpoint(
@@ -145,15 +139,14 @@
               EVP_HPKE_KEM_enc_len(EVP_HPKE_KEY_kem(server_hpke_key_.get()))) {
         return absl::OutOfRangeError("Not enough data to read header.");
       }
-      absl::StatusOr<ObliviousHttpRequest::Context> context =
+
+      QUICHE_ASSIGN_OR_RETURN(
+          ObliviousHttpRequest::Context context,
           ObliviousHttpRequest::DecodeEncapsulatedRequestHeader(
               reader, *server_hpke_key_, ohttp_key_config_,
-              ObliviousHttpHeaderKeyConfig::kChunkedOhttpRequestLabel);
-      if (!context.ok()) {
-        return context.status();
-      }
+              ObliviousHttpHeaderKeyConfig::kChunkedOhttpRequestLabel));
 
-      oblivious_http_request_context_ = std::move(*context);
+      oblivious_http_request_context_ = std::move(context);
       SaveCheckpoint(reader);
       request_current_section_ = RequestMessageSection::kChunk;
     }
@@ -174,18 +167,12 @@
             return absl::InternalError(
                 "HPKE context has not been derived from an encrypted request.");
           }
-          absl::StatusOr<std::string> decrypted_chunk =
-              ObliviousHttpRequest::DecryptChunk(
-                  *oblivious_http_request_context_, chunk,
-                  /*is_final_chunk=*/false);
-          if (!decrypted_chunk.ok()) {
-            return decrypted_chunk.status();
-          }
-          absl::Status handle_chunk_status =
-              chunk_handler_.OnDecryptedChunk(*decrypted_chunk);
-          if (!handle_chunk_status.ok()) {
-            return handle_chunk_status;
-          }
+          QUICHE_ASSIGN_OR_RETURN(std::string decrypted_chunk,
+                                  ObliviousHttpRequest::DecryptChunk(
+                                      *oblivious_http_request_context_, chunk,
+                                      /*is_final_chunk=*/false));
+          QUICHE_RETURN_IF_ERROR(
+              chunk_handler_.OnDecryptedChunk(decrypted_chunk));
         }
 
         SaveCheckpoint(reader);
@@ -202,22 +189,13 @@
         return absl::InternalError(
             "HPKE context has not been derived from an encrypted request.");
       }
-      absl::StatusOr<std::string> decrypted_chunk =
+      QUICHE_ASSIGN_OR_RETURN(
+          std::string decrypted_chunk,
           ObliviousHttpRequest::DecryptChunk(*oblivious_http_request_context_,
                                              reader.PeekRemainingPayload(),
-                                             /*is_final_chunk=*/true);
-      if (!decrypted_chunk.ok()) {
-        return decrypted_chunk.status();
-      }
-      absl::Status handle_chunk_status =
-          chunk_handler_.OnDecryptedChunk(*decrypted_chunk);
-      if (!handle_chunk_status.ok()) {
-        return handle_chunk_status;
-      }
-      handle_chunk_status = chunk_handler_.OnChunksDone();
-      if (!handle_chunk_status.ok()) {
-        return handle_chunk_status;
-      }
+                                             /*is_final_chunk=*/true));
+      QUICHE_RETURN_IF_ERROR(chunk_handler_.OnDecryptedChunk(decrypted_chunk));
+      QUICHE_RETURN_IF_ERROR(chunk_handler_.OnChunksDone());
     }
   }
   return absl::OkStatus();
@@ -277,32 +255,29 @@
   }
 
   if (!aead_context_data_.has_value()) {
-    absl::StatusOr<ObliviousHttpResponse::CommonAeadParamsResult> aead_params =
+    QUICHE_ASSIGN_OR_RETURN(
+        ObliviousHttpResponse::CommonAeadParamsResult aead_params,
         ObliviousHttpResponse::GetCommonAeadParams(
-            *oblivious_http_request_context_);
-    if (!aead_params.ok()) {
-      return aead_params.status();
-    }
+            *oblivious_http_request_context_));
 
     // secret_len represents max(Nn, Nk))
-    response_nonce_ = std::string(aead_params->secret_len, '\0');
+    response_nonce_ = std::string(aead_params.secret_len, '\0');
     quiche_random_->RandBytes(response_nonce_.data(), response_nonce_.size());
 
-    auto aead_context_data = ObliviousHttpResponse::GetAeadContextData(
-        *oblivious_http_request_context_, *aead_params,
-        ObliviousHttpHeaderKeyConfig::kChunkedOhttpResponseLabel,
-        response_nonce_);
-    if (!aead_context_data.ok()) {
-      return aead_context_data.status();
-    }
-    aead_context_data_.emplace(std::move(*aead_context_data));
+    QUICHE_ASSIGN_OR_RETURN(
+        ObliviousHttpResponse::AeadContextData aead_context_data,
+        ObliviousHttpResponse::GetAeadContextData(
+            *oblivious_http_request_context_, aead_params,
+            ObliviousHttpHeaderKeyConfig::kChunkedOhttpResponseLabel,
+            response_nonce_));
 
-    auto response_chunk_counter = ObliviousHttpResponse::ChunkCounter::Create(
-        aead_context_data_->aead_nonce);
-    if (!response_chunk_counter.ok()) {
-      return response_chunk_counter.status();
-    }
-    response_chunk_counter_.emplace(std::move(*response_chunk_counter));
+    aead_context_data_.emplace(std::move(aead_context_data));
+
+    QUICHE_ASSIGN_OR_RETURN(
+        ObliviousHttpResponse::ChunkCounter response_chunk_counter,
+        ObliviousHttpResponse::ChunkCounter::Create(
+            aead_context_data_->aead_nonce));
+    response_chunk_counter_.emplace(std::move(response_chunk_counter));
   }
 
   if (!response_chunk_counter_.has_value()) {
@@ -310,14 +285,12 @@
         "Response chunk counter has not been initialized.");
   }
 
-  absl::StatusOr<std::string> encrypted_data =
+  QUICHE_ASSIGN_OR_RETURN(
+      std::string encrypted_data,
       ObliviousHttpResponse::EncryptChunk(
           *oblivious_http_request_context_, *aead_context_data_,
           plaintext_payload, response_chunk_counter_->GetChunkNonce(),
-          is_final_chunk);
-  if (!encrypted_data.ok()) {
-    return encrypted_data.status();
-  }
+          is_final_chunk));
 
   absl::string_view maybe_nonce;
   if (response_current_section_ == ResponseMessageSection::kNonce) {
@@ -326,8 +299,8 @@
   }
 
   uint8_t chunk_var_int_length =
-      QuicheDataWriter::GetVarInt62Len(encrypted_data->size());
-  uint64_t chunk_var_int = encrypted_data->size();
+      QuicheDataWriter::GetVarInt62Len(encrypted_data.size());
+  uint64_t chunk_var_int = encrypted_data.size();
   if (is_final_chunk) {
     response_current_section_ = ResponseMessageSection::kEnd;
     chunk_var_int_length =
@@ -342,7 +315,7 @@
   }
 
   std::string response_buffer(
-      maybe_nonce.size() + chunk_var_int_length + encrypted_data->size(), '\0');
+      maybe_nonce.size() + chunk_var_int_length + encrypted_data.size(), '\0');
   QuicheDataWriter writer(response_buffer.size(), response_buffer.data());
 
   if (!writer.WriteStringPiece(maybe_nonce)) {
@@ -351,7 +324,7 @@
   if (!writer.WriteVarInt62(chunk_var_int)) {
     return absl::InternalError("Failed to write chunk to buffer.");
   }
-  if (!writer.WriteStringPiece(*encrypted_data)) {
+  if (!writer.WriteStringPiece(encrypted_data)) {
     return absl::InternalError("Failed to write encrypted data to buffer.");
   }
 
diff --git a/quiche/quic/masque/masque_ohttp_client.cc b/quiche/quic/masque/masque_ohttp_client.cc
index c5dd87d..72f2f1c 100644
--- a/quiche/quic/masque/masque_ohttp_client.cc
+++ b/quiche/quic/masque/masque_ohttp_client.cc
@@ -101,12 +101,13 @@
   request.headers[":authority"] = url.HostPort();
   request.headers[":path"] = url.path();
   request.headers["accept"] = "application/ohttp-keys";
-  absl::StatusOr<RequestId> request_id = connection_pool_.SendRequest(request);
-  if (!request_id.ok()) {
-    QUICHE_LOG(ERROR) << "Failed to send request: " << request_id.status();
-    return request_id.status();
-  }
-  key_fetch_request_id_ = *request_id;
+
+  QUICHE_ASSIGN_OR_RETURN(
+      key_fetch_request_id_, connection_pool_.SendRequest(request),
+      [](const absl::Status& status) {
+        QUICHE_LOG(ERROR) << "Failed to send request: " << status;
+        return status;
+      });
   return absl::OkStatus();
 }