Fixed smtp.lua loading.

Adjusted tftp module.
Added some comments.

1 files changed, 77 insertions, 57 deletions

diff --git a/etc/tftp.lua b/etc/tftp.lua
index 4113034..5238de0 100644
--- a/etc/tftp.lua
+++ b/etc/tftp.lua
@@ -2,13 +2,20 @@
 -- TFTP support for the Lua language
 -- LuaSocket toolkit.
 -- Author: Diego Nehab
--- Conforming to: RFC 783, LTN7
 -- RCS ID: $Id$
 -----------------------------------------------------------------------------
 
-local Public, Private = {}, {}
-local socket = _G[LUASOCKET_LIBNAME] -- get LuaSocket namespace
-socket.tftp = Public  -- create tftp sub namespace
+-----------------------------------------------------------------------------
+-- Load required files
+-----------------------------------------------------------------------------
+local socket = require("socket")
+local ltn12 = require("ltn12")
+local url = require("url")
+
+-----------------------------------------------------------------------------
+-- Setup namespace
+-----------------------------------------------------------------------------
+_LOADED["tftp"] = getfenv(1)
 
 -----------------------------------------------------------------------------
 -- Program constants
@@ -16,33 +23,33 @@ socket.tftp = Public  -- create tftp sub namespace
 local char = string.char
 local byte = string.byte
 
-Public.PORT = 69
-Private.OP_RRQ = 1
-Private.OP_WRQ = 2
-Private.OP_DATA = 3
-Private.OP_ACK = 4
-Private.OP_ERROR = 5
-Private.OP_INV = {"RRQ", "WRQ", "DATA", "ACK", "ERROR"}
+PORT = 69
+local OP_RRQ = 1
+local OP_WRQ = 2
+local OP_DATA = 3
+local OP_ACK = 4
+local OP_ERROR = 5
+local OP_INV = {"RRQ", "WRQ", "DATA", "ACK", "ERROR"}
 
 -----------------------------------------------------------------------------
 -- Packet creation functions
 -----------------------------------------------------------------------------
-function Private.RRQ(source, mode)
-        return char(0, Private.OP_RRQ) .. source .. char(0) .. mode .. char(0)
+local function RRQ(source, mode)
+        return char(0, OP_RRQ) .. source .. char(0) .. mode .. char(0)
 end
 
-function Private.WRQ(source, mode)
-        return char(0, Private.OP_RRQ) .. source .. char(0) .. mode .. char(0)
+local function WRQ(source, mode)
+        return char(0, OP_RRQ) .. source .. char(0) .. mode .. char(0)
 end
 
-function Private.ACK(block)
+local function ACK(block)
         local low, high
         low = math.mod(block, 256)
         high = (block - low)/256
-        return char(0, Private.OP_ACK, high, low) 
+        return char(0, OP_ACK, high, low) 
 end
 
-function Private.get_OP(dgram)
+local function get_OP(dgram)
     local op = byte(dgram, 1)*256 + byte(dgram, 2)
     return op
 end
@@ -50,13 +57,13 @@ end
 -----------------------------------------------------------------------------
 -- Packet analysis functions
 -----------------------------------------------------------------------------
-function Private.split_DATA(dgram)
+local function split_DATA(dgram)
         local block = byte(dgram, 3)*256 + byte(dgram, 4)
         local data = string.sub(dgram, 5)
         return block, data
 end
 
-function Private.get_ERROR(dgram)
+local function get_ERROR(dgram)
         local code = byte(dgram, 3)*256 + byte(dgram, 4)
         local msg
         _,_, msg = string.find(dgram, "(.*)\000", 5)
@@ -64,68 +71,81 @@ function Private.get_ERROR(dgram)
 end
 
 -----------------------------------------------------------------------------
--- Downloads and returns a file pointed to by url
+-- High level TFTP API
 -----------------------------------------------------------------------------
-function Public.get(url)
-    local parsed = socket.url.parse(url, {
-        host = "",
-        port = Public.PORT,
-        path ="/",
-        scheme = "tftp"
-    })
-    if parsed.scheme ~= "tftp" then
-        return nil, string.format("unknown scheme '%s'", parsed.scheme)
-    end
+local function tget(gett)
     local retries, dgram, sent, datahost, dataport, code
-    local cat = socket.concat.create()
     local last = 0
-        local udp, err = socket.udp()
-        if not udp then return nil, err end
+        local con = socket.try(socket.udp())
     -- convert from name to ip if needed
-        parsed.host = socket.dns.toip(parsed.host)
-        udp:settimeout(1)
+        gett.host = socket.try(socket.dns.toip(gett.host))
+        con:settimeout(1)
     -- first packet gives data host/port to be used for data transfers
     retries = 0
         repeat 
-                sent, err = udp:sendto(Private.RRQ(parsed.path, "octet"), 
-            parsed.host, parsed.port)
-                if err then return nil, err end
-                dgram, datahost, dataport = udp:receivefrom()
+                sent = socket.try(con:sendto(RRQ(gett.path, "octet"), 
+            gett.host, gett.port))
+                dgram, datahost, dataport = con:receivefrom()
         retries = retries + 1
         until dgram or datahost ~= "timeout" or retries > 5
-        if not dgram then return nil, datahost end
+        socket.try(dgram, datahost) 
     -- associate socket with data host/port
-        udp:setpeername(datahost, dataport)
+        socket.try(con:setpeername(datahost, dataport))
+    -- default sink
+    local sink = gett.sink or ltn12.sink.null()
     -- process all data packets
         while 1 do
         -- decode packet
-                code = Private.get_OP(dgram)
-                if code == Private.OP_ERROR then 
-            return nil, Private.get_ERROR(dgram) 
-        end
-                if code ~= Private.OP_DATA then 
-            return nil, "unhandled opcode " .. code 
-        end
+                code = get_OP(dgram)
+                socket.try(code ~= OP_ERROR, get_ERROR(dgram))
+        socket.try(code == OP_DATA, "unhandled opcode " .. code)
         -- get data packet parts
-                local block, data = Private.split_DATA(dgram)
+                local block, data = split_DATA(dgram)
         -- if not repeated, write
         if block == last+1 then
-                    cat:addstring(data)
+                    socket.try(sink(data))
             last = block
         end
         -- last packet brings less than 512 bytes of data
                 if string.len(data) < 512 then 
-            sent, err = udp:send(Private.ACK(block)) 
-            return cat:getresult()
+            socket.try(con:send(ACK(block)))
+            socket.try(con:close())
+            socket.try(sink(nil))
+            return 1
         end
         -- get the next packet
         retries = 0
                 repeat 
-                        sent, err = udp:send(Private.ACK(last))
-                        if err then return err end
-                        dgram, err = udp:receive()
+                        sent = socket.try(con:send(ACK(last)))
+                        dgram, err = con:receive()
             retries = retries + 1
                 until dgram or err ~= "timeout" or retries > 5
-                if not dgram then return err end
+                socket.try(dgram, err)
         end
 end
+
+local default = {
+    port = PORT,
+    path ="/",
+    scheme = "tftp"
+}
+
+local function parse(u)
+    local t = socket.try(url.parse(u, default))
+    socket.try(t.scheme == "tftp", "invalid scheme '" .. t.scheme .. "'")
+    socket.try(t.host, "invalid host")
+    return t
+end
+
+local function sget(u)
+    local gett = parse(u) 
+    local t = {}
+    gett.sink = ltn12.sink.table(t)
+    tget(gett)
+    return table.concat(t)
+end
+
+get = socket.protect(function(gett)
+    if type(gett) == "string" then return sget(gett)
+    else return tget(gett) end
+end)