sha3 1.0.5 → 2.1.0

data/ext/sha3/digest.c

@@ -1,10 +1,8 @@
-/* Copyright (c) 2012 - 2013 Johanns Gregorian <io+sha3@jsani.com> */
+#include "digest.h"
 
+#include "KeccakHash.h"
 #include "sha3.h"
 
-VALUE cSHA3Digest;
-VALUE eSHA3DigestError;
-
 /*
  * == Notes
  *
@@ -19,252 +17,811 @@ VALUE eSHA3DigestError;
  *
  */
 
-static void free_allox(MDX *mdx)
-{
-    if (mdx)
-    {
-        if (mdx->state)
-        {
-            free(mdx->state);
-        }
+/*** 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");
 
-        free(mdx);
+    /* 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");
+
+    if (NIL_P(_sha3_module)) {
+        // This is both a safeguard and a workaround for RDoc
+        _sha3_module = rb_define_module("SHA3");
     }
 
-    return;
+    /*
+     * Document-class: SHA3::Digest
+     *
+     * It is a subclass of the Digest::Class class, which provides a framework for
+     * creating and manipulating hash digest. Supported Algorithms are:
+     * - SHA3-224
+     * - SHA3-256
+     * - SHA3-384
+     * - SHA3-512
+     * - SHAKE128
+     * - SHAKE256
+     */
+    _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.
+     */
+    _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(_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");
 }
 
-static VALUE c_digest_alloc(VALUE klass)
-{
-    MDX *mdx;
-    VALUE obj;
+// 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!");
+    }
+}
 
-    mdx = (MDX *)malloc(sizeof(MDX));
-    if (!mdx)
-    {
-        rb_raise(eSHA3DigestError, "failed to allocate object memory");
+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);
+}
 
-    mdx->state = (Keccak_HashInstance *)malloc(sizeof(Keccak_HashInstance));
-    if (!mdx->state)
-    {
-        free_allox(mdx);
-        rb_raise(eSHA3DigestError, "failed to allocate state memory");
+int get_hashbit_length(VALUE obj, sha3_digest_algorithms* algorithm) {
+    if (TYPE(obj) == T_SYMBOL) {
+        ID symid = SYM2ID(obj);
+
+        if (symid == _sha3_224_id) {
+            *algorithm = SHA3_224;
+            return 224;
+        } else if (symid == _sha3_256_id) {
+            *algorithm = SHA3_256;
+            return 256;
+        } else if (symid == _sha3_384_id) {
+            *algorithm = SHA3_384;
+            return 384;
+        } else if (symid == _sha3_512_id) {
+            *algorithm = SHA3_512;
+            return 512;
+        } else if (symid == _shake_128_id) {
+            *algorithm = SHAKE_128;
+            return 128;
+        } else if (symid == _shake_256_id) {
+            *algorithm = SHAKE_256;
+            return 256;
+        }
+
+        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)");
     }
 
-    obj = Data_Wrap_Struct(klass, 0, free_allox, mdx);
+    rb_raise(_sha3_digest_error_class, "unknown type value");
+    return 0;  // Never reached, but silences compiler warnings
+}
 
-    memset(mdx->state, 0, sizeof(Keccak_HashInstance));
-    mdx->hashbitlen = 0;
+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(context);
+    }
+}
 
-    return obj;
+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 (context && context->state) {
+        size += sizeof(Keccak_HashInstance);
+    }
+
+    return size;
 }
 
