//--------------------------------------------------------------------------
// Copyright (C) 2015-2025 Cisco and/or its affiliates. All rights reserved.
//
// This program is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License Version 2 as published
// by the Free Software Foundation.  You may not use, modify or distribute
// this program under any other version of the GNU General Public License.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
//--------------------------------------------------------------------------

// hyperscan.cc author Russ Combs <rucombs@cisco.com>

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <algorithm>
#include <cassert>
#include <cstring>
#include <fstream>
#include <mutex>
#include <sstream>

#include <hs_compile.h>
#include <hs_runtime.h>

#include "framework/module.h"
#include "framework/mpse.h"
#include "hash/hashes.h"
#include "helpers/scratch_allocator.h"
#include "log/log_stats.h"
#include "log/messages.h"
#include "main/snort_config.h"
#include "main/thread.h"
#include "profiler/profiler.h"

using namespace snort;

static const char* s_name = "hyperscan";
static const char* s_help = "intel hyperscan-based MPSE with regex support";

//-------------------------------------------------------------------------
// pattern foo
//-------------------------------------------------------------------------

namespace
{
struct Pattern
{
    std::string pat;
    unsigned len;
    unsigned flags;
    bool no_case;
    bool negate;

    void* user;
    void* user_tree;
    void* user_list;

    Pattern(const uint8_t*, unsigned, const Mpse::PatternDescriptor&, void*);
    void escape(const uint8_t*, unsigned, bool);
};
}

Pattern::Pattern(
    const uint8_t* s, unsigned n, const Mpse::PatternDescriptor& d, void* u)
{
    escape(s, n, d.literal);

    len = n;
    no_case = d.no_case;
    negate = d.negated;
    flags = d.flags;
    user = u;
    user_tree = user_list = nullptr;

    if ( no_case )
        flags |= HS_FLAG_CASELESS;

    if ( !d.multi_match )
        flags |= HS_FLAG_SINGLEMATCH;
}

void Pattern::escape(const uint8_t* s, unsigned n, bool literal)
{
    for ( unsigned i = 0; i < n; ++i )
    {
        if ( !isprint(s[i]) )
        {
            char hex[5];
            snprintf(hex, sizeof(hex), "\\x%02X", s[i]);
            pat += hex;
        }
        else
        {
            const char* special = ".^$*+?()[]{}\\|";

            if ( literal and strchr(special, s[i]) )
                pat += '\\';

            pat += s[i];
        }
    }
}

static bool compare(const Pattern& a, const Pattern& b)
{
    if ( a.pat != b.pat )
        return a.pat < b.pat;

    return a.flags < b.flags;
}

typedef std::vector<Pattern> PatternVector;

//-------------------------------------------------------------------------
// scratch
//-------------------------------------------------------------------------

static std::mutex s_mutex;
static hs_scratch_t* s_scratch = nullptr;
static unsigned int scratch_index;
static ScratchAllocator* scratcher = nullptr;

static bool scratch_setup(SnortConfig* sc)
{
    for ( unsigned i = 0; i < sc->num_slots; ++i )
    {
        hs_scratch_t** ss = (hs_scratch_t**) &sc->state[i][scratch_index];
        hs_clone_scratch(s_scratch, ss);
    }

    return true;
}

static void scratch_cleanup(SnortConfig* sc)
{
    for ( unsigned i = 0; i < sc->num_slots; ++i )
    {
        hs_scratch_t* ss = (hs_scratch_t*)sc->state[i][scratch_index];
        hs_free_scratch(ss);
        sc->state[i][scratch_index] = nullptr;
    }
}

static void scratch_update(SnortConfig* sc)
{
    hs_scratch_t** ss = (hs_scratch_t**) &sc->state[get_instance_id()][scratch_index];

    if ( *ss == s_scratch )
        return;

    hs_free_scratch(*ss);
    *ss = nullptr;

    hs_clone_scratch(s_scratch, ss);
}

//-------------------------------------------------------------------------
// module
//-------------------------------------------------------------------------

struct FullCounts
{
    PegCount searches;
    PegCount matches;
    PegCount bytes;
};

static THREAD_LOCAL FullCounts hyper_counts;
static THREAD_LOCAL ProfileStats hyper_stats;

const PegInfo hyper_pegs[] =
{
    { CountType::SUM, "searches", "number of search attempts" },
    { CountType::SUM, "matches", "number of times a match was found" },
    { CountType::SUM, "bytes", "total bytes searched" },

    { CountType::END, nullptr, nullptr }
};

class HyperscanModule : public Module
{
public:
    HyperscanModule() : Module(s_name, s_help)
    {
        scratcher = new SimpleScratchAllocator(scratch_setup, scratch_cleanup, scratch_update);
        scratch_index = scratcher->get_id();
    }

