commit b6f01cf: [Rework] Pass upstream when sending TCP requests

Vsevolod Stakhov vsevolod at rspamd.com
Sat Jul 2 13:56:03 UTC 2022


Author: Vsevolod Stakhov
Date: 2022-07-02 14:53:24 +0100
URL: https://github.com/rspamd/rspamd/commit/b6f01cfb0c5fda0b92fc7d8e2e0f7c63d2a7b07f (HEAD -> master)

[Rework] Pass upstream when sending TCP requests

---
 lualib/lua_scanners/avast.lua        | 3 ++-
 lualib/lua_scanners/clamav.lua       | 6 ++----
 lualib/lua_scanners/dcc.lua          | 6 ++----
 lualib/lua_scanners/fprot.lua        | 5 ++---
 lualib/lua_scanners/icap.lua         | 7 ++-----
 lualib/lua_scanners/kaspersky_av.lua | 5 ++---
 lualib/lua_scanners/oletools.lua     | 7 ++-----
 lualib/lua_scanners/pyzor.lua        | 7 ++-----
 lualib/lua_scanners/razor.lua        | 8 ++------
 lualib/lua_scanners/savapi.lua       | 6 ++----
 lualib/lua_scanners/sophos.lua       | 6 ++----
 lualib/lua_scanners/spamassassin.lua | 7 ++-----
 12 files changed, 24 insertions(+), 49 deletions(-)

diff --git a/lualib/lua_scanners/avast.lua b/lualib/lua_scanners/avast.lua
index 6ecf6756f..5c54f8c40 100644
--- a/lualib/lua_scanners/avast.lua
+++ b/lualib/lua_scanners/avast.lua
@@ -93,6 +93,7 @@ local function avast_check(task, content, digest, rule, maybe_part)
       stop_pattern = CRLF,
       host = addr:to_string(),
       port = addr:get_port(),
+      upstream = upstream,
       timeout = rule.timeout,
       task = task
     }
@@ -155,6 +156,7 @@ local function avast_check(task, content, digest, rule, maybe_part)
 
       upstream = rule.upstreams:get_upstream_round_robin()
       addr = upstream:get_addr()
+      tcp_opts.upstream = upstream
       tcp_opts.callback = avast_helo_cb
 
       local is_succ, err = tcp.request(tcp_opts)
@@ -224,7 +226,6 @@ local function avast_check(task, content, digest, rule, maybe_part)
             end
           elseif beg == '200' then
             -- Final line
-            upstream:ok()
             if tcp_conn then
               tcp_conn:close()
               tcp_conn = nil
diff --git a/lualib/lua_scanners/clamav.lua b/lualib/lua_scanners/clamav.lua
index 48aaa66c4..f984864e7 100644
--- a/lualib/lua_scanners/clamav.lua
+++ b/lualib/lua_scanners/clamav.lua
@@ -91,9 +91,6 @@ local function clamav_check(task, content, digest, rule, maybe_part)
     local function clamav_callback(err, data)
       if err then
 
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
-
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
 
@@ -110,6 +107,7 @@ local function clamav_check(task, content, digest, rule, maybe_part)
             task = task,
             host = addr:to_string(),
             port = addr:get_port(),
+            upstream = upstream,
             timeout = rule['timeout'],
             callback = clamav_callback,
             data = { header, content, footer },
@@ -123,7 +121,6 @@ local function clamav_check(task, content, digest, rule, maybe_part)
         end
 
       else
-        upstream:ok()
         data = tostring(data)
         local cached
         lua_util.debugm(rule.name, task, '%s: got reply: %s',
@@ -172,6 +169,7 @@ local function clamav_check(task, content, digest, rule, maybe_part)
       port = addr:get_port(),
       timeout = rule['timeout'],
       callback = clamav_callback,
+      upstream = upstream,
       data = { header, content, footer },
       stop_pattern = '\0'
     })
diff --git a/lualib/lua_scanners/dcc.lua b/lualib/lua_scanners/dcc.lua
index 9c4619479..5cf0ef4a6 100644
--- a/lualib/lua_scanners/dcc.lua
+++ b/lualib/lua_scanners/dcc.lua
@@ -138,9 +138,6 @@ local function dcc_check(task, content, digest, rule)
     local function dcc_callback(err, data, conn)
 
       local function dcc_requery()
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
-
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
 
