Remove throw from inplaceEncrypt

Summary: Continuing the theme, this time with inplaceEncrypt. We have to catch the exceptions from Fizz for the bridge version.

Reviewed By: kvtsoy

Differential Revision: D75886829

fbshipit-source-id: f74d5ce73e242e20e4c62552748970b2d5d7cb79

quic/api/QuicTransportFunctions.cpp

@@ -302,8 +302,12 @@ continuousMemoryBuildScheduleEncrypt(
   connection.bufAccessor->trimStart(prevSize + headerLen);
   // buf and packetBuf is actually the same.
   auto buf = connection.bufAccessor->obtain();
-  auto packetBuf =
+  auto encryptResult =
       aead.inplaceEncrypt(std::move(buf), &packet->header, packetNum);
+  if (encryptResult.hasError()) {
+    return folly::makeUnexpected(encryptResult.error());
+  }
+  auto packetBuf = std::move(encryptResult.value());
   CHECK(packetBuf->headroom() == headerLen + prevSize);
   // Include header back.
   packetBuf->prepend(headerLen);
@@ -410,8 +414,12 @@ iobufChainBasedBuildScheduleEncrypt(
   bodyCursor.pull(unencrypted->writableData() + headerLen, bodyLen);
   unencrypted->advance(headerLen);
   unencrypted->append(bodyLen);
-  auto packetBuf =
+  auto encryptResult =
       aead.inplaceEncrypt(std::move(unencrypted), &packet->header, packetNum);
+  if (encryptResult.hasError()) {
+    return folly::makeUnexpected(encryptResult.error());
+  }
+  auto packetBuf = std::move(encryptResult.value());
   DCHECK(packetBuf->headroom() == headerLen);
   packetBuf->clear();
   auto headerCursor = Cursor(&packet->header);
@@ -1447,8 +1455,13 @@ void writeCloseCommon(
   auto packet = std::move(packetBuilder).buildPacket();
   CHECK_GE(packet.body.tailroom(), aead.getCipherOverhead());
   auto bufUniquePtr = packet.body.clone();
-  bufUniquePtr =
+  auto encryptResult =
       aead.inplaceEncrypt(std::move(bufUniquePtr), &packet.header, packetNum);
+  if (encryptResult.hasError()) {
+    LOG(ERROR) << "Error encrypting packet: " << encryptResult.error().message;
+    return;
+  }
+  bufUniquePtr = std::move(encryptResult.value());
   bufUniquePtr->coalesce();
   encryptPacketHeader(
       headerForm,

quic/codec/test/QuicPacketBuilderTest.cpp

@@ -39,7 +39,9 @@ BufPtr packetToBuf(
   }
   if (aead && !packet.header.empty()) {
     auto bodySize = body->computeChainDataLength();
-    body = aead->inplaceEncrypt(std::move(body), &packet.header, num);
+    auto result = aead->inplaceEncrypt(std::move(body), &packet.header, num);
+    CHECK(!result.hasError());
+    body = std::move(result.value());
     EXPECT_GT(body->computeChainDataLength(), bodySize);
   }
   if (body) {

quic/common/test/TestUtils.cpp

@@ -448,8 +448,13 @@ BufPtr packetToBufCleartext(
     body->appendToChain(folly::IOBuf::create(tagLen));
   }
   body->coalesce();
-  auto encryptedBody = cleartextCipher.inplaceEncrypt(
+  auto encryptResult = cleartextCipher.inplaceEncrypt(
       std::move(body), &packet.header, packetNum);
+  if (encryptResult.hasError()) {
+    throw std::runtime_error(
+        "Failed to encrypt packet: " + encryptResult.error().message);
+  }
+  auto encryptedBody = std::move(encryptResult.value());
   encryptedBody->coalesce();
   encryptPacketHeader(
       headerForm,

quic/dsr/backend/DSRPacketizer.cpp

@@ -88,8 +88,14 @@ bool PacketGroupWriter::writeSingleQuicPacket(
   // buildBuf's data starts from the body part of buildBuf.
   buildBuf->trimStart(prevSize_ + headerLen);
   // buildBuf and packetbuildBuf is actually the same.
-  auto packetbuildBuf =
+  auto encryptResult =
       aead.inplaceEncrypt(std::move(buildBuf), &packet.header, packetNum);
+  if (encryptResult.hasError()) {
+    throw QuicInternalException(
+        "DSR Send failed: Encryption error: " + encryptResult.error().message,
+        LocalErrorCode::INTERNAL_ERROR);
+  }
+  auto packetbuildBuf = std::move(encryptResult.value());
   CHECK_EQ(packetbuildBuf->headroom(), headerLen + prevSize_);
   // Include header back.
   packetbuildBuf->prepend(headerLen);

quic/fizz/handshake/FizzBridge.h

@@ -31,14 +31,20 @@ class FizzAead final : public Aead {
   Optional<TrafficKey> getKey() const override;
 
   /**
-   * Simply forward all calls to fizz::Aead.
+   * Forward calls to fizz::Aead, catching any exceptions and converting them to
+   * folly::Expected.
    */
-  std::unique_ptr<folly::IOBuf> inplaceEncrypt(
+  folly::Expected<std::unique_ptr<folly::IOBuf>, QuicError> inplaceEncrypt(
       std::unique_ptr<folly::IOBuf>&& plaintext,
       const folly::IOBuf* associatedData,
       uint64_t seqNum) const override {
+    try {
       return fizzAead->inplaceEncrypt(
           std::move(plaintext), associatedData, seqNum);
+    } catch (const std::exception& ex) {
+      return folly::makeUnexpected(
+          QuicError(TransportErrorCode::INTERNAL_ERROR, ex.what()));
+    }
   }
 
   std::unique_ptr<folly::IOBuf> decrypt(

quic/handshake/Aead.h

@@ -7,7 +7,9 @@
 
 #pragma once
 
+#include <folly/Expected.h>
 #include <folly/io/IOBuf.h>
+#include <quic/QuicException.h>
 #include <quic/common/Optional.h>
 
 namespace quic {
@@ -27,9 +29,12 @@ class Aead {
   virtual Optional<TrafficKey> getKey() const = 0;
 
   /**
-   * Encrypts plaintext inplace. Will throw on error.
+   * Encrypts plaintext inplace. Returns folly::Expected with the encrypted
+   * buffer or an error.
    */
-  virtual std::unique_ptr<folly::IOBuf> inplaceEncrypt(
+  [[nodiscard]] virtual folly::
+      Expected<std::unique_ptr<folly::IOBuf>, QuicError>
+      inplaceEncrypt(
           std::unique_ptr<folly::IOBuf>&& plaintext,
           const folly::IOBuf* associatedData,
           uint64_t seqNum) const = 0;

quic/handshake/BUCK

@@ -8,7 +8,9 @@ mvfst_cpp_library(
         "Aead.h",
     ],
     exported_deps = [
+        "//folly:expected",
         "//folly/io:iobuf",
+        "//quic:exception",
         "//quic/common:optional",
     ],
 )

quic/handshake/test/Mocks.h

@@ -54,14 +54,14 @@ class MockAead : public Aead {
 
   MOCK_METHOD(Optional<TrafficKey>, getKey, (), (const));
   MOCK_METHOD(
-      std::unique_ptr<folly::IOBuf>,
+      (folly::Expected<std::unique_ptr<folly::IOBuf>, QuicError>),
       _inplaceEncrypt,
       (std::unique_ptr<folly::IOBuf> & plaintext,
        const folly::IOBuf* associatedData,
        uint64_t seqNum),
       (const));
 
-  std::unique_ptr<folly::IOBuf> inplaceEncrypt(
+  folly::Expected<std::unique_ptr<folly::IOBuf>, QuicError> inplaceEncrypt(
       std::unique_ptr<folly::IOBuf>&& plaintext,
       const folly::IOBuf* associatedData,
       uint64_t seqNum) const override {
@@ -102,7 +102,9 @@ class MockAead : public Aead {
     using namespace testing;
     ON_CALL(*this, _inplaceEncrypt(_, _, _))
         .WillByDefault(InvokeWithoutArgs(
-            []() { return folly::IOBuf::copyBuffer("ciphertext"); }));
+            []() -> folly::Expected<std::unique_ptr<folly::IOBuf>, QuicError> {
+              return folly::IOBuf::copyBuffer("ciphertext");
+            }));
     ON_CALL(*this, _decrypt(_, _, _)).WillByDefault(InvokeWithoutArgs([]() {
       return folly::IOBuf::copyBuffer("plaintext");
     }));