1/* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */23#include "proxy.h"45enum mcp_ins_type {6    INS_REQ = 1,7    INS_RES,8};910enum mcp_ins_steptype {11    mcp_ins_step_none = 0,12    mcp_ins_step_sepkey,13    mcp_ins_step_keybegin,14    mcp_ins_step_keyis,15    mcp_ins_step_hasflag,16    mcp_ins_step_flagtoken,17    mcp_ins_step_flagint,18    mcp_ins_step_flagis,19    mcp_ins_step_final, // not used.20};2122// START STEP STRUCTS2324struct mcp_ins_sepkey {25    char sep;26    int pos;27    int mapref;28};2930struct mcp_ins_string {31    unsigned int str; // arena offset for match string.32    unsigned int len;33};3435struct mcp_ins_flag {36    uint64_t bit; // flag converted for bitmask test37    char f;38};3940// TODO: it might make more sense to flatten the structs into the ins_step41// struct. It wouldn't take much more space if we can be careful with42// alignment.43struct mcp_ins_flagstr {44    unsigned int str;45    unsigned int len;46    uint64_t bit; // flag bit47    char f;48};4950struct mcp_ins_step {51    enum mcp_ins_steptype type;52    union {53        struct mcp_ins_sepkey sepkey;54        struct mcp_ins_string string;55        struct mcp_ins_flag flag;56        struct mcp_ins_flagstr flagstr;57    } c;58};5960// END STEP STRUCTS6162struct mcp_inspector {63    enum mcp_ins_type type;64    int scount;65    unsigned int aused; // arena memory used66    unsigned int rcount; // number of results to expect67    char *arena; // string/data storage for steps68    struct mcp_ins_step steps[];69};7071// PRIVATE INTERFACE7273#define res_buf(r) (r->cresp ? r->cresp->iov[0].iov_base : r->buf)7475// COMMON ARG HANDLERS7677// multiple step types only take 'flag' as an argument.78static int mcp_inspector_flag_c_g(lua_State *L, int tidx) {79    if (lua_getfield(L, tidx, "flag") != LUA_TNIL) {80        size_t len = 0;81        const char *flag = lua_tolstring(L, -1, &len);82        if (len != 1) {83            proxy_lua_ferror(L, "inspector step %d: 'flag' must be a single character", tidx);84        }85        if (mcp_is_flag_invalid(flag[0])) {86            proxy_lua_ferror(L, "inspect step %d: 'flag' must be alphanumeric", tidx);87        }88    } else {89        proxy_lua_ferror(L, "inspector step %d: must provide 'flag' argument", tidx);90    }91    lua_pop(L, 1); // val or nil92    return 0;93}9495static int mcp_inspector_flag_i_g(lua_State *L, int tidx, int sc, struct mcp_inspector *ins) {96    struct mcp_ins_step *s = &ins->steps[sc];97    struct mcp_ins_flag *c = &s->c.flag;9899    if (lua_getfield(L, tidx, "flag") != LUA_TNIL) {100        const char *flag = lua_tostring(L, -1);101        c->f = flag[0];102        c->bit = (uint64_t)1 << (c->f - 65);103    }104    lua_pop(L, 1); // val or nil105106    return 0;107}108109static int mcp_inspector_string_c_g(lua_State *L, int tidx) {110    size_t len = 0;111112    if (lua_getfield(L, tidx, "str") != LUA_TNIL) {113        lua_tolstring(L, -1, &len);114        if (len < 1) {115            proxy_lua_ferror(L, "inspector step %d: 'str' must have nonzero length", tidx);116        }117    } else {118        proxy_lua_ferror(L, "inspector step %d: must provide 'str' argument", tidx);119    }120    lua_pop(L, 1); // val or nil121122    return len;123}124125static int mcp_inspector_string_i_g(lua_State *L, int tidx, int sc, struct mcp_inspector *ins) {126    struct mcp_ins_step *s = &ins->steps[sc];127    struct mcp_ins_string *c = &s->c.string;128    size_t len = 0;129130    // store our match string in the arena space that we reserved before.131    if (lua_getfield(L, tidx, "str") != LUA_TNIL) {132        const char *str = lua_tolstring(L, -1, &len);133        c->str = ins->aused;134        c->len = len;135        char *a = ins->arena + ins->aused;136        memcpy(a, str, len);137        ins->aused += len;138    }139    lua_pop(L, 1); // val or nil140141    return len;142}143144// END COMMMON ARG HANDLERS145146static int mcp_inspector_sepkey_c(lua_State *L, int tidx) {147    if (lua_getfield(L, tidx, "sep") != LUA_TNIL) {148        size_t len = 0;149        lua_tolstring(L, -1, &len);150        if (len != 1) {151            proxy_lua_ferror(L, "inspector step %d: separator must be one character", tidx);152        }153    }154    lua_pop(L, 1); // val or nil155156    if (lua_getfield(L, tidx, "pos") != LUA_TNIL) {157        luaL_checktype(L, -1, LUA_TNUMBER);158    }159    lua_pop(L, 1); // val or nil160161    if (lua_getfield(L, tidx, "map") != LUA_TNIL) {162        luaL_checktype(L, -1, LUA_TTABLE);163    }164    lua_pop(L, 1); // val or nil165166    return 0;167}168169// initializer. arguments already checked, so just fill out the slot.170static int mcp_inspector_sepkey_i(lua_State *L, int tidx, int sc, struct mcp_inspector *ins) {171    struct mcp_ins_step *s = &ins->steps[sc];172    struct mcp_ins_sepkey *c = &s->c.sepkey;173174    if (lua_getfield(L, tidx, "sep") != LUA_TNIL) {175        const char *sep = lua_tostring(L, -1);176        c->sep = sep[0];177    } else {178        // default separator179        c->sep = '/';180    }181    lua_pop(L, 1); // val or nil182183    if (lua_getfield(L, tidx, "pos") != LUA_TNIL) {184        c->pos = lua_tointeger(L, -1);185    } else {186        c->pos = 1;187    }188    lua_pop(L, 1);189190    if (lua_getfield(L, tidx, "map") != LUA_TNIL) {191        c->mapref = luaL_ref(L, LUA_REGISTRYINDEX);192    } else {193        c->mapref = 0;194        lua_pop(L, 1);195    }196    // ref was popped197198    return 0;199}200201// TODO: abstract out the token-position-finder202static int mcp_inspector_sepkey_r(lua_State *L, struct mcp_inspector *ins, struct mcp_ins_step *s, void *arg) {203    mcp_request_t *rq = arg;204    struct mcp_ins_sepkey *c = &s->c.sepkey;205206    const char *key = MCP_PARSER_KEY(&rq->pr);207    const char *end = key + rq->pr.klen;208    char sep = c->sep;209    int pos = c->pos;210211    // skip initial separators212    while (key != end) {213        if (*key == sep) {214            key++;215        } else {216            break;217        }218    }219    const char *token = key;220    int tlen = 0;221222    while (key != end) {223        if (*key == sep) {224            // measure token length and stop if at position.225            if (--pos == 0) {226                tlen = key - token;227                break;228            } else {229                // NOTE: this could point past the end of the key, but unless230                // it's the token we want we won't look at it.231                token = key+1;232            }233        }234        key++;235    }236237    // either the separator was never found, or we ended before finding238    // another one, which gives us an end token.239    if (pos == 1) {240        tlen = key - token;241    }242243    // now have *token and tlen244    if (tlen != 0) {245        if (c->mapref) {246            // look up this string against the map.247            // NOTE: this still ends up creating a garbage string. However,248            // since the map is internal we can optimize this later by moving249            // the map lookup op to C.250            lua_rawgeti(L, LUA_REGISTRYINDEX, c->mapref);251            lua_pushlstring(L, token, tlen);252            lua_rawget(L, -2); // pops string.253            lua_remove(L, -2); // removes map, shifts lookup result down.254            // stack should be clean: just the result.255        } else {256            // no map, return the actual token.257            lua_pushlstring(L, token, tlen);258        }259    } else {260        lua_pushnil(L); // not found.261    }262263    return 1;264}265266static int mcp_inspector_keybegin_r(lua_State *L, struct mcp_inspector *ins, struct mcp_ins_step *s, void *arg) {267    mcp_request_t *rq = arg;268    struct mcp_ins_string *c = &s->c.string;269270    const char *key = MCP_PARSER_KEY(&rq->pr);271    int klen = rq->pr.klen;272    const char *str = ins->arena + c->str;273274    if (c->len < klen && strncmp(key, str, c->len) == 0) {275        lua_pushboolean(L, 1);276    } else {277        lua_pushboolean(L, 0);278    }279280    return 1;281}282283static int mcp_inspector_keyis_r(lua_State *L, struct mcp_inspector *ins, struct mcp_ins_step *s, void *arg) {284    mcp_request_t *rq = arg;285    struct mcp_ins_string *c = &s->c.string;286287    const char *key = MCP_PARSER_KEY(&rq->pr);288    int klen = rq->pr.klen;289    const char *str = ins->arena + c->str;290291    if (c->len == klen && strncmp(key, str, c->len) == 0) {292        lua_pushboolean(L, 1);293    } else {294        lua_pushboolean(L, 0);295    }296297    return 1;298}299300static int mcp_inspector_hasflag_r(lua_State *L, struct mcp_inspector *ins, struct mcp_ins_step *s, void *arg) {301    struct mcp_ins_flag *c = &s->c.flag;302    if (ins->type == INS_REQ) {303        mcp_request_t *rq = arg;304        // requests should always be tokenized, so we can just check the bit.305        if (mcmc_token_has_flag_bit(&rq->pr.tok, c->bit) == MCMC_OK) {306            lua_pushboolean(L, 1);307        } else {308            lua_pushboolean(L, 0);309        }310    } else {311        mcp_resp_t *res = arg;312        if (res->resp.type == MCMC_RESP_META) {313            // result object may not be tokenized. this will do so if not314            // already. any future hits agains the same object will use the315            // cached tokenizer struct.316            mcmc_tokenize_res(res_buf(res), res->resp.reslen, &res->tok);317            if (mcmc_token_has_flag_bit(&res->tok, c->bit) == MCMC_OK) {318                lua_pushboolean(L, 1);319            } else {320                lua_pushboolean(L, 0);321            }322        } else {323            proxy_lua_error(L, "inspector error: response is not meta protocol");324        }325    }326    return 1;327}328329// This mirrors `bool, (str|nil) = r:flag_token("T")`330static int mcp_inspector_flagtoken_r(lua_State *L, struct mcp_inspector *ins, struct mcp_ins_step *s, void *arg) {331    struct mcp_ins_flag *c = &s->c.flag;332    if (ins->type == INS_REQ) {333        mcp_request_t *rq = arg;334335        if (mcmc_token_has_flag_bit(&rq->pr.tok, c->bit) == MCMC_OK) {336            lua_pushboolean(L, 1); // flag exists337            int tlen = 0;338            const char *tok = mcmc_token_get_flag(rq->pr.request, &rq->pr.tok, c->f, &tlen);339            lua_pushlstring(L, tok, tlen); // flag's token340            return 2;341        }342    } else {343        mcp_resp_t *res = arg;344        if (res->resp.type == MCMC_RESP_META) {345            mcmc_tokenize_res(res_buf(res), res->resp.reslen, &res->tok);346            if (mcmc_token_has_flag_bit(&res->tok, c->bit) == MCMC_OK) {347                lua_pushboolean(L, 1); // flag exists348                int tlen = 0;349                const char *tok = mcmc_token_get_flag(res_buf(res), &res->tok, c->f, &tlen);350                lua_pushlstring(L, tok, tlen); // flag's token351                return 2;352            }353        }354    }355    lua_pushboolean(L, 0);356    lua_pushnil(L);357358    return 2;359}360361// TODO: flaguint variant?362// still stuck as signed in lua but would reject signed tokens363static int mcp_inspector_flagint_r(lua_State *L, struct mcp_inspector *ins, struct mcp_ins_step *s, void *arg) {364    struct mcp_ins_flag *c = &s->c.flag;365    if (ins->type == INS_REQ) {366        mcp_request_t *rq = arg;367368        if (mcmc_token_has_flag_bit(&rq->pr.tok, c->bit) == MCMC_OK) {369            lua_pushboolean(L, 1); // flag exists370            int64_t tok = 0;371            int status = mcmc_token_get_flag_64(rq->pr.request, &rq->pr.tok, c->f, &tok);372            if (status == MCMC_OK) {373                lua_pushinteger(L, tok);374            } else {375                // Potential error messages?376                lua_pushnil(L);377            }378            return 2;379        }380    } else {381        mcp_resp_t *res = arg;382        if (res->resp.type == MCMC_RESP_META) {383            mcmc_tokenize_res(res_buf(res), res->resp.reslen, &res->tok);384            if (mcmc_token_has_flag_bit(&res->tok, c->bit) == MCMC_OK) {385                lua_pushboolean(L, 1); // flag exists386                int64_t tok = 0;387                if (mcmc_token_get_flag_64(res_buf(res), &res->tok, c->f, &tok) == MCMC_OK) {388                    lua_pushinteger(L, tok);389                } else {390                    lua_pushnil(L); // token couldn't be converted391                }392                return 2;393            }394        }395    }396    lua_pushboolean(L, 0);397    lua_pushnil(L);398399    return 2;400}401402static int mcp_inspector_flagstr_c(lua_State *L, int tidx) {403    mcp_inspector_flag_c_g(L, tidx);404    int size = mcp_inspector_string_c_g(L, tidx);405    return size;406}407408static int mcp_inspector_flagstr_i(lua_State *L, int tidx, int sc, struct mcp_inspector *ins) {409    // TODO: if we never use mcp_ins_step we can remove it and just pass parts410    // of the relevant structs down into these functions.411    struct mcp_ins_step *s = &ins->steps[sc];412    struct mcp_ins_flagstr *c = &s->c.flagstr;413    size_t len = 0;414415    if (lua_getfield(L, tidx, "flag") != LUA_TNIL) {416        const char *flag = lua_tostring(L, -1);417        c->f = flag[0];418        c->bit = (uint64_t)1 << (c->f - 65);419    }420    lua_pop(L, 1); // val or nil421422    if (lua_getfield(L, tidx, "str") != LUA_TNIL) {423        const char *str = lua_tolstring(L, -1, &len);424        c->str = ins->aused;425        c->len = len;426        char *a = ins->arena + ins->aused;427        memcpy(a, str, len);428        ins->aused += len;429    }430    lua_pop(L, 1); // val or nil431432    return len;433}434435static int mcp_inspector_flagis_r(lua_State *L, struct mcp_inspector *ins, struct mcp_ins_step *s, void *arg) {436    struct mcp_ins_flagstr *c = &s->c.flagstr;437    const char *str = ins->arena + c->str;438    if (ins->type == INS_REQ) {439        mcp_request_t *rq = arg;440441        if (mcmc_token_has_flag_bit(&rq->pr.tok, c->bit) == MCMC_OK) {442            lua_pushboolean(L, 1); // flag exists443            int tlen = -1;444            const char *tok = mcmc_token_get_flag(rq->pr.request, &rq->pr.tok, c->f, &tlen);445446            if (tlen == c->len && strncmp(tok, str, c->len) == 0) {447                lua_pushboolean(L, 1);448            } else {449                lua_pushboolean(L, 0);450            }451            return 2;452        }453    } else {454        mcp_resp_t *res = arg;455        if (res->resp.type == MCMC_RESP_META) {456            mcmc_tokenize_res(res_buf(res), res->resp.reslen, &res->tok);457            if (mcmc_token_has_flag_bit(&res->tok, c->bit) == MCMC_OK) {458                lua_pushboolean(L, 1); // flag exists459                int tlen = 0;460                const char *tok = mcmc_token_get_flag(res_buf(res), &res->tok, c->f, &tlen);461                if (tlen == c->len && strncmp(tok, str, c->len) == 0) {462                    lua_pushboolean(L, 1);463                } else {464                    lua_pushboolean(L, 0);465                }466                return 2;467            }468        }469    }470    lua_pushboolean(L, 0);471    lua_pushnil(L);472473    return 2;474}475476// END STEPS477478typedef int (*mcp_ins_c)(lua_State *L, int tidx);479typedef int (*mcp_ins_i)(lua_State *L, int tidx, int sc, struct mcp_inspector *ins);480typedef int (*mcp_ins_r)(lua_State *L, struct mcp_inspector *ins, struct mcp_ins_step *s, void *arg);481482struct mcp_ins_entry {483    const char *s; // string name484    mcp_ins_c c;485    mcp_ins_i i;486    mcp_ins_r r;487    unsigned int t; // allowed object types488    int n; // number of results to expect489};490491static const struct mcp_ins_entry mcp_ins_entries[] = {492    [mcp_ins_step_none] = {NULL, NULL, NULL, NULL, 0, 0},493    [mcp_ins_step_sepkey] = {"sepkey", mcp_inspector_sepkey_c, mcp_inspector_sepkey_i, mcp_inspector_sepkey_r, INS_REQ, 1},494    [mcp_ins_step_keybegin] = {"keybegin", mcp_inspector_string_c_g, mcp_inspector_string_i_g, mcp_inspector_keybegin_r, INS_REQ, 1},495    [mcp_ins_step_keyis] = {"keyis", mcp_inspector_string_c_g, mcp_inspector_string_i_g, mcp_inspector_keyis_r, INS_REQ, 1},496    [mcp_ins_step_hasflag] = {"hasflag", mcp_inspector_flag_c_g, mcp_inspector_flag_i_g, mcp_inspector_hasflag_r, INS_REQ|INS_RES, 1},497    [mcp_ins_step_flagtoken] = {"flagtoken", mcp_inspector_flag_c_g, mcp_inspector_flag_i_g, mcp_inspector_flagtoken_r, INS_REQ|INS_RES, 2},498    [mcp_ins_step_flagint] = {"flagint", mcp_inspector_flag_c_g, mcp_inspector_flag_i_g, mcp_inspector_flagint_r, INS_REQ|INS_RES, 2},499    [mcp_ins_step_flagis] = {"flagis", mcp_inspector_flagstr_c, mcp_inspector_flagstr_i, mcp_inspector_flagis_r, INS_REQ|INS_RES, 2},500};501502// call with type string on top503static enum mcp_ins_steptype mcp_inspector_steptype(lua_State *L) {504    const char *type = luaL_checkstring(L, -1);505    for (int x = 0; x < mcp_ins_step_final; x++) {506        const struct mcp_ins_entry *e = &mcp_ins_entries[x];507        if (e->s && strcmp(type, e->s) == 0) {508            return x;509        }510    }511    return mcp_ins_step_none;512}513514// - arguments given as list of tables:515//   { t = "type", arg = "bar", etc },516//   { etc }517//   - can take table-of-tables via: mcp.req_inspector_new(table.unpack(args))518// NOTES:519// - can we inline necessary strings/etc via extra allocated memory?520// - can we get mcp.inspector metatable into the upvalue of the _call and _gc521// funcs for fast-compare?522static int mcp_inspector_new(lua_State *L, enum mcp_ins_type type) {523    int argc = lua_gettop(L);524    size_t size = 0;525    int scount = 0;526527    // loop argument tables once for validation and pre-calculations.528    for (int x = 1; x <= argc; x++) {529        luaL_checktype(L, x, LUA_TTABLE);530        if (lua_getfield(L, x, "t") != LUA_TNIL) {531            enum mcp_ins_steptype st = mcp_inspector_steptype(L);532            const struct mcp_ins_entry *e = &mcp_ins_entries[st];533            if (!(e->t & type)) {534                proxy_lua_ferror(L, "inspector step %d: step incompatible with inspector type", x);535            }536            if ((st == mcp_ins_step_none) || e->c == NULL) {537                proxy_lua_ferror(L, "inspector step %d: unknown step type", x);538            }539            size += e->c(L, x);540        }541        lua_pop(L, 1); // drop 't' or nil542        scount++;543    }544545    // we now know the size and number of steps. allocate some flat memory.546547    // TODO: we need memory for steps + arbitrary step data. (ie; string stems548    // and the like)549    // - now: single extra malloc, divvy out the buffer as requested550    // - later: if alignment of the flexible step array can be reliably551    // determined (C11 alignas or etc), inline memory can be used instead.552    size_t extsize = sizeof(struct mcp_ins_step) * scount;553    struct mcp_inspector *ins = lua_newuserdatauv(L, sizeof(*ins) + extsize, 1);554    memset(ins, 0, sizeof(*ins));555556    ins->arena = malloc(size);557    if (ins->arena == NULL) {558        proxy_lua_error(L, "mcp.req_inspector_new: failed to allocate memory");559    }560561    luaL_setmetatable(L, "mcp.inspector");562    switch (type) {563        case INS_REQ:564            luaL_getmetatable(L, "mcp.request");565            break;566        case INS_RES:567            luaL_getmetatable(L, "mcp.response");568            break;569    }570    // set metatable to the upvalue for a fast comparison during __call571    lua_setiuservalue(L, -2, 1);572    ins->type = type;573574    // loop the arg tables again to fill in the steps575    // skip checks since we did that during the first loop.576    scount = 0;577    for (int x = 1; x <= argc; x++) {578        if (lua_getfield(L, x, "t") != LUA_TNIL) {579            enum mcp_ins_steptype st = mcp_inspector_steptype(L);580            ins->steps[scount].type = st;581            mcp_ins_entries[st].i(L, x, scount, ins);582            ins->rcount += mcp_ins_entries[st].n;583        }584        lua_pop(L, 1); // drop t or nil585        scount++;586    }587588    if (size != ins->aused) {589        proxy_lua_error(L, "inspector failed to properly initialize, memory not filled correctly");590    }591    ins->scount = scount;592593    return 1;594}595596static int mcp_ins_run(lua_State *L, struct mcp_inspector *ins, void *arg) {597    int ret = 0;598599    for (int x = 0; x < ins->scount; x++) {600        struct mcp_ins_step *s = &ins->steps[x];601        assert(s->type != mcp_ins_step_none);602        ret += mcp_ins_entries[s->type].r(L, ins, s, arg);603    }604605    return ret;606}607608// PUBLIC INTERFACE609610int mcplib_req_inspector_new(lua_State *L) {611    return mcp_inspector_new(L, INS_REQ);612}613614int mcplib_res_inspector_new(lua_State *L) {615    return mcp_inspector_new(L, INS_RES);616}617618// walk each step and free references/memory/etc619int mcplib_inspector_gc(lua_State *L) {620    struct mcp_inspector *ins = lua_touserdata(L, 1);621622    if (ins->arena) {623        free(ins->arena);624        ins->arena = NULL;625    }626627    for (int x = 0; x < ins->scount; x++) {628        struct mcp_ins_step *s = &ins->steps[x];629        switch (s->type) {630            case mcp_ins_step_sepkey:631                if (s->c.sepkey.mapref) {632                    luaL_unref(L, LUA_REGISTRYINDEX, s->c.sepkey.mapref);633                    s->c.sepkey.mapref = 0;634                }635                break;636            case mcp_ins_step_keybegin:637            case mcp_ins_step_keyis:638            case mcp_ins_step_hasflag:639            case mcp_ins_step_flagtoken:640            case mcp_ins_step_flagint:641            case mcp_ins_step_flagis:642            case mcp_ins_step_none:643            case mcp_ins_step_final:644                break;645        }646    }647648    return 0;649}650651// - iterate steps, call function callbacks with as context arg652// TODO:653// - second arg _may_ be a table: in which case we fill the results into this654//   table rather than return them directly.655//   - do this via a different run function that pops each step result?656int mcplib_inspector_call(lua_State *L) {657    // since we're here from a __call, assume the type is correct.658    struct mcp_inspector *ins = lua_touserdata(L, 1);659    luaL_checktype(L, 2, LUA_TUSERDATA);660    if (lua_checkstack(L, ins->rcount) == 0) {661        proxy_lua_error(L, "inspector ran out of stack space for results");662    }663664    // luaL_checkudata() is slow. Trying a new method here where we pull the665    // metatable from a reference then compare it against the meta table of666    // the argument object.667    lua_getmetatable(L, 2); // put arg metatable on stack668    lua_getiuservalue(L, 1, 1); // put stashed metatable on stack669    luaL_argcheck(L, lua_rawequal(L, -1, -2), 2,670            "invalid argument to inspector object");671    lua_pop(L, 2); // toss both metatables672673    // we're valid now. run the steps674    void *arg = lua_touserdata(L, 2);675    return mcp_ins_run(L, ins, arg);676}