    ~HyperscanModule() override
    {
        delete scratcher;
        hs_free_scratch(s_scratch);
        s_scratch = nullptr;
    }
    ProfileStats* get_profile() const override
    { return &hyper_stats; }

    const PegInfo* get_pegs() const override
    { return hyper_pegs; }

    PegCount* get_counts() const override
    { return (PegCount*)&hyper_counts; }

    Usage get_usage() const override
    { return GLOBAL; }
};

//-------------------------------------------------------------------------
// mpse
//-------------------------------------------------------------------------

class HyperscanMpse : public Mpse
{
public:
    HyperscanMpse(const MpseAgent* a)
        : Mpse("hyperscan")
    {
        agent = a;
        ++instances;
    }

    ~HyperscanMpse() override
    {
        if ( hs_db )
            hs_free_database(hs_db);

        if ( agent )
            user_dtor();
    }

    int add_pattern(
        const uint8_t* pat, unsigned len, const PatternDescriptor& desc, void* user) override
    {
        Pattern p(pat, len, desc, user);
        pvector.emplace_back(p);
        ++patterns;
        return 0;
    }

    int prep_patterns(SnortConfig*) override;
    void reuse_search() override;

    int search(const uint8_t*, int, MpseMatch, void*, int*) override;

    int get_pattern_count() const override
    { return pvector.size(); }

    int match(unsigned id, unsigned long long to, MpseMatch match_cb, void* match_ctx);

    static int match(
        unsigned id, unsigned long long from, unsigned long long to,
        unsigned flags, void*);

    int serialize(uint8_t*&, size_t&) const override;
    bool deserialize(const uint8_t*, size_t) override;

    void get_hash(std::string&) override;

private:
    void user_ctor(SnortConfig*);
    void user_dtor();

    const MpseAgent* agent;
    PatternVector pvector;

    hs_database_t* hs_db = nullptr;
    bool compiled = false;

public:
    static uint64_t instances;
    static uint64_t patterns;
};

uint64_t HyperscanMpse::instances = 0;
uint64_t HyperscanMpse::patterns = 0;

int HyperscanMpse::serialize(uint8_t*& buf, size_t& sz) const
{
    if ( !compiled )
        return 0;

    return (hs_db and (hs_serialize_database(hs_db, (char**)&buf, &sz) == HS_SUCCESS) and buf) ? 1 : -1;
}

bool HyperscanMpse::deserialize(const uint8_t* buf, size_t sz)
{
    if ( hs_error_t err = hs_deserialize_database((const char*)buf, sz, &hs_db) )
    {
        ParseWarning(WARN_RULES, "can't deserialize hyperscan database (%d)", err);
        return false;
    }

    if ( hs_error_t err = hs_alloc_scratch(hs_db, &s_scratch) )
    {
        ParseError("can't allocate search scratch space (%d)", err);
        return false;
    }

    return true;
}

// other mpse have direct access to their fsm match states and populate
// user list and tree with each pattern that leads to the same match state.
// with hyperscan we don't have internal fsm knowledge and each fast
// pattern is considered to be in a distinct match state.  the resulting
// detection option trees for hyperscan are thus just single option chains.

void HyperscanMpse::user_ctor(SnortConfig* sc)
{
    for ( auto& p : pvector )
    {
        if ( p.user )
        {
            if ( p.negate )
                agent->negate_list(p.user, &p.user_list);
            else
                agent->build_tree(sc, p.user, &p.user_tree);
        }
        agent->build_tree(sc, nullptr, &p.user_tree);
    }
}

void HyperscanMpse::user_dtor()
{
    for ( auto& p : pvector )
    {
        if ( p.user )
            agent->user_free(p.user);

        if ( p.user_list )
            agent->list_free(&p.user_list);

        if ( p.user_tree )
            agent->tree_free(&p.user_tree);
    }
}

int HyperscanMpse::prep_patterns(SnortConfig* sc)
{
    if ( hs_db )
    {
        if ( agent )
            user_ctor(sc);

        return 0;
    }

    if ( pvector.empty() )
        return -1;

    if ( hs_valid_platform() != HS_SUCCESS )
    {
        ParseError("This host does not support Hyperscan.");
        return -1;
    }

    // sort for consistent serialization
    std::sort(pvector.begin(), pvector.end(), compare);

    hs_compile_error_t* errptr = nullptr;
    std::vector<const char*> pats;
    std::vector<unsigned> flags;
    std::vector<unsigned> ids;

    unsigned id = 0;

    for ( auto& p : pvector )
    {
        pats.emplace_back(p.pat.c_str());
        flags.emplace_back(p.flags);
        ids.emplace_back(id++);
    }

#ifdef REG_TEST
    hs_platform_info_t info = { HS_TUNE_FAMILY_SNB, 0, 0, 0 };
    auto pinfo = &info;
#else
    hs_platform_info_t* pinfo = nullptr;
#endif

    if ( hs_compile_multi(&pats[0], &flags[0], &ids[0], pvector.size(), HS_MODE_BLOCK,
        pinfo, &hs_db, &errptr) or !hs_db )
    {
        ParseError("can't compile hyperscan pattern database: %s (%d) - '%s'",
            errptr->message, errptr->expression,
            errptr->expression >= 0 ? pats[errptr->expression] : "");
        hs_free_compile_error(errptr);
        return -2;
    }

    compiled = true;

    if ( agent )
        user_ctor(sc);

    std::lock_guard<std::mutex> lock(s_mutex);

    if ( hs_error_t err = hs_alloc_scratch(hs_db, &s_scratch) )
    {
        ParseError("can't allocate search scratch space (%d)", err);
        return -3;
    }

    return 0;
}

