#ifdef WIN32
#pragma warning(disable:4996)
#include <windows.h>
#else
#include <pthread.h>
#endif
#include "test_common.h"

// Количество приватных объектов на токене не должно превышать 256,
// количество публичных объектов на токене тоже не должно превышать 256,
// количество одновременно открытых сессий не должно превышать 256,
// а количество запускаемых потоков ограничено операционной системой,
// поэтому задавать здесь произвольно большие значения не рекомендуется.
#define THREADCNT    4  // количество потоков (максимум зависит от системы)
#define NUMSESS      3  // количество сессий в потоке
#define NUMOBJ       5  // максимальное количество объектов в сессии

CK_RV
open_session_and_login(CK_SLOT_ID SlotID, CK_SESSION_HANDLE *phSession)
{
    CK_FLAGS          flags;
    CK_RV             rc;

    flags = CKF_SERIAL_SESSION | CKF_RW_SESSION;
    rc    = funcs->C_OpenSession(
        SlotID, flags, NULL, NULL, phSession);

    rc |= funcs->C_Login(*phSession, CKU_USER,
                        user_pin, (CK_ULONG)strlen((char *)user_pin));

    return rc;
}

int create_token_objects(int id, int is)
{
    CK_FLAGS          flags;
    CK_SESSION_HANDLE h_session;
    CK_RV             rc;
    static CK_BYTE           ltrue      = TRUE;
    static CK_BYTE           lfalse     = FALSE;
    static CK_OBJECT_CLASS   data_class = CKO_DATA;
    static CK_UTF8CHAR       label[32];
    static CK_UTF8CHAR           *app = (CK_UTF8CHAR *)"pkcs11_library";
	CK_ULONG app_len = (CK_ULONG)strlen((char *)app) + 1;
    static CK_BYTE           data_value[]
        = "AaBbCcDdEeFfGgHhIiJjKkLlMmNnOoPpQqRrSsTtUuVvWwXxYyZz";
    CK_ATTRIBUTE      data_attribs[] =
    {
        { CKA_CLASS,       &data_class, sizeof(data_class) },
        { CKA_APPLICATION, NULL,        0    },	// 1
        { CKA_TOKEN,       &ltrue,      sizeof(ltrue)      },
        { CKA_VALUE,       &data_value, sizeof(data_value) },
        { CKA_PRIVATE,     &lfalse,     sizeof(lfalse)     }
    };
    CK_ATTRIBUTE      label_attribs[] =
    {
        { CKA_LABEL, label, 0 },
    };
    CK_OBJECT_HANDLE  obj_list[NUMOBJ];
    CK_ULONG          i     = 0;
    CK_ULONG          count = NUMOBJ;

	data_attribs[1].pValue = app;
	data_attribs[1].ulValueLen = app_len;

    flags = CKF_SERIAL_SESSION | CKF_RW_SESSION;
    rc    = funcs->C_OpenSession(
        SlotId, flags, NULL, NULL, &h_session);
    if (rc != CKR_OK)
    {
        show_error("   C_OpenSession", rc);
        goto done;
    }
    fprintf(stderr, "Thread %d: Session %d opened\n", id, is);

    // В каждой сессии мы выполняем операции создания,
    // модификации и уничтожения случайного количества
    // объектов токена, чтобы создать конкурентную обстановку
    // между потоками.


    rc = funcs->C_GenerateRandom(h_session,
                                 (CK_BYTE *) &count, sizeof(count));
    if (rc != CKR_OK)
    {
        show_error("   C_GenerateRandom", rc);
        goto done;
    }
    count = count % NUMOBJ + 1;

    for (i = 0; i < count; i++)
    {
        rc = funcs->C_CreateObject(h_session,
                                   data_attribs, sizeof(data_attribs) / sizeof(CK_ATTRIBUTE),
                                   &obj_list[i]);
        if (rc != CKR_OK)
        {
            show_error("   C_CreateObject", rc);
            goto done;
        }

        sprintf((char *)label, "data_%d_%d_%ld", id, is, i + 1);
        label_attribs[0].ulValueLen = (CK_ULONG)strlen((char *)label) + 1;

        rc = funcs->C_SetAttributeValue(h_session, obj_list[i],
                                        label_attribs, sizeof(label_attribs) / sizeof(CK_ATTRIBUTE));
        if (rc != CKR_OK)
        {
            show_error("   C_SetAttributeValue", rc);
            goto done;
        }

        fprintf(stderr, "Object %s created\n", label);
    }
    fprintf(stderr,
            "Thread %d: %ld objects of session %d created\n",
            id, count, is);

    for (i = 0; i < count; i++)
    {
        rc = funcs->C_DestroyObject(h_session, obj_list[i]);
        if (rc != CKR_OK)
        {
            show_error("   C_DestroyObject", rc);
            goto done;
        }
        sprintf((char *)label, "data_%d_%d_%ld", id, is, i + 1);
        fprintf(stderr, "Object %s destroyed\n", label);
    }
    fprintf(stderr,
            "Thread %d: %ld objects of %d destroyed\n",
            id, count, is);

 done:

    funcs->C_CloseSession(h_session);

    fprintf(stderr, "Session %d closed\n", id);

    return rc;
}

