ta trieu rsa

#include <tee_internal_api.h>
#include <string.h>
#include <stdio.h>
#include "TA.h"

// Biến toàn cục lưu giữ Handle của cặp khóa RSA bên trong TA
static TEE_ObjectHandle rsa_keypair = TEE_HANDLE_NULL;

TEE_Result TA_CreateEntryPoint(void)
{
    TEE_Result res;
    // 1. Cấp phát một Transient Object cho RSA Keypair
    res = TEE_AllocateTransientObject(TEE_TYPE_RSA_KEYPAIR, 2048, &rsa_keypair);
    if (res != TEE_SUCCESS) return res;

    // 2. Tạo ngẫu nhiên cặp khóa RSA-2048
    res = TEE_GenerateKey(rsa_keypair, 2048, NULL, 0);
    if (res != TEE_SUCCESS) {
        TEE_FreeTransientObject(rsa_keypair);
        rsa_keypair = TEE_HANDLE_NULL;
    }
    return res;
}

void TA_DestroyEntryPoint(void)
{
    if (rsa_keypair != TEE_HANDLE_NULL) {
        TEE_FreeTransientObject(rsa_keypair);
    }
}

TEE_Result TA_OpenSessionEntryPoint(uint32_t paramTypes,
    TEE_Param params[4], void **sessionContext)
{
    (void) paramTypes;
    (void) params;
    (void) sessionContext;

    return TEE_SUCCESS;
}

void TA_CloseSessionEntryPoint(void *sessionContext)
{
    (void) sessionContext;
}

static void create_obj_id(PasswordEntry *e, char *obj_id, size_t *len){
	// Ensure buffer is null-terminated
	obj_id[0] = '\0';
	
	// Ensure we don't overflow the buffer
	size_t package_len = strlen(e->package_name);
	size_t user_len = strlen(e->user_id);
	size_t max_len = MAX_SIZE * 2 - 1; // Leave room for null terminator
	
	// Truncate if strings are too long to fit with separator
	if (package_len + user_len + 2 > max_len) { // +2 for '|' and '\0'
		if (package_len > MAX_SIZE - 1) package_len = MAX_SIZE - 1;
		if (user_len > MAX_SIZE - 1) user_len = MAX_SIZE - 1;
		// Ensure total fits
		if (package_len + user_len + 2 > max_len) {
			// Reduce the longer one
			if (package_len > user_len) {
				package_len = max_len - user_len - 2;
			} else {
				user_len = max_len - package_len - 2;
			}
		}
	}
	
	snprintf(obj_id, MAX_SIZE * 2, "%.*s|%.*s", 
	         (int)package_len, e->package_name,
	         (int)user_len, e->user_id);
	obj_id[MAX_SIZE * 2 - 1] = '\0'; // Ensure null termination
	*len = strlen(obj_id); 
}

