diff --git a/nselib/brute.lua b/nselib/brute.lua index d1e8fcdb5..f50f19986 100644 --- a/nselib/brute.lua +++ b/nselib/brute.lua @@ -1465,6 +1465,17 @@ Iterators = { } +-- These functions all return a boolean and an error (or result) +-- and should all be wrapped in order to check status of the engine. +checkwrap = { + connect = true, + send = true, + receive = true, + receive_lines = true, + receive_buf = true, + receive_bytes = true, +} + -- A socket wrapper class. -- Instances of this class can be treated as regular sockets. -- This wrapper is used to relay connection errors to the corresponding Engine @@ -1476,20 +1487,29 @@ BruteSocket = { } setmetatable(o, self) - self.__index = function (table, key) - if self[key] then - return self[key] - elseif o.socket[key] then - if type(o.socket[key]) == "function" then - return function (self, ...) - return o.socket[key](o.socket, ...) - end - else - return o.socket[key] + self.__index = function (instance, key) + local f = rawget(self, key) + if f then + -- BruteSocket function + return f + else + -- something provided by NSE socket + f = instance.socket[key] + end + -- Check if it should be wrapped with a checkStatus call + if checkwrap[key] then + return function(s, ...) + local status, err = f(instance.socket, ...) + instance:checkStatus(status, err) + return status, err + end + elseif type(f) == "function" then + -- not wrapped? call the function on the underlying socket + return function (s, ...) + return f(instance.socket, ...) end end - - return nil + return f end o.socket = nmap.new_socket() @@ -1521,31 +1541,6 @@ BruteSocket = { thread_data.con_error_reason = err end end, - - connect = function (self, host, port) - local status, err = self.socket:connect(host, port) - self:checkStatus(status, err) - - return status, err - end, - - send = function (self, data) - local status, err = self.socket:send(data) - self:checkStatus(status, err) - - return status, err - end, - - receive = function (self) - local status, data = self.socket:receive() - self:checkStatus(status, data) - - return status, data - end, - - close = function (self) - self.socket:close() - end, } function new_socket ()