// SPDX-License-Identifier: (LGPL-2.1 OR BSD-2-Clause)
/* Copyright (c) 2024, Oracle and/or its affiliates. */

#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif

#ifdef __KERNEL__
#include <linux/bpf.h>
#include <linux/bsearch.h>
#include <linux/btf.h>
#include <linux/sort.h>
#include <linux/string.h>
#include <linux/bpf_verifier.h>

#define btf_type_by_id                          (struct btf_type *)btf_type_by_id
#define btf__type_cnt                           btf_nr_types
#define btf__base_btf                           btf_base_btf
#define btf__name_by_offset                     btf_name_by_offset
#define btf__str_by_offset                      btf_str_by_offset
#define btf_kflag                               btf_type_kflag

#define calloc(nmemb, sz)                       kvcalloc(nmemb, sz, GFP_KERNEL | __GFP_NOWARN)
#define free(ptr)                               kvfree(ptr)
#define qsort(base, num, sz, cmp)               sort(base, num, sz, cmp, NULL)

#else

#include "btf.h"
#include "bpf.h"
#include "libbpf.h"
#include "libbpf_internal.h"

#endif /* __KERNEL__ */

struct btf;

struct btf_relocate {
        struct btf *btf;
        const struct btf *base_btf;
        const struct btf *dist_base_btf;
        unsigned int nr_base_types;
        unsigned int nr_split_types;
        unsigned int nr_dist_base_types;
        int dist_str_len;
        int base_str_len;
        __u32 *id_map;
        __u32 *str_map;
};

/* Set temporarily in relocation id_map if distilled base struct/union is
 * embedded in a split BTF struct/union; in such a case, size information must
 * match between distilled base BTF and base BTF representation of type.
 */
#define BTF_IS_EMBEDDED ((__u32)-1)

/* <name, size, id> triple used in sorting/searching distilled base BTF. */
struct btf_name_info {
        const char *name;
        /* set when search requires a size match */
        bool needs_size: 1;
        unsigned int size: 31;
        __u32 id;
};

static int btf_relocate_rewrite_type_id(struct btf_relocate *r, __u32 i)
{
        struct btf_type *t = btf_type_by_id(r->btf, i);
        struct btf_field_iter it;
        __u32 *id;
        int err;

        err = btf_field_iter_init(&it, t, BTF_FIELD_ITER_IDS);
        if (err)
                return err;

        while ((id = btf_field_iter_next(&it)))
                *id = r->id_map[*id];
        return 0;
}

/* Simple string comparison used for sorting within BTF, since all distilled
 * types are named.  If strings match, and size is non-zero for both elements
 * fall back to using size for ordering.
 */
static int cmp_btf_name_size(const void *n1, const void *n2)
{
        const struct btf_name_info *ni1 = n1;
        const struct btf_name_info *ni2 = n2;
        int name_diff = strcmp(ni1->name, ni2->name);

        if (!name_diff && ni1->needs_size && ni2->needs_size)
                return ni2->size - ni1->size;
        return name_diff;
}

/* Binary search with a small twist; find leftmost element that matches
 * so that we can then iterate through all exact matches.  So for example
 * searching { "a", "bb", "bb", "c" }  we would always match on the
 * leftmost "bb".
 */
static struct btf_name_info *search_btf_name_size(struct btf_name_info *key,
                                                  struct btf_name_info *vals,
                                                  int nelems)
{
        struct btf_name_info *ret = NULL;
        int high = nelems - 1;
        int low = 0;

        while (low <= high) {
                int mid = (low + high)/2;
                struct btf_name_info *val = &vals[mid];
                int diff = cmp_btf_name_size(key, val);

                if (diff == 0)
                        ret = val;
                /* even if found, keep searching for leftmost match */
                if (diff <= 0)
                        high = mid - 1;
                else
                        low = mid + 1;
        }
        return ret;
}

/* If a member of a split BTF struct/union refers to a base BTF
 * struct/union, mark that struct/union id temporarily in the id_map
 * with BTF_IS_EMBEDDED.  Members can be const/restrict/volatile/typedef
 * reference types, but if a pointer is encountered, the type is no longer
 * considered embedded.
 */
