sha3 2.0.0 → 2.1.0

data/ext/sha3/digest.c

@@ -1,5 +1,8 @@
 #include "digest.h"
 
+#include "KeccakHash.h"
+#include "sha3.h"
+
 /*
  * == Notes
  *
@@ -7,181 +10,255 @@
  *    | .alloc() ->
  *    | .new() ->
  *    | .update() ->
- *    | .digest or .hexdigest or .inspect -> (Instance.digest or .hexdigest())
- * ->
+ *    | .digest or .hexdigest or .inspect -> (Instance.digest or .hexdigest()) ->
  *    --| .alloc() ->
  *      | .copy() ->
  *      | .finish() ->
  *
  */
 
-/*
- * SHA3 module
- */
+/*** Types and structs  ***/
+
+typedef enum { SHA3_224 = 0, SHA3_256, SHA3_384, SHA3_512, SHAKE_128, SHAKE_256 } sha3_digest_algorithms;
+
+typedef struct {
+    Keccak_HashInstance* state;
+    int hashbitlen;
+    sha3_digest_algorithms algorithm;
+} sha3_digest_context_t;
+
+typedef HashReturn (*keccak_init_func)(Keccak_HashInstance*);
+
+/*** Function prototypes ***/
+
+static int compare_contexts(const sha3_digest_context_t*, const sha3_digest_context_t*);
+static inline void get_sha3_digest_context(VALUE, sha3_digest_context_t**);
+static inline void safe_get_sha3_digest_context(VALUE, sha3_digest_context_t**);
+
+static int get_hashbit_length(VALUE, sha3_digest_algorithms*);
+static HashReturn keccak_hash_initialize(sha3_digest_context_t*);
+
+static void sha3_digest_free_context(void*);
+static size_t sha3_digest_context_size(const void*);
+
+/* Allocation and initialization */
+static VALUE rb_sha3_digest_alloc(VALUE);
+static VALUE rb_sha3_digest_init(int, VALUE*, VALUE);
+static VALUE rb_sha3_digest_copy(VALUE, VALUE);
+
+/* Core digest operations */
+static VALUE rb_sha3_digest_finish(int, VALUE*, VALUE);
+static VALUE rb_sha3_digest_reset(VALUE);
+static VALUE rb_sha3_digest_update(VALUE, VALUE);
+
+/* Digest properties */
+static VALUE rb_sha3_digest_block_length(VALUE);
+static VALUE rb_sha3_digest_length(VALUE);
+static VALUE rb_sha3_digest_name(VALUE);
+
+/* Output methods */
+static VALUE rb_sha3_digest_digest(int, VALUE*, VALUE);
+static VALUE rb_sha3_digest_hexdigest(int, VALUE*, VALUE);
+static VALUE rb_sha3_digest_hex_squeeze(VALUE, VALUE);
+static VALUE rb_sha3_digest_squeeze(VALUE, VALUE);
+static VALUE rb_sha3_digest_self_digest(VALUE, VALUE, VALUE);
+static VALUE rb_sha3_digest_self_hexdigest(VALUE, VALUE, VALUE);
+
+/*** Globals variables  ***/
+
+VALUE _sha3_digest_class;
+VALUE _sha3_digest_error_class;
+
+/* Define the ID variables */
+static ID _sha3_224_id;
+static ID _sha3_256_id;
+static ID _sha3_384_id;
+static ID _sha3_512_id;
+static ID _shake_128_id;
+static ID _shake_256_id;
+
+/* TypedData structure for sha3_digest_context_t */
+const rb_data_type_t sha3_digest_data_type_t = {"SHA3::Digest",
+                                                {
+                                                    NULL,
+                                                    sha3_digest_free_context,
+                                                    sha3_digest_context_size,
+                                                },
+                                                NULL,
+                                                NULL,
+                                                RUBY_TYPED_FREE_IMMEDIATELY};
+
 void Init_sha3_digest(void) {
     rb_require("digest");
 
     /* Initialize static symbol IDs for faster lookup in get_hlen() */
-    sha3_224_id = rb_intern("sha3_224");
-    sha3_256_id = rb_intern("sha3_256");
-    sha3_384_id = rb_intern("sha3_384");
-    sha3_512_id = rb_intern("sha3_512");
-    shake_128_id = rb_intern("shake_128");
-    shake_256_id = rb_intern("shake_256");
+    _sha3_224_id = rb_intern("sha3_224");
+    _sha3_256_id = rb_intern("sha3_256");
+    _sha3_384_id = rb_intern("sha3_384");
+    _sha3_512_id = rb_intern("sha3_512");
+    _shake_128_id = rb_intern("shake_128");
+    _shake_256_id = rb_intern("shake_256");
 
-    /*
-     * Document-module: SHA3
-     *
-     * This hosts the SHA3::Digest classes.
-     */
-    sha3_module = rb_define_module("SHA3");
+    if (NIL_P(_sha3_module)) {
+        // This is both a safeguard and a workaround for RDoc
+        _sha3_module = rb_define_module("SHA3");
+    }
 
     /*
      * Document-class: SHA3::Digest
      *
      * It is a subclass of the Digest::Class class, which provides a framework for
-     * creating and manipulating hash digests.
+     * creating and manipulating hash digest. Supported Algorithms are:
+     * - SHA3-224
+     * - SHA3-256
+     * - SHA3-384
+     * - SHA3-512
+     * - SHAKE128
+     * - SHAKE256
      */
-    digest_class = rb_define_class_under(sha3_module, "Digest", rb_path2class("Digest::Class"));
-
-    /*
-     * Default-const: SHA3::VERSION
-     *
-     * It is the version of the SHA3 module.
-     */
-    rb_define_const(sha3_module, "VERSION", rb_str_new2("2.0.0"));
+    _sha3_digest_class = rb_define_class_under(_sha3_module, "Digest", rb_path2class("Digest::Class"));
 
     /*
      * Document-class: SHA3::Digest::DigestError
      *
+     * All SHA3::Digest methods raise this exception on error.
+     *
      * It is a subclass of the StandardError class -- see the Ruby documentation
      * for more information.
      */
-    digest_error_class = rb_define_class_under(digest_class, "DigestError", rb_eStandardError);
-
-    rb_define_alloc_func(digest_class, rb_digest_alloc);
-    rb_define_method(digest_class, "initialize", rb_digest_init, -1);
-    rb_define_method(digest_class, "update", rb_digest_update, 1);
-    rb_define_method(digest_class, "reset", rb_digest_reset, 0);
-    rb_define_method(digest_class, "initialize_copy", rb_digest_copy, 1);
-    rb_define_method(digest_class, "digest_length", rb_digest_length, 0);
-    rb_define_method(digest_class, "block_length", rb_digest_block_length, 0);
-    rb_define_method(digest_class, "name", rb_digest_name, 0);
-    rb_define_method(digest_class, "squeeze", rb_digest_squeeze, 1);
-    rb_define_method(digest_class, "hex_squeeze", rb_digest_hex_squeeze, 1);
-    rb_define_method(digest_class, "digest", rb_digest_digest, -1);
-    rb_define_method(digest_class, "hexdigest", rb_digest_hexdigest, -1);
-    rb_define_private_method(digest_class, "finish", rb_digest_finish, -1);
+    _sha3_digest_error_class = rb_define_class_under(_sha3_digest_class, "DigestError", rb_eStandardError);
+
+    rb_define_alloc_func(_sha3_digest_class, rb_sha3_digest_alloc);
+    rb_define_method(_sha3_digest_class, "initialize", rb_sha3_digest_init, -1);
+    rb_define_method(_sha3_digest_class, "update", rb_sha3_digest_update, 1);
+    rb_define_method(_sha3_digest_class, "reset", rb_sha3_digest_reset, 0);
+    rb_define_method(_sha3_digest_class, "initialize_copy", rb_sha3_digest_copy, 1);
+    rb_define_method(_sha3_digest_class, "digest_length", rb_sha3_digest_length, 0);
+    rb_define_method(_sha3_digest_class, "block_length", rb_sha3_digest_block_length, 0);
+    rb_define_method(_sha3_digest_class, "name", rb_sha3_digest_name, 0);
+
+    rb_define_method(_sha3_digest_class, "squeeze", rb_sha3_digest_squeeze, 1);
+    rb_define_method(_sha3_digest_class, "hex_squeeze", rb_sha3_digest_hex_squeeze, 1);
+    rb_define_method(_sha3_digest_class, "digest", rb_sha3_digest_digest, -1);
+    rb_define_method(_sha3_digest_class, "hexdigest", rb_sha3_digest_hexdigest, -1);
+
+    rb_define_private_method(_sha3_digest_class, "finish", rb_sha3_digest_finish, -1);
 
     /* Define the class method self.digest */
-    rb_define_singleton_method(digest_class, "digest", rb_digest_self_digest, 2);
-    rb_define_singleton_method(digest_class, "hexdigest", rb_digest_self_hexdigest, 2);
-    rb_define_alias(digest_class, "<<", "update");
+    rb_define_singleton_method(_sha3_digest_class, "digest", rb_sha3_digest_self_digest, 2);
+    rb_define_singleton_method(_sha3_digest_class, "hexdigest", rb_sha3_digest_self_hexdigest, 2);
+
+    rb_define_alias(_sha3_digest_class, "<<", "update");
 }
 
