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

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

static VALUE contrast_protect_fork(const int argc, const VALUE *argv,
                                   const VALUE self) {
    VALUE ret;
    if (rb_block_given_p()) {
        /* We call our hook, but it's a little complicated.
         * We wrap the fork block with our own lambda in 
         * order to instrument it.  There are no public
         * methods in the Ruby C API to set the prevailing
         * block, so we have to use rb_funcall_with_block.
         * Also, rb_funcall_with_block does a public call,
         * and our method is private.
         * So we (as a hack) temporarily set it to public.
         */
        VALUE wrapper;
        wrapper =
            rb_funcall_with_block(kernel_protect, rb_sym_protect_kernel_wrapper, 0,
                                   NULL, rb_block_proc());
        rb_funcall(rb_mKernel, rb_intern("public"), 1, ID2SYM(rb_sym_protect_kernel_fork));
        ret = rb_funcall_with_block(self, rb_sym_protect_kernel_fork, argc,
                                    argv, wrapper);
        rb_funcall(rb_mKernel, rb_intern("private"), 1, ID2SYM(rb_sym_protect_kernel_fork));
    } else {
        ret = rb_funcall2(self, rb_sym_protect_kernel_fork, argc, argv);
    }
    return ret;
}

void Init_cs__protect_kernel(void) {
    VALUE core_protect = rb_define_module_under(core_extensions, "Protect");
    kernel_protect = rb_define_module_under(core_protect, "Kernel");
    rb_sym_protect_kernel_wrapper = rb_intern("build_wrapper");

    rb_sym_protect_kernel_fork = contrast_register_patch("Kernel",
                                                         "fork",
                                                         &contrast_protect_fork);

    rb_sym_protect_kernel_fork = contrast_register_singleton_patch("Kernel",
                                                                   "fork",
                                                                   &contrast_protect_fork);

}