Implement an explicit inplace encrypt in Aead.

Summary: This is useful when you want to ensure that the IOBuf you pass in is encrypted inplace, as opposed to potentially creating a new one.

Reviewed By: yangchi

Differential Revision: D21135253

fbshipit-source-id: 89b6e718fc8da1324685c390c721a564bb77d01d

quic/api/QuicTransportFunctions.cpp

@@ -204,8 +204,8 @@ DataPathResult iobufChainBasedBuildScheduleEncrypt(
   bodyCursor.pull(unencrypted->writableData() + headerLen, bodyLen);
   unencrypted->advance(headerLen);
   unencrypted->append(bodyLen);
-  auto packetBuf =
+  auto packetBuf = aead.inplaceEncrypt(
-      aead.encrypt(std::move(unencrypted), packet->header.get(), packetNum);
+      std::move(unencrypted), packet->header.get(), packetNum);
   DCHECK(packetBuf->headroom() == headerLen);
   packetBuf->clear();
   auto headerCursor = folly::io::Cursor(packet->header.get());
@@ -832,8 +832,8 @@ void writeCloseCommon(
   }
   auto packet = std::move(packetBuilder).buildPacket();
   packet.header->coalesce();
-  auto body =
+  auto body = aead.inplaceEncrypt(
-      aead.encrypt(std::move(packet.body), packet.header.get(), packetNum);
+      std::move(packet.body), packet.header.get(), packetNum);
   body->coalesce();
   encryptPacketHeader(
       headerForm,

quic/api/test/QuicTransportFunctionsTest.cpp

@@ -1402,7 +1402,7 @@ TEST_F(QuicTransportFunctionsTest, WriteProbingOldData) {
   writeDataToQuicStream(*stream, buf->clone(), true);
 
   folly::IOBuf pktBodyCaptured;
-  EXPECT_CALL(*capturingAead, _encrypt(_, _, _))
+  EXPECT_CALL(*capturingAead, _inplaceEncrypt(_, _, _))
       .WillRepeatedly(Invoke([&](auto& buf, auto, auto) {
         if (buf) {
           pktBodyCaptured.prependChain(buf->clone());
@@ -1418,7 +1418,7 @@ TEST_F(QuicTransportFunctionsTest, WriteProbingOldData) {
   // Now we have no new data, let's probe again, and verify the same old data is
   // sent.
   folly::IOBuf secondBodyCaptured;
-  EXPECT_CALL(*capturingAead, _encrypt(_, _, _))
+  EXPECT_CALL(*capturingAead, _inplaceEncrypt(_, _, _))
       .WillRepeatedly(Invoke([&](auto& buf, auto, auto) {
         if (buf) {
           secondBodyCaptured.prependChain(buf->clone());

quic/api/test/QuicTransportTest.cpp

@@ -186,7 +186,7 @@ class QuicTransportTest : public Test {
     // a cipher.
     auto aead = std::make_unique<NiceMock<MockAead>>();
     aead_ = aead.get();
-    EXPECT_CALL(*aead_, _encrypt(_, _, _))
+    EXPECT_CALL(*aead_, _inplaceEncrypt(_, _, _))
         .WillRepeatedly(
             Invoke([&](auto& buf, auto, auto) { return buf->clone(); }));
     EXPECT_CALL(*aead_, _decrypt(_, _, _))

quic/client/handshake/ClientHandshake.cpp

@@ -163,12 +163,11 @@ bool ClientHandshake::verifyRetryIntegrityTag(
       retryPacket.header.getToken().size());
 
   pseudoRetryPacket.coalesce();
-  std::unique_ptr<folly::IOBuf> emptyPlaintext =
-      std::make_unique<folly::IOBuf>();
 
   auto retryCipher = getRetryPacketCipher();
-  auto expectedIntegrityTag =
+  auto emptyPlaintext = folly::IOBuf::create(retryCipher->getCipherOverhead());
-      retryCipher->encrypt(std::move(emptyPlaintext), &pseudoRetryPacket, 0);
+  auto expectedIntegrityTag = retryCipher->inplaceEncrypt(
+      std::move(emptyPlaintext), &pseudoRetryPacket, 0);
 
   return folly::IOBufEqualTo()(
       *expectedIntegrityTag, *retryPacket.integrityTag);

quic/codec/test/QuicPacketBuilderTest.cpp

@@ -75,7 +75,7 @@ Buf packetToBuf(
   }
   if (aead && packet.header) {
     auto bodySize = body->computeChainDataLength();
-    body = aead->encrypt(std::move(body), packet.header.get(), num);
+    body = aead->inplaceEncrypt(std::move(body), packet.header.get(), num);
     EXPECT_GT(body->computeChainDataLength(), bodySize);
   }
   if (body) {

quic/common/test/TestUtils.cpp

@@ -265,10 +265,10 @@ template <class T>
 std::unique_ptr<T> createNoOpAeadImpl() {
   // Fake that the handshake has already occured
   auto aead = std::make_unique<NiceMock<T>>();
-  ON_CALL(*aead, _encrypt(_, _, _))
+  ON_CALL(*aead, _inplaceEncrypt(_, _, _))
       .WillByDefault(Invoke([&](auto& buf, auto, auto) {
         if (buf) {
-          return buf->clone();
+          return std::move(buf);
         } else {
           return folly::IOBuf::create(0);
         }
@@ -426,12 +426,20 @@ Buf packetToBufCleartext(
   auto packetBuf = packet.header->clone();
   Buf body;
   if (packet.body) {
+    packet.body->coalesce();
     body = packet.body->clone();
+  } else {
+    body = folly::IOBuf::create(0);
   }
   auto headerForm = packet.packet.header.getHeaderForm();
   packet.header->coalesce();
-  auto encryptedBody =
+  auto tagLen = cleartextCipher.getCipherOverhead();
-      cleartextCipher.encrypt(std::move(body), packet.header.get(), packetNum);
+  if (body->tailroom() < tagLen) {
+    body->prependChain(folly::IOBuf::create(tagLen));
+  }
+  body->coalesce();
+  auto encryptedBody = cleartextCipher.inplaceEncrypt(
+      std::move(body), packet.header.get(), packetNum);
   encryptedBody->coalesce();
   encryptPacketHeader(
       headerForm,

quic/fizz/handshake/FizzBridge.h

@@ -32,11 +32,12 @@ class FizzAead final : public Aead {
   /**
    * Simply forward all calls to fizz::Aead.
    */
-  std::unique_ptr<folly::IOBuf> encrypt(
+  std::unique_ptr<folly::IOBuf> inplaceEncrypt(
       std::unique_ptr<folly::IOBuf>&& plaintext,
       const folly::IOBuf* associatedData,
       uint64_t seqNum) const override {
-    return fizzAead->encrypt(std::move(plaintext), associatedData, seqNum);
+    return fizzAead->inplaceEncrypt(
+        std::move(plaintext), associatedData, seqNum);
   }
   std::unique_ptr<folly::IOBuf> decrypt(
       std::unique_ptr<folly::IOBuf>&& ciphertext,

quic/handshake/Aead.h

@@ -26,9 +26,9 @@ class Aead {
   virtual ~Aead() = default;
 
   /**
-   * Encrypts plaintext. Will throw on error.
+   * Encrypts plaintext inplace. Will throw on error.
    */
-  virtual std::unique_ptr<folly::IOBuf> encrypt(
+  virtual std::unique_ptr<folly::IOBuf> inplaceEncrypt(
       std::unique_ptr<folly::IOBuf>&& plaintext,
       const folly::IOBuf* associatedData,
       uint64_t seqNum) const = 0;

quic/handshake/test/Mocks.h

@@ -32,16 +32,16 @@ class MockAead : public Aead {
   MOCK_CONST_METHOD0(getCipherOverhead, size_t());
 
   MOCK_CONST_METHOD3(
-      _encrypt,
+      _inplaceEncrypt,
       std::unique_ptr<folly::IOBuf>(
           std::unique_ptr<folly::IOBuf>& plaintext,
           const folly::IOBuf* associatedData,
           uint64_t seqNum));
-  std::unique_ptr<folly::IOBuf> encrypt(
+  std::unique_ptr<folly::IOBuf> inplaceEncrypt(
       std::unique_ptr<folly::IOBuf>&& plaintext,
       const folly::IOBuf* associatedData,
       uint64_t seqNum) const override {
-    return _encrypt(plaintext, associatedData, seqNum);
+    return _inplaceEncrypt(plaintext, associatedData, seqNum);
   }
 
   MOCK_CONST_METHOD3(
@@ -71,9 +71,9 @@ class MockAead : public Aead {
   }
 
   void setDefaults() {
-    ON_CALL(*this, _encrypt(_, _, _)).WillByDefault(InvokeWithoutArgs([]() {
+    ON_CALL(*this, _inplaceEncrypt(_, _, _))
-      return folly::IOBuf::copyBuffer("ciphertext");
+        .WillByDefault(InvokeWithoutArgs(
-    }));
+            []() { return folly::IOBuf::copyBuffer("ciphertext"); }));
     ON_CALL(*this, _decrypt(_, _, _)).WillByDefault(InvokeWithoutArgs([]() {
       return folly::IOBuf::copyBuffer("plaintext");
     }));