-int get_hlen(VALUE obj, algorithm_type* algorithm) {
+// Static inline functions replacing macros
+static inline void get_sha3_digest_context(VALUE obj, sha3_digest_context_t** context) {
+    TypedData_Get_Struct((obj), sha3_digest_context_t, &sha3_digest_data_type_t, (*context));
+    if (!(*context)) {
+        rb_raise(rb_eRuntimeError, "Digest data not initialized!");
+    }
+}
+
+static inline void safe_get_sha3_digest_context(VALUE obj, sha3_digest_context_t** context) {
+    if (!rb_obj_is_kind_of(obj, _sha3_digest_class)) {
+        rb_raise(rb_eTypeError, "wrong argument (%s)! (expected %s)", rb_obj_classname(obj),
+                 rb_class2name(_sha3_digest_class));
+    }
+    get_sha3_digest_context(obj, context);
+}
+
+int get_hashbit_length(VALUE obj, sha3_digest_algorithms* algorithm) {
     if (TYPE(obj) == T_SYMBOL) {
         ID symid = SYM2ID(obj);
 
-        if (symid == sha3_224_id) {
+        if (symid == _sha3_224_id) {
             *algorithm = SHA3_224;
             return 224;
-        } else if (symid == sha3_256_id) {
+        } else if (symid == _sha3_256_id) {
             *algorithm = SHA3_256;
             return 256;
-        } else if (symid == sha3_384_id) {
+        } else if (symid == _sha3_384_id) {
             *algorithm = SHA3_384;
             return 384;
-        } else if (symid == sha3_512_id) {
+        } else if (symid == _sha3_512_id) {
             *algorithm = SHA3_512;
             return 512;
-        } else if (symid == shake_128_id) {
+        } else if (symid == _shake_128_id) {
             *algorithm = SHAKE_128;
             return 128;
-        } else if (symid == shake_256_id) {
+        } else if (symid == _shake_256_id) {
             *algorithm = SHAKE_256;
             return 256;
         }
 
-        rb_raise(digest_error_class,
+        rb_raise(_sha3_digest_error_class,
                  "invalid hash algorithm symbol (should be: :sha3_224, "
                  ":sha3_256, :sha3_384, :sha3_512, :shake_128, or :shake_256)");
     }
 
-    rb_raise(digest_error_class, "unknown type value");
+    rb_raise(_sha3_digest_error_class, "unknown type value");
     return 0;  // Never reached, but silences compiler warnings
 }
 
-static void mdx_free(void* ptr) {
-    MDX* mdx = (MDX*)ptr;
-    if (mdx) {
-        if (mdx->state) {
-            free(mdx->state);
+static void sha3_digest_free_context(void* ptr) {
+    sha3_digest_context_t* context = (sha3_digest_context_t*)ptr;
+    if (context) {
+        if (context->state) {
+            free(context->state);
         }
-        free(mdx);
+        free(context);
     }
 }
 
-static size_t mdx_memsize(const void* ptr) {
-    const MDX* mdx = (const MDX*)ptr;
-    size_t size = sizeof(MDX);
+static size_t sha3_digest_context_size(const void* ptr) {
+    const sha3_digest_context_t* context = (const sha3_digest_context_t*)ptr;
+    size_t size = sizeof(sha3_digest_context_t);
 
-    if (mdx && mdx->state) {
+    if (context && context->state) {
         size += sizeof(Keccak_HashInstance);
     }
 
     return size;
 }
 
-/* TypedData structure for MDX */
-const rb_data_type_t mdx_type = {"SHA3::Digest",
-                                 {
-                                     NULL,
-                                     mdx_free,
-                                     mdx_memsize,
-                                 },
-                                 NULL,
-                                 NULL,
-                                 RUBY_TYPED_FREE_IMMEDIATELY};
-
-static VALUE rb_digest_alloc(VALUE klass) {
-    MDX* mdx = (MDX*)malloc(sizeof(MDX));
-    if (!mdx) {
-        rb_raise(digest_error_class, "failed to allocate object memory");
-    }
-
-    mdx->state = (Keccak_HashInstance*)calloc(1, sizeof(Keccak_HashInstance));
-    if (!mdx->state) {
-        mdx_free(mdx);
-        rb_raise(digest_error_class, "failed to allocate state memory");
-    }
-
-    VALUE obj = TypedData_Wrap_Struct(klass, &mdx_type, mdx);
-    mdx->hashbitlen = 0;
-    mdx->algorithm = SHA3_256;  // Default algorithm
-
-    return obj;
-}
-
-HashReturn keccak_hash_initialize(MDX* mdx) {
-    switch (mdx->algorithm) {
+static HashReturn keccak_hash_initialize(sha3_digest_context_t* context) {
+    switch (context->algorithm) {
         case SHA3_224:
-            return Keccak_HashInitialize_SHA3_224(mdx->state);
+            return Keccak_HashInitialize_SHA3_224(context->state);
         case SHA3_256:
-            return Keccak_HashInitialize_SHA3_256(mdx->state);
+            return Keccak_HashInitialize_SHA3_256(context->state);
         case SHA3_384:
-            return Keccak_HashInitialize_SHA3_384(mdx->state);
+            return Keccak_HashInitialize_SHA3_384(context->state);
         case SHA3_512:
-            return Keccak_HashInitialize_SHA3_512(mdx->state);
+            return Keccak_HashInitialize_SHA3_512(context->state);
         case SHAKE_128:
-            return Keccak_HashInitialize_SHAKE128(mdx->state);
+            return Keccak_HashInitialize_SHAKE128(context->state);
         case SHAKE_256:
-            return Keccak_HashInitialize_SHAKE256(mdx->state);
+            return Keccak_HashInitialize_SHAKE256(context->state);
     }
 
     return KECCAK_FAIL;
 }
 
+static VALUE rb_sha3_digest_alloc(VALUE klass) {
+    sha3_digest_context_t* context = (sha3_digest_context_t*)malloc(sizeof(sha3_digest_context_t));
+    if (!context) {
+        rb_raise(_sha3_digest_error_class, "failed to allocate object memory");
+    }
+
+    context->state = (Keccak_HashInstance*)calloc(1, sizeof(Keccak_HashInstance));
+    if (!context->state) {
+        sha3_digest_free_context(context);
+        rb_raise(_sha3_digest_error_class, "failed to allocate state memory");
+    }
+
+    VALUE obj = TypedData_Wrap_Struct(klass, &sha3_digest_data_type_t, context);
+    context->hashbitlen = 0;
+    context->algorithm = SHA3_256;  // Default algorithm
+
+    return obj;
+}
+
 /*
  * :call-seq:
  *   ::new() -> instance
@@ -206,26 +283,26 @@ HashReturn keccak_hash_initialize(MDX* mdx) {
  *   SHA3::Digest.new(:sha3_256)
  *   SHA3::Digest.new(:shake_128, "initial data")
  */
-static VALUE rb_digest_init(int argc, VALUE* argv, VALUE self) {
-    MDX* mdx;
+static VALUE rb_sha3_digest_init(int argc, VALUE* argv, VALUE self) {
+    sha3_digest_context_t* context;
     VALUE hlen, data;
 
     rb_scan_args(argc, argv, "02", &hlen, &data);
-    get_mdx(self, &mdx);
+    get_sha3_digest_context(self, &context);
 
     if (NIL_P(hlen)) {
-        mdx->algorithm = SHA3_256;
-        mdx->hashbitlen = 256;
+        context->algorithm = SHA3_256;
+        context->hashbitlen = 256;
     } else {
-        mdx->hashbitlen = get_hlen(hlen, &mdx->algorithm);
+        context->hashbitlen = get_hashbit_length(hlen, &context->algorithm);
     }
 
-    if (keccak_hash_initialize(mdx) != KECCAK_SUCCESS) {
-        rb_raise(digest_error_class, "failed to initialize algorithm state");
+    if (keccak_hash_initialize(context) != KECCAK_SUCCESS) {
+        rb_raise(_sha3_digest_error_class, "failed to initialize algorithm state");
     }
 
     if (!NIL_P(data)) {
-        return rb_digest_update(self, data);
+        return rb_sha3_digest_update(self, data);
     }
 
     return self;
@@ -244,12 +321,12 @@ static VALUE rb_digest_init(int argc, VALUE* argv, VALUE self) {
  *   digest.update("more data")
  *   digest << "more data"  # alias for update
  */
-static VALUE rb_digest_update(VALUE self, VALUE data) {
-    MDX* mdx;
+static VALUE rb_sha3_digest_update(VALUE self, VALUE data) {
+    sha3_digest_context_t* context;
     BitLength dlen;
 
     StringValue(data);
-    get_mdx(self, &mdx);
+    get_sha3_digest_context(self, &context);
 
     // Check for empty data
     if (RSTRING_LEN(data) == 0) {
@@ -258,13 +335,13 @@ static VALUE rb_digest_update(VALUE self, VALUE data) {
 
     // Check for NULL data pointer
     if (RSTRING_PTR(data) == NULL) {
-        rb_raise(digest_error_class, "cannot update with NULL data");
+        rb_raise(_sha3_digest_error_class, "cannot update with NULL data");
     }
 
     dlen = (RSTRING_LEN(data) * 8);
 
-    if (Keccak_HashUpdate(mdx->state, (BitSequence*)RSTRING_PTR(data), dlen) != KECCAK_SUCCESS) {
-        rb_raise(digest_error_class, "failed to update hash data");
+    if (Keccak_HashUpdate(context->state, (BitSequence*)RSTRING_PTR(data), dlen) != KECCAK_SUCCESS) {
+        rb_raise(_sha3_digest_error_class, "failed to update hash data");
     }
 
     return self;
@@ -279,41 +356,41 @@ static VALUE rb_digest_update(VALUE self, VALUE data) {
  * = example
  *   digest.reset
  */
-static VALUE rb_digest_reset(VALUE self) {
-    MDX* mdx;
-    get_mdx(self, &mdx);
+static VALUE rb_sha3_digest_reset(VALUE self) {
+    sha3_digest_context_t* context;
+    get_sha3_digest_context(self, &context);
 
-    memset(mdx->state, 0, sizeof(Keccak_HashInstance));
+    memset(context->state, 0, sizeof(Keccak_HashInstance));
 
-    if (keccak_hash_initialize(mdx) != KECCAK_SUCCESS) {
-        rb_raise(digest_error_class, "failed to reset internal state");
+    if (keccak_hash_initialize(context) != KECCAK_SUCCESS) {
+        rb_raise(_sha3_digest_error_class, "failed to reset internal state");
     }
 
     return self;
 }
 
-static int cmp_states(const MDX* mdx1, const MDX* mdx2) {
+static int compare_contexts(const sha3_digest_context_t* context1, const sha3_digest_context_t* context2) {
     // First check the hashbitlen and algorithm
-    if (mdx1->hashbitlen != mdx2->hashbitlen || mdx1->algorithm != mdx2->algorithm) {
+    if (context1->hashbitlen != context2->hashbitlen || context1->algorithm != context2->algorithm) {
         return 0;
     }
 
     // Compare the internal state structure
-    if (memcmp(&(mdx1->state->sponge.state), &(mdx2->state->sponge.state),
-               sizeof(mdx1->state->sponge.state)) != 0) {
+    if (memcmp(&(context1->state->sponge.state), &(context2->state->sponge.state),
+               sizeof(context1->state->sponge.state)) != 0) {
         return 0;
     }
 
     // Compare sponge parameters
-    if ((mdx1->state->sponge.rate != mdx2->state->sponge.rate) ||
-        (mdx1->state->sponge.byteIOIndex != mdx2->state->sponge.byteIOIndex) ||
-        (mdx1->state->sponge.squeezing != mdx2->state->sponge.squeezing)) {
+    if ((context1->state->sponge.rate != context2->state->sponge.rate) ||
+        (context1->state->sponge.byteIOIndex != context2->state->sponge.byteIOIndex) ||
+        (context1->state->sponge.squeezing != context2->state->sponge.squeezing)) {
         return 0;
     }
 
     // Compare hash-specific parameters
-    if ((mdx1->state->fixedOutputLength != mdx2->state->fixedOutputLength) ||
-        (mdx1->state->delimitedSuffix != mdx2->state->delimitedSuffix)) {
+    if ((context1->state->fixedOutputLength != context2->state->fixedOutputLength) ||
+        (context1->state->delimitedSuffix != context2->state->delimitedSuffix)) {
         return 0;
     }
 
@@ -333,23 +410,29 @@ static int cmp_states(const MDX* mdx1, const MDX* mdx2) {
  * = example
  *   new_digest = digest.dup
  */
-static VALUE rb_digest_copy(VALUE self, VALUE obj) {
-    MDX *mdx1, *mdx2;
+static VALUE rb_sha3_digest_copy(VALUE self, VALUE other) {
+    sha3_digest_context_t* context;
+    sha3_digest_context_t* other_context;
 
     rb_check_frozen(self);
-    if (self == obj) {
+    if (self == other) {
         return self;
     }
 
-    get_mdx(self, &mdx1);
-    safe_get_mdx(obj, &mdx2);
+    if (!rb_obj_is_kind_of(other, _sha3_digest_class)) {
+        rb_raise(rb_eTypeError, "wrong argument (%s)! (expected %s)", rb_obj_classname(other),
+                 rb_class2name(_sha3_digest_class));
+    }
 
-    memcpy(mdx1->state, mdx2->state, sizeof(Keccak_HashInstance));
-    mdx1->hashbitlen = mdx2->hashbitlen;
-    mdx1->algorithm = mdx2->algorithm;
+    safe_get_sha3_digest_context(other, &other_context);
+    get_sha3_digest_context(self, &context);
 
-    if (!cmp_states(mdx1, mdx2)) {
-        rb_raise(digest_error_class, "failed to copy state");
+    context->hashbitlen = other_context->hashbitlen;
+    context->algorithm = other_context->algorithm;
+    memcpy(context->state, other_context->state, sizeof(Keccak_HashInstance));
+
+    if (!compare_contexts(context, other_context)) {
+        rb_raise(_sha3_digest_error_class, "failed to copy state");
     }
 
     return self;
@@ -364,11 +447,11 @@ static VALUE rb_digest_copy(VALUE self, VALUE obj) {
  * = example
  *   digest.length  #=> 32 for SHA3-256
  */
-static VALUE rb_digest_length(VALUE self) {
-    MDX* mdx;
-    get_mdx(self, &mdx);
+static VALUE rb_sha3_digest_length(VALUE self) {
+    sha3_digest_context_t* context;
+    get_sha3_digest_context(self, &context);
 
-    return ULL2NUM(mdx->hashbitlen / 8);
+    return ULL2NUM(context->hashbitlen / 8);
 }
 
 /*
@@ -380,11 +463,11 @@ static VALUE rb_digest_length(VALUE self) {
  * = example
  *   digest.block_length
  */
-static VALUE rb_digest_block_length(VALUE self) {
-    MDX* mdx;
-    get_mdx(self, &mdx);
+static VALUE rb_sha3_digest_block_length(VALUE self) {
+    sha3_digest_context_t* context;
+    get_sha3_digest_context(self, &context);
 
-    return ULL2NUM(200 - (2 * (mdx->hashbitlen / 8)));
+    return ULL2NUM(200 - (2 * (context->hashbitlen / 8)));
 }
 
 /*
@@ -396,11 +479,11 @@ static VALUE rb_digest_block_length(VALUE self) {
  * = example
  *   digest.name  #=> "SHA3-256"
  */
-static VALUE rb_digest_name(VALUE self) {
-    MDX* mdx;
-    get_mdx(self, &mdx);
+static VALUE rb_sha3_digest_name(VALUE self) {
+    sha3_digest_context_t* context;
+    get_sha3_digest_context(self, &context);
 
-    switch (mdx->algorithm) {
+    switch (context->algorithm) {
         case SHA3_224:
             return rb_str_new2("SHA3-224");
         case SHA3_256:
@@ -414,7 +497,7 @@ static VALUE rb_digest_name(VALUE self) {
         case SHAKE_256:
             return rb_str_new2("SHAKE256");
         default:
-            rb_raise(digest_error_class, "unknown algorithm");
+            rb_raise(_sha3_digest_error_class, "unknown algorithm");
     }
 }
 
@@ -431,17 +514,17 @@ static VALUE rb_digest_name(VALUE self) {
  *   digest.finish
  *   digest.finish("final chunk")
  */
-static VALUE rb_digest_finish(int argc, VALUE* argv, VALUE self) {
-    MDX* mdx;
+static VALUE rb_sha3_digest_finish(int argc, VALUE* argv, VALUE self) {
+    sha3_digest_context_t* context;
     VALUE str;
     int digest_bytes;
 
     rb_scan_args(argc, argv, "01", &str);
-    get_mdx(self, &mdx);
+    get_sha3_digest_context(self, &context);
 
     // For both SHA3 and SHAKE algorithms, use the security strength (hashbitlen)
     // as the default output length
-    digest_bytes = mdx->hashbitlen / 8;
+    digest_bytes = context->hashbitlen / 8;
 
     if (NIL_P(str)) {
         str = rb_str_new(0, digest_bytes);
@@ -450,8 +533,8 @@ static VALUE rb_digest_finish(int argc, VALUE* argv, VALUE self) {
         rb_str_resize(str, digest_bytes);
     }
 
-    if (Keccak_HashFinal(mdx->state, (BitSequence*)RSTRING_PTR(str)) != KECCAK_SUCCESS) {
-        rb_raise(digest_error_class, "failed to finalize digest");
+    if (Keccak_HashFinal(context->state, (BitSequence*)RSTRING_PTR(str)) != KECCAK_SUCCESS) {
+        rb_raise(_sha3_digest_error_class, "failed to finalize digest");
     }
 
     return str;
@@ -462,6 +545,7 @@ static VALUE rb_digest_finish(int argc, VALUE* argv, VALUE self) {
  *   squeeze(length) -> String
  *
  * Returns the squeezed output as a binary string. Only available for SHAKE algorithms.
+ * This method creates a copy of the current instance to preserve the original state.
  *
  * +length+::
  *   The length in bytes of the output to squeeze.
@@ -469,8 +553,8 @@ static VALUE rb_digest_finish(int argc, VALUE* argv, VALUE self) {
  * = example
  *   digest.squeeze(32)  # Get 32 bytes of output
  */
-static VALUE rb_digest_squeeze(VALUE self, VALUE length) {
-    MDX* mdx;
+static VALUE rb_sha3_digest_squeeze(VALUE self, VALUE length) {
+    sha3_digest_context_t* context;
     VALUE str, copy;
     int output_bytes;
 
@@ -478,38 +562,35 @@ static VALUE rb_digest_squeeze(VALUE self, VALUE length) {
     output_bytes = NUM2INT(length);
 
     if (output_bytes <= 0) {
-        rb_raise(digest_error_class, "output length must be positive");
+        rb_raise(_sha3_digest_error_class, "output length must be positive");
     }
 
-    get_mdx(self, &mdx);
+    get_sha3_digest_context(self, &context);
 
     // Only SHAKE algorithms support arbitrary-length output
-    if (mdx->algorithm != SHAKE_128 && mdx->algorithm != SHAKE_256) {
-        rb_raise(digest_error_class, "squeeze is only supported for SHAKE algorithms");
+    if (context->algorithm != SHAKE_128 && context->algorithm != SHAKE_256) {
+        rb_raise(_sha3_digest_error_class, "squeeze is only supported for SHAKE algorithms");
     }
 
     // Create a copy of the digest object to avoid modifying the original
     copy = rb_obj_clone(self);
 
-    // Get the MDX struct from the copy
-    MDX* mdx_copy;
-    get_mdx(copy, &mdx_copy);
+    // Get the sha3_digest_context_t struct from the copy
+    sha3_digest_context_t* context_copy;
+    get_sha3_digest_context(copy, &context_copy);
 
     str = rb_str_new(0, output_bytes);
 
     // Finalize the hash on the copy
-    if (Keccak_HashFinal(mdx_copy->state, NULL) != KECCAK_SUCCESS) {
-        rb_raise(digest_error_class, "failed to finalize digest");
+    if (Keccak_HashFinal(context_copy->state, NULL) != KECCAK_SUCCESS) {
+        rb_raise(_sha3_digest_error_class, "failed to finalize digest");
     }
 
     // Then squeeze out the desired number of bytes
-    if (Keccak_HashSqueeze(mdx_copy->state, (BitSequence*)RSTRING_PTR(str), output_bytes * 8) !=
-        KECCAK_SUCCESS) {
-        rb_raise(digest_error_class, "failed to squeeze output");
+    if (Keccak_HashSqueeze(context_copy->state, (BitSequence*)RSTRING_PTR(str), output_bytes * 8) != KECCAK_SUCCESS) {
+        rb_raise(_sha3_digest_error_class, "failed to squeeze output");
     }
 
-    // NOTE: We don't need the copy anymore...Ruby's GC will handle freeing it
-
     return str;
 }
 
@@ -526,17 +607,11 @@ static VALUE rb_digest_squeeze(VALUE self, VALUE length) {
  * = example
  *   digest.hex_squeeze(32)  # Get 64 hex characters (32 bytes)
  */
-static VALUE rb_digest_hex_squeeze(VALUE self, VALUE length) {
-    VALUE bin_str, result_array;
-
+static VALUE rb_sha3_digest_hex_squeeze(VALUE self, VALUE length) {
     // Get the binary output using the existing squeeze function
-    bin_str = rb_digest_squeeze(self, length);
-
+    VALUE bin_str = rb_sha3_digest_squeeze(self, length);
     // Use Ruby's built-in unpack method to convert to hex
-    result_array = rb_funcall(bin_str, rb_intern("unpack"), 1, rb_str_new2("H*"));
-
-    // Extract the first element from the array
-    return rb_ary_entry(result_array, 0);
+    return rb_funcall(bin_str, rb_intern("unpack1"), 1, rb_str_new2("H*"));
 }
 
 /*
@@ -561,11 +636,11 @@ static VALUE rb_digest_hex_squeeze(VALUE self, VALUE length) {
  *   digest.digest(12)  # For SHAKE algorithms
  *   digest.digest(12, 'compute me')  # For SHAKE algorithms
  */
-static VALUE rb_digest_digest(int argc, VALUE* argv, VALUE self) {
-    MDX* mdx;
-    get_mdx(self, &mdx);
+static VALUE rb_sha3_digest_digest(int argc, VALUE* argv, VALUE self) {
+    sha3_digest_context_t* context;
+    get_sha3_digest_context(self, &context);
 
-    if (mdx->algorithm != SHAKE_128 && mdx->algorithm != SHAKE_256) {
+    if (context->algorithm != SHAKE_128 && context->algorithm != SHAKE_256) {
         return rb_call_super(argc, argv);
     }
 
@@ -574,15 +649,18 @@ static VALUE rb_digest_digest(int argc, VALUE* argv, VALUE self) {
 
     // For SHAKE algorithms
     if (NIL_P(length)) {
-        rb_raise(digest_error_class, "output length must be specified for SHAKE algorithms");
+        rb_raise(_sha3_digest_error_class, "output length must be specified for SHAKE algorithms");
     }
 
+    // Add type checking for length
+    Check_Type(length, T_FIXNUM);
+
     // If data is provided, update the state before squeezing
     if (!NIL_P(data)) {
-        rb_digest_update(self, data);
+        rb_sha3_digest_update(self, data);
     }
 
-    return rb_digest_squeeze(self, length);
+    return rb_sha3_digest_squeeze(self, length);
 }
 
 /*
@@ -607,11 +685,11 @@ static VALUE rb_digest_digest(int argc, VALUE* argv, VALUE self) {
  *   digest.hexdigest(12)  # For SHAKE algorithms
  *   digest.hexdigest(12, 'compute me')  # For SHAKE algorithms
  */
-static VALUE rb_digest_hexdigest(int argc, VALUE* argv, VALUE self) {
-    MDX* mdx;
-    get_mdx(self, &mdx);
+static VALUE rb_sha3_digest_hexdigest(int argc, VALUE* argv, VALUE self) {
+    sha3_digest_context_t* context;
+    get_sha3_digest_context(self, &context);
 
-    if (mdx->algorithm != SHAKE_128 && mdx->algorithm != SHAKE_256) {
+    if (context->algorithm != SHAKE_128 && context->algorithm != SHAKE_256) {
         return rb_call_super(argc, argv);
     }
 
@@ -619,15 +697,18 @@ static VALUE rb_digest_hexdigest(int argc, VALUE* argv, VALUE self) {
     rb_scan_args(argc, argv, "02", &length, &data);
 
     if (NIL_P(length)) {
-        rb_raise(digest_error_class, "output length must be specified for SHAKE algorithms");
+        rb_raise(_sha3_digest_error_class, "output length must be specified for SHAKE algorithms");
     }
 
+    // Add type checking for length
+    Check_Type(length, T_FIXNUM);
+
     // If data is provided, update the state before squeezing
     if (!NIL_P(data)) {
-        rb_digest_update(self, data);
+        rb_sha3_digest_update(self, data);
     }
 
-    return rb_digest_hex_squeeze(self, length);
+    return rb_sha3_digest_hex_squeeze(self, length);
 }
 
 /*
@@ -656,25 +737,27 @@ static VALUE rb_digest_hexdigest(int argc, VALUE* argv, VALUE self) {
  * This method defaults to squeezing 16 bytes for SHAKE128 and 32 bytes for SHAKE256.
  * To squeeze a different length, use #squeeze instance method.
  */
-static VALUE rb_digest_self_digest(VALUE klass, VALUE name, VALUE data) {
+static VALUE rb_sha3_digest_self_digest(VALUE klass, VALUE name, VALUE data) {
     VALUE args[2];
-    algorithm_type algorithm;
+
+    // Need to add type checking for the data parameter
+    StringValue(data);
 
     /* For SHAKE algorithms, we need to handle them differently */
     if (TYPE(name) == T_SYMBOL) {
         ID symid = SYM2ID(name);
-        if (symid == shake_128_id || symid == shake_256_id) {
+        if (symid == _shake_128_id || symid == _shake_256_id) {
             /* Create a new digest instance with the specified algorithm */
             VALUE digest = rb_class_new_instance(1, &name, klass);
 
             /* Update it with the data */
-            rb_digest_update(digest, data);
+            rb_sha3_digest_update(digest, data);
 
             /* For SHAKE algorithms, use a default output length based on the security strength */
-            int output_length = (symid == shake_128_id) ? 16 : 32; /* 128/8 or 256/8 */
+            int output_length = (symid == _shake_128_id) ? 16 : 32; /* 128/8 or 256/8 */
 
             /* Return the squeezed output */
-            return rb_digest_squeeze(digest, INT2NUM(output_length));
+            return rb_sha3_digest_squeeze(digest, INT2NUM(output_length));
         }
     }
 
@@ -712,25 +795,27 @@ static VALUE rb_digest_self_digest(VALUE klass, VALUE name, VALUE data) {
  * This method defaults to squeezing 16 bytes for SHAKE128 and 32 bytes for SHAKE256.
  * To squeeze a different length, use #hex_squeeze instance method.
  */
-static VALUE rb_digest_self_hexdigest(VALUE klass, VALUE name, VALUE data) {
+static VALUE rb_sha3_digest_self_hexdigest(VALUE klass, VALUE name, VALUE data) {
     VALUE args[2];
-    algorithm_type algorithm;
+
+    // Need to add type checking for the data parameter
+    StringValue(data);
 
     /* For SHAKE algorithms, we need to handle them differently */
     if (TYPE(name) == T_SYMBOL) {
         ID symid = SYM2ID(name);
-        if (symid == shake_128_id || symid == shake_256_id) {
+        if (symid == _shake_128_id || symid == _shake_256_id) {
             /* Create a new digest instance with the specified algorithm */
             VALUE digest = rb_class_new_instance(1, &name, klass);
 
             /* Update it with the data */
-            rb_digest_update(digest, data);
+            rb_sha3_digest_update(digest, data);
 
             /* For SHAKE algorithms, use a default output length based on the security strength */
-            int output_length = (symid == shake_128_id) ? 16 : 32; /* 128/8 or 256/8 */
+            int output_length = (symid == _shake_128_id) ? 16 : 32; /* 128/8 or 256/8 */
 
             /* Return the hexadecimal representation of the squeezed output */
-            return rb_digest_hex_squeeze(digest, INT2NUM(output_length));
+            return rb_sha3_digest_hex_squeeze(digest, INT2NUM(output_length));
         }
     }