TEE_Result TA_InvokeCommandEntryPoint(void *sessionContext,
    uint32_t commandID, uint32_t paramTypes, TEE_Param params[4])
{
    (void) sessionContext;
    (void) commandID;

    if (paramTypes != TEE_PARAM_TYPES(
            TEE_PARAM_TYPE_MEMREF_INOUT,
            TEE_PARAM_TYPE_NONE,
            TEE_PARAM_TYPE_NONE,
            TEE_PARAM_TYPE_NONE)) {
        /* Bad parameter types */
        return TEE_ERROR_BAD_PARAMETERS;
    }

	PasswordEntry *entry = (PasswordEntry *)params[0].memref.buffer;
	uint32_t entry_size = params[0].memref.size;
	TEE_ObjectHandle object;
	TEE_Result res;
	char obj_id[MAX_SIZE * 2];
	size_t obj_id_len;

	switch (commandID)
	{
	case CMD_GET_PUBLIC_KEY:
        // Kiểm tra tham số đầu ra (Memref để chứa Modulus N)
        if (paramTypes != TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_OUTPUT, TEE_PARAM_TYPE_NONE, TEE_PARAM_TYPE_NONE, TEE_PARAM_TYPE_NONE)) {
            return TEE_ERROR_BAD_PARAMETERS;
        }
        if (params[0].memref.size < RSA_2048_KEY_SIZE) return TEE_ERROR_SHORT_BUFFER;

        // Trích xuất Modulus (N) từ cặp khóa RSA để gửi cho CA
        uint32_t exp_size = params[0].memref.size;
        res = TEE_GetObjectAttributeBuffer(rsa_keypair, TEE_ATTR_RSA_MODULUS, params[0].memref.buffer, &exp_size);
        params[0].memref.size = exp_size;
        return res;
        
	case CMD_SAVE_PASSWORD:
		if (paramTypes != TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INOUT, TEE_PARAM_TYPE_NONE, TEE_PARAM_TYPE_NONE, TEE_PARAM_TYPE_NONE)) {
            return TEE_ERROR_BAD_PARAMETERS;
        }

        // Nhận struct đã bị mã hóa từ Shared Memory
        EncryptedPasswordEntry *enc_entry = (EncryptedPasswordEntry *)params[0].memref.buffer;
        
        // Tạo một cấu trúc PasswordEntry (Plaintext) trên Stack của TA để chuẩn bị giải mã vào đó
        PasswordEntry plain_entry;
        memset(&plain_entry, 0, sizeof(PasswordEntry));
        memcpy(plain_entry.package_name, enc_entry->package_name, MAX_SIZE);
        memcpy(plain_entry.user_id, enc_entry->user_id, MAX_SIZE);

        // Tiến hành giải mã RSA dữ liệu password
        TEE_OperationHandle op;
        res = TEE_AllocateOperation(&op, TEE_ALG_RSAES_PKCS1_V1_5, TEE_MODE_DECRYPT, 2048);
        if (res != TEE_SUCCESS) return res;

        res = TEE_SetOperationKey(op, rsa_keypair);
        if (res != TEE_SUCCESS) { TEE_FreeOperation(op); return res; }

        size_t decrypted_len = MAX_SIZE - 1;
        res = TEE_AsymmetricDecrypt(op, NULL, 0, 
                                    enc_entry->encrypted_password, RSA_2048_KEY_SIZE, 
                                    plain_entry.password, &decrypted_len);
        TEE_FreeOperation(op);
        
        if (res != TEE_SUCCESS) {
            printf("TA Decryption failed: 0x%x\n", res);
            return res;
        }
        plain_entry.password[decrypted_len] = '\0'; // Đảm bảo null-terminated

        // Tạo Object ID dựa trên dữ liệu đã giải mã
        create_obj_id(&plain_entry, obj_id, &obj_id_len);

        // Lưu bản rõ (plain_entry) vào Secure Storage của TEE
        res = TEE_CreatePersistentObject(
                TEE_STORAGE_PRIVATE, 
                obj_id, obj_id_len, 
                TEE_DATA_FLAG_ACCESS_WRITE | TEE_DATA_FLAG_OVERWRITE,
                TEE_HANDLE_NULL,
                &plain_entry, sizeof(PasswordEntry),
                &object
            );
        if (res == TEE_SUCCESS) {
            TEE_CloseObject(object);
        }
        return res;

	case CMD_CALL_PASSWORD:
		create_obj_id(entry, obj_id, &obj_id_len);
		/**
		TEE_Result TEE_OpenPersistentObject(
			 uint32_t storageID,
			[in(objectIDLength)] void* objectID, size_t objectIDLen,
			 uint32_t flags,
			[out] TEE_ObjectHandle* object);
		**/
		res = TEE_OpenPersistentObject(
				TEE_STORAGE_PRIVATE,
				obj_id, obj_id_len,
				TEE_DATA_FLAG_ACCESS_READ,
				&object
			);
		if(res != TEE_SUCCESS){
			printf("open persistent object fail\n");
			return res;
		} else {
			printf("open persistent object success\n");
		}

		// Use stack buffer instead of struct to avoid memory access issues
		char buffer[MAX_SIZE * 3]; // Enough space for the PasswordEntry struct
		if (entry_size > sizeof(buffer)) {
			TEE_CloseObject(object);
			return TEE_ERROR_OVERFLOW;
		}

		size_t count;
		/**
		TEE_Result TEE_ReadObjectData(
			 TEE_ObjectHandle object,
			 [out] void* buffer,
			 size_t size,
			 [out] size_t* count 
		**/
		res = TEE_ReadObjectData(object, buffer, entry_size, &count);
		if(res != TEE_SUCCESS) {
			printf("read object data fail: 0x%x\n", res);
		} else {
			printf("read object data success\n");
			// Copy data back to entry struct
			memcpy(entry, buffer, count);
		}
		TEE_CloseObject(object);
		return res;

	default:
		return TEE_ERROR_NOT_SUPPORTED;
	}
}