Raise FrozenError when modifying a frozen digest

Author: wtn

ext/digest/digest.c

@@ -266,7 +266,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 +315,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 +691,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 +715,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 +744,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;

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_raise(FrozenError) { md.new }
+  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_raise(FrozenError) { md.new }
+  end
+
   def test_s_file
     Tempfile.create("test_digest_file") { |tmpfile|
       str = Data1