static int btf_mark_embedded_composite_type_ids(struct btf_relocate *r, __u32 i)
{
        struct btf_type *t = btf_type_by_id(r->btf, i);
        struct btf_field_iter it;
        __u32 *id;
        int err;

        if (!btf_is_composite(t))
                return 0;

        err = btf_field_iter_init(&it, t, BTF_FIELD_ITER_IDS);
        if (err)
                return err;

        while ((id = btf_field_iter_next(&it))) {
                __u32 next_id = *id;

                while (next_id) {
                        t = btf_type_by_id(r->btf, next_id);
                        switch (btf_kind(t)) {
                        case BTF_KIND_CONST:
                        case BTF_KIND_RESTRICT:
                        case BTF_KIND_VOLATILE:
                        case BTF_KIND_TYPEDEF:
                        case BTF_KIND_TYPE_TAG:
                                next_id = t->type;
                                break;
                        case BTF_KIND_ARRAY: {
                                struct btf_array *a = btf_array(t);

                                next_id = a->type;
                                break;
                        }
                        case BTF_KIND_STRUCT:
                        case BTF_KIND_UNION:
                                if (next_id < r->nr_dist_base_types)
                                        r->id_map[next_id] = BTF_IS_EMBEDDED;
                                next_id = 0;
                                break;
                        default:
                                next_id = 0;
                                break;
                        }
                }
        }

        return 0;
}

/* Build a map from distilled base BTF ids to base BTF ids. To do so, iterate
 * through base BTF looking up distilled type (using binary search) equivalents.
 */
