/* Copyright (c) 2022 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>

/*
 * Marshal#load source:
 * static VALUE
 * marshal_load(int argc, VALUE *argv)
 *  {
 *     VALUE port, proc;
 *
 *     rb_check_arity(argc, 1, 2);
 *     port = argv[0];
 *     proc = argc > 1 ? argv[1] : Qnil;
 *     return rb_marshal_load_with_proc(port, proc);
 * }
 */
static VALUE contrast_marshal_module_prepend(const int argc, const VALUE *argv) {
    VALUE result;
    VALUE source_string;

    /* Our patches need to apply only 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");

    /* check if prepended called form the after_load_patcher */
    VALUE marshal = rb_const_get(rb_cObject, rb_intern("Marshal"));
    VALUE marshal_load = ID2SYM(rb_intern("load"));
    VALUE is_prepended = contrast_check_prepended(marshal, marshal_load, Qfalse);

    if (is_prepended == Qtrue) {
        contrast_register_singleton_prepend_patch("Marshal", "load",
                                              &contrast_marshal_module_prepend);
    }
}