-static VALUE c_digest_update(VALUE, VALUE);
-
-HashReturn c_keccak_hash_initialize(MDX *mdx)
-{
-    HashReturn r = KECCAK_FAIL;
-
-    switch (mdx->hashbitlen)
-    {
-    case 224:
-        r = Keccak_HashInitialize_SHA3_224(mdx->state);
-        break;
-    case 256:
-        r = Keccak_HashInitialize_SHA3_256(mdx->state);
-        break;
-    case 384:
-        r = Keccak_HashInitialize_SHA3_384(mdx->state);
-        break;
-    case 512:
-        r = Keccak_HashInitialize_SHA3_512(mdx->state);
-        break;
-    }
-
-    return r;
+static HashReturn keccak_hash_initialize(sha3_digest_context_t* context) {
+    switch (context->algorithm) {
+        case SHA3_224:
+            return Keccak_HashInitialize_SHA3_224(context->state);
+        case SHA3_256:
+            return Keccak_HashInitialize_SHA3_256(context->state);
+        case SHA3_384:
+            return Keccak_HashInitialize_SHA3_384(context->state);
+        case SHA3_512:
+            return Keccak_HashInitialize_SHA3_512(context->state);
+        case SHAKE_128:
+            return Keccak_HashInitialize_SHAKE128(context->state);
+        case SHAKE_256:
+            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;
 }
 
-// SHA3::Digest.new(type, [data]) -> self
-static VALUE c_digest_init(int argc, VALUE *argv, VALUE self)
-{
-    MDX *mdx;
+/*
+ * :call-seq:
+ *   ::new() -> instance
+ *   ::new([algorithm], [message]) -> instance
+ *
+ * Creates a new digest object.
+ *
+ * +algorithm+::
+ *   _optional_ The algorithm to use.
+ *   Valid algorithms are:
+ *   - :sha3_224
+ *   - :sha3_256
+ *   - :sha3_384
+ *   - :sha3_512
+ *   - :shake_128
+ *   - :shake_256
+ *
+ * +message+::
+ *   _optional_ The message to hash.
+ *
+ * = example
+ *   SHA3::Digest.new(:sha3_256)
+ *   SHA3::Digest.new(:shake_128, "initial data")
+ */
+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);
-    GETMDX(self, mdx);
+    get_sha3_digest_context(self, &context);
 
-    if (!NIL_P(hlen))
-    {
-        mdx->hashbitlen = get_hlen(hlen);
-    }
-    else
-    {
-        mdx->hashbitlen = 256;
+    if (NIL_P(hlen)) {
+        context->algorithm = SHA3_256;
+        context->hashbitlen = 256;
+    } else {
+        context->hashbitlen = get_hashbit_length(hlen, &context->algorithm);
     }
 
-    if (c_keccak_hash_initialize(mdx) != KECCAK_SUCCESS)
-    {
-        rb_raise(eSHA3DigestError, "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 c_digest_update(self, data);
+    if (!NIL_P(data)) {
+        return rb_sha3_digest_update(self, data);
     }
 
     return self;
 }
 
-// SHA3::Digest.update(data) -> self
-static VALUE c_digest_update(VALUE self, VALUE data)
-{
-    MDX *mdx;
+/*
+ * :call-seq:
+ *   update(string) -> digest
+ *
+ * Updates the digest with the given string.
+ *
+ * +string+::
+ *   The string to update the digest with.
+ *
+ * = example
+ *   digest.update("more data")
+ *   digest << "more data"  # alias for update
+ */
+static VALUE rb_sha3_digest_update(VALUE self, VALUE data) {
+    sha3_digest_context_t* context;
     BitLength dlen;
 
     StringValue(data);
-    GETMDX(self, mdx);
+    get_sha3_digest_context(self, &context);
+
+    // Check for empty data
+    if (RSTRING_LEN(data) == 0) {
+        return self;
+    }
+
+    // Check for NULL data pointer
+    if (RSTRING_PTR(data) == NULL) {
+        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(eSHA3DigestError, "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;
 }
 
-// SHA3::Digest.reset() -> self
-static VALUE c_digest_reset(VALUE self)
-{
-    MDX *mdx;
-
-    GETMDX(self, mdx);
+/*
+ * :call-seq:
+ *   reset -> digest
+ *
+ * Resets the digest to its initial state.
+ *
+ * = example
+ *   digest.reset
+ */
+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 (c_keccak_hash_initialize(mdx) != KECCAK_SUCCESS)
-    {
-        rb_raise(eSHA3DigestError, "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(MDX *mdx1, MDX *mdx2)
-{
-    return (
-        (mdx1->hashbitlen == mdx2->hashbitlen) &&
-        (strcmp((const char *)mdx1->state->sponge.state, (const char *)mdx2->state->sponge.state) == 0) &&
-        (mdx1->state->sponge.rate == mdx2->state->sponge.rate) &&
-        (mdx1->state->sponge.byteIOIndex == mdx2->state->sponge.byteIOIndex) &&
-        (mdx1->state->sponge.squeezing == mdx2->state->sponge.squeezing) &&
-        (mdx1->state->fixedOutputLength == mdx2->state->fixedOutputLength) &&
-        (mdx1->state->delimitedSuffix == mdx2->state->delimitedSuffix));
+static int compare_contexts(const sha3_digest_context_t* context1, const sha3_digest_context_t* context2) {
+    // First check the hashbitlen and algorithm
+    if (context1->hashbitlen != context2->hashbitlen || context1->algorithm != context2->algorithm) {
+        return 0;
+    }
+
+    // Compare the internal state structure
+    if (memcmp(&(context1->state->sponge.state), &(context2->state->sponge.state),
+               sizeof(context1->state->sponge.state)) != 0) {
+        return 0;
+    }
+
+    // Compare sponge parameters
+    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 ((context1->state->fixedOutputLength != context2->state->fixedOutputLength) ||
+        (context1->state->delimitedSuffix != context2->state->delimitedSuffix)) {
+        return 0;
+    }
+
+    // All comparisons passed
+    return 1;
 }
 
-// SHA3::Digest.copy(obj) -> self
-static VALUE c_digest_copy(VALUE self, VALUE obj)
-{
-    MDX *mdx1, *mdx2;
+/*
+ * :call-seq:
+ *   initialize_copy(other) -> digest
+ *
+ * Initializes the digest with the state of another digest.
+ *
+ * +other+::
+ *   The digest to copy the state from.
+ *
+ * = example
+ *   new_digest = digest.dup
+ */
+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;
     }
 
-    GETMDX(self, mdx1);
-    SAFEGETMDX(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;
+    safe_get_sha3_digest_context(other, &other_context);
+    get_sha3_digest_context(self, &context);
 
-    // Fetch the data again to make sure it was copied
-    GETMDX(self, mdx1);
-    SAFEGETMDX(obj, mdx2);
+    context->hashbitlen = other_context->hashbitlen;
+    context->algorithm = other_context->algorithm;
+    memcpy(context->state, other_context->state, sizeof(Keccak_HashInstance));
 
-    if (!cmp_states(mdx1, mdx2))
-    {
-        rb_raise(eSHA3DigestError, "failed to copy state");
+    if (!compare_contexts(context, other_context)) {
+        rb_raise(_sha3_digest_error_class, "failed to copy state");
     }
 
     return self;
 }
 
-// SHA3::Digest.digest_length -> Integer
-static VALUE c_digest_length(VALUE self)
-{
-    MDX *mdx;
-    GETMDX(self, mdx);
+/*
+ * :call-seq:
+ *   length -> Integer
+ *
+ * Returns the length of the digest in bytes.
+ *
+ * = example
+ *   digest.length  #=> 32 for SHA3-256
+ */
+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);
 }
 
-// SHA3::Digest.block_length -> Integer
-static VALUE c_digest_block_length(VALUE self)
-{
-    MDX *mdx;
-    GETMDX(self, mdx);
+/*
+ * :call-seq:
+ *   block_length -> Integer
+ *
+ * Returns the block length of the algorithm in bytes.
+ *
+ * = example
+ *   digest.block_length
+ */
+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)));
 }
 
-// SHA3::Digest.name -> String
-static VALUE c_digest_name(VALUE self)
-{
-    return rb_str_new2("SHA3");
+/*
+ * :call-seq:
+ *   name -> String
+ *
+ * Returns the name of the algorithm.
+ *
+ * = example
+ *   digest.name  #=> "SHA3-256"
+ */
+static VALUE rb_sha3_digest_name(VALUE self) {
+    sha3_digest_context_t* context;
+    get_sha3_digest_context(self, &context);
+
+    switch (context->algorithm) {
+        case SHA3_224:
+            return rb_str_new2("SHA3-224");
+        case SHA3_256:
+            return rb_str_new2("SHA3-256");
+        case SHA3_384:
+            return rb_str_new2("SHA3-384");
+        case SHA3_512:
+            return rb_str_new2("SHA3-512");
+        case SHAKE_128:
+            return rb_str_new2("SHAKE128");
+        case SHAKE_256:
+            return rb_str_new2("SHAKE256");
+        default:
+            rb_raise(_sha3_digest_error_class, "unknown algorithm");
+    }
 }
 
-// SHA3::Digest.finish() -> String
-static VALUE c_digest_finish(int argc, VALUE *argv, VALUE self)
-{
-    MDX *mdx;
+/*
+ * :call-seq:
+ *   finish([message]) -> String
+ *
+ * Returns the final digest as a binary string.
+ *
+ * +message+::
+ *   _optional_ Update state with additional data before finalizing.
+ *
+ * = example
+ *   digest.finish
+ *   digest.finish("final chunk")
+ */
+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);
-    GETMDX(self, mdx);
+    get_sha3_digest_context(self, &context);
 
-    if (NIL_P(str))
-    {
-        str = rb_str_new(0, mdx->hashbitlen / 8);
-    }
-    else
-    {
+    // For both SHA3 and SHAKE algorithms, use the security strength (hashbitlen)
+    // as the default output length
+    digest_bytes = context->hashbitlen / 8;
+
+    if (NIL_P(str)) {
+        str = rb_str_new(0, digest_bytes);
+    } else {
         StringValue(str);
-        rb_str_resize(str, mdx->hashbitlen / 8);
+        rb_str_resize(str, digest_bytes);
     }
 
-    if (Keccak_HashFinal(mdx->state, (BitSequence *)RSTRING_PTR(str)) != KECCAK_SUCCESS)
-    {
-        rb_raise(eSHA3DigestError, "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;
 }
 
-void Init_sha3_n_digest()
-{
-    rb_require("digest");
+/*
+ * :call-seq:
+ *   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.
+ *
+ * = example
+ *   digest.squeeze(32)  # Get 32 bytes of output
+ */
+static VALUE rb_sha3_digest_squeeze(VALUE self, VALUE length) {
+    sha3_digest_context_t* context;
+    VALUE str, copy;
+    int output_bytes;
+
+    Check_Type(length, T_FIXNUM);
+    output_bytes = NUM2INT(length);
+
+    if (output_bytes <= 0) {
+        rb_raise(_sha3_digest_error_class, "output length must be positive");
+    }
+
+    get_sha3_digest_context(self, &context);
+
+    // Only SHAKE algorithms support arbitrary-length output
+    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 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(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(context_copy->state, (BitSequence*)RSTRING_PTR(str), output_bytes * 8) != KECCAK_SUCCESS) {
+        rb_raise(_sha3_digest_error_class, "failed to squeeze output");
+    }
+
+    return str;
+}
+
+/*
+ * :call-seq:
+ *   hex_squeeze(length) -> String
+ *
+ * Returns the hexadecimal representation of the squeezed output. Only available for SHAKE
+ * algorithms.
+ *
+ * +length+::
+ *   The length in bytes of the output to squeeze.
+ *
+ * = example
+ *   digest.hex_squeeze(32)  # Get 64 hex characters (32 bytes)
+ */
+static VALUE rb_sha3_digest_hex_squeeze(VALUE self, VALUE length) {
+    // Get the binary output using the existing squeeze function
+    VALUE bin_str = rb_sha3_digest_squeeze(self, length);
+    // Use Ruby's built-in unpack method to convert to hex
+    return rb_funcall(bin_str, rb_intern("unpack1"), 1, rb_str_new2("H*"));
+}
+
+/*
+ * :call-seq:
+ *   digest() -> string
+ *   digest([data]) -> string
+ *   digest(length) -> string
+ *   digest(length, data) -> string
+ *
+ * Returns the binary representation of the digest.
+ *
+ * +length+::
+ *   The length of the output to squeeze when using SHAKE algorithms.
+ *   This parameter is required for SHAKE algorithms.
+ *
+ * +data+::
+ *   _optional_ Update state with additional data before returning digest.
+ *
+ * = example
+ *   digest.digest()
+ *   digest.digest('compute me')
+ *   digest.digest(12)  # For SHAKE algorithms
+ *   digest.digest(12, 'compute me')  # For SHAKE algorithms
+ */
+static VALUE rb_sha3_digest_digest(int argc, VALUE* argv, VALUE self) {
+    sha3_digest_context_t* context;
+    get_sha3_digest_context(self, &context);
+
+    if (context->algorithm != SHAKE_128 && context->algorithm != SHAKE_256) {
+        return rb_call_super(argc, argv);
+    }
+
+    VALUE length, data;
+    rb_scan_args(argc, argv, "02", &length, &data);
+
+    // For SHAKE algorithms
+    if (NIL_P(length)) {
+        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_sha3_digest_update(self, data);
+    }
+
+    return rb_sha3_digest_squeeze(self, length);
+}
+
+/*
+ * :call-seq:
+ *   hexdigest() -> string
+ *   hexdigest([data]) -> string
+ *   hexdigest(length) -> string
+ *   hexdigest(length, data) -> string
+ *
+ * Returns the hexadecimal representation of the digest.
+ *
+ * +length+::
+ *   The length of the output to squeeze when using SHAKE algorithms.
+ *   This parameter is required for SHAKE algorithms.
+ *
+ * +data+::
+ *   _optional_ Update state with additional data before returning digest.
+ *
+ * = example
+ *   digest.hexdigest()
+ *   digest.hexdigest('compute me')
+ *   digest.hexdigest(12)  # For SHAKE algorithms
+ *   digest.hexdigest(12, 'compute me')  # For SHAKE algorithms
+ */
+static VALUE rb_sha3_digest_hexdigest(int argc, VALUE* argv, VALUE self) {
+    sha3_digest_context_t* context;
+    get_sha3_digest_context(self, &context);
+
+    if (context->algorithm != SHAKE_128 && context->algorithm != SHAKE_256) {
+        return rb_call_super(argc, argv);
+    }
+
+    VALUE length, data;
+    rb_scan_args(argc, argv, "02", &length, &data);
+
+    if (NIL_P(length)) {
+        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_sha3_digest_update(self, data);
+    }
+
+    return rb_sha3_digest_hex_squeeze(self, length);
+}
+
+/*
+ * :call-seq:
+ *   SHA3::Digest.digest(name, data) -> string
+ *
+ * Returns the binary digest of the given +data+ using the algorithm specified by +name+.
+ *
+ * +name+::
+ *   The hash algorithm to use (as a Symbol).
+ *   Valid algorithms are:
+ *   - :sha3_224
+ *   - :sha3_256
+ *   - :sha3_384
+ *   - :sha3_512
+ *   - :shake_128
+ *   - :shake_256
+ *
+ * +data+::
+ *   The data to hash.
+ *
+ * = example
+ *   SHA3::Digest.digest(:sha3_256, "data to hash")
+ *
+ * = note
+ * 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_sha3_digest_self_digest(VALUE klass, VALUE name, VALUE data) {
+    VALUE args[2];
+
+    // 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) {
+            /* Create a new digest instance with the specified algorithm */
+            VALUE digest = rb_class_new_instance(1, &name, klass);
+
+            /* Update it with the 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 */
+
+            /* Return the squeezed output */
+            return rb_sha3_digest_squeeze(digest, INT2NUM(output_length));
+        }
+    }
+
+    /* Call the superclass method with arguments in reverse order */
+    args[0] = data;
+    args[1] = name;
+
+    return rb_call_super(2, args);
+}
+
+/*
+ * :call-seq:
+ *   SHA3::Digest.hexdigest(name, data) -> string
+ *
+ * Returns the hexadecimal representation of the given +data+ using the algorithm specified by
+ * +name+.
+ *
+ * +name+::
+ *   The hash algorithm to use (as a Symbol).
+ *   Valid algorithms are:
+ *   - :sha3_224
+ *   - :sha3_256
+ *   - :sha3_384
+ *   - :sha3_512
+ *   - :shake_128
+ *   - :shake_256
+ *
+ * +data+::
+ *   The data to hash.
+ *
+ * = example
+ *   SHA3::Digest.hexdigest(:sha3_256, "data to hash")
+ *
+ * = note
+ * 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_sha3_digest_self_hexdigest(VALUE klass, VALUE name, VALUE data) {
+    VALUE args[2];
+
+    // 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) {
+            /* Create a new digest instance with the specified algorithm */
+            VALUE digest = rb_class_new_instance(1, &name, klass);
+
+            /* Update it with the 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 */
+
+            /* Return the hexadecimal representation of the squeezed output */
+            return rb_sha3_digest_hex_squeeze(digest, INT2NUM(output_length));
+        }
+    }
+
+    /* Call the superclass method with arguments in reverse order */
+    args[0] = data;
+    args[1] = name;
 
-    /* SHA3::Digest (class) */
-    cSHA3Digest = rb_define_class_under(mSHA3, "Digest", rb_path2class("Digest::Class"));
-    /* SHA3::Digest::DigestError (class) */
-    eSHA3DigestError = rb_define_class_under(cSHA3Digest, "DigestError", rb_eStandardError);
-
-    // SHA3::Digest (class) methods
-    rb_define_alloc_func(cSHA3Digest, c_digest_alloc);
-    rb_define_method(cSHA3Digest, "initialize", c_digest_init, -1);
-    rb_define_method(cSHA3Digest, "update", c_digest_update, 1);
-    rb_define_method(cSHA3Digest, "reset", c_digest_reset, 0);
-    rb_define_method(cSHA3Digest, "initialize_copy", c_digest_copy, 1);
-    rb_define_method(cSHA3Digest, "digest_length", c_digest_length, 0);
-    rb_define_method(cSHA3Digest, "block_length", c_digest_block_length, 0);
-    rb_define_method(cSHA3Digest, "name", c_digest_name, 0);
-    rb_define_private_method(cSHA3Digest, "finish", c_digest_finish, -1);
-
-    rb_define_alias(cSHA3Digest, "<<", "update");
-
-    return;
+    return rb_call_super(2, args);
 }