LibCrypto+LibTLS+Kernel: Switch the Cipher::Mode interface to use Span

This shaves 2.5 more runtime seconds off 'disasm /bin/id', and makes the
Mode<T> interface a lot more allocation-friendly.

7 files changed, 87 insertions, 64 deletions

diff --git a/Libraries/LibCrypto/Authentication/HMAC.h b/Libraries/LibCrypto/Authentication/HMAC.h
index fbe1a47eb2..8b0ba343cd 100644
--- a/Libraries/LibCrypto/Authentication/HMAC.h
+++ b/Libraries/LibCrypto/Authentication/HMAC.h
@@ -68,8 +68,11 @@ public:
         m_inner_hasher.update(message, length);
     }
 
+    TagType process(const ReadonlyBytes& span) { return process(span.data(), span.size()); }
     TagType process(const ByteBuffer& buffer) { return process(buffer.data(), buffer.size()); }
     TagType process(const StringView& string) { return process((const u8*)string.characters_without_null_termination(), string.length()); }
+
+    void update(const ReadonlyBytes& span) { return update(span.data(), span.size()); }
     void update(const ByteBuffer& buffer) { return update(buffer.data(), buffer.size()); }
     void update(const StringView& string) { return update((const u8*)string.characters_without_null_termination(), string.length()); }
 
diff --git a/Libraries/LibCrypto/Cipher/AES.cpp b/Libraries/LibCrypto/Cipher/AES.cpp
index 3118ae000a..919fa1e7f5 100644
--- a/Libraries/LibCrypto/Cipher/AES.cpp
+++ b/Libraries/LibCrypto/Cipher/AES.cpp
@@ -396,13 +396,11 @@ void AESCipher::decrypt_block(const AESCipherBlock& in, AESCipherBlock& out)
     // clang-format on
 }
 
-void AESCipherBlock::overwrite(const ByteBuffer& buffer)
+void AESCipherBlock::overwrite(const ReadonlyBytes& span)
 {
-    overwrite(buffer.data(), buffer.size());
-}
+    auto data = span.data();
+    auto length = span.size();
 
