Raise FrozenError when modifying a frozen digest

Author: wtn

ext/digest/digest.c

@@ -14,13 +14,15 @@
 ************************************************/
 
 #include "digest.h"
+#include <ruby/version.h>
 
 static VALUE rb_mDigest;
 static VALUE rb_mDigest_Instance;
 static VALUE rb_cDigest_Class;
 static VALUE rb_cDigest_Base;
 
 static ID id_reset, id_update, id_finish, id_digest, id_hexdigest, id_digest_length;
+static ID id_clone, id_freeze;
 static ID id_metadata;
 
 RUBY_EXTERN void Init_digest_base(void);
@@ -217,12 +219,23 @@ rb_digest_instance_reset(VALUE self)
  *     digest_obj.new -> another_digest_obj
  *
  * Returns a new, initialized copy of the digest object.  Equivalent
- * to digest_obj.clone().reset().
+ * to digest_obj.clone(freeze: false).reset().
  */
 static VALUE
 rb_digest_instance_new(VALUE self)
 {
-    VALUE clone = rb_obj_clone(self);
+    VALUE clone;
+
+#if RUBY_API_VERSION_MAJOR >= 3
+    /* dup() would drop the singleton class, and the copy must not stay frozen */
+    VALUE kwargs = rb_hash_new();
+    rb_hash_aset(kwargs, ID2SYM(id_freeze), Qfalse);
+    clone = rb_funcallv_kw(self, id_clone, 1, &kwargs, RB_PASS_KEYWORDS);
+#else
+    /* clone(freeze:) needs Ruby 3.0; dup() loses the singleton class but resets */
+    clone = OBJ_FROZEN(self) ? rb_obj_dup(self) : rb_obj_clone(self);
+#endif
+
     rb_funcall(clone, id_reset, 0);
     return clone;
 }
@@ -266,7 +279,11 @@ rb_digest_instance_digest(int argc, VALUE *argv, VALUE self)
 static VALUE
 rb_digest_instance_digest_bang(VALUE self)
 {
-    VALUE value = rb_funcall(self, id_finish, 0);
+    VALUE value;
+
+    rb_check_frozen(self);
+
+    value = rb_funcall(self, id_finish, 0);
     rb_funcall(self, id_reset, 0);
 
     return value;
@@ -311,7 +328,11 @@ rb_digest_instance_hexdigest(int argc, VALUE *argv, VALUE self)
 static VALUE
 rb_digest_instance_hexdigest_bang(VALUE self)
 {
-    VALUE value = rb_funcall(self, id_finish, 0);
+    VALUE value;
+
+    rb_check_frozen(self);
+
+    value = rb_funcall(self, id_finish, 0);
     rb_funcall(self, id_reset, 0);
 
     return hexencode_str_new(value);
@@ -683,6 +704,8 @@ rb_digest_base_reset(VALUE self)
     rb_digest_metadata_t *algo;
     void *pctx;
 
+    rb_check_frozen(self);
+
     algo = get_digest_obj_metadata(self);
 
     TypedData_Get_Struct(self, void, &digest_type, pctx);
@@ -705,6 +728,8 @@ rb_digest_base_update(VALUE self, VALUE str)
     rb_digest_metadata_t *algo;
     void *pctx;
 
+    rb_check_frozen(self);
+
     algo = get_digest_obj_metadata(self);
 
     TypedData_Get_Struct(self, void, &digest_type, pctx);
@@ -732,6 +757,7 @@ rb_digest_base_finish(VALUE self)
     algo->finish_func(pctx, (unsigned char *)RSTRING_PTR(str));
 
     /* avoid potential coredump caused by use of a finished context */
+    /* not frozen-checked: #digest and #hexdigest call this on a clone */
     algo_init(algo, pctx);
 
     return str;
@@ -777,6 +803,8 @@ Init_digest(void)
     id_digest          = rb_intern("digest");
     id_hexdigest       = rb_intern("hexdigest");
     id_digest_length   = rb_intern("digest_length");
+    id_clone           = rb_intern("clone");
+    id_freeze          = rb_intern("freeze");
     id_metadata        = rb_id_metadata();
     InitVM(digest);
 }

lib/digest/sha2.rb

@@ -77,6 +77,7 @@ def initialize(bitlen = 256)
     #
     # Reset the digest to the initial state and return self.
     def reset
+      raise FrozenError, "can't modify frozen #{self.class}" if frozen?
       @sha2.reset
       self
     end
@@ -87,6 +88,7 @@ def reset
     #
     # Update the digest using a given _string_ and return self.
     def update(str)
+      raise FrozenError, "can't modify frozen #{self.class}" if frozen?
       @sha2.update(str)
       self
     end

test/digest/test_digest.rb

@@ -81,6 +81,42 @@ def test_eq
     assert_equal(md1, md2, self.class::ALGO)
   end
 
+  def test_frozen
+    md = self.class::ALGO.new
+    md << "a"
+    md.freeze
+
+    assert_raise(FrozenError) { md.update("b") }
+    assert_raise(FrozenError) { md << "b" }
+    assert_raise(FrozenError) { md.reset }
+    assert_raise(FrozenError) { md.digest! }
+    assert_raise(FrozenError) { md.hexdigest! }
+    assert_raise(FrozenError) { md.base64digest! }
+    assert_raise(FrozenError) { md.digest("b") }
+    assert_raise(FrozenError) { md.hexdigest("b") }
+
+    assert_equal(self.class::ALGO.hexdigest("a"), md.hexdigest)
+  end
+
+  def test_new_keeps_singleton
+    md = self.class::ALGO.new
+    md.extend(Module.new { def extended_marker; end })
+    def md.singleton_marker; end
+
+    assert_respond_to(md.new, :extended_marker)
+    assert_respond_to(md.new, :singleton_marker)
+  end
+
+  def test_frozen_copy
+    md = self.class::ALGO.new
+    md << "a"
+    md.freeze
+
+    assert_equal(self.class::ALGO.hexdigest("a"), md.clone.hexdigest)
+    assert_equal(self.class::ALGO.hexdigest("a"), md.dup.hexdigest)
+    assert_equal(self.class::ALGO.hexdigest(""), md.new.hexdigest)
+  end
+
   def test_s_file
     Tempfile.create("test_digest_file", mode: File::BINARY) { |tmpfile|
       str = "hello, world.\r\n"
@@ -174,6 +210,32 @@ class TestSHA512 < Test::Unit::TestCase
 
   class TestSHA2 < Test::Unit::TestCase
 
+  def test_frozen
+    md = Digest::SHA2.new
+    md << "a"
+    md.freeze
+
+    assert_raise(FrozenError) { md.update("b") }
+    assert_raise(FrozenError) { md.reset }
+    assert_raise(FrozenError) { md.digest! }
+    assert_raise(FrozenError) { md.hexdigest! }
+    assert_raise(FrozenError) { md.base64digest! }
+    assert_raise(FrozenError) { md.digest("b") }
+    assert_raise(FrozenError) { md.hexdigest("b") }
+
+    assert_equal(Digest::SHA256.hexdigest("a"), md.hexdigest)
+  end
+
+  def test_frozen_copy
+    md = Digest::SHA2.new
+    md << "a"
+    md.freeze
+
+    assert_equal(Digest::SHA256.hexdigest("a"), md.clone.hexdigest)
+    assert_equal(Digest::SHA256.hexdigest("a"), md.dup.hexdigest)
+    assert_equal(Digest::SHA256.hexdigest(""), md.new.hexdigest)
+  end
+
   def test_s_file
     Tempfile.create("test_digest_file") { |tmpfile|
       str = Data1