5 files changed, 185 insertions, 152 deletions

diff --git a/src/except.c b/src/except.c
index 0482e1c..53a65ac 100644
--- a/src/except.c
+++ b/src/except.c
@@ -14,17 +14,20 @@
 \*=========================================================================*/
 static int global_try(lua_State *L);
 static int global_protect(lua_State *L);
+static int global_newtry(lua_State *L);
 static int protected(lua_State *L);
+static int finalize(lua_State *L);
 
 /* except functions */
 static luaL_reg func[] = {
-    {"try",      global_try},
-    {"protect",  global_protect},
-    {NULL,       NULL}
+    {"try",       global_try},
+    {"newtry",    global_newtry},
+    {"protect",   global_protect},
+    {NULL,        NULL}
 };
 
 /*-------------------------------------------------------------------------*\
-* Exception handling: try method
+* Try method
 \*-------------------------------------------------------------------------*/
 static int global_try(lua_State *L) {
     if (lua_isnil(L, 1) || (lua_isboolean(L, 1) && !lua_toboolean(L, 1))) {
@@ -35,7 +38,25 @@ static int global_try(lua_State *L) {
 }
 
 /*-------------------------------------------------------------------------*\
-* Exception handling: protect factory
+* Finalizer factory
+\*-------------------------------------------------------------------------*/
+static int finalize(lua_State *L) {
+    if (lua_isnil(L, 1) || (lua_isboolean(L, 1) && !lua_toboolean(L, 1))) {
+        lua_pushvalue(L, lua_upvalueindex(1));
+        lua_pcall(L, 0, 0, 0);
+        lua_settop(L, 2);
+        lua_error(L);
+        return 0;
+    } else return lua_gettop(L);
+}
+
+static int global_newtry(lua_State *L) {
+    lua_pushcclosure(L, finalize, 1);
+    return 1;
+}
+
+/*-------------------------------------------------------------------------*\
+* Protect factory
 \*-------------------------------------------------------------------------*/
 static int protected(lua_State *L) {
     lua_pushvalue(L, lua_upvalueindex(1));
@@ -48,7 +69,6 @@ static int protected(lua_State *L) {
 }
 
 static int global_protect(lua_State *L) {
-    lua_insert(L, 1);
     lua_pushcclosure(L, protected, 1);
     return 1;
 }
diff --git a/src/ftp.lua b/src/ftp.lua
index 3868158..ad02c72 100644
--- a/src/ftp.lua
+++ b/src/ftp.lua
@@ -32,143 +32,153 @@ local metat = { __index = {} }
 
 function open(server, port)
     local tp = socket.try(tp.connect(server, port or PORT, TIMEOUT))
-    return setmetatable({tp = tp}, metat)
+    local f = { tp = tp }
+    -- make sure everything gets closed in an exception
+    f.try = socket.newtry(function() 
+        tp:close() 
+        if f.data then f.data:close() end
+        if f.server then f.server:close() end
+    end)
+    return setmetatable(f, metat)
 end
 
-local function port(portt)
-    return portt.server:accept()
+function metat.__index:portconnect()
+    self.try(self.server:settimeout(TIMEOUT))
+    self.data = self.try(self.server:accept())
+    self.try(self.data:settimeout(TIMEOUT))
 end
 
-local function pasv(pasvt)
-    local data = socket.try(socket.tcp())
-    socket.try(data:settimeout(TIMEOUT))
-    socket.try(data:connect(pasvt.ip, pasvt.port))
-    return data
+function metat.__index:pasvconnect()
+    self.data = self.try(socket.tcp())
+    self.try(self.data:settimeout(TIMEOUT))
+    self.try(self.data:connect(self.pasvt.ip, self.pasvt.port))
 end
 
 function metat.__index:login(user, password)
-    socket.try(self.tp:command("user", user or USER))
-    local code, reply = socket.try(self.tp:check{"2..", 331})
+    self.try(self.tp:command("user", user or USER))
+    local code, reply = self.try(self.tp:check{"2..", 331})
     if code == 331 then
-        socket.try(self.tp:command("pass", password or PASSWORD))
-        socket.try(self.tp:check("2.."))
+        self.try(self.tp:command("pass", password or PASSWORD))
+        self.try(self.tp:check("2.."))
     end
     return 1
 end
 
 function metat.__index:pasv()
-    socket.try(self.tp:command("pasv"))
-    local code, reply = socket.try(self.tp:check("2.."))
+    self.try(self.tp:command("pasv"))
+    local code, reply = self.try(self.tp:check("2.."))
     local pattern = "(%d+)%D(%d+)%D(%d+)%D(%d+)%D(%d+)%D(%d+)"
     local a, b, c, d, p1, p2 = socket.skip(2, string.find(reply, pattern))
-    socket.try(a and b and c and d and p1 and p2, reply)
+    self.try(a and b and c and d and p1 and p2, reply)
     self.pasvt = { 
         ip = string.format("%d.%d.%d.%d", a, b, c, d), 
         port = p1*256 + p2
     }
-    if self.portt then 
-        self.portt.server:close()
-        self.portt = nil
+    if self.server then 
+        self.server:close()
+        self.server = nil
     end
     return self.pasvt.ip, self.pasvt.port 
 end
 
 function metat.__index:port(ip, port)
     self.pasvt = nil
-    local server
     if not ip then 
-        ip, port = socket.try(self.tp:getcontrol():getsockname())
-        server = socket.try(socket.bind(ip, 0))
-        ip, port = socket.try(server:getsockname())
-        socket.try(server:settimeout(TIMEOUT))
+        ip, port = self.try(self.tp:getcontrol():getsockname())
+        self.server = self.try(socket.bind(ip, 0))
+        ip, port = self.try(self.server:getsockname())
+        self.try(server:settimeout(TIMEOUT))
     end
     local pl = math.mod(port, 256)
     local ph = (port - pl)/256
     local arg = string.gsub(string.format("%s,%d,%d", ip, ph, pl), "%.", ",")
-    socket.try(self.tp:command("port", arg))
-    socket.try(self.tp:check("2.."))
-    self.portt = server and {ip = ip, port = port, server = server}
+    self.try(self.tp:command("port", arg))
+    self.try(self.tp:check("2.."))
     return 1
 end
 
 function metat.__index:send(sendt)
-    local data
-    socket.try(self.pasvt or self.portt, "need port or pasv first")
-    if self.pasvt then data = socket.try(pasv(self.pasvt)) end
+    self.try(self.pasvt or self.server, "need port or pasv first")
+    -- if there is a pasvt table, we already sent a PASV command 
+    -- we just get the data connection into self.data
+    if self.pasvt then self:pasvconnect() end
+    -- get the transfer argument and command 
     local argument = sendt.argument or string.gsub(sendt.path, "^/", "")
     if argument == "" then argument = nil end
-    local command =  sendt.command or "stor"
-    socket.try(self.tp:command(command, argument))
-    local code, reply = socket.try(self.tp:check{"2..", "1.."})
-    if self.portt then data = socket.try(port(self.portt)) end
+    local command = sendt.command or "stor"
+    -- send the transfer command and check the reply
+    self.try(self.tp:command(command, argument))
+    local code, reply = self.try(self.tp:check{"2..", "1.."})
+    -- if there is not a a pasvt table, then there is a server
+    -- and we already sent a PORT command
+    if not self.pasvt then self:portconnect() end
+    -- get the sink, source and step for the transfer
     local step = sendt.step or ltn12.pump.step
     local checkstep = function(src, snk)
+        -- check status in control connection while downloading
         local readyt = socket.select(readt, nil, 0)
-        if readyt[tp] then
-            code, reply = self.tp:check("2..")
-            if not code then 
-                data:close()
-                return nil, reply 
-            end
-        end
-        local ret, err = step(src, snk)
-        if err then data:close() end
-        return ret, err
+        if readyt[tp] then self.try(self.tp:check("2..")) end
+        return step(src, snk)
     end
-    local sink = socket.sink("close-when-done", data)
-    socket.try(ltn12.pump.all(sendt.source, sink, checkstep))
-    if string.find(code, "1..") then socket.try(self.tp:check("2..")) end
+    local sink = socket.sink("close-when-done", self.data)
+    -- transfer all data and check error
+    self.try(ltn12.pump.all(sendt.source, sink, checkstep))
+    if string.find(code, "1..") then self.try(self.tp:check("2..")) end
+    -- done with data connection
+    self.data:close()
+    self.data = nil
     return 1
 end
 
 function metat.__index:receive(recvt)
-    local data
-    socket.try(self.pasvt or self.portt, "need port or pasv first")
-    if self.pasvt then data = socket.try(pasv(self.pasvt)) end
+    self.try(self.pasvt or self.server, "need port or pasv first")
+    if self.pasvt then self:pasvconnect() end
     local argument = recvt.argument or string.gsub(recvt.path, "^/", "")
     if argument == "" then argument = nil end
-    local command =  recvt.command or "retr"
-    socket.try(self.tp:command(command, argument))
-    local code = socket.try(self.tp:check{"1..", "2.."})
-    if self.portt then data = socket.try(port(self.portt)) end
-    local source = socket.source("until-closed", data)
+    local command = recvt.command or "retr"
+    self.try(self.tp:command(command, argument))
+    local code = self.try(self.tp:check{"1..", "2.."})
+    if not self.pasvt then self:portconnect() end
+    local source = socket.source("until-closed", self.data)
     local step = recvt.step or ltn12.pump.step
-    local checkstep = function(src, snk)
-        local ret, err = step(src, snk)
-        if err then data:close() end
-        return ret, err
-    end
-    socket.try(ltn12.pump.all(source, recvt.sink, checkstep))
-    if string.find(code, "1..") then socket.try(self.tp:check("2..")) end
+    self.try(ltn12.pump.all(source, recvt.sink, step))
+    if string.find(code, "1..") then self.try(self.tp:check("2..")) end
+    self.data:close()
+    self.data = nil
     return 1
 end
 
 function metat.__index:cwd(dir)
-    socket.try(self.tp:command("cwd", dir))
-    socket.try(self.tp:check(250))
+    self.try(self.tp:command("cwd", dir))
+    self.try(self.tp:check(250))
     return 1
 end
 
 function metat.__index:type(type)
-    socket.try(self.tp:command("type", type))
-    socket.try(self.tp:check(200))
+    self.try(self.tp:command("type", type))
+    self.try(self.tp:check(200))
     return 1
 end
 
 function metat.__index:greet()
-    local code = socket.try(self.tp:check{"1..", "2.."})
-    if string.find(code, "1..") then socket.try(self.tp:check("2..")) end
+    local code = self.try(self.tp:check{"1..", "2.."})
+    if string.find(code, "1..") then self.try(self.tp:check("2..")) end
     return 1
 end
 
 function metat.__index:quit()
-    socket.try(self.tp:command("quit"))
-    socket.try(self.tp:check("2.."))
+    self.try(self.tp:command("quit"))
+    self.try(self.tp:check("2.."))
     return 1
 end
 
 function metat.__index:close()
-    socket.try(self.tp:close())
+    self.tp:close()
+    if self.data then self.data:close() end
+    if self.server then self.server:close() end
+    self.tp = nil
+    self.data = nil
+    self.server = nil
     return 1
 end
 
@@ -176,14 +186,14 @@ end
 -- High level FTP API
 -----------------------------------------------------------------------------
 local function tput(putt)
-    local con = open(putt.host, putt.port)
-    con:greet()
-    con:login(putt.user, putt.password)
-    if putt.type then con:type(putt.type) end
-    con:pasv()
-    con:send(putt)
-    con:quit()
-    return con:close()
+    local f = open(putt.host, putt.port)
+    f:greet()
+    f:login(putt.user, putt.password)
+    if putt.type then f:type(putt.type) end
+    f:pasv()
+    f:send(putt)
+    f:quit()
+    return f:close()
 end
 
 local default = {
@@ -216,14 +226,14 @@ put = socket.protect(function(putt, body)
 end)
 
 local function tget(gett)
-    local con = open(gett.host, gett.port)
-    con:greet()
-    con:login(gett.user, gett.password)
-    if gett.type then con:type(gett.type) end
-    con:pasv()
-    con:receive(gett)
-    con:quit()
-    return con:close()
+    local f = open(gett.host, gett.port)
+    f:greet()
+    f:login(gett.user, gett.password)
+    if gett.type then f:type(gett.type) end
+    f:pasv()
+    f:receive(gett)
+    f:quit()
+    return f:close()
 end
 
 local function sget(u)
diff --git a/src/http.lua b/src/http.lua
index 3bd4d6a..d8889e1 100644
--- a/src/http.lua
+++ b/src/http.lua
@@ -31,24 +31,25 @@ BLOCKSIZE = 2048
 local metat = { __index = {} }
 
 function open(host, port)
-    local con = socket.try(socket.tcp()) 
-    socket.try(con:settimeout(TIMEOUT))
-    port = port or PORT
-    socket.try(con:connect(host, port))
-    return setmetatable({ con = con }, metat)
+    local c = socket.try(socket.tcp()) 
+    -- make sure the connection gets closed on exception
+    local try = socket.newtry(function() c:close() end)
+    try(c:settimeout(TIMEOUT))
+    try(c:connect(host, port or PORT))
+    return setmetatable({ c = c, try = try }, metat)
 end
 
 function metat.__index:sendrequestline(method, uri)
     local reqline = string.format("%s %s HTTP/1.1\r\n", method or "GET", uri) 
-    return socket.try(self.con:send(reqline))
+    return self.try(self.c:send(reqline))
 end
 
 function metat.__index:sendheaders(headers)
     for i, v in pairs(headers) do
-        socket.try(self.con:send(i .. ": " .. v .. "\r\n"))
+        self.try(self.c:send(i .. ": " .. v .. "\r\n"))
     end
     -- mark end of request headers
-    socket.try(self.con:send("\r\n"))
+    self.try(self.c:send("\r\n"))
     return 1
 end
 
@@ -59,32 +60,32 @@ function metat.__index:sendbody(headers, source, step)
     local mode
     if headers["content-length"] then mode = "keep-open"
     else mode = "http-chunked" end
-    return socket.try(ltn12.pump.all(source, socket.sink(mode, self.con), step))
+    return self.try(ltn12.pump.all(source, socket.sink(mode, self.c), step))
 end
 
 function metat.__index:receivestatusline()
-    local status = socket.try(self.con:receive())
+    local status = self.try(self.c:receive())
     local code = socket.skip(2, string.find(status, "HTTP/%d*%.%d* (%d%d%d)"))
-    return socket.try(tonumber(code), status)
+    return self.try(tonumber(code), status)
 end
 
 function metat.__index:receiveheaders()
     local line, name, value
     local headers = {}
     -- get first line
-    line = socket.try(self.con:receive())
+    line = self.try(self.c:receive())
     -- headers go until a blank line is found
     while line ~= "" do
         -- get field-name and value
         name, value = socket.skip(2, string.find(line, "^(.-):%s*(.*)"))
-        socket.try(name and value, "malformed reponse headers")
+        self.try(name and value, "malformed reponse headers")
         name = string.lower(name)
         -- get next line (value might be folded)
-        line  = socket.try(self.con:receive())
+        line  = self.try(self.c:receive())
         -- unfold any folded values
         while string.find(line, "^%s") do
             value = value .. line
-            line = socket.try(self.con:receive())
+            line = self.try(self.c:receive())
         end
         -- save pair in table
         if headers[name] then headers[name] = headers[name] .. ", " .. value
@@ -102,12 +103,12 @@ function metat.__index:receivebody(headers, sink, step)
     if TE and TE ~= "identity" then mode = "http-chunked"
     elseif tonumber(headers["content-length"]) then mode = "by-length"
     else mode = "default" end
-    return socket.try(ltn12.pump.all(socket.source(mode, self.con, length), 
+    return self.try(ltn12.pump.all(socket.source(mode, self.c, length), 
         sink, step))
 end
 
 function metat.__index:close()
-    return self.con:close()
+    return self.c:close()
 end
 
 -----------------------------------------------------------------------------
@@ -204,23 +205,23 @@ end
 
 function trequest(reqt)
     reqt = adjustrequest(reqt)
-    local con = open(reqt.host, reqt.port)
-    con:sendrequestline(reqt.method, reqt.uri)
-    con:sendheaders(reqt.headers)
-    con:sendbody(reqt.headers, reqt.source, reqt.step)
+    local h = open(reqt.host, reqt.port)
+    h:sendrequestline(reqt.method, reqt.uri)
+    h:sendheaders(reqt.headers)
+    h:sendbody(reqt.headers, reqt.source, reqt.step)
     local code, headers, status
-    code, status = con:receivestatusline()
-    headers = con:receiveheaders()
+    code, status = h:receivestatusline()
+    headers = h:receiveheaders()
     if shouldredirect(reqt, code) then 
-        con:close()
+        h:close()
         return tredirect(reqt, headers)
     elseif shouldauthorize(reqt, code) then 
-        con:close()
+        h:close()
         return tauthorize(reqt)
     elseif shouldreceivebody(reqt, code) then
-        con:receivebody(headers, reqt.sink, reqt.step)
+        h:receivebody(headers, reqt.sink, reqt.step)
     end
-    con:close()
+    h:close()
     return 1, code, headers, status
 end
 
diff --git a/src/smtp.lua b/src/smtp.lua
index 1708053..d6357d2 100644
--- a/src/smtp.lua
+++ b/src/smtp.lua
@@ -31,51 +31,51 @@ ZONE = "-0000"
 local metat = { __index = {} }
 
 function metat.__index:greet(domain)
-    socket.try(self.tp:check("2.."))
-    socket.try(self.tp:command("EHLO", domain or DOMAIN))
-    return socket.skip(1, socket.try(self.tp:check("2..")))
+    self.try(self.tp:check("2.."))
+    self.try(self.tp:command("EHLO", domain or DOMAIN))
+    return socket.skip(1, self.try(self.tp:check("2..")))
 end 
 
 function metat.__index:mail(from)
-    socket.try(self.tp:command("MAIL", "FROM:" .. from))
-    return socket.try(self.tp:check("2.."))
+    self.try(self.tp:command("MAIL", "FROM:" .. from))
+    return self.try(self.tp:check("2.."))
 end 
 
 function metat.__index:rcpt(to)
-    socket.try(self.tp:command("RCPT", "TO:" .. to))
-    return socket.try(self.tp:check("2.."))
+    self.try(self.tp:command("RCPT", "TO:" .. to))
+    return self.try(self.tp:check("2.."))
 end
 
 function metat.__index:data(src, step)
-    socket.try(self.tp:command("DATA"))
-    socket.try(self.tp:check("3.."))
-    socket.try(self.tp:source(src, step))
-    socket.try(self.tp:send("\r\n.\r\n"))
-    return socket.try(self.tp:check("2.."))
+    self.try(self.tp:command("DATA"))
+    self.try(self.tp:check("3.."))
+    self.try(self.tp:source(src, step))
+    self.try(self.tp:send("\r\n.\r\n"))
+    return self.try(self.tp:check("2.."))
 end
 
 function metat.__index:quit()
-    socket.try(self.tp:command("QUIT"))
-    return socket.try(self.tp:check("2.."))
+    self.try(self.tp:command("QUIT"))
+    return self.try(self.tp:check("2.."))
 end
 
 function metat.__index:close()
-    return socket.try(self.tp:close())
+    return self.try(self.tp:close())
 end
 
 function metat.__index:login(user, password)
-    socket.try(self.tp:command("AUTH", "LOGIN"))
-    socket.try(self.tp:check("3.."))
-    socket.try(self.tp:command(mime.b64(user)))
-    socket.try(self.tp:check("3.."))
-    socket.try(self.tp:command(mime.b64(password)))
-    return socket.try(self.tp:check("2.."))
+    self.try(self.tp:command("AUTH", "LOGIN"))
+    self.try(self.tp:check("3.."))
+    self.try(self.tp:command(mime.b64(user)))
+    self.try(self.tp:check("3.."))
+    self.try(self.tp:command(mime.b64(password)))
+    return self.try(self.tp:check("2.."))
 end
 
 function metat.__index:plain(user, password)
     local auth = "PLAIN " .. mime.b64("\0" .. user .. "\0" .. password)
-    socket.try(self.tp:command("AUTH", auth))
-    return socket.try(self.tp:check("2.."))
+    self.try(self.tp:command("AUTH", auth))
+    return self.try(self.tp:check("2.."))
 end
 
 function metat.__index:auth(user, password, ext)
@@ -85,7 +85,7 @@ function metat.__index:auth(user, password, ext)
     elseif string.find(ext, "AUTH[^\n]+PLAIN") then
         return self:plain(user, password)
     else
-        socket.try(nil, "authentication not supported")
+        self.try(nil, "authentication not supported")
     end
 end
 
@@ -104,7 +104,9 @@ end
 
 function open(server, port)
     local tp = socket.try(tp.connect(server or SERVER, port or PORT, TIMEOUT))
-    return setmetatable({tp = tp}, metat)
+    -- make sure tp is closed if we get an exception
+    local try = socket.newtry(function() tp:close() end)
+    return setmetatable({ tp = tp, try = try}, metat)
 end
 
 ---------------------------------------------------------------------------
@@ -222,10 +224,10 @@ end
 -- High level SMTP API
 -----------------------------------------------------------------------------
 send = socket.protect(function(mailt)
-    local con = open(mailt.server, mailt.port)
-    local ext = con:greet(mailt.domain)
-    con:auth(mailt.user, mailt.password, ext)
-    con:send(mailt)
-    con:quit()
-    return con:close()
+    local s = open(mailt.server, mailt.port)
+    local ext = s:greet(mailt.domain)
+    s:auth(mailt.user, mailt.password, ext)
+    s:send(mailt)
+    s:quit()
+    return s:close()
 end)
diff --git a/src/socket.lua b/src/socket.lua
index f73d167..0a681bf 100644
--- a/src/socket.lua
+++ b/src/socket.lua
@@ -156,8 +156,8 @@ socket.sourcet["http-chunked"] = function(sock)
             else
                 -- get chunk and skip terminating CRLF
                 local chunk, err = sock:receive(size)
-                if err or socket.skip(2, sock:receive()) then return nil, err 
-                else return chunk end
+                if chunk then sock:receive() end 
+                return chunk, err
             end
         end
     })