@@ -158,6 +155,7 @@ local function dcc_check(task, content, digest, rule)
             host = addr:to_string(),
             port = addr:get_port(),
             timeout = rule.timeout or 2.0,
+            upstream = upstream,
             shutdown = true,
             data = request_data,
             callback = dcc_callback,
@@ -178,7 +176,6 @@ local function dcc_check(task, content, digest, rule)
 
       else
         -- Parse the response
-        if upstream then upstream:ok() end
         local _,_,result,disposition,header = tostring(data):find("(.-)\n(.-)\n(.-)$")
         lua_util.debugm(rule.name, task, 'DCC result=%1 disposition=%2 header="%3"',
             result, disposition, header)
@@ -287,6 +284,7 @@ local function dcc_check(task, content, digest, rule)
       port = addr:get_port(),
       timeout = rule.timeout or 2.0,
       shutdown = true,
+      upstream = upstream,
       data = request_data,
       callback = dcc_callback,
       body_max = 999999,
diff --git a/lualib/lua_scanners/fprot.lua b/lualib/lua_scanners/fprot.lua
index 35c4c943d..d2153f7ac 100644
--- a/lualib/lua_scanners/fprot.lua
+++ b/lualib/lua_scanners/fprot.lua
@@ -91,9 +91,6 @@ local function fprot_check(task, content, digest, rule, maybe_part)
 
     local function fprot_callback(err, data)
       if err then
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
-
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
 
@@ -110,6 +107,7 @@ local function fprot_check(task, content, digest, rule, maybe_part)
             task = task,
             host = addr:to_string(),
             port = addr:get_port(),
+            upstream = upstream,
             timeout = rule['timeout'],
             callback = fprot_callback,
             data = { header, content, footer },
@@ -155,6 +153,7 @@ local function fprot_check(task, content, digest, rule, maybe_part)
       task = task,
       host = addr:to_string(),
       port = addr:get_port(),
+      upstream = upstream,
       timeout = rule['timeout'],
       callback = fprot_callback,
       data = { header, content, footer },
diff --git a/lualib/lua_scanners/icap.lua b/lualib/lua_scanners/icap.lua
index 1d783e834..ba571b13d 100644
--- a/lualib/lua_scanners/icap.lua
+++ b/lualib/lua_scanners/icap.lua
@@ -207,9 +207,6 @@ local function icap_check(task, content, digest, rule, maybe_part)
     local function icap_callback(err, conn)
 
       local function icap_requery(err_m, info)
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
-
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
 
@@ -230,6 +227,7 @@ local function icap_check(task, content, digest, rule, maybe_part)
           tcp_options.port = addr:get_port()
           tcp_options.callback = icap_callback
           tcp_options.data = options_request
+          tcp_options.upstream = upstream
 
           tcp.request(tcp_options)
 
@@ -649,8 +647,6 @@ local function icap_check(task, content, digest, rule, maybe_part)
       if err or conn == nil then
         icap_requery(err, "options_request")
       else
-        -- set upstream ok
-        if upstream then upstream:ok() end
         conn:add_read(icap_r_options_cb, '\r\n\r\n')
       end
     end
@@ -671,6 +667,7 @@ local function icap_check(task, content, digest, rule, maybe_part)
 
     tcp_options.host = addr:to_string()
     tcp_options.port = addr:get_port()
+    tcp_options.upstream = upstream
 
     tcp.request(tcp_options)
   end
diff --git a/lualib/lua_scanners/kaspersky_av.lua b/lualib/lua_scanners/kaspersky_av.lua
index fc395f363..69aaad7e5 100644
--- a/lualib/lua_scanners/kaspersky_av.lua
+++ b/lualib/lua_scanners/kaspersky_av.lua
@@ -110,8 +110,6 @@ local function kaspersky_check(task, content, digest, rule, maybe_part)
 
     local function kaspersky_callback(err, data)
       if err then
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
 
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
@@ -129,6 +127,7 @@ local function kaspersky_check(task, content, digest, rule, maybe_part)
             task = task,
             host = addr:to_string(),
             port = addr:get_port(),
+            upstream = upstream,
             timeout = rule['timeout'],
             callback = kaspersky_callback,
             data = { clamav_compat_cmd },
@@ -144,7 +143,6 @@ local function kaspersky_check(task, content, digest, rule, maybe_part)
         end
 
       else
