/* **********************************************************
 * Copyright (c) 2005-2006 VMware, Inc.  All rights reserved.
 * **********************************************************/

/*
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * * Redistributions of source code must retain the above copyright notice,
 *   this list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright notice,
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 *
 * * Neither the name of VMware, Inc. nor the names of its contributors may be
 *   used to endorse or promote products derived from this software without
 *   specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL VMWARE, INC. OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
 * DAMAGE.
 */

/* DR core / shared library interface unit tests */
/* executed as part of the unit tests for processes.c */

/* need to have a core build anyway for these tests to work */
#include "win32/events.h"

/* nudge tends to hang/timeout if you nudge right after the
 *  process starts. so in order to let the basic tests pass,
 *  they all sleep for at least this long before trying to
 *  nudge.
 * of course this should be fixed and then we should create
 *  stress tests to address this specifically. */
#define NUDGE_LET_PROCESS_START_WAIT 500

#define TEST_TIMEOUT 2000
#define LAUNCH_TIMEOUT 1000

#define WAIT_FOR_APP(hProc) \
    DO_ASSERT(WAIT_OBJECT_0 == WaitForSingleObject(hProc, TEST_TIMEOUT * 2))

#define LONG_WAIT_FOR_APP(hProc) \
    DO_ASSERT(WAIT_OBJECT_0 == WaitForSingleObject(hProc, TEST_TIMEOUT * 20))

#define DO_TEST(name, appstr, block) DO_TEST_HP(name, appstr, TRUE, block)

#define TESTER_1_BLOCK                                                                   \
    "BEGIN_BLOCK\n"                                                                      \
    "APP_NAME=tester_1.exe\n"                                                            \
    "DYNAMORIO_OPTIONS=-kill_thread -kill_thread_max 1000 -report_max 0 -dumpcore_mask " \
    "0xff\n"                                                                             \
    "END_BLOCK\n"

#define TESTER_2_BLOCK                        \
    "BEGIN_BLOCK\n"                           \
    "APP_NAME=tester_2.exe\n"                 \
    "DYNAMORIO_OPTIONS=-dumpcore_mask 0xff\n" \
    "END_BLOCK\n"

#define TESTER_1_HOT_PATCH_BLOCK                           \
    "BEGIN_BLOCK\n"                                        \
    "APP_NAME=tester_1.exe\n"                              \
    "DYNAMORIO_OPTIONS=-kill_thread -dumpcore_mask 0xff\n" \
    "DYNAMORIO_HOT_PATCH_MODES=\\conf\\test-modes.cfg\n"   \
    "BEGIN_MP_MODES\n"                                     \
    "1\n"                                                  \
    "TEST.000A:2\n"                                        \
    "END_MP_MODES\n"                                       \
    "END_BLOCK\n"

#define TESTER_1_HOT_PATCH_DETECT_BLOCK                    \
    "BEGIN_BLOCK\n"                                        \
    "APP_NAME=tester_1.exe\n"                              \
    "DYNAMORIO_OPTIONS=-kill_thread -dumpcore_mask 0xff\n" \
    "DYNAMORIO_HOT_PATCH_MODES=\\conf\\test-modes.cfg\n"   \
    "BEGIN_MP_MODES\n"                                     \
    "1\n"                                                  \
    "TEST.000A:1\n"                                        \
    "END_MP_MODES\n"                                       \
    "END_BLOCK\n"

/* simple detach */
DO_TEST(detach, TESTER_1_BLOCK, {
    UINT pid;
    LAUNCH_APP(L"tester_1.exe 5000", &pid);
    /* FIXME: if this isn't here, the first run (right after
     *  building the tests) always fails. */
    Sleep(LAUNCH_TIMEOUT * 2);
    VERIFY_UNDER_DR(pid);
    CHECKED_OPERATION(detach(pid, TRUE, DETACH_RECOMMENDED_TIMEOUT));
    VERIFY_NOT_UNDER_DR(pid);
    TERMINATE_PROCESS(pid);
});

/* pending restart */
DO_TEST(pending_restart, TESTER_1_BLOCK, {
    UINT pid;
    LAUNCH_APP(L"tester_1.exe 5000", &pid);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid);
    CHECKED_OPERATION(detach(pid, TRUE, DETACH_RECOMMENDED_TIMEOUT));
    VERIFY_NOT_UNDER_DR(pid);
    DO_ASSERT(is_process_pending_restart(pid));
    DO_ASSERT(is_any_process_pending_restart());
    TERMINATE_PROCESS(pid);
});

