Merge pull request #1654 from cloudflare/jsnell/crypto-memory-tracker

cloudflare · Feb 10, 2024 · 18885cb · 18885cb
2 parents 15ec0a0 + e70e54c
commit 18885cb
Show file tree

Hide file tree

Showing 9 changed files with 87 additions and 0 deletions.
diff --git a/src/workerd/api/crypto-impl-aes.c++ b/src/workerd/api/crypto-impl-aes.c++
@@ -123,6 +123,13 @@ protected:
            CRYPTO_memcmp(keyData.begin(), other.begin(), keyData.size()) == 0;
   }
 
+  kj::StringPtr jsgGetMemoryName() const override { return "AesKeyBase"_kjc; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(AesKeyBase); }
+  void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const override {
+    tracker.trackFieldWithSize("keyData", keyData.size());
+    tracker.trackField("keyAlgorithm", keyAlgorithm);
+  }
+
 private:
   CryptoKey::AlgorithmVariant getAlgorithm(jsg::Lock& js) const override final { return keyAlgorithm; }
 

diff --git a/src/workerd/api/crypto-impl-asymmetric.c++ b/src/workerd/api/crypto-impl-asymmetric.c++
@@ -341,6 +341,10 @@ public:
     return false;
   }
 
+  kj::StringPtr jsgGetMemoryName() const override { return "AsymmetricKey"; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(AsymmetricKey); }
+  void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const override {}
+
 private:
   virtual SubtleCrypto::JsonWebKey exportJwk() const = 0;
   virtual kj::Array<kj::byte> exportRaw() const = 0;
@@ -533,6 +537,13 @@ public:
     : AsymmetricKey(kj::mv(keyData), keyType, extractable, usages),
       keyAlgorithm(kj::mv(keyAlgorithm)) {}
 
+  kj::StringPtr jsgGetMemoryName() const override { return "AsymmetricKey"; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(AsymmetricKey); }
+  void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const override {
+    AsymmetricKey::jsgGetMemoryInfo(tracker);
+    tracker.trackField("keyAlgorithm", keyAlgorithm);
+  }
+
 protected:
   CryptoKey::RsaKeyAlgorithm keyAlgorithm;
 
@@ -1527,6 +1538,13 @@ public:
       SubtleCrypto::GenerateKeyAlgorithm&& algorithm, bool extractable,
       CryptoKeyUsageSet privateKeyUsages, CryptoKeyUsageSet publicKeyUsages);
 
+  kj::StringPtr jsgGetMemoryName() const override { return "EllipticKey"; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(EllipticKey); }
+  void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const override {
+    AsymmetricKey::jsgGetMemoryInfo(tracker);
+    tracker.trackField("keyAlgorithm", keyAlgorithm);
+  }
+
 private:
   static kj::Array<kj::byte> bigNumToPaddedArray(const BIGNUM& n, size_t paddedLength) {
     kj::Vector<kj::byte> result(paddedLength);
@@ -2134,6 +2152,12 @@ public:
     return CryptoKey::AsymmetricKeyDetails {};
   }
 
+  kj::StringPtr jsgGetMemoryName() const override { return "EdDsaKey"; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(EdDsaKey); }
+  void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const override {
+    AsymmetricKey::jsgGetMemoryInfo(tracker);
+  }
+
 private:
   kj::StringPtr keyAlgorithm;
 

diff --git a/src/workerd/api/crypto-impl-hkdf.c++ b/src/workerd/api/crypto-impl-hkdf.c++
@@ -16,6 +16,13 @@ public:
       : CryptoKey::Impl(extractable, usages),
         keyData(kj::mv(keyData)), keyAlgorithm(kj::mv(keyAlgorithm)) {}
 
+  kj::StringPtr jsgGetMemoryName() const override { return "HkdfKey"; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(HkdfKey); }
+  void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const override {
+    tracker.trackFieldWithSize("keyData", keyData.size());
+    tracker.trackField("keyAlgorithm", keyAlgorithm);
+  }
+
 private:
   kj::Array<kj::byte> deriveBits(
       jsg::Lock& js, SubtleCrypto::DeriveKeyAlgorithm&& algorithm,

diff --git a/src/workerd/api/crypto-impl-hmac.c++ b/src/workerd/api/crypto-impl-hmac.c++
@@ -18,6 +18,13 @@ public:
       : CryptoKey::Impl(extractable, usages),
         keyData(kj::mv(keyData)), keyAlgorithm(kj::mv(keyAlgorithm)) {}
 
+  kj::StringPtr jsgGetMemoryName() const override { return "HmacKey"; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(HmacKey); }
+  void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const override {
+    tracker.trackFieldWithSize("keyData", keyData.size());
+    tracker.trackField("keyAlgorithm", keyAlgorithm);
+  }
+
 private:
   kj::Array<kj::byte> sign(
       SubtleCrypto::SignAlgorithm&& algorithm,

diff --git a/src/workerd/api/crypto-impl-pbkdf2.c++ b/src/workerd/api/crypto-impl-pbkdf2.c++
@@ -15,6 +15,13 @@ public:
       : CryptoKey::Impl(extractable, usages),
         keyData(kj::mv(keyData)), keyAlgorithm(kj::mv(keyAlgorithm)) {}
 
+  kj::StringPtr jsgGetMemoryName() const override { return "Pbkdf2Key"; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(Pbkdf2Key); }
+  void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const override {
+    tracker.trackFieldWithSize("keyData", keyData.size());
+    tracker.trackField("keyAlgorithm", keyAlgorithm);
+  }
+
 private:
   kj::Array<kj::byte> deriveBits(
       jsg::Lock& js, SubtleCrypto::DeriveKeyAlgorithm&& algorithm,

diff --git a/src/workerd/api/crypto-impl.c++ b/src/workerd/api/crypto-impl.c++
@@ -220,4 +220,7 @@ void checkPbkdfLimits(jsg::Lock& js, size_t iterations) {
   }
 }
 
+void CryptoKey::visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+  tracker.trackField("impl", impl);
+}
 }  // namespace workerd::api
