proxy: iterate modified request handling

have a bug where updating a token and then requesting it again returns
the previous token. Also have other branches which require use of the
flattened request.

Removes the extra allocation space for lua request objects as we're not
flattening into the end of the memory.

I was originally doing this using a lot of lua but just copying the
string a few times has some better properties:
1) should actually be faster with less lua + fewer allocations
2) can be optimized to do minimal copying (avoid keys, append new flags,
   etc)

1 files changed, 73 insertions, 100 deletions

diff --git a/proxy_request.c b/proxy_request.c
index e133240..cd520b0 100644
--- a/proxy_request.c
+++ b/proxy_request.c
@@ -440,7 +440,7 @@ int process_request(mcp_parser_t *pr, const char *command, size_t cmdlen) {
 // FIXME (v2): any reason to pass in command/cmdlen separately?
 mcp_request_t *mcp_new_request(lua_State *L, mcp_parser_t *pr, const char *command, size_t cmdlen) {
     // reserving an upvalue for key.
-    mcp_request_t *rq = lua_newuserdatauv(L, sizeof(mcp_request_t) + MCP_REQUEST_MAXLEN * 2 + KEY_MAX_LENGTH, 1);
+    mcp_request_t *rq = lua_newuserdatauv(L, sizeof(mcp_request_t) + MCP_REQUEST_MAXLEN + KEY_MAX_LENGTH, 1);
     // TODO (v2): memset only the non-data part? as the rest gets memcpy'd
     // over.
     memset(rq, 0, sizeof(mcp_request_t));
@@ -458,63 +458,68 @@ mcp_request_t *mcp_new_request(lua_State *L, mcp_parser_t *pr, const char *comma
     return rq;
 }
 
-// TODO (v2):
-// if modified, this will re-serialize every time it's accessed.
-// a simple opt could copy back over the original space
-// a "better" one could A/B the request ptr and clear the modified state
-// each time it gets serialized.
-void mcp_request_attach(lua_State *L, mcp_request_t *rq, io_pending_proxy_t *p) {
+// Replaces a token inside a request and re-parses.
+// Note that this has some optimization opportunities. Delaying until
+// required.
+// We should not guarantee order when updating meta flags, which would allow
+// blanking tokens and appending new ones.
+// TODO (v2): function doesn't allow appending.
+// TODO (v2): much of the length is the key, avoid copying it.
+int mcp_request_render(mcp_request_t *rq, int idx, const char *tok, size_t len) {
+    char temp[MCP_REQUEST_MAXLEN];
+    char *p = temp;
     mcp_parser_t *pr = &rq->pr;
-    char *r = (char *) pr->request;
-    size_t len = pr->reqlen;
 
-    // one or more of the tokens were changed
-    if (rq->was_modified) {
-        assert(rq->tokent_ref);
-        // option table to top of stack.
-        lua_rawgeti(L, LUA_REGISTRYINDEX, rq->tokent_ref);
-
-        // space was reserved in case of modification.
-        char *nr = rq->request + MCP_REQUEST_MAXLEN;
-        r = nr;
-        char *or = NULL;
-
-        for (int x = 0; x < pr->ntokens; x++) {
-            const char *newtok = NULL;
-            size_t newlen = 0;
-            if (x != 0 && x != pr->keytoken) {
-                int type = lua_rawgeti(L, -1, x+1);
-                if (type != LUA_TNIL) {
-                    newtok = lua_tolstring(L, -1, &newlen);
-                    memcpy(nr, newtok, newlen);
-                    nr += newlen;
-                }
-                lua_pop(L, 1);
-            }
+    if (pr->reqlen + len > MCP_REQUEST_MAXLEN) {
+        return -1;
+    }
+    // Cannot add/append tokens yet.
+    if (idx >= pr->ntokens) {
+        return -1;
+    }
 
-            if (newtok == NULL) {
-                // TODO (v2): if we add an extra "end" token that's just reqlen we can
-                // memcpy... however most args are short and that may not be worth
-                // it.
-                or = rq->request + pr->tokens[x];
-                // will walk past the end without the \r test.
-                // if we add the end token trick this can be changed.
-                while (*or != ' ' && *or != '\r' && *or != '\n') {
-                    *nr = *or;
-                    nr++;
-                    or++;
-                }
-            }
-            *nr = ' ';
-            nr++;
+    memcpy(p, pr->request, pr->tokens[idx]);
+    p += pr->tokens[idx];
+
+    memcpy(p, tok, len);
+    p += len;
+
+    // Add a space and copy more tokens if there were more.
+    if (idx+1 < pr->ntokens) {
+        if (len != 0) {
+            // Only pre-space if not deleting the token.
+            *p = ' ';
+            p++;
         }
-        // tag the end bits.
-        memcpy(nr-1, "\r\n\0", 3);
-        nr++;
+        memcpy(p, &pr->request[pr->tokens[idx+1]], pr->tokens[pr->ntokens] - pr->tokens[idx+1]);
+        p += pr->tokens[pr->ntokens] - pr->tokens[idx+1];
+    }
+
+    memcpy(p, "\r\n\0", 3);
+    p += 2;
 
-        len = nr - (rq->request + MCP_REQUEST_MAXLEN);
-        lua_pop(L, 1); // pop the table
+    memcpy(rq->request, temp, p - temp);
+
+    // Hold the vlen/vbuf and restore after re-parsing. Since we can only edit
+    // the command line, not the value here, we would otherwise allow sending
+    // arbitrary memory over the network if someone modifies a SET.
+    void *vbuf = pr->vbuf;
+    int vlen = pr->vlen;
+
+    memset(pr, 0, sizeof(mcp_parser_t)); // TODO: required?
+    int ret = process_request(pr, rq->request, p - temp);
+    if (ret != 0) {
+        return ret;
     }
+    pr->vbuf = vbuf;
+    pr->vlen = vlen;
+    return 0;
+}
+
+void mcp_request_attach(lua_State *L, mcp_request_t *rq, io_pending_proxy_t *p) {
+    mcp_parser_t *pr = &rq->pr;
+    char *r = (char *) pr->request;
+    size_t len = pr->reqlen;
 
     // The stringified request. This is also referencing into the coroutine
     // stack, which should be safe from gc.
@@ -528,7 +533,6 @@ void mcp_request_attach(lua_State *L, mcp_request_t *rq, io_pending_proxy_t *p)
         p->iovcnt = 2;
         p->iovbytes += pr->vlen;
     }
-
 }
 
 // second argument is optional, for building set requests.
@@ -649,43 +653,23 @@ int mcplib_request_token(lua_State *L) {
         return 1;
     }
 
-    // we hold overwritten or parsed tokens in a lua table.
-    if (rq->tokent_ref == 0) {
-        // create a presized table that can hold our tokens.
-        lua_createtable(L, rq->pr.ntokens, 0);
-        // duplicate value to set back
-        lua_pushvalue(L, -1);
-        rq->tokent_ref = luaL_ref(L, LUA_REGISTRYINDEX);
-    } else {
-        lua_rawgeti(L, LUA_REGISTRYINDEX, rq->tokent_ref);
-    }
-    // top of stack should be token table.
-
     size_t vlen = 0;
     if (argc > 2) {
         // overwriting a token.
-        luaL_checklstring(L, 3, &vlen);
-        lua_pushvalue(L, 3); // copy to top of stack
-        lua_rawseti(L, -2, token);
-        rq->was_modified = true;
+        size_t newlen = 0;
+        const char *newtok = lua_tolstring(L, 3, &newlen);
+        if (mcp_request_render(rq, token-1, newtok, newlen) != 0) {
+            proxy_lua_error(L, "token(): request malformed after edit");
+            return 0;
+        }
         return 0;
     } else {
         // fetching a token.
-        if (lua_rawgeti(L, -1, token) != LUA_TSTRING) {
-            lua_pop(L, 1); // got a nil, drop it.
-
-            // token not uploaded yet. find the len.
-            const char *start = rq->pr.request + rq->pr.tokens[token-1];
-            vlen = _process_token_len(&rq->pr, token-1);
+        const char *start = rq->pr.request + rq->pr.tokens[token-1];
+        vlen = _process_token_len(&rq->pr, token-1);
 
-            P_DEBUG("%s: pushing token of len: %lu\n", __func__, vlen);
-            lua_pushlstring(L, start, vlen);
-            lua_pushvalue(L, -1); // copy
-
-            lua_rawseti(L, -3, token); // pops copy.
-        }
-
-        // return fetched token or copy of new token.
+        P_DEBUG("%s: pushing token of len: %lu\n", __func__, vlen);
+        lua_pushlstring(L, start, vlen);
         return 1;
     }
 
@@ -743,12 +727,6 @@ int mcplib_request_flag_token(lua_State *L) {
     }
     if (lua_isstring(L, 3)) {
         // overwriting a flag/token with the third argument.
-        // NOTE: semi duplicated from mcplib_request_token()
-        if (rq->tokent_ref == 0) {
-            // create a presized table that can hold our tokens.
-            lua_createtable(L, rq->pr.ntokens, 0);
-            rq->tokent_ref = luaL_ref(L, LUA_REGISTRYINDEX);
-        }
         replace = true;
     }
     uint64_t flagbit = (uint64_t)1 << (flagstr[0] - 65);
@@ -770,14 +748,12 @@ int mcplib_request_flag_token(lua_State *L) {
 
                 // Have something to replace the flag/token with.
                 if (replace) {
-                    // table is top of stack.
-                    lua_rawgeti(L, LUA_REGISTRYINDEX, rq->tokent_ref);
-                    // need to copy the string to top of stack.
-                    lua_pushvalue(L, 3);
-
-                    lua_rawseti(L, -2, x+1); // pops copy
-                    lua_pop(L, 1); // pop table
-                    rq->was_modified = true;
+                    size_t newlen = 0;
+                    const char *newtok = lua_tolstring(L, 3, &newlen);
+                    if (mcp_request_render(rq, x, newtok, newlen) != 0) {
+                        proxy_lua_error(L, "token(): request malformed after edit");
+                        return 0;
+                    }
                 }
                 break;
             }
@@ -798,9 +774,6 @@ int mcplib_request_gc(lua_State *L) {
         free(rq->pr.vbuf);
     }
 
-    if (rq->tokent_ref != 0) {
-        luaL_unref(L, LUA_REGISTRYINDEX, rq->tokent_ref);
-    }
     return 0;
 }