/* detach_exe */
DO_TEST(detach_exe, TESTER_1_BLOCK TESTER_2_BLOCK, {
    UINT pid1;
    UINT pid2;
    UINT pid3;
    LAUNCH_APP(L"tester_1.exe", &pid1);
    LAUNCH_APP(L"tester_2.exe", &pid2);
    LAUNCH_APP(L"tester_1.exe", &pid3);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid1);
    VERIFY_UNDER_DR(pid2);
    VERIFY_UNDER_DR(pid3);
    CHECKED_OPERATION(detach_exe(L"tester_1.exe", DETACH_RECOMMENDED_TIMEOUT));
    VERIFY_NOT_UNDER_DR(pid1);
    VERIFY_UNDER_DR(pid2);
    VERIFY_NOT_UNDER_DR(pid3);
    DO_ASSERT(is_process_pending_restart(pid1));
    DO_ASSERT(!is_process_pending_restart(pid2));
    DO_ASSERT(is_any_process_pending_restart());
    TERMINATE_PROCESS(pid1);
    TERMINATE_PROCESS(pid2);
    TERMINATE_PROCESS(pid3);
});

/* detach_all */
DO_TEST(detach_all, TESTER_1_BLOCK TESTER_2_BLOCK, {
    UINT pid1;
    UINT pid2;
    UINT pid3;
    LAUNCH_APP(L"tester_1.exe", &pid1);
    LAUNCH_APP(L"tester_2.exe", &pid2);
    LAUNCH_APP(L"tester_1.exe", &pid3);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid1);
    VERIFY_UNDER_DR(pid2);
    VERIFY_UNDER_DR(pid3);
    CHECKED_OPERATION(detach_all(DETACH_RECOMMENDED_TIMEOUT));
    VERIFY_NOT_UNDER_DR(pid1);
    VERIFY_NOT_UNDER_DR(pid2);
    VERIFY_NOT_UNDER_DR(pid3);
    DO_ASSERT(is_process_pending_restart(pid1));
    DO_ASSERT(is_any_process_pending_restart());
    TERMINATE_PROCESS(pid1);
    TERMINATE_PROCESS(pid2);
    TERMINATE_PROCESS(pid3);
});

/* consistency_detach */
DO_TEST(consistency_detach, TESTER_1_BLOCK TESTER_2_BLOCK, {
    UINT pid1;
    UINT pid2;
    UINT pid3;
    LAUNCH_APP(L"tester_1.exe", &pid1);
    LAUNCH_APP(L"tester_2.exe", &pid2);
    LAUNCH_APP(L"tester_1.exe", &pid3);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid1);
    VERIFY_UNDER_DR(pid2);
    VERIFY_UNDER_DR(pid3);
    CHECKED_OPERATION(load_test_config(TESTER_2_BLOCK, FALSE));
    CHECKED_OPERATION(consistency_detach(DETACH_RECOMMENDED_TIMEOUT));
    VERIFY_NOT_UNDER_DR(pid1);
    VERIFY_UNDER_DR(pid2);
    VERIFY_NOT_UNDER_DR(pid3);
    TERMINATE_PROCESS(pid1);
    TERMINATE_PROCESS(pid2);
    TERMINATE_PROCESS(pid3);
});

/* check start/stop events */
DO_TEST(start_stop_event, TESTER_1_BLOCK, {
    UINT pid;
    HANDLE hProc;
    LAUNCH_APP_HANDLE(L"tester_1.exe 10", &pid, &hProc);
    WAIT_FOR_APP(hProc);
    DO_ASSERT(
        check_for_event(MSG_INFO_PROCESS_START, L"tester_1.exe", pid, NULL, NULL, 0));
    DO_ASSERT(
        check_for_event(MSG_INFO_PROCESS_STOP, L"tester_1.exe", pid, NULL, NULL, 0));
});

