// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2026, Oracle and/or its affiliates. */
#include <test_progs.h>
#include <linux/btf.h>
#include "bpf/libbpf_internal.h"
#include "../test_btf.h"
#include "kfree_skb.skel.h"

#define TYPE_LEN        (sizeof(struct btf_type) + sizeof(__u32))
#define MAX_NR_LAYOUT   2
#define LAYOUT_LEN      (sizeof(struct btf_layout) * MAX_NR_LAYOUT)
#define STR_LEN         sizeof("\0int")

struct layout_btf {
        struct btf_header hdr;
        __u32 types[TYPE_LEN/sizeof(__u32)];
        struct btf_layout layout[MAX_NR_LAYOUT];
        char strs[STR_LEN];
};

static const struct layout_btf layout_btf = {
        .hdr = {
                .magic          = BTF_MAGIC,
                .version        = BTF_VERSION,
                .hdr_len        = sizeof(struct btf_header),
                .type_off       = 0,
                .type_len       = TYPE_LEN,
                .str_off        = TYPE_LEN + LAYOUT_LEN,
                .str_len        = STR_LEN,
                .layout_off     = TYPE_LEN,
                .layout_len     = LAYOUT_LEN,
        },
        .types = {
                BTF_TYPE_INT_ENC(1, BTF_INT_SIGNED, 0, 32, 4),
        },
        .layout = {
                { .info_sz = 0,          .elem_sz = 0, .flags = 0 },
                { .info_sz = sizeof(__u32), .elem_sz = 0, .flags = 0 },
        },
        .strs = "\0int",
};

void test_btf_sanitize_layout(void)
{
        struct btf *orig = NULL, *sanitized = NULL;
        struct kern_feature_cache *cache = NULL;
        struct kfree_skb *skel = NULL;
        const struct btf_header *hdr;
        const void *raw;
        __u32 raw_sz;

        skel = kfree_skb__open();
        if (!ASSERT_OK_PTR(skel, "kfree_skb_skel"))
                return;
        orig = btf__new(&layout_btf, sizeof(layout_btf));
        if (!ASSERT_OK_PTR(orig, "btf_new_layout"))
                goto out;
        raw = btf__raw_data(orig, &raw_sz);
        if (!ASSERT_OK_PTR(raw, "btf__raw_data_orig"))
                goto out;
        hdr = (struct btf_header *)raw;
        ASSERT_EQ(hdr->layout_off, TYPE_LEN, "layout_off_nonzero");
        ASSERT_EQ(hdr->layout_len, LAYOUT_LEN, "layout_len_nonzero");

        cache = calloc(1, sizeof(*cache));
        if (!ASSERT_OK_PTR(cache, "alloc_feat_cache"))
                goto out;
        for (int i = 0; i < __FEAT_CNT; i++)
                cache->res[i] = FEAT_SUPPORTED;
        cache->res[FEAT_BTF_LAYOUT] = FEAT_MISSING;

        bpf_object_set_feat_cache(skel->obj, cache);

        if (!ASSERT_FALSE(kernel_supports(skel->obj, FEAT_BTF_LAYOUT), "layout_feature_missing"))
                goto out;
        if (!ASSERT_TRUE(kernel_supports(skel->obj, FEAT_BTF_FUNC), "other_feature_allowed"))
                goto out;

        sanitized = bpf_object__sanitize_btf(skel->obj, orig);
        if (!ASSERT_OK_PTR(sanitized, "bpf_object__sanitize_btf"))
                goto out;

        raw = btf__raw_data(sanitized, &raw_sz);
        if (!ASSERT_OK_PTR(raw, "btf__raw_data_sanitized"))
                goto out;
        hdr = (struct btf_header *)raw;
        ASSERT_EQ(hdr->layout_off, 0, "layout_off_zero");
        ASSERT_EQ(hdr->layout_len, 0, "layout_len_zero");
        ASSERT_EQ(hdr->str_off, TYPE_LEN, "strs_after_types");
        ASSERT_EQ(hdr->str_len, STR_LEN, "strs_len_unchanged");
        ASSERT_EQ(raw_sz, hdr->hdr_len + hdr->type_len + hdr->str_len, "btf_raw_sz_reduced");
out:
        /* This will free the cache we allocated above */
        kfree_skb__destroy(skel);
        btf__free(sanitized);
        btf__free(orig);
}