Shootout Echo

wiki

-- 2003-09 (Lua 5.0)
-- lua -lposix -lsocket echo.lua [number]

data = 'Hello there sailor\n'
dlen = string.len(data)

-- this documents the ways a tcp receive() can go wrong.  it should
-- not be called if the 2nd result of calling receive() == nil, but it
-- guards against the possibility even so.
function socket_receive_error(err)
   if err == nil then
      print'socket_receive_error(): err == nil, oops!'
   elseif err == 'closed' then
      print'<socket>:receive(): connection closed'
      return
   elseif err == 'timeout' then
      print'<socket>:receive(): connection timed out'
   else
      error('<socket>:receive(): error: '..err)
   end
end

function server_sock()
   local server,err = socket.bind('127.0.0.1',0)
   assert(server,err)
   return server
end

function get_port(sock)
   local _,port = sock:getsockname()
   return port
end

function client_sock(port)
   local client,err = socket.connect('127.0.0.1',port)
   assert(client,err)
   return client
end

function echo_client(n,port)
   local s = client_sock(port)
   for i=1,n do
      s:send(data)
      local d,err = s:receive(dlen)
      if err then
         socket_receive_error(err)
         break
      end
      if d ~= data then error('server reply: '..d) end
   end
   s:close()
end

function echo_server(n)
   local ssock = server_sock()
   local pid = posix.fork()
   if pid == 0 then
      echo_client(n,get_port(ssock)) -- child is client
   else
      local n = 0
      local client,err = ssock:accept()
      while true do
         local data,err = client:receive(dlen)
         if err then
            socket_receive_error(err)
            break
         end
         n = n + string.len(data)
         client:send(data)
      end
      print('server received '..n..' bytes')
   end
end

-- insist on a proper argument if one's present, no silent corrections
if arg and arg[1] then
   n = tonumber(arg[1])
   if not n then error('arg must be a number') end
else
   n = 1
end
echo_server(n)

--SteveElkins