-        upstream:ok()
         data = tostring(data)
         local cached
         lua_util.debugm(rule.name, task,
@@ -174,6 +172,7 @@ local function kaspersky_check(task, content, digest, rule, maybe_part)
       task = task,
       host = addr:to_string(),
       port = addr:get_port(),
+      upstream = upstream,
       timeout = rule['timeout'],
       callback = kaspersky_callback,
       data = { clamav_compat_cmd },
diff --git a/lualib/lua_scanners/oletools.lua b/lualib/lua_scanners/oletools.lua
index 04cdfad3e..8513699a0 100644
--- a/lualib/lua_scanners/oletools.lua
+++ b/lualib/lua_scanners/oletools.lua
@@ -98,8 +98,6 @@ local function oletools_check(task, content, digest, rule, maybe_part)
     local function oletools_callback(err, data, conn)
 
       local function oletools_requery(error)
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
 
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
@@ -117,6 +115,7 @@ local function oletools_check(task, content, digest, rule, maybe_part)
             task = task,
             host = addr:to_string(),
             port = addr:get_port(),
+            upstream = upstream,
             timeout = rule.timeout,
             shutdown = true,
             data = { protocol, content },
@@ -136,9 +135,6 @@ local function oletools_check(task, content, digest, rule, maybe_part)
         oletools_requery(err)
 
       else
-        -- Parse the response
-        if upstream then upstream:ok() end
-
         json_response = json_response .. tostring(data)
 
 	      if not string.find(json_response, '\t\n\n\t') and #data == 8192 then
@@ -346,6 +342,7 @@ local function oletools_check(task, content, digest, rule, maybe_part)
       task = task,
       host = addr:to_string(),
       port = addr:get_port(),
+      upstream = upstream,
       timeout = rule.timeout,
       shutdown = true,
       data = { protocol, content },
diff --git a/lualib/lua_scanners/pyzor.lua b/lualib/lua_scanners/pyzor.lua
index ed0e3d16b..78250a3ad 100644
--- a/lualib/lua_scanners/pyzor.lua
+++ b/lualib/lua_scanners/pyzor.lua
@@ -85,9 +85,6 @@ local function pyzor_check(task, content, digest, rule)
 
       if err then
 
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
-
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
 
@@ -104,6 +101,7 @@ local function pyzor_check(task, content, digest, rule)
             task = task,
             host = addr:to_string(),
             port = addr:get_port(),
+            upstream = upstream,
             timeout = rule['timeout'],
             shutdown = true,
             data = content,
@@ -116,8 +114,6 @@ local function pyzor_check(task, content, digest, rule)
               'failed to scan and retransmits exceed')
         end
       else
-        -- Parse the response
-        if upstream then upstream:ok() end
         -- pyzor output is unicode (\x09 -> tab, \0a -> newline)
         --   public.pyzor.org:24441  (200, 'OK')     21285091   206759
         --   server:port             Code  Diag      Count      WL-Count
@@ -187,6 +183,7 @@ local function pyzor_check(task, content, digest, rule)
       task = task,
       host = addr:to_string(),
       port = addr:get_port(),
+      upstream = upstream,
       timeout = rule.timeout,
       shutdown = true,
       data = content,
diff --git a/lualib/lua_scanners/razor.lua b/lualib/lua_scanners/razor.lua
index 62b1e06a1..7de4c84eb 100644
--- a/lualib/lua_scanners/razor.lua
+++ b/lualib/lua_scanners/razor.lua
@@ -90,9 +90,6 @@ local function razor_check(task, content, digest, rule)
     local function razor_callback(err, data, conn)
 
       local function razor_requery()
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
-
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
 
@@ -112,6 +109,7 @@ local function razor_check(task, content, digest, rule)
             task = task,
             host = addr:to_string(),
             port = addr:get_port(),
+            upstream = upstream,
             timeout = rule.timeout or 2.0,
             shutdown = true,
             data = content,
@@ -129,9 +127,6 @@ local function razor_check(task, content, digest, rule)
         razor_requery()
 
       else