static int btf_relocate_map_distilled_base(struct btf_relocate *r)
{
        struct btf_name_info *info, *info_end;
        struct btf_type *base_t, *dist_t;
        __u8 *base_name_cnt = NULL;
        int err = 0;
        __u32 id;

        /* generate a sort index array of name/type ids sorted by name for
         * distilled base BTF to speed name-based lookups.
         */
        info = calloc(r->nr_dist_base_types, sizeof(*info));
        if (!info) {
                err = -ENOMEM;
                goto done;
        }
        info_end = info + r->nr_dist_base_types;
        for (id = 0; id < r->nr_dist_base_types; id++) {
                dist_t = btf_type_by_id(r->dist_base_btf, id);
                info[id].name = btf__name_by_offset(r->dist_base_btf, dist_t->name_off);
                info[id].id = id;
                info[id].size = dist_t->size;
                info[id].needs_size = true;
        }
        qsort(info, r->nr_dist_base_types, sizeof(*info), cmp_btf_name_size);

        /* Mark distilled base struct/union members of split BTF structs/unions
         * in id_map with BTF_IS_EMBEDDED; this signals that these types
         * need to match both name and size, otherwise embedding the base
         * struct/union in the split type is invalid.
         */
        for (id = r->nr_dist_base_types; id < r->nr_dist_base_types + r->nr_split_types; id++) {
                err = btf_mark_embedded_composite_type_ids(r, id);
                if (err)
                        goto done;
        }

        /* Collect name counts for composite types in base BTF.  If multiple
         * instances of a struct/union of the same name exist, we need to use
         * size to determine which to map to since name alone is ambiguous.
         */
        base_name_cnt = calloc(r->base_str_len, sizeof(*base_name_cnt));
        if (!base_name_cnt) {
                err = -ENOMEM;
                goto done;
        }
        for (id = 1; id < r->nr_base_types; id++) {
                base_t = btf_type_by_id(r->base_btf, id);
                if (!btf_is_composite(base_t) || !base_t->name_off)
                        continue;
                if (base_name_cnt[base_t->name_off] < 255)
                        base_name_cnt[base_t->name_off]++;
        }

        /* Now search base BTF for matching distilled base BTF types. */
        for (id = 1; id < r->nr_base_types; id++) {
                struct btf_name_info *dist_info, base_info = {};
                int dist_kind, base_kind;

                base_t = btf_type_by_id(r->base_btf, id);
                /* distilled base consists of named types only. */
                if (!base_t->name_off)
                        continue;
                base_kind = btf_kind(base_t);
                base_info.id = id;
                base_info.name = btf__name_by_offset(r->base_btf, base_t->name_off);
                switch (base_kind) {
                case BTF_KIND_INT:
                case BTF_KIND_FLOAT:
                case BTF_KIND_ENUM:
                case BTF_KIND_ENUM64:
                        /* These types should match both name and size */
                        base_info.needs_size = true;
                        base_info.size = base_t->size;
                        break;
                case BTF_KIND_FWD:
                        /* No size considerations for fwds. */
                        break;
                case BTF_KIND_STRUCT:
                case BTF_KIND_UNION:
                        /* Size only needs to be used for struct/union if there
                         * are multiple types in base BTF with the same name.
                         * If there are multiple _distilled_ types with the same
                         * name (a very unlikely scenario), that doesn't matter
                         * unless corresponding _base_ types to match them are
                         * missing.
                         */
                        base_info.needs_size = base_name_cnt[base_t->name_off] > 1;
                        base_info.size = base_t->size;
                        break;
                default:
                        continue;
                }
                /* iterate over all matching distilled base types */
                for (dist_info = search_btf_name_size(&base_info, info, r->nr_dist_base_types);
                     dist_info != NULL && dist_info < info_end &&
                     cmp_btf_name_size(&base_info, dist_info) == 0;
                     dist_info++) {
                        if (!dist_info->id || dist_info->id >= r->nr_dist_base_types) {
                                pr_warn("base BTF id [%d] maps to invalid distilled base BTF id [%d]\n",
                                        id, dist_info->id);
                                err = -EINVAL;
                                goto done;
                        }
                        dist_t = btf_type_by_id(r->dist_base_btf, dist_info->id);
                        dist_kind = btf_kind(dist_t);

                        /* Validate that the found distilled type is compatible.
                         * Do not error out on mismatch as another match may
                         * occur for an identically-named type.
                         */
                        switch (dist_kind) {
                        case BTF_KIND_FWD:
                                switch (base_kind) {
                                case BTF_KIND_FWD:
                                        if (btf_kflag(dist_t) != btf_kflag(base_t))
                                                continue;
                                        break;
                                case BTF_KIND_STRUCT:
                                        if (btf_kflag(base_t))
                                                continue;
                                        break;
                                case BTF_KIND_UNION:
                                        if (!btf_kflag(base_t))
                                                continue;
                                        break;
                                default:
                                        continue;
                                }
                                break;
                        case BTF_KIND_INT:
                                if (dist_kind != base_kind ||
                                    btf_int_encoding(base_t) != btf_int_encoding(dist_t))
                                        continue;
                                break;
                        case BTF_KIND_FLOAT:
                                if (dist_kind != base_kind)
                                        continue;
                                break;
                        case BTF_KIND_ENUM:
                                /* ENUM and ENUM64 are encoded as sized ENUM in
                                 * distilled base BTF.
                                 */
                                if (base_kind != dist_kind && base_kind != BTF_KIND_ENUM64)
                                        continue;
                                break;
                        case BTF_KIND_STRUCT:
                        case BTF_KIND_UNION:
                                /* size verification is required for embedded
                                 * struct/unions.
                                 */
                                if (r->id_map[dist_info->id] == BTF_IS_EMBEDDED &&
                                    base_t->size != dist_t->size)
                                        continue;
                                break;
                        default:
                                continue;
                        }
                        if (r->id_map[dist_info->id] &&
                            r->id_map[dist_info->id] != BTF_IS_EMBEDDED) {
                                /* we already have a match; this tells us that
                                 * multiple base types of the same name
                                 * have the same size, since for cases where
                                 * multiple types have the same name we match
                                 * on name and size.  In this case, we have
                                 * no way of determining which to relocate
                                 * to in base BTF, so error out.
                                 */
                                pr_warn("distilled base BTF type '%s' [%u], size %u has multiple candidates of the same size (ids [%u, %u]) in base BTF\n",
                                        base_info.name, dist_info->id,
                                        base_t->size, id, r->id_map[dist_info->id]);
                                err = -EINVAL;
                                goto done;
                        }
                        /* map id and name */
                        r->id_map[dist_info->id] = id;
                        r->str_map[dist_t->name_off] = base_t->name_off;
                }
        }
        /* ensure all distilled BTF ids now have a mapping... */
        for (id = 1; id < r->nr_dist_base_types; id++) {
                const char *name;

                if (r->id_map[id] && r->id_map[id] != BTF_IS_EMBEDDED)
                        continue;
                dist_t = btf_type_by_id(r->dist_base_btf, id);
                name = btf__name_by_offset(r->dist_base_btf, dist_t->name_off);
                pr_warn("distilled base BTF type '%s' [%d] is not mapped to base BTF id\n",
                        name, id);
                err = -EINVAL;
                break;
        }
done:
        free(base_name_cnt);
        free(info);
        return err;
}