/* check detach event */
DO_TEST(detach_event, TESTER_1_BLOCK, {
    UINT pid;
    LAUNCH_APP(L"tester_1.exe 2000", &pid);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid);
    CHECKED_OPERATION(detach(pid, TRUE, DETACH_RECOMMENDED_TIMEOUT));
    VERIFY_NOT_UNDER_DR(pid);
    Sleep(100);
    DO_ASSERT(
        check_for_event(MSG_INFO_PROCESS_START, L"tester_1.exe", pid, NULL, NULL, 0));
    DO_ASSERT(check_for_event(MSG_INFO_DETACHING, L"tester_1.exe", pid, NULL, NULL, 0));
    TERMINATE_PROCESS(pid);
});

/* check attack event */
DO_TEST(attack_event, TESTER_1_BLOCK, {
    UINT pid;
    HANDLE hProc;
    WCHAR s3[MAX_PATH];
    WCHAR s4[MAX_PATH];
    LAUNCH_APP_HANDLE(L"tester_1.exe 10 10 1", &pid, &hProc);
    WAIT_FOR_APP(hProc);
    DO_ASSERT(check_for_event(MSG_SEC_VIOLATION_THREAD, L"tester_1.exe", pid, s3, s4,
                              MAX_PATH));
    /* make sure threat id looks ok */
    DO_ASSERT(11 == wcslen(s3));
});

/* check forensics event/file */
DO_TEST(forensics_file, TESTER_1_BLOCK, {
    UINT pid;
    HANDLE hProc;
    WCHAR s3[MAX_PATH];
    WCHAR s4[MAX_PATH];
    LAUNCH_APP_HANDLE(L"tester_1.exe 10 10 1", &pid, &hProc);
    WAIT_FOR_APP(hProc);
    DO_ASSERT(
        check_for_event(MSG_SEC_VIOLATION_THREAD, L"tester_1.exe", pid, NULL, NULL, 0));
    DO_ASSERT(check_for_event(MSG_SEC_FORENSICS, L"tester_1.exe", pid, s3, s4, MAX_PATH));
    DO_ASSERT(file_exists(s3));
});

/* stress forensics event/file */
DO_TEST(forensics_stress, TESTER_1_BLOCK, {
    UINT pid;
    HANDLE hProc;
    WCHAR s3[MAX_PATH];
    WCHAR s4[MAX_PATH];
    int i;
    LAUNCH_APP_HANDLE(L"tester_1.exe 10 10 0 100", &pid, &hProc);
    LONG_WAIT_FOR_APP(hProc);
    for (i = 0; i < 100; i++) {
        DO_ASSERT(check_for_event(MSG_SEC_VIOLATION_THREAD, L"tester_1.exe", pid, NULL,
                                  NULL, 0));
        DO_ASSERT(
            check_for_event(MSG_SEC_FORENSICS, L"tester_1.exe", pid, s3, s4, MAX_PATH));
        DO_ASSERT(file_exists(s3));
    }
    DO_ASSERT(
        check_for_event(MSG_INFO_PROCESS_STOP, L"tester_1.exe", pid, NULL, NULL, 0));
});

/* check to make sure nudge doesn't leave the code lying around */
DO_TEST(check_nudge, TESTER_1_BLOCK, {
    UINT pid;
    WCHAR nudge_code_buf[MAX_PATH];

    LAUNCH_APP(L"tester_1.exe", &pid);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid);
    Sleep(NUDGE_LET_PROCESS_START_WAIT);
    CHECKED_OPERATION(hotp_notify_modes_update(pid, TRUE, TEST_TIMEOUT));
    DO_ASSERT(ERROR_SUCCESS !=
              get_config_parameter(L_PRODUCT_NAME, FALSE, L_DYNAMORIO_VAR_NUDGE,
                                   nudge_code_buf, MAX_PATH));
    TERMINATE_PROCESS(pid);
});

/* simple test for hotpatching */
DO_TEST(hotp_protect, TESTER_1_HOT_PATCH_BLOCK, {
    UINT pid;
    char fc[MAX_PATH];
    LAUNCH_APP_AND_WAIT(L"tester_1.exe 100", &pid);
    CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
    DO_ASSERT(0 == strncmp(fc, "10", 2));
});

/* detect should report 00 */
DO_TEST(hotp_detect, TESTER_1_HOT_PATCH_DETECT_BLOCK, {
    UINT pid;
    char fc[MAX_PATH];
    LAUNCH_APP_AND_WAIT(L"tester_1.exe 100", &pid);
    CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
    DO_ASSERT(0 == strncmp(fc, "00", 2));
});

