// SPDX-License-Identifier: GPL-2.0
// Copyright(c) 2018 Linus Torvalds. All rights reserved.
// Copyright(c) 2018 Alexei Starovoitov. All rights reserved.
// Copyright(c) 2018 Intel Corporation. All rights reserved.

#ifndef _LINUX_NOSPEC_H
#define _LINUX_NOSPEC_H

#include <linux/compiler.h>
#include <linux/types.h>
#include <asm/barrier.h>

struct task_struct;

#ifndef barrier_nospec
__respectre_inline void barrier_nospec(void)
{
}
#endif

#ifndef barrier_ssb
__respectre_inline void barrier_ssb(void)
{
    barrier_nospec();
}
#endif

/**
 * array_index_mask_nospec() - generate a ~0 mask when index < size, 0 otherwise
 * @index: array element index
 * @size: number of elements in array
 *
 * When @index is out of bounds (@index >= @size), the sign bit will be
 * set.  Extend the sign bit to all bits and invert, giving a result of
 * zero for an out of bounds index, or ~0 if within bounds [0, @size).
 */
#ifndef array_index_mask_nospec
__respectre_inline unsigned long array_index_mask_nospec(unsigned long index,
                            unsigned long size)
{
    /*
     * Always calculate and emit the mask even if the compiler
     * thinks the mask is not needed. The compiler does not take
     * into account the value of @index under speculation.
     */
    OPTIMIZER_HIDE_VAR(index);
    return 0 - (index < size);
}
#endif

#ifndef array_index_mask_nospec_extra
__respectre_inline unsigned long array_index_mask_nospec_extra(unsigned long index,
                            unsigned long size)
{
    if (__builtin_constant_p(size))
        return size ? array_index_mask_nospec(index, size) : ~0UL;

    /*
     * Always calculate and emit the mask even if the compiler
     * thinks the mask is not needed. The compiler does not take
     * into account the value of @index under speculation.
     */
    OPTIMIZER_HIDE_VAR(index);
    return 0 - (index < size) - !size;
}
#endif

#if BITS_PER_LONG == 32
#ifndef array_index_mask_nospec_ll
__respectre_inline unsigned long long array_index_mask_nospec_ll(unsigned long long index,
                            unsigned long long size)
{
    /*
     * Always calculate and emit the mask even if the compiler
     * thinks the mask is not needed. The compiler does not take
     * into account the value of @index under speculation.
     */
    OPTIMIZER_HIDE_VAR(index);
    return 0 - (index < size);
}
#endif

#ifndef array_index_mask_nospec_extra_ll
__respectre_inline unsigned long long array_index_mask_nospec_extra_ll(unsigned long long index,
                            unsigned long long size)
{
    if (__builtin_constant_p(size))
        return size ? array_index_mask_nospec_ll(index, size) : ~0ULL;

    /*
     * Always calculate and emit the mask even if the compiler
     * thinks the mask is not needed. The compiler does not take
     * into account the value of @index under speculation.
     */
    OPTIMIZER_HIDE_VAR(index);
    return 0 - (index < size) - !size;
}
#endif
#else /* BITS_PER_LONG == 64 */
extern unsigned long long array_index_mask_nospec_ll(unsigned long long index,
        unsigned long long size) __compiletime_error("Can't happen");
extern unsigned long long array_index_mask_nospec_extra_ll(unsigned long long index,
        unsigned long long size) __compiletime_error("Can't happen");
#endif

/*
 * array_index_nospec - sanitize an array index after a bounds check
 *
 * For a code sequence like:
 *
 *     if (index < size) {
 *         index = array_index_nospec(index, size);
 *         val = array[index];
 *     }
 *
 * ...if the CPU speculates past the bounds check then
 * array_index_nospec() will clamp the index within the range of [0, size).
 */
#define array_index_nospec_type(itype, stype, offset)                    \
__respectre_inline __maybe_used itype array_index_nospec_##itype(itype _index, stype size, bool extra)\
{                                            \
    stype index = _index;                                \
                                            \
    if (__builtin_constant_p(_index))                        \
        return index;                                \
                                            \
    if (!__builtin_constant_p(size))                        \
        barrier_ssb();                                \
                                            \
    if (extra) {                                    \
        if (sizeof(stype) > BITS_PER_LONG/8)                    \
            index &= array_index_mask_nospec_extra_ll(index + offset, size + offset);\
        else                                    \
            index &= array_index_mask_nospec_extra(index + offset, size + offset);\
    } else {                                    \
        if (sizeof(stype) > BITS_PER_LONG/8)                    \
            index &= array_index_mask_nospec_ll(index + offset, size + offset);\
        else                                    \
            index &= array_index_mask_nospec(index + offset, size + offset);\
    }                                        \
                                            \
    return index;                                    \
}

#define array_index_nospec_type_x(ibits, stype)                        \
    array_index_nospec_type(s##ibits, stype, (((u##ibits)(~0ULL)) >> 1) + 1)    \
    array_index_nospec_type(u##ibits, stype, 0)

array_index_nospec_type_x( 8, unsigned long)
array_index_nospec_type_x(16, unsigned long)
array_index_nospec_type_x(32, unsigned long)
#if BITS_PER_LONG == 32
array_index_nospec_type_x(64, unsigned long long)
#else
array_index_nospec_type_x(64, unsigned long)
#endif

#undef array_index_nospec_type_x

#define choose_array_index_nospec_x(type,sign,bits,index,size)        \
    __builtin_choose_expr(__builtin_types_compatible_p(typeof(index), type), \
                   array_index_nospec_##sign##bits((index), (size), false),

#define choose_array_index_nospec(itype,bits,index,size)    \
    choose_array_index_nospec_x(signed itype,s,bits,(index),(size))    \
    choose_array_index_nospec_x(unsigned itype,u,bits,(index),(size))

#define array_index_nospec(index, size)                    \
    choose_array_index_nospec(char,            8,(index),(size))    \
    choose_array_index_nospec(short,          16,(index),(size))    \
    choose_array_index_nospec(int,            32,(index),(size))    \
    choose_array_index_nospec(long,BITS_PER_LONG,(index),(size))    \
    (void)0))))))))

#define access_ok_nospec(addr, size) ({            \
    unsigned long limit = TASK_SIZE_MAX;        \
/*    if (!__builtin_constant_p(size))        \
        size = array_index_nospec(size, limit);    \
*/    addr = (typeof(addr))array_index_nospec((unsigned long)addr, limit - size + 1); \
})

/* Speculation control prctl */
int arch_prctl_spec_ctrl_get(struct task_struct *task, unsigned long which);
int arch_prctl_spec_ctrl_set(struct task_struct *task, unsigned long which,
                 unsigned long ctrl);
/* Speculation control for seccomp enforced mitigation */
void arch_seccomp_spec_mitigate(struct task_struct *task);

#endif /* _LINUX_NOSPEC_H */