Added forward server.

samples/forward.lua

@@ -0,0 +1,229 @@
+-- load our favourite library
+local socket = require"socket"
+-- timeout before an inactive thread is kicked
+local TIMEOUT = 10
+-- local address to bind to
+local ihost, iport = arg[1] or "localhost", arg[2] or 8080
+-- address to forward all data to
+local ohost, oport = arg[3] or "localhost", arg[4] or 3128
+
+-- creates a new set data structure
+function newset()
+    local reverse = {}
+    local set = {}
+    return setmetatable(set, {__index = {
+        insert = function(set, value)
+            if not reverse[value] then
+                table.insert(set, value)
+                reverse[value] = table.getn(set)
+            end
+        end,
+        remove = function(set, value)
+            local index = reverse[value]
+            if index then
+                reverse[value] = nil
+                local top = table.remove(set)
+                if top ~= value then 
+                    reverse[top] = index
+                    set[index] = top
+                end 
+            end
+        end
+    }})
+end
+
+local receiving = newset()
+local sending = newset() 
+local context = {}
+
+-- starts a non-blocking connect
+function nconnect(host, port)
+    local peer, err = socket.tcp()
+    if not peer then return nil, err end
+    peer:settimeout(0)
+    local ret, err = peer:connect(host, port)
+    if ret then return peer end
+    if err ~= "timeout" then 
+        peer:close()
+        return nil, err 
+    end
+    return peer
+end
+
+-- gets rid of a client
+function kick(who)
+if who == server then error("FUDEU") end
+    if context[who] then
+        sending:remove(who)
+        receiving:remove(who)
+        context[who] = nil
+        who:close()
+    end
+end
+
+-- decides what to do with a thread based on coroutine return
+function route(who, status, what)
+print(who, status, what)
+    if status and what then
+        if what == "receiving" then receiving:insert(who) end
+        if what == "sending" then sending:insert(who) end
+    else kick(who) end
+end
+
+-- loops accepting connections and creating new threads to deal with them
+function accept(server)
+    while true do
+print(server, "accepting a new client")
+        -- accept a new connection and start a new coroutine to deal with it
+        local client = server:accept()
+        if client then
+            -- start a new connection, non-blockingly, to the forwarding address
+            local peer = nconnect(ohost, oport)
+            if peer then
+                context[client] = {
+                    last = socket.gettime(),
+                    thread = coroutine.create(forward),
+                    peer = peer,
+                }
+                -- make sure peer will be tested for writing in the next select
+                -- round, which means the connection attempt has finished
+                sending:insert(peer)
+                context[peer] = {
+                    peer = client,
+                    thread = coroutine.create(check),
+                    last = socket.gettime()
+                }
+                -- put both in non-blocking mode
+                client:settimeout(0)
+                peer:settimeout(0)
+            else 
+                -- otherwise just dump the client
+                client:close() 
+            end
+        end
+        -- tell scheduler we are done for now
+        coroutine.yield("receiving")
+    end
+end
+
+-- forwards all data arriving to the appropriate peer
+function forward(who)
+    while true do
+print(who, "getting data")
+        -- try to read as much as possible
+        local data, rec_err, partial = who:receive("*a")
+        -- if we had an error other than timeout, abort
+        if rec_err and rec_err ~= "timeout" then return error(rec_err) end
+        -- if we got a timeout, we probably have partial results to send
+        data = data or partial
+print(who, " got ", string.len(data))
+        -- renew our timestamp so scheduler sees we are active
+        context[who].last = socket.gettime()
+        -- forward what we got right away
+        local peer = context[who].peer
+        while true do
+            -- tell scheduler we need to wait until we can send something
+            coroutine.yield("sending")
+            local ret, snd_err
+            local start = 0
+print(who, "sending data")
+            ret, snd_err, start = peer:send(data, start+1)
+            if ret then break 
+            elseif snd_err ~= "timeout" then return error(snd_err) end
+            -- renew our timestamp so scheduler sees we are active
+            context[who].last = socket.gettime()
+        end
+        -- if we are done receiving, we are done with this side of the
+        -- connection
+        if not rec_err then return nil end
+        -- otherwise tell schedule we have to wait for more data to arrive
+        coroutine.yield("receiving")
+    end
+end
+
+-- checks if a connection completed successfully and if it did, starts
+-- forwarding all data
+function check(who)
+    local ret, err = who:connected()
+    if ret then 
+print(who, "connection completed")
+        receiving:insert(context[who].peer)
+        context[who].last = socket.gettime()
+print(who, "yielding until there is input data")
+        coroutine.yield("receiving")
+        return forward(who)
+    else return error(err) end
+end
+
+-- initializes the forward server
+function init()
+    -- socket sets to test for events 
+    -- create our server socket
+    server = assert(socket.bind(ihost, iport))
+    server:settimeout(0.1) -- we don't want to be killed by bad luck
+    -- we initially
+    receiving:insert(server)
+    context[server] = { thread = coroutine.create(accept) }
+end
+
+-- loop waiting until something happens, restarting the thread to deal with
+-- what happened, and routing it to wait until something else happens
+function go()
+    while true  do
+print("will select for readability")
+for i,v in ipairs(receiving) do
+    print(i, v)
+end
+print("will select for writability")
+for i,v in ipairs(sending) do
+    print(i, v)
+end
+        -- check which sockets are interesting and act on them
+        readable, writable = socket.select(receiving, sending, 3)
+print("returned as readable")
+for i,v in ipairs(readable) do
+    print(i, v)
+end
+print("returned as writable")
+for i,v in ipairs(writable) do
+    print(i, v)
+end
+        -- for all readable connections, resume its thread and route it 
+        for _, who in ipairs(readable) do
+            receiving:remove(who)
+            if context[who] then
+                route(who, coroutine.resume(context[who].thread, who))
+            end
+        end
+        -- for all writable connections, do the same
+        for _, who in ipairs(writable) do
+            sending:remove(who)
+            if context[who] then
+                route(who, coroutine.resume(context[who].thread, who))
+            end
+        end
+        -- put all inactive threads in death row
+        local now = socket.gettime()
+        local deathrow
+        for who, data in pairs(context) do
+            if data.last then
+print("hung for" , now - data.last, who)
+                if  now - data.last > TIMEOUT then
+                    -- only create table if someone is doomed
+                    deathrow = deathrow or {} 
+                    deathrow[who] = true
+                end
+            end
+        end
+        -- finally kick everyone in deathrow
+        if deathrow then
+print("in death row")
+for i,v in pairs(deathrow) do
+    print(i, v)
+end
+            for who in pairs(deathrow) do kick(who) end
+        end
+    end
+end
+
+go(init())