/* basic nudge test */
DO_TEST_HP(hotp_defs_nudge, TESTER_1_BLOCK, FALSE, /* don't load the modes file! */
           {
               UINT pid;
               char fc[MAX_PATH];
               HANDLE hProc;
               hotp_policy_status_table_t *status_table = NULL;

               /* first launch app w/o hotpatch */
               LAUNCH_APP_AND_WAIT(L"tester_1.exe 10", &pid);
               CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
               DO_ASSERT(0 == strncmp(fc, "00", 2));

               /* now, same thing with longer wait */
               LAUNCH_APP_HANDLE(L"tester_1.exe 2000", &pid, &hProc);
               Sleep(LAUNCH_TIMEOUT);
               VERIFY_UNDER_DR(pid);

               /* make sure nothing's there */
               DO_ASSERT(ERROR_DRMARKER_ERROR == get_hotp_status(pid, &status_table));

               /* load the new config -- this time with hot patching*/
               CHECKED_OPERATION(load_test_config(TESTER_1_HOT_PATCH_BLOCK, TRUE));

               /* and nudge */
               Sleep(NUDGE_LET_PROCESS_START_WAIT);
               CHECKED_OPERATION(hotp_notify_defs_update(pid, TRUE, TEST_TIMEOUT));
               VERIFY_UNDER_DR(pid);
               WAIT_FOR_APP(hProc);

               CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
               DO_ASSERT(0 == strncmp(fc, "10", 2));
           });

DO_TEST(hotp_modes_nudge, TESTER_1_HOT_PATCH_BLOCK, {
    UINT pid;
    char fc[MAX_PATH];
    HANDLE hProc;

    /* first launch app w/hotpatch protect */
    LAUNCH_APP_AND_WAIT(L"tester_1.exe 10", &pid);
    CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
    DO_ASSERT(0 == strncmp(fc, "10", 2));

    /* now, same thing with longer wait */
    LAUNCH_APP_HANDLE(L"tester_1.exe 2000", &pid, &hProc);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid);

    /* load the new config */
    CHECKED_OPERATION(load_test_config(TESTER_1_HOT_PATCH_DETECT_BLOCK, TRUE));

    /* and do a modes nudge */
    Sleep(NUDGE_LET_PROCESS_START_WAIT);
    CHECKED_OPERATION(hotp_notify_modes_update(pid, TRUE, TEST_TIMEOUT));
    VERIFY_UNDER_DR(pid);
    WAIT_FOR_APP(hProc);

    CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
    DO_ASSERT(0 == strncmp(fc, "00", 2));
});

/* nudge twice to make sure we don't die */
DO_TEST_HP(hotp_nudge_twice, TESTER_1_BLOCK, FALSE, {
    UINT pid;
    char fc[MAX_PATH];
    HANDLE hProc;

    /* first launch app w/o hotpatch */
    LAUNCH_APP_AND_WAIT(L"tester_1.exe 10", &pid);
    CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
    DO_ASSERT(0 == strncmp(fc, "00", 2));

    LAUNCH_APP_HANDLE(L"tester_1.exe 2000", &pid, &hProc);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid);

    /* load the new config */
    CHECKED_OPERATION(load_test_config(TESTER_1_HOT_PATCH_BLOCK, TRUE));
    Sleep(NUDGE_LET_PROCESS_START_WAIT);
    CHECKED_OPERATION(hotp_notify_defs_update(pid, TRUE, TEST_TIMEOUT));
    VERIFY_UNDER_DR(pid);

    /* load the old config back */
    CHECKED_OPERATION(load_test_config(TESTER_1_HOT_PATCH_DETECT_BLOCK, TRUE));
    CHECKED_OPERATION(hotp_notify_defs_update(pid, TRUE, TEST_TIMEOUT));
    VERIFY_UNDER_DR(pid);
    WAIT_FOR_APP(hProc);

    CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
    DO_ASSERT(0 == strncmp(fc, "00", 2));
});