/* distilled base should only have named int/float/enum/fwd/struct/union types. */
static int btf_relocate_validate_distilled_base(struct btf_relocate *r)
{
        unsigned int i;

        for (i = 1; i < r->nr_dist_base_types; i++) {
                struct btf_type *t = btf_type_by_id(r->dist_base_btf, i);
                int kind = btf_kind(t);

                switch (kind) {
                case BTF_KIND_INT:
                case BTF_KIND_FLOAT:
                case BTF_KIND_ENUM:
                case BTF_KIND_STRUCT:
                case BTF_KIND_UNION:
                case BTF_KIND_FWD:
                        if (t->name_off)
                                break;
                        pr_warn("type [%d], kind [%d] is invalid for distilled base BTF; it is anonymous\n",
                                i, kind);
                        return -EINVAL;
                default:
                        pr_warn("type [%d] in distilled based BTF has unexpected kind [%d]\n",
                                i, kind);
                        return -EINVAL;
                }
        }
        return 0;
}

static int btf_relocate_rewrite_strs(struct btf_relocate *r, __u32 i)
{
        struct btf_type *t = btf_type_by_id(r->btf, i);
        struct btf_field_iter it;
        __u32 *str_off;
        int off, err;

        err = btf_field_iter_init(&it, t, BTF_FIELD_ITER_STRS);
        if (err)
                return err;

        while ((str_off = btf_field_iter_next(&it))) {
                if (!*str_off)
                        continue;
                if (*str_off >= r->dist_str_len) {
                        *str_off += r->base_str_len - r->dist_str_len;
                } else {
                        off = r->str_map[*str_off];
                        if (!off) {
                                pr_warn("string '%s' [offset %u] is not mapped to base BTF\n",
                                        btf__str_by_offset(r->btf, off), *str_off);
                                return -ENOENT;
                        }
                        *str_off = off;
                }
        }
        return 0;
}

/* If successful, output of relocation is updated BTF with base BTF pointing
 * at base_btf, and type ids, strings adjusted accordingly.
 */
int btf_relocate(struct btf *btf, const struct btf *base_btf, __u32 **id_map)
{
        unsigned int nr_types = btf__type_cnt(btf);
        const struct btf_header *dist_base_hdr;
        const struct btf_header *base_hdr;
        struct btf_relocate r = {};
        int err = 0;
        __u32 id, i;

        r.dist_base_btf = btf__base_btf(btf);
        if (!base_btf || r.dist_base_btf == base_btf)
                return -EINVAL;

        r.nr_dist_base_types = btf__type_cnt(r.dist_base_btf);
        r.nr_base_types = btf__type_cnt(base_btf);
        r.nr_split_types = nr_types - r.nr_dist_base_types;
        r.btf = btf;
        r.base_btf = base_btf;

        r.id_map = calloc(nr_types, sizeof(*r.id_map));
        r.str_map = calloc(btf_header(r.dist_base_btf)->str_len, sizeof(*r.str_map));
        dist_base_hdr = btf_header(r.dist_base_btf);
        base_hdr = btf_header(r.base_btf);
        r.dist_str_len = dist_base_hdr->str_len;
        r.base_str_len = base_hdr->str_len;
        if (!r.id_map || !r.str_map) {
                err = -ENOMEM;
                goto err_out;
        }

        err = btf_relocate_validate_distilled_base(&r);
        if (err)
                goto err_out;

        /* Split BTF ids need to be adjusted as base and distilled base
         * have different numbers of types, changing the start id of split
         * BTF.
         */
        for (id = r.nr_dist_base_types; id < nr_types; id++)
                r.id_map[id] = id + r.nr_base_types - r.nr_dist_base_types;

        /* Build a map from distilled base ids to actual base BTF ids; it is used
         * to update split BTF id references.  Also build a str_map mapping from
         * distilled base BTF names to base BTF names.
         */
        err = btf_relocate_map_distilled_base(&r);
        if (err)
                goto err_out;

        /* Next, rewrite type ids in split BTF, replacing split ids with updated
         * ids based on number of types in base BTF, and base ids with
         * relocated ids from base_btf.
         */
        for (i = 0, id = r.nr_dist_base_types; i < r.nr_split_types; i++, id++) {
                err = btf_relocate_rewrite_type_id(&r, id);
                if (err)
                        goto err_out;
        }
        /* String offsets now need to be updated using the str_map. */
        for (i = 0; i < r.nr_split_types; i++) {
                err = btf_relocate_rewrite_strs(&r, i + r.nr_dist_base_types);
                if (err)
                        goto err_out;
        }
        /* Finally reset base BTF to be base_btf */
        btf_set_base_btf(btf, base_btf);

        if (id_map) {
                *id_map = r.id_map;
                r.id_map = NULL;
        }
err_out:
        free(r.id_map);
        free(r.str_map);
        return err;
}