void HyperscanMpse::get_hash(std::string& hash)
{
    if ( !hs_db )
        std::sort(pvector.begin(), pvector.end(), compare);

    std::stringstream ss;

    for ( const auto& p : pvector )
        ss << p.pat << p.flags;

    std::string str = ss.str();
    uint8_t buf[MD5_HASH_SIZE];

    md5((const uint8_t*)str.c_str(), str.size(), buf);
    hash.assign((const char*)buf, sizeof(buf));
}

void HyperscanMpse::reuse_search()
{
    if ( pvector.empty() )
        return;

    if ( hs_error_t err = hs_alloc_scratch(hs_db, &s_scratch) )
        ErrorMessage("can't allocate search scratch space (%d)", err);
}

int HyperscanMpse::match(unsigned id, unsigned long long to, MpseMatch match_cb, void* match_ctx)
{
    assert(id < pvector.size());
    Pattern& p = pvector[id];
    hyper_counts.matches++;
    return match_cb(p.user, p.user_tree, (int)to, match_ctx, p.user_list);
}

namespace
{
struct ScanContext
{
    HyperscanMpse* mpse;
    MpseMatch match_cb;
    void* match_ctx;
    int nfound = 0;

    ScanContext(HyperscanMpse* m, MpseMatch cb, void* ctx)
    { mpse = m; match_cb = cb; match_ctx = ctx; }

};
}

int HyperscanMpse::match(
    unsigned id, unsigned long long /*from*/, unsigned long long to,
    unsigned /*flags*/, void* pv)
{
    ScanContext* scan = (ScanContext*)pv;
    scan->nfound++;
    return  scan->mpse->match(id, to, scan->match_cb, scan->match_ctx);
}

int HyperscanMpse::search(
    const uint8_t* buf, int n, MpseMatch mf, void* pv, int* current_state)
{
    Profile profile(hyper_stats);  // cppcheck-suppress unreadVariable

    *current_state = 0;
    ScanContext scan(this, mf, pv);

    hs_scratch_t* ss =
        (hs_scratch_t*)SnortConfig::get_conf()->state[get_instance_id()][scratch_index];

    // scratch is null for the degenerate case w/o patterns
    assert(!hs_db or ss);

    hyper_counts.searches++;
    hyper_counts.bytes += n;

    hs_scan(hs_db, (const char*)buf, n, 0, ss, HyperscanMpse::match, &scan);

    return scan.nfound;
}

//-------------------------------------------------------------------------
// api
//-------------------------------------------------------------------------

static Module* mod_ctor()
{ return new HyperscanModule; }

static void mod_dtor(Module* p)
{ delete p; }

static Mpse* hs_ctor(const SnortConfig*, class Module*, const MpseAgent* a)
{
    return new HyperscanMpse(a);
}

static void hs_dtor(Mpse* p)
{
    delete p;
}

static void hs_init()
{
    HyperscanMpse::instances = 0;
    HyperscanMpse::patterns = 0;
}

static void hs_print()
{
    LogCount("instances", HyperscanMpse::instances);
    LogCount("patterns", HyperscanMpse::patterns);
}

static const MpseApi hs_api =
{
    {
        PT_SEARCH_ENGINE,
        sizeof(MpseApi),
        SEAPI_VERSION,
        1,
        API_RESERVED,
        API_OPTIONS,
        s_name,
        s_help,
        mod_ctor,
        mod_dtor
    },
    MPSE_REGEX | MPSE_MTBLD,
    nullptr,  // activate
    nullptr,  // setup
    nullptr,  // start
    nullptr,  // stop
    hs_ctor,
    hs_dtor,
    hs_init,
    hs_print,
    nullptr,
};

#ifdef BUILDING_SO
SO_PUBLIC const BaseApi* snort_plugins[] =
#else
const BaseApi* se_hyperscan[] =
#endif
{
    &hs_api.base,
    nullptr
};