-void AESCipherBlock::overwrite(const u8* data, size_t length)
-{
     ASSERT(length <= m_data.size());
     m_data.overwrite(0, data, length);
     if (length < m_data.size()) {
diff --git a/Libraries/LibCrypto/Cipher/AES.h b/Libraries/LibCrypto/Cipher/AES.h
index 7b20091c3c..91b1e3bf3a 100644
--- a/Libraries/LibCrypto/Cipher/AES.h
+++ b/Libraries/LibCrypto/Cipher/AES.h
@@ -46,7 +46,7 @@ public:
     AESCipherBlock(const u8* data, size_t length, PaddingMode mode = PaddingMode::CMS)
         : AESCipherBlock(mode)
     {
-        overwrite(data, length);
+        CipherBlock::overwrite(data, length);
     }
 
     static size_t block_size() { return BlockSizeInBits / 8; };
@@ -54,8 +54,9 @@ public:
     virtual ByteBuffer get() const override { return m_data; };
     virtual const ByteBuffer& data() const override { return m_data; };
 
-    virtual void overwrite(const ByteBuffer&) override;
-    virtual void overwrite(const u8* data, size_t length) override;
+    virtual void overwrite(const ReadonlyBytes&) override;
+    virtual void overwrite(const ByteBuffer& buffer) override { overwrite(buffer.span()); }
+    virtual void overwrite(const u8* data, size_t size) override { overwrite({ data, size }); }
 
     virtual void apply_initialization_vector(const u8* ivec) override
     {
diff --git a/Libraries/LibCrypto/Cipher/Cipher.h b/Libraries/LibCrypto/Cipher/Cipher.h
index f9d59b6b2b..5712a43fd3 100644
--- a/Libraries/LibCrypto/Cipher/Cipher.h
+++ b/Libraries/LibCrypto/Cipher/Cipher.h
@@ -64,12 +64,14 @@ public:
     virtual ByteBuffer get() const = 0;
     virtual const ByteBuffer& data() const = 0;
 
-    virtual void overwrite(const ByteBuffer&) = 0;
-    virtual void overwrite(const u8*, size_t) = 0;
+    virtual void overwrite(const ReadonlyBytes&) = 0;
+    virtual void overwrite(const ByteBuffer& buffer) { overwrite(buffer.span()); }
+    virtual void overwrite(const u8* data, size_t size) { overwrite({ data, size }); }
 
     virtual void apply_initialization_vector(const u8* ivec) = 0;
 
     PaddingMode padding_mode() const { return m_padding_mode; }
+    void set_padding_mode(PaddingMode mode) { m_padding_mode = mode; }
 
     template<typename T>
     void put(size_t offset, T value)
diff --git a/Libraries/LibCrypto/Cipher/Mode/CBC.h b/Libraries/LibCrypto/Cipher/Mode/CBC.h
index db6880de54..6083cfae70 100644
--- a/Libraries/LibCrypto/Cipher/Mode/CBC.h
+++ b/Libraries/LibCrypto/Cipher/Mode/CBC.h
@@ -27,8 +27,8 @@
 #pragma once
 
 #include <AK/String.h>
-#include <AK/StringView.h>
 #include <AK/StringBuilder.h>
+#include <AK/StringView.h>
 #include <LibCrypto/Cipher/Mode/Mode.h>
 
 namespace Crypto {
@@ -56,44 +56,48 @@ public:
 
     virtual size_t IV_length() const override { return IVSizeInBits / 8; }
 
-    virtual Optional<ByteBuffer> encrypt(const ByteBuffer& in, ByteBuffer& out, Optional<ByteBuffer> ivec = {}) override
+    virtual void encrypt(const ReadonlyBytes& in, Bytes& out, const Bytes& ivec = {}, Bytes* ivec_out = nullptr) override
     {
         auto length = in.size();
         if (length == 0)
-            return {};
+            return;
 
         auto& cipher = this->cipher();
 
         // FIXME: We should have two of these encrypt/decrypt functions that
         //        we SFINAE out based on whether the Cipher mode needs an ivec
-        ASSERT(ivec.has_value());
-        const auto* iv = ivec.value().data();
+        ASSERT(!ivec.is_empty());
+        const auto* iv = ivec.data();
 
-        typename T::BlockType block { cipher.padding_mode() };
+        m_cipher_block.set_padding_mode(cipher.padding_mode());
         size_t offset { 0 };
         auto block_size = cipher.block_size();
 
         while (length >= block_size) {
-            block.overwrite(in.slice_view(offset, block_size));
-            block.apply_initialization_vector(iv);
-            cipher.encrypt_block(block, block);
-            out.overwrite(offset, block.get().data(), block_size);
-            iv = out.offset_pointer(offset);
+            m_cipher_block.overwrite(in.slice(offset, block_size));
+            m_cipher_block.apply_initialization_vector(iv);
+            cipher.encrypt_block(m_cipher_block, m_cipher_block);
+            ASSERT(offset + block_size <= out.size());
+            __builtin_memcpy(out.offset(offset), m_cipher_block.get().data(), block_size);
+            iv = out.offset(offset);
             length -= block_size;
             offset += block_size;
         }
 
         if (length > 0) {
-            block.overwrite(in.slice_view(offset, length));
-            block.apply_initialization_vector(iv);
-            cipher.encrypt_block(block, block);
-            out.overwrite(offset, block.get().data(), block_size);
-            iv = out.offset_pointer(offset);
+            m_cipher_block.overwrite(in.slice(offset, length));
+            m_cipher_block.apply_initialization_vector(iv);
+            cipher.encrypt_block(m_cipher_block, m_cipher_block);
+            ASSERT(offset + block_size <= out.size());
+            __builtin_memcpy(out.offset(offset), m_cipher_block.get().data(), block_size);
+            iv = out.offset(offset);
         }
 
-        return ByteBuffer::copy(iv, block_size);
+        if (ivec_out)
+            __builtin_memcpy(ivec_out->data(), iv, min(IV_length(), ivec_out->size()));
     }
-    virtual void decrypt(const ByteBuffer& in, ByteBuffer& out, Optional<ByteBuffer> ivec = {}) override
+
+    virtual void decrypt(const ReadonlyBytes& in, Bytes& out, const Bytes& ivec = {}) override
     {
         auto length = in.size();
         if (length == 0)
@@ -101,8 +105,8 @@ public:
 
         auto& cipher = this->cipher();
 
-        ASSERT(ivec.has_value());
-        const auto* iv = ivec.value().data();
+        ASSERT(!ivec.is_empty());
+        const auto* iv = ivec.data();
 
         auto block_size = cipher.block_size();
 
@@ -110,23 +114,27 @@ public:
         // FIXME (ponder): Should we simply decrypt as much as we can?
         ASSERT(length % block_size == 0);
 
-        typename T::BlockType block { cipher.padding_mode() };
+        m_cipher_block.set_padding_mode(cipher.padding_mode());
         size_t offset { 0 };
 
         while (length > 0) {
-            auto* slice = in.offset_pointer(offset);
-            block.overwrite(slice, block_size);
-            cipher.decrypt_block(block, block);
-            block.apply_initialization_vector(iv);
-            auto decrypted = block.get();
-            out.overwrite(offset, decrypted.data(), decrypted.size());
+            auto* slice = in.offset(offset);
+            m_cipher_block.overwrite(slice, block_size);
+            cipher.decrypt_block(m_cipher_block, m_cipher_block);
+            m_cipher_block.apply_initialization_vector(iv);
+            auto decrypted = m_cipher_block.get();
+            ASSERT(offset + decrypted.size() <= out.size());
+            __builtin_memcpy(out.offset(offset), decrypted.data(), decrypted.size());
             iv = slice;
             length -= block_size;
             offset += block_size;
         }
-        out.trim(offset);
+        out = out.slice(0, offset);
         this->prune_padding(out);
     }
+
+private:
+    typename T::BlockType m_cipher_block {};
 };
 
 }
diff --git a/Libraries/LibCrypto/Cipher/Mode/CTR.h b/Libraries/LibCrypto/Cipher/Mode/CTR.h
index fb2ab93569..9799044f1e 100644
--- a/Libraries/LibCrypto/Cipher/Mode/CTR.h
+++ b/Libraries/LibCrypto/Cipher/Mode/CTR.h
@@ -98,8 +98,8 @@ public:
     // Must intercept `Intent`, because AES must always be set to
     // Encryption, even when decrypting AES-CTR.
     // TODO: How to deal with ciphers that take different arguments?
-    template<typename... Args>
-    explicit constexpr CTR<T>(const ByteBuffer& user_key, size_t key_bits, Intent = Intent::Encryption, Args... args)
+    template<typename KeyType, typename... Args>
+    explicit constexpr CTR<T>(const KeyType& user_key, size_t key_bits, Intent = Intent::Encryption, Args... args)
         : Mode<T>(user_key, key_bits, args...)
     {
     }
@@ -114,26 +114,29 @@ public:
 
     virtual size_t IV_length() const override { return IVSizeInBits / 8; }
 
-    virtual Optional<ByteBuffer> encrypt(const ByteBuffer& in, ByteBuffer& out, Optional<ByteBuffer> ivec = {}) override
+    virtual void encrypt(const ReadonlyBytes& in, Bytes& out, const Bytes& ivec = {}, Bytes* ivec_out = nullptr) override
     {
         // Our interpretation of "ivec" is what AES-CTR
         // would define as nonce + IV + 4 zero bytes.
-        return this->encrypt_or_stream(&in, out, ivec);
+        this->encrypt_or_stream(&in, out, ivec, ivec_out);
     }
 
-    Optional<ByteBuffer> key_stream(ByteBuffer& out, Optional<ByteBuffer> ivec = {})
+    void key_stream(Bytes& out, const Bytes& ivec = {}, Bytes* ivec_out = nullptr)
     {
-        return this->encrypt_or_stream(nullptr, out, ivec);
+        this->encrypt_or_stream(nullptr, out, ivec, ivec_out);
     }
 
-    virtual void decrypt(const ByteBuffer& in, ByteBuffer& out, Optional<ByteBuffer> ivec = {}) override
+    virtual void decrypt(const ReadonlyBytes& in, Bytes& out, const Bytes& ivec = {}) override
     {
         // XOR (and thus CTR) is the most symmetric mode.
-        (void)this->encrypt(in, out, ivec);
+        this->encrypt(in, out, ivec);
     }
 
 private:
-    static void increment_inplace(ByteBuffer& in)
+    u8 m_ivec_storage[IVSizeInBits / 8];
+    typename T::BlockType m_cipher_block {};
+
+    static void increment_inplace(Bytes& in)
     {
         for (size_t i = in.size(); i > 0;) {
             --i;
@@ -146,14 +149,14 @@ private:
         }
     }
 
-    Optional<ByteBuffer> encrypt_or_stream(const ByteBuffer* in, ByteBuffer& out, Optional<ByteBuffer> ivec)
+    void encrypt_or_stream(const ReadonlyBytes* in, Bytes& out, const Bytes& ivec, Bytes* ivec_out = nullptr)
     {
         size_t length;
         if (in) {
             ASSERT(in->size() <= out.size());
             length = in->size();
             if (length == 0)
-                return {};
+                return;
         } else {
             length = out.size();
         }
@@ -162,29 +165,36 @@ private:
 
         // FIXME: We should have two of these encrypt/decrypt functions that
         //        we SFINAE out based on whether the Cipher mode needs an ivec
-        ASSERT(ivec.has_value());
-        auto iv = ivec.value();
+        ASSERT(!ivec.is_empty());
+        ASSERT(ivec.size() >= IV_length());
+
+        m_cipher_block.set_padding_mode(cipher.padding_mode());
+
+        __builtin_memcpy(m_ivec_storage, ivec.data(), IV_length());
+        Bytes iv { m_ivec_storage, IV_length() };
 
-        typename T::BlockType block { cipher.padding_mode() };
         size_t offset { 0 };
         auto block_size = cipher.block_size();
 
         while (length > 0) {
-            block.overwrite(iv.slice_view(0, block_size));
+            m_cipher_block.overwrite(iv.slice(0, block_size));
 
-            cipher.encrypt_block(block, block);
+            cipher.encrypt_block(m_cipher_block, m_cipher_block);
             if (in) {
-                block.apply_initialization_vector(in->data() + offset);
+                m_cipher_block.apply_initialization_vector(in->data() + offset);
             }
             auto write_size = min(block_size, length);
-            out.overwrite(offset, block.get().data(), write_size);
+
+            ASSERT(offset + write_size <= out.size());
+            __builtin_memcpy(out.offset(offset), m_cipher_block.get().data(), write_size);
 
             increment_inplace(iv);
             length -= write_size;
             offset += write_size;
         }
 
-        return iv;
+        if (ivec_out)
+            __builtin_memcpy(ivec_out->data(), iv.data(), min(ivec_out->size(), IV_length()));
     }
 };
 
diff --git a/Libraries/LibCrypto/Cipher/Mode/Mode.h b/Libraries/LibCrypto/Cipher/Mode/Mode.h
index cf54ea6d5f..65b65e4c1c 100644
--- a/Libraries/LibCrypto/Cipher/Mode/Mode.h
+++ b/Libraries/LibCrypto/Cipher/Mode/Mode.h
@@ -27,6 +27,8 @@
 #pragma once
 
 #include <AK/ByteBuffer.h>
+#include <AK/Span.h>
+#include <AK/StdLibExtras.h>
 #include <LibCrypto/Cipher/Cipher.h>
 
 namespace Crypto {
@@ -37,9 +39,8 @@ class Mode {
 public:
     virtual ~Mode() { }
 
-    // FIXME: Somehow communicate that encrypt returns the last initialization vector (if the mode supports it)
-    virtual Optional<ByteBuffer> encrypt(const ByteBuffer& in, ByteBuffer& out, Optional<ByteBuffer> ivec = {}) = 0;
-    virtual void decrypt(const ByteBuffer& in, ByteBuffer& out, Optional<ByteBuffer> ivec = {}) = 0;
+    virtual void encrypt(const ReadonlyBytes& in, Bytes& out, const Bytes& ivec = {}, Bytes* ivec_out = nullptr) = 0;
+    virtual void decrypt(const ReadonlyBytes& in, Bytes& out, const Bytes& ivec = {}) = 0;
 
     virtual size_t IV_length() const = 0;
 
@@ -58,7 +59,7 @@ public:
     T& cipher() { return m_cipher; }
 
 protected:
-    virtual void prune_padding(ByteBuffer& data)
+    virtual void prune_padding(Bytes& data)
     {
         auto size = data.size();
         switch (m_cipher.padding_mode()) {
@@ -74,7 +75,7 @@ protected:
                     return;
                 }
             }
-            data.trim(size - maybe_padding_length);
+            data = data.slice(0, size - maybe_padding_length);
             break;
         }
         case PaddingMode::RFC5246: {
@@ -86,13 +87,13 @@ protected:
                     return;
                 }
             }
-            data.trim(size - maybe_padding_length - 1);
+            data = data.slice(0, size - maybe_padding_length - 1);
             break;
         }
         case PaddingMode::Null: {
             while (data[size - 1] == 0)
                 --size;
-            data.trim(size);
+            data = data.slice(0, size);
             break;
         }
         default: