/* SPDX-License-Identifier: BSD-3-Clause
 * Copyright (c) 2020 BIFIT <mitigator@bifit.com>
 */
#include <mitigator_bpf.h>

struct Challenge {
    char signature[4];
    uint32_t cookie;
};

struct Response {
    char signature[4];
    uint32_t cookie;
    uint32_t hash1;
    uint32_t hash2;
};

LOCAL enum Result
process_mcr_request(Context ctx, uint8_t* data) {
    struct Challenge* challenge = (struct Challenge*)data;

    struct Flow flow;
    packet_flow(ctx, &flow);

    challenge->signature[0] = 'M';
    challenge->signature[1] = 'C';
    challenge->signature[2] = 'R';
    challenge->signature[3] = 'C';
    challenge->cookie = cookie_make(ctx, &flow);

    set_packet_length(ctx, sizeof(struct Challenge));
    return RESULT_BACK;
}

LOCAL enum Result
process_mcr_response(Context ctx, const uint8_t* data) {
    static const size_t KEY_SIZE = 8;

    const struct Response* response = (const struct Response*)data;

    struct Flow flow;
    packet_flow(ctx, &flow);

    if (!cookie_check(ctx, &flow, response->cookie)) {
        return RESULT_DROP;
    }

    const uint8_t* key = parameters_get(ctx);

    uint32_t cookie_hash = hash_crc32_u32(response->cookie, 0);
    uint32_t hash1 = hash_crc32_data(key, key + KEY_SIZE, cookie_hash);
    uint32_t hash2 = hash_crc32_u32(response->cookie, hash1);
    if ((response->hash1 == hash1) && (response->hash2 == hash2)) {
        table_put(ctx, flow.src_ip.v4, 0);
    }
    return RESULT_DROP;
}

LOCAL enum Result
process_any(Context ctx) {
    struct IpHeader* ip = packet_network_header(ctx);
    struct TableRecord record = {};
    return table_get(ctx, ip->ip_src, &record) ? RESULT_PASS : RESULT_DROP;
}

LOCAL enum Result
process_udp(Context ctx) {
    static const uint8_t MCRHELLO[] = {'M', 'C', 'R', 'H', '3', '1', '1', '0'};
    static const uint8_t MCRR[] = {'M', 'C', 'R', 'R'};

    uint16_t length = 0;
    uint8_t* data = packet_transport_payload(ctx, &length);

    if (length != 400) {
        return process_any(ctx);
    }

    if (*(uint64_t*)data == *(const uint64_t*)MCRHELLO) {
        return process_mcr_request(ctx, data);
    }
    if (*(uint32_t*)data == *(const uint32_t*)MCRR) {
        return process_mcr_response(ctx, data);
    }

    return process_any(ctx);
}

ENTRYPOINT enum Result
mcr(Context ctx) {
    if (packet_transport_proto(ctx) == IP_PROTO_UDP) {
        return process_udp(ctx);
    }
    return process_any(ctx);
}

PROGRAM_DISPLAY_ID("MCR UDP (example)")