/* tests state of the patches */
DO_TEST(hotp_protect_status, TESTER_1_HOT_PATCH_BLOCK, {
    UINT pid;
    HANDLE hProc;
    hotp_policy_status_table_t *status_table = NULL;

    LAUNCH_APP_HANDLE(L"tester_1.exe 10 2500", &pid, &hProc);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid);
    Sleep(500);
    CHECKED_OPERATION(get_hotp_status(pid, &status_table));
    DO_ASSERT(status_table != NULL && status_table->num_policies > 0);
    if (status_table != NULL && status_table->num_policies > 0) {
        UINT j;
        for (j = 0; j < status_table->num_policies; j++) {
            if (0 ==
                    strcmp(status_table->policy_status_array[j].policy_id, "TEST.000A") ||
                0 ==
                    strcmp(status_table->policy_status_array[j].policy_id, "TEST.000B")) {
                DO_ASSERT(status_table->policy_status_array[j].inject_status ==
                          HOTP_INJECT_PROTECT);
            }
        }
        free_hotp_status_table(status_table);
    }
    WAIT_FOR_APP(hProc);
});

/* tests state of the patches */
DO_TEST(hotp_pending_status, TESTER_1_HOT_PATCH_BLOCK, {
    UINT pid;
    HANDLE hProc;
    hotp_policy_status_table_t *status_table = NULL;

    LAUNCH_APP_HANDLE(L"tester_1.exe 2500", &pid, &hProc);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid);
    Sleep(200);
    CHECKED_OPERATION(get_hotp_status(pid, &status_table));
    DO_ASSERT(status_table != NULL && status_table->num_policies > 0);
    if (status_table != NULL && status_table->num_policies > 0) {
        UINT j;
        for (j = 0; j < status_table->num_policies; j++) {
            if (0 ==
                    strcmp(status_table->policy_status_array[j].policy_id, "TEST.000A") ||
                0 ==
                    strcmp(status_table->policy_status_array[j].policy_id, "TEST.000B")) {
                DO_ASSERT(status_table->policy_status_array[j].inject_status ==
                          HOTP_INJECT_PENDING);
            }
        }
        free_hotp_status_table(status_table);
    }
    WAIT_FOR_APP(hProc);
});

/* sanity check to make sure we have different state */
DO_TEST(hotp_detect_status, TESTER_1_HOT_PATCH_DETECT_BLOCK, {
    UINT pid;
    HANDLE hProc;
    hotp_policy_status_table_t *status_table = NULL;

    LAUNCH_APP_HANDLE(L"tester_1.exe 10 2500", &pid, &hProc);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid);
    Sleep(500);
    CHECKED_OPERATION(get_hotp_status(pid, &status_table));
    DO_ASSERT(status_table != NULL && status_table->num_policies > 0);
    if (status_table != NULL && status_table->num_policies > 0) {
        UINT j;
        for (j = 0; j < status_table->num_policies; j++) {
            if (0 ==
                    strcmp(status_table->policy_status_array[j].policy_id, "TEST.000A") ||
                0 ==
                    strcmp(status_table->policy_status_array[j].policy_id, "TEST.000B")) {
                DO_ASSERT(status_table->policy_status_array[j].inject_status ==
                          HOTP_INJECT_DETECT);
            }
        }
        free_hotp_status_table(status_table);
    }
    WAIT_FOR_APP(hProc);
});

DO_TEST(hotp_modes_nudge_all, TESTER_1_HOT_PATCH_BLOCK, {
    UINT pid;
    char fc[MAX_PATH];
    HANDLE hProc;

    /* first launch app w/hotpatch protect */
    LAUNCH_APP_AND_WAIT(L"tester_1.exe 10", &pid);
    CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
    DO_ASSERT(0 == strncmp(fc, "10", 2));

    /* now, same thing with longer wait */
    LAUNCH_APP_HANDLE(L"tester_1.exe 2000", &pid, &hProc);
    Sleep(LAUNCH_TIMEOUT);
    VERIFY_UNDER_DR(pid);

    /* load the new config */
    CHECKED_OPERATION(load_test_config(TESTER_1_HOT_PATCH_DETECT_BLOCK, TRUE));

    /* and do a modes nudge */
    Sleep(NUDGE_LET_PROCESS_START_WAIT);
    CHECKED_OPERATION(hotp_notify_all_modes_update(TEST_TIMEOUT));
    VERIFY_UNDER_DR(pid);
    WAIT_FOR_APP(hProc);

    CHECKED_OPERATION(read_file_contents(L"tester.out", fc, MAX_PATH, NULL));
    DO_ASSERT(0 == strncmp(fc, "00", 2));
});