// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */

#include <vmlinux.h>
#include <stdbool.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>
#include "bpf_misc.h"

#define ENOENT 2
#define EEXIST 17

char _license[] SEC("license") = "GPL";

int err;

struct elem {
        char arr[128];
        int val;
};

struct {
        __uint(type, BPF_MAP_TYPE_RHASH);
        __uint(map_flags, BPF_F_NO_PREALLOC);
        __uint(max_entries, 128);
        __type(key, int);
        __type(value, struct elem);
} rhmap SEC(".maps");

SEC("syscall")
int test_rhash_lookup_update(void *ctx)
{
        int key = 5;
        struct elem empty = {.val = 3, .arr = {0}};
        struct elem *e;

        err = 1;
        e = bpf_map_lookup_elem(&rhmap, &key);
        if (e)
                return 1;

        err = bpf_map_update_elem(&rhmap, &key, &empty, BPF_NOEXIST);
        if (err)
                return 1;

        e = bpf_map_lookup_elem(&rhmap, &key);
        if (!e || e->val != empty.val) {
                err = 2;
                return 2;
        }

        err = 0;
        return 0;
}

SEC("syscall")
int test_rhash_update_delete(void *ctx)
{
        int key = 6;
        struct elem empty = {.val = 4, .arr = {0}};
        struct elem *e;

        err = 1;
        e = bpf_map_lookup_elem(&rhmap, &key);
        if (e)
                return 1;

        err = bpf_map_update_elem(&rhmap, &key, &empty, BPF_NOEXIST);
        if (err)
                return 2;

        err = bpf_map_delete_elem(&rhmap, &key);
        if (err)
                return 3;

        e = bpf_map_lookup_elem(&rhmap, &key);
        if (e) {
                err = 4;
                return 4;
        }

        err = 0;
        return 0;
}

SEC("syscall")
int test_rhash_update_elements(void *ctx)
{
        int key = 0;
        struct elem empty = {.val = 4, .arr = {0}};
        struct elem *e;
        int i;

        err = 1;

        for (i = 0; i < 128; ++i) {
                key = i;
                e = bpf_map_lookup_elem(&rhmap, &key);
                if (e)
                        return 1;

                empty.val = key;
                err = bpf_map_update_elem(&rhmap, &key, &empty, BPF_NOEXIST);
                if (err)
                        return 2;

                e = bpf_map_lookup_elem(&rhmap, &key);
                if (!e || e->val != key) {
                        err = 4;
                        return 4;
                }
        }

        for (i = 0; i < 128; ++i) {
                key = i;
                err = bpf_map_delete_elem(&rhmap, &key);
                if (err)
                        return 3;

                e = bpf_map_lookup_elem(&rhmap, &key);
                if (e) {
                        err = 5;
                        return 5;
                }
        }

        err = 0;
        return 0;
}

SEC("syscall")
int test_rhash_update_exist(void *ctx)
{
        int key = 10;
        struct elem val1 = {.val = 100, .arr = {0}};
        struct elem val2 = {.val = 200, .arr = {0}};
        struct elem *e;
        int ret;

        err = 1;

        /* BPF_EXIST on non-existent key should fail with -ENOENT */
        ret = bpf_map_update_elem(&rhmap, &key, &val1, BPF_EXIST);
        if (ret != -ENOENT)
                return 1;

        /* Insert element first */
        ret = bpf_map_update_elem(&rhmap, &key, &val1, BPF_NOEXIST);
        if (ret)
                return 2;

        /* Verify initial value */
        e = bpf_map_lookup_elem(&rhmap, &key);
        if (!e || e->val != 100)
                return 3;

        /* BPF_EXIST on existing key should succeed and update value */
        ret = bpf_map_update_elem(&rhmap, &key, &val2, BPF_EXIST);
        if (ret)
                return 4;

        /* Verify value was updated */
        e = bpf_map_lookup_elem(&rhmap, &key);
        if (!e || e->val != 200)
                return 5;

        /* Cleanup */
        bpf_map_delete_elem(&rhmap, &key);
        err = 0;
        return 0;
}

SEC("syscall")
int test_rhash_update_any(void *ctx)
{
        int key = 11;
        struct elem val1 = {.val = 111, .arr = {0}};
        struct elem val2 = {.val = 222, .arr = {0}};
        struct elem *e;
        int ret;

        err = 1;

        /* BPF_ANY on non-existent key should insert */
        ret = bpf_map_update_elem(&rhmap, &key, &val1, BPF_ANY);
        if (ret)
                return 1;

        e = bpf_map_lookup_elem(&rhmap, &key);
        if (!e || e->val != 111)
                return 2;

        /* BPF_ANY on existing key should update */
        ret = bpf_map_update_elem(&rhmap, &key, &val2, BPF_ANY);
        if (ret)
                return 3;

        e = bpf_map_lookup_elem(&rhmap, &key);
        if (!e || e->val != 222)
                return 4;

        /* Cleanup */
        bpf_map_delete_elem(&rhmap, &key);
        err = 0;
        return 0;
}

SEC("syscall")
int test_rhash_noexist_duplicate(void *ctx)
{
        int key = 12;
        struct elem val = {.val = 600, .arr = {0}};
        int ret;

        err = 1;

        /* Insert element */
        ret = bpf_map_update_elem(&rhmap, &key, &val, BPF_NOEXIST);
        if (ret)
                return 1;

        /* Try to insert again with BPF_NOEXIST - should fail with -EEXIST */
        ret = bpf_map_update_elem(&rhmap, &key, &val, BPF_NOEXIST);
        if (ret != -EEXIST)
                return 2;

        /* Cleanup */
        bpf_map_delete_elem(&rhmap, &key);
        err = 0;
        return 0;
}

SEC("syscall")
int test_rhash_delete_nonexistent(void *ctx)
{
        int key = 99999;
        int ret;

        err = 1;

        /* Delete non-existent key should return -ENOENT */
        ret = bpf_map_delete_elem(&rhmap, &key);
        if (ret != -ENOENT)
                return 1;

        err = 0;
        return 0;
}