/* Actual create_context() + actual policy; identity/vendor/callback doubles. */
#include <assert.h>
#include <stdio.h>
#include <string.h>
#include "ssh_protocol_policy.h"
#include "context_constants.h"

typedef int esp_err_t;
enum { ESP_OK = 0, ESP_FAIL = -1, ESP_ERR_NO_MEM = 0x101 };
static WOLFSSH_CTX candidate, *s_context;
static unsigned copies, news, imports, wipes, setters, callbacks, frees;
static unsigned fail_setter;
static int copy_error, allocation_fail, import_error, setter_error;
static unsigned char *identity;
static size_t identity_capacity;

static void unpublished(void)
{
    assert(s_context == NULL);
}
static void wiped(void)
{
    assert(wipes == 1);
    assert(identity_capacity == SSH_SECURITY_PRIVATE_KEY_DER_CAPACITY);
    /* Only called while the extracted create_context stack frame is alive. */
    for (size_t i = 0; i < identity_capacity; ++i) assert(identity[i] == 0);
}
static esp_err_t ssh_security_copy_private_key(unsigned char *out, size_t capacity,
                                               size_t *length)
{
    unpublished();
    assert(++copies == 1);
    assert(capacity == SSH_SECURITY_PRIVATE_KEY_DER_CAPACITY);
    for (size_t i = 0; i < capacity; ++i) assert(out[i] == 0);
    identity = out;
    identity_capacity = capacity;
    /* Mark the unused tail too, so a short wipe cannot pass this test. */
    memset(out, 0x6d, capacity);
    *length = 31;
    return copy_error;
}
static void secure_wipe(void *ptr, size_t length)
{
    unpublished();
    assert(ptr == identity && length == identity_capacity);
    assert(wipes++ == 0);
    memset(ptr, 0, length);
}
static WOLFSSH_CTX *wolfSSH_CTX_new(int endpoint, void *heap)
{
    unpublished();
    assert(copies == 1 && wipes == 0 && imports == 0);
    assert(endpoint == WOLFSSH_ENDPOINT_SERVER && heap == NULL);
    assert(++news == 1);
    return allocation_fail ? NULL : &candidate;
}
static int wolfSSH_CTX_UsePrivateKey_buffer(WOLFSSH_CTX *ctx,
        const unsigned char *key, word32 length, int format)
{
    unpublished();
    assert(ctx == &candidate && news == 1 && wipes == 0);
    assert(++imports == 1);
    assert(key == identity && length == 31 && format == WOLFSSH_FORMAT_ASN1);
    for (size_t i = 0; i < identity_capacity; ++i) assert(key[i] == 0x6d);
    return import_error;
}
static void wolfSSH_CTX_free(WOLFSSH_CTX *ctx)
{
    unpublished();
    wiped();
    assert(ctx == &candidate && news == 1 && callbacks == 0);
    assert(++frees == 1);
}
static int set_list(WOLFSSH_CTX *ctx, const char *list, const char **field,
                    unsigned step)
{
    unpublished();
    wiped();
    assert(imports == 1 && import_error == 0 && frees == 0 && callbacks == 0);
    assert(ctx == &candidate && ++setters == step);
    assert(list != NULL && list[0] != '\0');
    if (step == fail_setter) return setter_error;
    *field = list;
    return WS_SUCCESS;
}
#define SETTER(name, field, step) \
    int wolfSSH_CTX_SetAlgoList##name(WOLFSSH_CTX *ctx, const char *list) \
    { return set_list(ctx, list, &ctx->field, step); }
SETTER(Kex, algoListKex, 1)
SETTER(Key, algoListKey, 2)
SETTER(Cipher, algoListCipher, 3)
SETTER(Mac, algoListMac, 4)
SETTER(KeyAccepted, algoListKeyAccepted, 5)

static void bounded_ssh_receive(void) {}
static void authenticate_user(void) {}
static void allowed_auth_types(void) {}
static void authentication_result(void) {}
static void reject_keyboard_auth(void) {}
static void accept_shell(void) {}
static void reject_channel_request(void) {}
static int callback(WOLFSSH_CTX *ctx, void (*actual)(void),
                    void (*expected)(void), unsigned step)
{
    unpublished();
    wiped();
    assert(ctx == &candidate && setters == 5 && fail_setter == 0 && frees == 0);
    assert(++callbacks == step && actual == expected);
    return WS_SUCCESS;
}
#define CALLBACK(name, expected, step) \
    static int name(WOLFSSH_CTX *ctx, void (*cb)(void)) \
    { return callback(ctx, cb, expected, step); }
CALLBACK(wolfSSH_SetIORecv, bounded_ssh_receive, 1)
CALLBACK(wolfSSH_SetUserAuth, authenticate_user, 2)
CALLBACK(wolfSSH_SetUserAuthTypes, allowed_auth_types, 3)
CALLBACK(wolfSSH_SetUserAuthResult, authentication_result, 4)
CALLBACK(wolfSSH_SetKeyboardAuthPrompts, reject_keyboard_auth, 5)
CALLBACK(wolfSSH_CTX_SetChannelReqShellCb, accept_shell, 6)
CALLBACK(wolfSSH_CTX_SetChannelReqExecCb, reject_channel_request, 7)
CALLBACK(wolfSSH_CTX_SetChannelReqSubsysCb, reject_channel_request, 8)

#include "context_actual.c"

static void reset(void)
{
    memset(&candidate, 0, sizeof(candidate));
    s_context = NULL;
    copies = news = imports = wipes = setters = callbacks = frees = 0;
    fail_setter = 0;
    copy_error = allocation_fail = import_error = setter_error = 0;
    identity = NULL;
    identity_capacity = 0;
}
static void failed(int expected)
{
    assert(create_context() == expected);
    assert(s_context == NULL && callbacks == 0 && wipes == 1 && copies == 1);
    /* identity points at a retired stack frame now: never inspect it here. */
    identity = NULL;
}
int main(void)
{
    unsigned cases = 0;
    reset();
    copy_error = 0x4321;
    failed(copy_error);
    assert(news == 0 && imports == 0 && setters == 0 && frees == 0);
    ++cases;
    reset();
    allocation_fail = 1;
    failed(ESP_ERR_NO_MEM);
    assert(news == 1 && imports == 0 && setters == 0 && frees == 0);
    ++cases;
    for (unsigned sign = 0; sign < 2; ++sign) {
        reset();
        import_error = sign ? 7001 : -7001;
        failed(ESP_FAIL);
        assert(news == 1 && imports == 1 && setters == 0 && frees == 1);
        ++cases;
    }
    for (unsigned step = 1; step <= 5; ++step) {
        for (unsigned sign = 0; sign < 2; ++sign) {
            reset();
            fail_setter = step;
            setter_error = sign ? (int)(8000 + step) : -(int)(8000 + step);
            failed(ESP_FAIL);
            assert(news == 1 && imports == 1 && setters == step && frees == 1);
            ++cases;
        }
    }
    reset();
    assert(create_context() == ESP_OK);
    identity = NULL;
    assert(s_context == &candidate);
    assert(copies == 1 && news == 1 && imports == 1 && wipes == 1);
    assert(setters == 5 && callbacks == 8 && frees == 0);
    ++cases;
    assert(cases == 15);
    puts("PASS: actual create_context + policy: 15 cases, full identity wipe before policy/callbacks/free, exact cleanup, publication only after success");
    return 0;
}