diff --git a/src/workerd/api/crypto-impl.h b/src/workerd/api/crypto-impl.h
@@ -217,6 +217,10 @@ class CryptoKey::Impl {
   virtual bool equals(const Impl& other) const = 0;
   virtual bool equals(const kj::Array<kj::byte>& other) const;
 
+  virtual kj::StringPtr jsgGetMemoryName() const { return "CryptoKey::Impl"; }
+  virtual size_t jsgGetMemorySelfSize() const { return sizeof(Impl); }
+  virtual void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {}
+
 private:
   const bool extractable;
   const CryptoKeyUsageSet usages;

diff --git a/src/workerd/api/crypto.h b/src/workerd/api/crypto.h
@@ -125,6 +125,7 @@ class CryptoKey: public jsg::Object {
   struct KeyAlgorithm {
     kj::StringPtr name;
     JSG_STRUCT(name);
+    JSG_MEMORY_INFO(KeyAlgorithm) {}
   };
 
   struct AesKeyAlgorithm {
@@ -135,6 +136,8 @@ class CryptoKey: public jsg::Object {
     uint16_t length;
 
     JSG_STRUCT(name, length);
+
+    JSG_MEMORY_INFO(AesKeyAlgorithm) {}
   };
 
   struct HmacKeyAlgorithm {
@@ -151,6 +154,7 @@ class CryptoKey: public jsg::Object {
     uint16_t length;
 
     JSG_STRUCT(name, hash, length);
+    JSG_MEMORY_INFO(HmacKeyAlgorithm) {}
   };
 
   using BigInteger = kj::Array<kj::byte>;
@@ -192,6 +196,8 @@ class CryptoKey: public jsg::Object {
     }
 
     JSG_STRUCT(name, modulusLength, publicExponent, hash);
+
+    JSG_MEMORY_INFO(RsaKeyAlgorithm) {}
   };
 
   struct EllipticKeyAlgorithm {
@@ -202,6 +208,8 @@ class CryptoKey: public jsg::Object {
     kj::StringPtr namedCurve;
 
     JSG_STRUCT(name, namedCurve);
+
+    JSG_MEMORY_INFO(EllipticKeyAlgorithm) {}
   };
 
   // Catch-all that can be used for extension algorithms. Combines fields of several known types.
@@ -276,6 +284,8 @@ class CryptoKey: public jsg::Object {
   // that a timing-safe comparison is used for the key material.
   bool operator==(const CryptoKey& other) const;
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const;
+
 private:
   kj::Own<Impl> impl;
 
@@ -650,6 +660,10 @@ class DigestStream: public WritableStream {
     JSG_TS_OVERRIDE(extends WritableStream<ArrayBuffer | ArrayBufferView>);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("promise", promise);
+  }
+
 private:
   jsg::MemoizedIdentity<jsg::Promise<kj::Array<kj::byte>>> promise;
 
@@ -699,6 +713,14 @@ class Crypto: public jsg::Object {
     });
   }
 
+  void visitForGc(jsg::GcVisitor& visitor) {
+    visitor.visit(subtle);
+  }
+
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("subtle", subtle);
+  }
+
 private:
   jsg::Ref<SubtleCrypto> subtle = jsg::alloc<SubtleCrypto>();
 };

diff --git a/src/workerd/api/node/crypto-keys.c++ b/src/workerd/api/node/crypto-keys.c++
@@ -50,6 +50,12 @@ public:
     return kj::heapArray(keyData.asPtr());
   }
 
+  kj::StringPtr jsgGetMemoryName() const override { return "SecretKey"; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(SecretKey); }
+  void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const override {
+    tracker.trackFieldWithSize("keyData", keyData.size());
+  }
+
 private:
   ZeroOnFree keyData;
 };