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

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

static VALUE contrast_marshal_module_load(const int argc,
                                                 const VALUE *argv) {
    VALUE result;
    VALUE source_string;

    // Our patches only need only apply in the case where there was valid input.
    if (argc >= 1) {
        source_string = argv[0];
    } else {
        source_string = Qnil;
    }

    // Run our protect code ahead of the original method
    if (source_string != Qnil) {
        rb_funcall(marshal_propagator, rb_sym_protect_marshal_load, 1, source_string);
    }

    // Invoke the original method
    result = rb_call_super(argc, argv);

    // Run our assess code after the original method
    if (source_string != Qnil) {
        VALUE tracked =
            rb_funcall(properties_hash, rb_sym_hash_tracked, 1, source_string);

        // Assuming the source is tracked and needs assess checks
        if (tracked == Qtrue) {
            VALUE skip =
                rb_funcall(contrast_patcher(), rb_sym_skip_assess_analysis, 0);
            // And Assess is enabled and applies to this request
            if (skip == Qfalse) {
                rb_funcall(marshal_propagator, rb_sym_assess_marshal_load, 2,
                           source_string, result);
            }
        }
    }
    return result;
}

void Init_cs__assess_marshal_module(void) {
    // Contrast::Agent::Assess::Tracker::PROPERTIES_HASH
    VALUE tracker = rb_define_class_under(assess, "Tracker", rb_cObject);
    properties_hash = rb_const_get(tracker, rb_intern("PROPERTIES_HASH"));
    marshal_propagator =
        rb_define_class_under(core_assess, "MarshalPropagator", rb_cObject);
    rb_sym_assess_marshal_load = rb_intern("cs__load_assess");
    rb_sym_protect_marshal_load = rb_intern("cs__load_protect");

    contrast_register_singleton_prepend_patch(
        "Marshal", "load", &contrast_marshal_module_load);
}