CK_RV test_gost3410(int id, int is)
{
    CK_RV                  rc = CKR_OK;
    CK_BYTE                value[256];
    CK_ULONG               len;
    CK_SESSION_HANDLE      sess;
    CK_OBJECT_HANDLE       pub_key            = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE       priv_key           = CK_INVALID_HANDLE;
//    CK_OBJECT_HANDLE       new_key            = CK_INVALID_HANDLE;
    CK_MECHANISM           mechanism_desc     = { CKM_GOSTR3410, NULL, 0 };
    CK_MECHANISM_PTR       mechanism          = &mechanism_desc;
    CK_MECHANISM           mechanism_gen_desc = { CKM_GOSTR3410_KEY_PAIR_GEN, NULL, 0 };
    CK_MECHANISM_PTR       mechanism_gen      = &mechanism_gen_desc;

    static CK_BBOOL        ltrue       = CK_TRUE;
//    static CK_BBOOL        lfalse      = CK_FALSE;
//    static CK_OBJECT_CLASS oclass      = CKO_SECRET_KEY;
//    static CK_OBJECT_CLASS oclass_pub  = CKO_PUBLIC_KEY;
//    static CK_OBJECT_CLASS oclass_priv = CKO_PRIVATE_KEY;
//    static CK_KEY_TYPE     key_type    = CKK_GOSTR3410;

    static CK_BYTE         gostR3410params[] = {
        0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x23, 0x01
    };
    static CK_BYTE         gostR3411params[] = {
        0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1e, 0x01
    };
    static CK_BYTE         data[] = {
        0xee, 0x8c, 0xd9, 0x40, 0x23, 0xdd, 0x9a, 0xf8,
        0x17, 0xdd, 0xe8, 0x1f, 0x02, 0x25, 0x5b, 0xb3,
        0xaa, 0x6b, 0xd3, 0x21, 0x09, 0x41, 0x95, 0x7e,
        0xea, 0x7e, 0x0e, 0x3c, 0x35, 0x9f, 0x04, 0x5a
    };
//    static CK_BYTE         pub_key_value[64];
//    static CK_UTF8CHAR     *label         = "Yet Another Keypair";
    CK_ATTRIBUTE           pub_template[] = {
        { CKA_TOKEN,           &ltrue,          sizeof(ltrue)           },
        { CKA_GOSTR3410PARAMS, gostR3410params, sizeof(gostR3410params) },
        { CKA_GOSTR3411PARAMS, gostR3411params, sizeof(gostR3411params) },
        { CKA_VERIFY,          &ltrue,          sizeof(CK_BBOOL)        },
    };
    CK_ATTRIBUTE           priv_template[] = {
        { CKA_TOKEN,   &ltrue, sizeof(ltrue)    },
        { CKA_PRIVATE, &ltrue, sizeof(ltrue)    },
        { CKA_SIGN,    &ltrue, sizeof(CK_BBOOL) },
    };

    rc = funcs->C_OpenSession(
        SlotId, CKF_RW_SESSION | CKF_SERIAL_SESSION, NULL, NULL, &sess);
    if (rc != CKR_OK)
    {
        show_error("   C_OpenSession", rc);
        return rc;
    }
    fprintf(stderr, "Thread %d: Session %d opened\n", id, is);
    rc = funcs->C_GenerateKeyPair(sess, mechanism_gen,
                                  pub_template, sizeof(pub_template) / sizeof(CK_ATTRIBUTE),
                                  priv_template, sizeof(priv_template) / sizeof(CK_ATTRIBUTE),
                                  &pub_key, &priv_key);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_GenerateKeyPair failed: 0x%lx\n", rc);
        return rc;
    }
    fprintf(stderr, "C_GenerateKeyPair OK\n");

    rc = funcs->C_SignInit(sess, mechanism, priv_key);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_SignInit failed: 0x%lx\n", rc);
        return rc;
    }

    len = sizeof(value);
    rc  = funcs->C_Sign(sess, data, sizeof(data), value, &len);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_Sign failed: 0x%lx\n", rc);
        return rc;
    }
    fprintf(stderr, "C_Sign OK\n");

    printf("SIGNATURE:\n");
    print_bytes(value, len);

    rc = funcs->C_VerifyInit(sess, mechanism, pub_key);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_VerifyInit failed: 0x%lx\n", rc);
        return rc;
    }

    rc = funcs->C_Verify(sess, data, sizeof(data), value, len);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_Verify failed: 0x%lx\n", rc);
        return rc;
    }
    fprintf(stderr, "C_Verify OK\n");

    rc = funcs->C_DestroyObject(sess, priv_key);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_DestroyObject failed: 0x%lx\n", rc);
        return rc;
    }
    rc = funcs->C_DestroyObject(sess, pub_key);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_DestroyObject failed: 0x%lx\n", rc);
        return rc;
    }

    funcs->C_CloseSession(sess);

    fprintf(stderr, "Session %d closed\n", id);

    return rc;
}

