commit dd4b3cd: [Minor] Allow to limit number of files returned by archive:get_files

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2020-04-28 15:44:18 +0100
URL: https://github.com/rspamd/rspamd/commit/dd4b3cdeef50038332ccd28c95a9882f17f7777c

[Minor] Allow to limit number of files returned by archive:get_files

---
 src/lua/lua_task.c | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 270d5ec06..1ca96ed76 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -6609,13 +6609,21 @@ lua_archive_get_files (lua_State *L)
 {
 	LUA_TRACE_POINT;
 	struct rspamd_archive *arch = lua_check_archive (L);
-	guint i;
+	guint i, max_files = 0;
 	struct rspamd_archive_file *f;
 
 	if (arch != NULL) {
-		lua_createtable (L, arch->files->len, 0);
+		if (lua_isnumber (L, 2)) {
+			max_files = lua_tointeger (L, 2);
+			max_files = MIN (arch->files->len, max_files);
+		}
+		else {
+			max_files = arch->files->len;
+		}
+
+		lua_createtable (L, max_files, 0);
 
-		for (i = 0; i < arch->files->len; i ++) {
+		for (i = 0; i < max_files; i ++) {
 			f = g_ptr_array_index (arch->files, i);
 
 			lua_pushlstring (L, f->fname->str, f->fname->len);
@@ -6634,13 +6642,21 @@ lua_archive_get_files_full (lua_State *L)
 {
 	LUA_TRACE_POINT;
 	struct rspamd_archive *arch = lua_check_archive (L);
-	guint i;
+	guint i, max_files = 0;
 	struct rspamd_archive_file *f;
 
 	if (arch != NULL) {
-		lua_createtable (L, arch->files->len, 0);
+		if (lua_isnumber (L, 2)) {
+			max_files = lua_tointeger (L, 2);
+			max_files = MIN (arch->files->len, max_files);
+		}
+		else {
+			max_files = arch->files->len;
+		}
+
+		lua_createtable (L, max_files, 0);
 
-		for (i = 0; i < arch->files->len; i ++) {
+		for (i = 0; i < max_files; i ++) {
 			f = g_ptr_array_index (arch->files, i);
 
 			lua_createtable (L, 0, 4);