From b6574dce52f857720ac789282565b6b9c95d341a Mon Sep 17 00:00:00 2001
From: Datong Sun <datong.sun@konghq.com>
Date: Thu, 21 Nov 2024 04:19:57 -0800
Subject: [PATCH] fix(prefix): fix bug when buffer realloc occurs during prefix
 lookup

Prefix based lookup may need to reallocate the key/value buffer when the
buffer size exceeds the default (currently 1MB).

However, because we use the `ops[0]` and `ops[1]` passed in is used for both
providing the `after` and `prefix`, in case of `NGX_AGAIN`, these will
be overwritten by the `ngx_lua_resty_lmdb_ffi_prefix()` function to the
first and second key respectively, which breaks the retry and causing
the second call to `ngx_lua_resty_lmdb_ffi_prefix()` to return empty.

This commit fixes the above issue by always resetting `ops[0]` and
`ops[1]` between retries, therefore the above issue no longer occurs.

Also, this commit fixes a use-after-free bug for `db_drop()` operation when
map is full.
---
 lib/resty/lmdb/prefix.lua            | 16 +++---
 src/ngx_lua_resty_lmdb_transaction.c |  3 +-
 t/10-prefix.t                        | 75 ++++++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 9 deletions(-)

diff --git a/lib/resty/lmdb/prefix.lua b/lib/resty/lmdb/prefix.lua
index 989530b..b25cf17 100644
--- a/lib/resty/lmdb/prefix.lua
+++ b/lib/resty/lmdb/prefix.lua
@@ -36,14 +36,6 @@ function _M.page(start, prefix, db, page_size)
     local value_buf_size = get_string_buf_size()
     local ops = ffi_new("ngx_lua_resty_lmdb_operation_t[?]", page_size)
 
-    ops[0].opcode = C.NGX_LMDB_OP_PREFIX
-    ops[0].key.data = start
-    ops[0].key.len = #start
-
-    ops[1].opcode = C.NGX_LMDB_OP_PREFIX
-    ops[1].key.data = prefix
-    ops[1].key.len = #prefix
-
     local dbi, err = get_dbi(false, db or DEFAULT_DB)
     if err then
         return nil, "unable to open DB for access: " .. err
@@ -55,6 +47,14 @@ function _M.page(start, prefix, db, page_size)
     ops[0].dbi = dbi
 
 ::again::
+    ops[0].opcode = C.NGX_LMDB_OP_PREFIX
+    ops[0].key.data = start
+    ops[0].key.len = #start
+
+    ops[1].opcode = C.NGX_LMDB_OP_PREFIX
+    ops[1].key.data = prefix
+    ops[1].key.len = #prefix
+
     local buf = get_string_buf(value_buf_size, false)
     local ret = C.ngx_lua_resty_lmdb_ffi_prefix(ops, page_size,
                     buf, value_buf_size, err_ptr)
diff --git a/src/ngx_lua_resty_lmdb_transaction.c b/src/ngx_lua_resty_lmdb_transaction.c
index ad892b3..5665b1e 100644
--- a/src/ngx_lua_resty_lmdb_transaction.c
+++ b/src/ngx_lua_resty_lmdb_transaction.c
@@ -148,7 +148,8 @@ int ngx_lua_resty_lmdb_ffi_execute(ngx_lua_resty_lmdb_operation_t *ops,
         rc = mdb_txn_commit(txn);
         if (rc != 0) {
             *err = mdb_strerror(rc);
-            goto err;
+
+            return NGX_ERROR;
         }
 
     } else {
diff --git a/t/10-prefix.t b/t/10-prefix.t
index 8ca2352..0d44370 100644
--- a/t/10-prefix.t
+++ b/t/10-prefix.t
@@ -353,3 +353,78 @@ false.../lua-resty-lmdb/lua-resty-lmdb/lib/resty/lmdb/prefix.lua:34: 'page_size'
 [error]
 [warn]
 [crit]
+
+
+
+=== TEST 9: prefix.page() operation with large page size [KAG-5874]
+--- http_config eval: $::HttpConfig
+--- main_config eval: $::MainConfig
+--- config
+    location = /t {
+        content_by_lua_block {
+            local l = require("resty.lmdb")
+
+            ngx.say(l.db_drop(true))
+
+            ngx.say(l.set("test", "value"))
+
+            local inserted = { test = "value" }
+
+            for i = 1, 2048 do
+                -- string.rep with 120 makes sure each page goes just over the 1MB
+                -- default buffer size and triggers a realloc
+                assert(l.set(string.rep("test", 120) .. i, string.rep("value", 120)))
+                inserted[string.rep("test", 120) .. i] = string.rep("value", 120)
+            end
+
+            ngx.say(l.set("u", "value4"))
+            ngx.say(l.set("u1", "value5"))
+
+            local p = require("resty.lmdb.prefix")
+
+            local res, err = p.page("test", "test", nil, 1000)
+            if not res then
+                ngx.say("page errored: ", err)
+            end
+
+            ngx.say("FIRST PAGE")
+            for _, pair in ipairs(res) do
+                inserted[pair.key] = nil
+            end
+
+            res, err = p.page(res[#res].key .. "\x00", "test", nil, 1000)
+            if not res then
+                ngx.say("page errored: ", err)
+            end
+            ngx.say("SECOND PAGE")
+            for _i, pair in ipairs(res) do
+                inserted[pair.key] = nil
+            end
+
+            res, err = p.page(res[#res].key .. "\x00", "test", nil, 1000)
+            if not res then
+                ngx.say("page errored: ", err)
+            end
+            ngx.say("THIRD PAGE")
+            for _i, pair in ipairs(res) do
+                inserted[pair.key] = nil
+            end
+
+            ngx.say(next(inserted))
+        }
+    }
+--- request
+GET /t
+--- response_body
+true
+true
+true
+true
+FIRST PAGE
+SECOND PAGE
+THIRD PAGE
+nil
+--- no_error_log
+[error]
+[warn]
+[crit]