#ifdef WIN32
DWORD WINAPI
#else
int
#endif
thread_func(void *thid)
{
    int   i  = 0;
    CK_RV rv = CKR_OK;
    int   id = *(int *) thid;

    fprintf(stderr, "thread_func %d started\n", id);
    while (i++ < NUMSESS)
    {
        fprintf(stderr,
                "create_token_objects %d, session %d\n", id, i);

        rv = create_token_objects(id, i);
        if (rv != CKR_OK)
        {
            fprintf(stderr,
                    "create_token_objects %d, session %d failed\n",
                    id, i);
            break;
        }

        fprintf(stderr,
                "create_token_objects %d, session %d finished\n",
                id, i);

        fprintf(stderr,
                "generate key pair %d, session %d\n", id, i);
        rv = test_gost3410(id, i);
        if (rv != CKR_OK)
        {
            fprintf(stderr,
                    "create_token_objects %d, session %d failed\n",
                    id, i);
            break;
        }
        fprintf(stderr,
                "generate key pair %d, session %d finished\n",
                id, i);

    }

    fprintf(stderr, "thread_func %d finished\n", id);
    return rv;
}

int
main(int argc, char **argv)
{
    int       i;
	CK_RV     rc;
	CK_SESSION_HANDLE hSession = CK_INVALID_HANDLE;
#ifdef WIN32
    HANDLE    id[THREADCNT];
    DWORD     rv;
#else
    pthread_t id[THREADCNT];
#endif
    int       thid[THREADCNT];

    for (i = 1; i < argc; i++)
    {
        if (strcmp(argv[i], "-api") == 0)
        {
            i++;
            api_path = argv[i];
        }
        else if (strcmp(argv[i], "-user_pin") == 0)
        {
            i++;
            user_pin = (CK_UTF8CHAR *)argv[i];
        }
        else if (strcmp(argv[i], "-h") == 0)
        {
            printf(
                "usage:  %s [-api <path>][-user_pin <PIN>] [-h]\n\n",
                argv[0]);
            return 0;
        }
    }

    fprintf(stderr, "==>> init...\n");
    // Use PKCS#11 library native multithreading locks.
    rc = init(api_path, &funcs, pinit_args_os_locking);
    // Use application-supplied multithreading callbacks for locking.
    //rc = init(api_path, &funcs, pinit_args_app_locking);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "==>> init failed\n\n");
        return rc;
    }
    fprintf(stderr, "==>> init Ok\n\n");

    fprintf(stderr, "==>> get_slot_list...\n");
    rc = get_slot_list(funcs,
                       &SlotList,
                       &SlotCount);
    if (rc != CKR_OK)
    {
        fprintf(stderr,
                "==>> get_slot_list failed, rc = 0x%lx\n\n", rc);
        return rc;
    }
    fprintf(stderr, "==>> get_slot_list Ok\n\n");

    fprintf(stderr, "==>> find_slot_with_token...\n");
    rc = find_slot_with_token(funcs,
                              SlotList,
                              SlotCount,
                              &SlotId);
    if (rc != CKR_OK)
    {
        fprintf(stderr,
                "==>> find_slot_with_token failed, rc = 0x%lx\n\n",
                rc);
        return rc;
    }
    fprintf(stderr, "==>> find_slot_with_token Ok\n\n");

    rc = open_session_and_login(SlotId, &hSession);
    if (rc != CKR_OK)
    {
        fprintf(stderr,
                "==>> open_session_and_login failed, rc = 0x%lx\n\n",
                rc);
        return rc;
    }
    fprintf(stderr, "==>> open_session_and_login Ok\n\n");

    for (i = 0; i < THREADCNT; i++)
    {
        thid[i] = i + 1;
        fprintf(stderr, "Creating thread %d \n", thid[i]);
#ifdef WIN32
        id[i] = CreateThread(NULL, 0,
                             (LPTHREAD_START_ROUTINE) thread_func,
                             (void *) &(thid[i]), 0, NULL);
#else
        pthread_create(&id[i], NULL,
                       (void*(*)(void *))thread_func, (void *) &(thid[i]));
#endif
    }

    fprintf(stderr, "Waiting for all application threads\n");
#ifdef WIN32
    rv = WaitForMultipleObjects(THREADCNT, id, TRUE, INFINITE);
    if (rv == WAIT_TIMEOUT)
    {
        fprintf(stderr, "Stop waiting by timeout\n");
    }
    else
    {
        fprintf(stderr, "All threads completed\n");
    }
#else
    for (i = 0; i < THREADCNT; i++)
    {
        if (thid[i])
        {
            printf("Joining thread %d\n", thid[i]);
            pthread_join(id[i], NULL);
            thid[i] = 0;
        }
    }
#endif

	funcs->C_Logout(hSession);
	funcs->C_CloseSession(hSession);

    fprintf(stderr,
            "Finalizing library...\n");
    funcs->C_Finalize(NULL);
	if (SlotList){
		free(SlotList);
	}
#ifdef WIN32
    FreeLibrary(dll);
#else
    dlclose(dll);
#endif
    fprintf(stderr, "SUCCESS\n");
    return 0;
}