-        -- Parse the response
-        if upstream then upstream:ok() end
-
         --[[
         @todo: Razorsocket currently only returns ham or spam. When the wrapper is fixed we should add dynamic scores here.
         Maybe check spamassassin implementation.
@@ -163,6 +158,7 @@ local function razor_check(task, content, digest, rule)
       task = task,
       host = addr:to_string(),
       port = addr:get_port(),
+      upstream = upstream,
       timeout = rule.timeout or 2.0,
       shutdown = true,
       data = content,
diff --git a/lualib/lua_scanners/savapi.lua b/lualib/lua_scanners/savapi.lua
index 2933aa7b4..430009df3 100644
--- a/lualib/lua_scanners/savapi.lua
+++ b/lualib/lua_scanners/savapi.lua
@@ -198,9 +198,6 @@ local function savapi_check(task, content, digest, rule)
     local function savapi_callback_init(err, data, conn)
       if err then
 
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
-
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
 
@@ -217,6 +214,7 @@ local function savapi_check(task, content, digest, rule)
             task = task,
             host = addr:to_string(),
             port = addr:get_port(),
+            upstream = upstream,
             timeout = rule['timeout'],
             callback = savapi_callback_init,
             stop_pattern = {'\n'},
@@ -226,7 +224,6 @@ local function savapi_check(task, content, digest, rule)
           common.yield_result(task, rule, 'failed to scan and retransmits exceed', 0.0, 'fail')
         end
       else
-        upstream:ok()
         local result = tostring(data)
 
         -- 100 SAVAPI:4.0 greeting
@@ -240,6 +237,7 @@ local function savapi_check(task, content, digest, rule)
       task = task,
       host = addr:to_string(),
       port = addr:get_port(),
+      upstream = upstream,
       timeout = rule['timeout'],
       callback = savapi_callback_init,
       stop_pattern = {'\n'},
diff --git a/lualib/lua_scanners/sophos.lua b/lualib/lua_scanners/sophos.lua
index b82a71025..d9b64f1a8 100644
--- a/lualib/lua_scanners/sophos.lua
+++ b/lualib/lua_scanners/sophos.lua
@@ -90,9 +90,6 @@ local function sophos_check(task, content, digest, rule, maybe_part)
     local function sophos_callback(err, data, conn)
 
       if err then
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
-
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
 
@@ -109,6 +106,7 @@ local function sophos_check(task, content, digest, rule, maybe_part)
             task = task,
             host = addr:to_string(),
             port = addr:get_port(),
+            upstream = upstream,
             timeout = rule['timeout'],
             callback = sophos_callback,
             data = { protocol, streamsize, content, bye }
@@ -119,7 +117,6 @@ local function sophos_check(task, content, digest, rule, maybe_part)
               0.0, 'fail', maybe_part)
         end
       else
-        upstream:ok()
         data = tostring(data)
         lua_util.debugm(rule.name, task,
             '%s [%s]: got reply: %s', rule['symbol'], rule['type'], data)
@@ -170,6 +167,7 @@ local function sophos_check(task, content, digest, rule, maybe_part)
       task = task,
       host = addr:to_string(),
       port = addr:get_port(),
+      upstream = upstream,
       timeout = rule['timeout'],
       callback = sophos_callback,
       data = { protocol, streamsize, content, bye }
diff --git a/lualib/lua_scanners/spamassassin.lua b/lualib/lua_scanners/spamassassin.lua
index c50c477a4..c4fd0d2e5 100644
--- a/lualib/lua_scanners/spamassassin.lua
+++ b/lualib/lua_scanners/spamassassin.lua
@@ -103,8 +103,6 @@ local function spamassassin_check(task, content, digest, rule)
     local function spamassassin_callback(err, data)
 
       local function spamassassin_requery(error)
-        -- set current upstream to fail because an error occurred
-        upstream:fail()
 
         -- retry with another upstream until retransmits exceeds
         if retransmits > 0 then
@@ -125,6 +123,7 @@ local function spamassassin_check(task, content, digest, rule)
             task = task,
             host = addr:to_string(),
             port = addr:get_port(),
+            upstream = upstream,
             timeout = rule['timeout'],
             data = request_data,
             callback = spamassassin_callback,
@@ -141,9 +140,6 @@ local function spamassassin_check(task, content, digest, rule)
         spamassassin_requery(err)
 
       else
-        -- Parse the response
-        if upstream then upstream:ok() end
-
         --lua_util.debugm(rule.N, task, '%s: returned result: %s', rule.log_prefix, data)
 
         --[[
@@ -193,6 +189,7 @@ local function spamassassin_check(task, content, digest, rule)
       task = task,
       host = addr:to_string(),
       port = addr:get_port(),
+      upstream = upstream,
       timeout = rule['timeout'],
       data = request_data,
       callback = spamassassin_callback,


More information about the Commits mailing list