luasocket/samples/forward.lua
2005-04-20 18:57:47 +00:00

225 lines
6.9 KiB
Lua

-- load our favourite library
local socket = require"socket"
-- 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
-- timeout before an inactive thread is kicked
local TIMEOUT = 10
-- set of connections waiting to receive data
local receiving = newset()
-- set of sockets waiting to send data
local sending = newset()
-- context for connections and servers
local context = {}
function wait(who, what)
if what == "input" then receiving:insert(who)
else sending:insert(who) end
context[who].last = socket.gettime()
coroutine.yield()
end
-- initializes the forward server
function init()
if table.getn(arg) < 1 then
print("Usage")
print(" lua forward.lua <iport:ohost:oport> ...")
os.exit(1)
end
-- for each tunnel, start a new server socket
for i, v in ipairs(arg) do
-- capture forwarding parameters
local iport, ohost, oport =
socket.skip(2, string.find(v, "([^:]+):([^:]+):([^:]+)"))
assert(iport, "invalid arguments")
-- create our server socket
local server = assert(socket.bind("*", iport))
server:settimeout(0) -- we don't want to be killed by bad luck
-- make sure server is tested for readability
receiving:insert(server)
-- add server context
context[server] = {
thread = coroutine.create(accept),
ohost = ohost,
oport = oport
}
end
end
-- starts a connection in a non-blocking way
function connect(who, host, port)
who:settimeout(0)
print("trying to connect peer", who, host, port)
local ret, err = who:connect(host, port)
if not ret and err == "timeout" then
print("got timeout, will wait", who)
wait(who, "output")
ret, err = who:connect(host, port)
print("connection results arrived", who, ret, err)
end
if not ret then
print("connection failed", who, err)
kick(who)
kick(context[who].peer)
else
return forward(who)
end
end
-- gets rid of a client
function kick(who)
if who and context[who] then
sending:remove(who)
receiving:remove(who)
context[who] = nil
who:close()
end
end
-- loops accepting connections and creating new threads to deal with them
function accept(server)
while true do
-- accept a new connection and start a new coroutine to deal with it
local client = server:accept()
print("accepted ", client)
if client then
-- create contexts for client and peer.
local peer, err = socket.tcp()
if peer then
context[client] = {
last = socket.gettime(),
-- client goes straight to forwarding loop
thread = coroutine.create(forward),
peer = peer,
}
context[peer] = {
last = socket.gettime(),
peer = client,
-- peer first tries to connect to forwarding address
thread = coroutine.create(connect),
last = socket.gettime()
}
-- resume peer and client so they can do their thing
local ohost = context[server].ohost
local oport = context[server].oport
coroutine.resume(context[peer].thread, peer, ohost, oport)
coroutine.resume(context[client].thread, client)
else
print(err)
client:close()
end
end
-- tell scheduler we are done for now
wait(server, "input")
end
end
-- forwards all data arriving to the appropriate peer
function forward(who)
print("starting to foward", who)
who:settimeout(0)
while true do
-- wait until we have something to read
wait(who, "input")
-- 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 kick(who) end
-- if we got a timeout, we probably have partial results to send
data = data or partial
-- 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
wait(who, "output")
local ret, snd_err
local start = 0
ret, snd_err, start = peer:send(data, start+1)
if ret then break
elseif snd_err ~= "timeout" then return kick(who) end
end
-- if we are done receiving, we are done
if not rec_err then
kick(who)
kick(peer)
end
end
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 reading")
for i,v in ipairs(receiving) do
print(i, v)
end
print("will select for sending")
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("was readable")
for i,v in ipairs(readable) do
print(i, v)
end
print("was writable")
for i,v in ipairs(writable) do
print(i, v)
end
-- for all readable connections, resume its thread
for _, who in ipairs(readable) do
receiving:remove(who)
coroutine.resume(context[who].thread, who)
end
-- for all writable connections, do the same
for _, who in ipairs(writable) do
sending:remove(who)
coroutine.resume(context[who].thread, who)
end
-- put all inactive threads in death row
local now = socket.gettime()
local deathrow
for who, data in pairs(context) do
if data.peer then
if now - data.last > TIMEOUT then
-- only create table if at least one is doomed
deathrow = deathrow or {}
deathrow[who] = true
end
end
end
-- finally kick everyone in deathrow
if deathrow then
for who in pairs(deathrow) do kick(who) end
end
end
end
init()
go()