#include "wvds_crypto.h"
#include <string.h>
 
int encrypt_message(
    const uint8_t* plaintext, size_t pt_len,
    const char* aad, size_t aad_len,
    uint8_t* ciphertext, size_t* ct_len,
    uint8_t nonce[12], uint8_t tag[16]
) {
    uint8_t request[65536];
    size_t request_len = sizeof(request);
 
    // 1. Build request
    int rc = wvds_build_aes_encrypt_request(
        request, &request_len,
        1,                      // Key-ID (must exist)
        aad, aad_len,          // Additional Authenticated Data
        plaintext, pt_len      // Data to encrypt
    );
    if (rc != 0) return rc;
 
    // 2. Send request to Crypto Service
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    // 3. Parse response
    rc = wvds_parse_aes_encrypt_response(
        g_shared_memory, g_response_len,
        nonce, tag,            // Will be filled
        ciphertext, ct_len     // Ciphertext + length
    );
 
    return rc;
}

int decrypt_message(
    const uint8_t* ciphertext, size_t ct_len,
    const uint8_t nonce[12], const uint8_t tag[16],
    const char* aad, size_t aad_len,
    uint8_t* plaintext, size_t* pt_len
) {
    uint8_t request[65536];
    size_t request_len = sizeof(request);
 
    // 1. Build request
    int rc = wvds_build_aes_decrypt_request(
        request, &request_len,
        1,                      // Key-ID
        nonce, tag,            // Received from Encrypt
        aad, aad_len,          // Must be identical!
        ciphertext, ct_len
    );
    if (rc != 0) return rc;
 
    // 2. Send request
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    // 3. Parse response
    rc = wvds_parse_aes_decrypt_response(
        g_shared_memory, g_response_len,
        plaintext, pt_len
    );
 
    // rc == 6 means: Tag verification failed!
    return rc;
}

int sign_message(
    const uint8_t* message, size_t msg_len,
    uint8_t* signature, size_t* sig_len
) {
    uint8_t request[65536];
    size_t request_len = sizeof(request);
 
    int rc = wvds_build_mldsa_sign_request(
        request, &request_len,
        1,                      // Private Key-ID
        message, msg_len
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    rc = wvds_parse_mldsa_sign_response(
        g_shared_memory, g_response_len,
        signature, sig_len      // ML-DSA-65: 3293 bytes
    );
 
    return rc;
}

int verify_signature(
    const uint8_t* message, size_t msg_len,
    const uint8_t* signature, size_t sig_len,
    int* valid                  // 1 = valid, 0 = invalid
) {
    uint8_t request[65536];
    size_t request_len = sizeof(request);
 
    int rc = wvds_build_mldsa_verify_request(
        request, &request_len,
        2,                      // Public Key-ID
        message, msg_len,
        signature, sig_len
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    rc = wvds_parse_mldsa_verify_response(
        g_shared_memory, g_response_len,
        valid
    );
 
    return rc;
}

// Server side: KeyGen + Decaps
int server_key_exchange(
    uint8_t public_key[1184],   // ML-KEM-768 Public Key
    size_t* pk_len
) {
    uint8_t request[256];
    size_t request_len = sizeof(request);
 
    // KeyGen Request
    int rc = wvds_build_mlkem_keygen_request(
        request, &request_len,
        3                       // Key-ID for new key pair
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    // Extract public key
    rc = wvds_parse_mlkem_keygen_response(
        g_shared_memory, g_response_len,
        public_key, pk_len      // Private Key remains in service
    );
 
    return rc;
}
 
int server_decaps(
    const uint8_t* ciphertext, size_t ct_len,
    uint8_t shared_secret[32]
) {
    uint8_t request[2048];
    size_t request_len = sizeof(request);
 
    int rc = wvds_build_mlkem_decaps_request(
        request, &request_len,
        3,                      // Key-ID from KeyGen
        ciphertext, ct_len
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    rc = wvds_parse_mlkem_decaps_response(
        g_shared_memory, g_response_len,
        shared_secret           // 32-byte Shared Secret
    );
 
    return rc;
}
 
// Client side: Encaps
int client_encaps(
    const uint8_t* public_key, size_t pk_len,
    uint8_t* ciphertext, size_t* ct_len,
    uint8_t shared_secret[32]
) {
    uint8_t request[2048];
    size_t request_len = sizeof(request);
 
    int rc = wvds_build_mlkem_encaps_request(
        request, &request_len,
        public_key, pk_len
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    rc = wvds_parse_mlkem_encaps_response(
        g_shared_memory, g_response_len,
        ciphertext, ct_len,     // For server
        shared_secret           // Identical to server after Decaps
    );
 
    return rc;
}

typedef struct {
    uint32_t sensor_id;
    uint32_t timestamp;
    float temperature;
    float humidity;
} SensorData;
 
int send_sensor_data(const SensorData* data) {
    // 1. Serialize
    uint8_t plaintext[sizeof(SensorData)];
    memcpy(plaintext, data, sizeof(SensorData));
 
    // 2. AAD: Sensor-ID + Timestamp (authenticated, not encrypted)
    char aad[64];
    snprintf(aad, sizeof(aad), "sensor:%u:ts:%u",
             data->sensor_id, data->timestamp);
 
    // 3. Encrypt
    uint8_t ciphertext[1024], nonce[12], tag[16];
    size_t ct_len;
 
    int rc = encrypt_message(
        plaintext, sizeof(plaintext),
        aad, strlen(aad),
        ciphertext, &ct_len,
        nonce, tag
    );
    if (rc != 0) return rc;
 
    // 4. Sign
    uint8_t to_sign[2048];
    size_t to_sign_len = 0;
    memcpy(to_sign + to_sign_len, nonce, 12); to_sign_len += 12;
    memcpy(to_sign + to_sign_len, tag, 16); to_sign_len += 16;
    memcpy(to_sign + to_sign_len, ciphertext, ct_len); to_sign_len += ct_len;
 
    uint8_t signature[4096];
    size_t sig_len;
 
    rc = sign_message(to_sign, to_sign_len, signature, &sig_len);
    if (rc != 0) return rc;
 
    // 5. Build packet and send
    // ... (network code) ...
 
    return 0;
}

< Installation | Next: Protocol Specification >