/* Copyright (c) 2022 Contrast Security, Inc.  See
 * https://www.contrastsecurity.com/enduser-terms-0317a for more details. */

#include "cs__assess_hash.h"
#include "../cs__common/cs__common.h"
#include <ruby.h>

/* Hashes can be constructed thusly):
 *   irb(main):001:0> Hash[:a, :b]
 *   => {:a=>:b}
 *
 * This method instruments that unique bracket-construction style
 * of initializing a hash.
 */
static VALUE contrast_assess_hash_bracket_constructor(const int argc,
                                                      VALUE *argv,
                                                      const VALUE hash) {
    VALUE result;

    /* Array of Arrays: Hash[ [ [key, value], ... ] ]   -> new_hash */
    if (RB_TYPE_P(argv[0], T_ARRAY)) {
        int i;
        for (i = 0; i < argc; i++) {
            argv[i] = rb_funcall(hash_propagator,
                                 rb_sym_assess_hash_dup_and_freeze, 1, argv[i]);
        }
        /* Hash[ key, value, ... ]         -> new_hash */
    } else if (argc > 1) {
        int i;
        for (i = 0; i < argc; i += 2) {
            argv[i] = rb_funcall(hash_propagator,
                                 rb_sym_assess_hash_dup_and_freeze, 1, argv[i]);
        }
    }

    const VALUE *argv_final = argv;
    /* unhandled case - shouldn't need it since issue is only unfrozen
     * String keys
     * # Hash[ object ]                  -> new_hash
     */
    result = rb_funcall2(hash, rb_sym_assess_hash_bracket_constructor, argc,
                         argv_final);

    return result;
}

/* Hashes, when keyed with a string, will dup & freeze that string.
 * This is resource-efficient, but inconvenient for instrumentation.
 */
static VALUE contrast_assess_hash_bracket_set(const int argc, VALUE *argv,
                                              const VALUE hash) {
    VALUE result;
    /* Current name (assess_hash_bracket_set).
     * It doesn't set anything on the hash.
     * It takes the arg that /would/ have been the key, and preemptively
     * calls #dup and then #freeze, and then gives you that key.
     *
     * We intentionally don't enter Contrast scope for this patch.
     * #dup instruments the string, and #freeze gets the hash to accept
     * the key directly, without calling its own #dup/#freeze.
     * (That naturally happens in C-land, our instrumentation is in Ruby,
     * so our patches to #dup don't take effect within Hash#[]= unless we
     * specifically do this instrumentation.
     * We haven't revisited this approach since we started more extensively
     * hooking public C functions.)
     */
    if (argc > 0) {
        argv[0] = rb_funcall(hash_propagator, rb_sym_assess_hash_dup_and_freeze,
                             1, argv[0]);
    }
    /* This is the underlying assignment, w/ our instrumented key. */
    result = rb_funcall2(hash, rb_sym_assess_hash_bracket_equals, argc, argv);

    return result;
}

void Init_cs__assess_hash(void) {
    hash_propagator =
        rb_define_class_under(core_assess, "HashPropagator", rb_cObject);
    rb_sym_assess_hash_dup_and_freeze = rb_intern("cs__duplicate_and_freeze");

    VALUE hash_class = rb_define_class("Hash", rb_cObject);

    rb_sym_assess_hash_bracket_constructor = contrast_register_singleton_patch(
        "Hash", "[]", contrast_assess_hash_bracket_constructor);

    rb_sym_assess_hash_bracket_equals = contrast_register_patch(
        "Hash", "[]=", contrast_assess_hash_bracket_set);
}