[osinetwork] TCP

network/bin/nc.lua

@@ -6,21 +6,31 @@ local os     = require("os")
 local socket = require("socket")
 
 
-local args, opts        = shell.parse(...)
-local udpSocket, reason = socket.udp()
+local args, opts = shell.parse(...)
+---@type TCPSocket|UDPSocket
+local localSocket
+---@type TCPSocket|nil
+local clientSocket
+---@type thread
 local listenerThread
 
 
 opts.p = tonumber(opts.p) or 0
 opts.b = opts.b or "0.0.0.0"
 
----@param listenedSocket UDPSocket
+---@param listenedSocket UDPSocket|TCPSocket
 local function listenSocket(listenedSocket)
     checkArg(1, listenedSocket, 'table')
     while true do
-        local datagram = listenedSocket:recieve()
-        if (datagram) then
-            term.write(datagram)
+        if (listenedSocket:instanceOf(socket.tcp)) then
+            ---@cast listenedSocket TCPSocket
+            if (listenedSocket:getState() ~= "ESTABLISHED") then
+                return false
+            end
+        end
+        local data = listenedSocket:recieve()
+        if (data) then
+            term.write(data)
         end
         os.sleep()
     end
@@ -37,8 +47,18 @@ local function help()
     print("\t nc -u 192.168.1.1 9999")
 end
 
-event.listen("interrupted", function(...)
-    if (udpSocket) then udpSocket:close() end
+---read stdin indefenitely and send what's read through the socket
+---@param outSocket TCPSocket|UDPSocket
+local function readUserInput(outSocket)
+    repeat
+        local msg = term.read()
+        if (msg) then outSocket:send(msg .. "\n") end
+    until not msg or listenerThread:status() == "dead"
+end
+
+local function exit()
+    if (localSocket) then localSocket:close() end
+    if (clientSocket) then clientSocket:close() end
     if (listenerThread) then
         if (not listenerThread:join(3)) then
             listenerThread:kill()
@@ -46,63 +66,62 @@ event.listen("interrupted", function(...)
     end
     return false
 end
-)
+
+event.listen("interrupted", exit)
+
 if (opts.h or opts.help) then
     help()
     os.exit()
 elseif (opts.l and opts.u and (tonumber(args[1]) or opts.p)) then --listen UDP
-    assert(udpSocket:setsockname("*", tonumber(args[1]) or opts.p))
-    --udpSocket:setCallback(listenSocket)
-    print(string.format("Listening on %s:%d", udpSocket:getsockname()))
-    listenerThread = thread.create(listenSocket, udpSocket)
+    localSocket = socket.udp()
+    assert(localSocket:setsockname("*", tonumber(args[1]) or opts.p))
+    print(string.format("Listening on %s:%d", localSocket:getsockname()))
+    listenerThread = thread.create(listenSocket, localSocket)
     while true do
         --no remote addr/port. We cannot send msgs
         os.sleep()
     end
-    udpSocket:close()
+    localSocket:close()
 elseif (opts.u) then --connect UDP
-    assert(udpSocket:setsockname(opts.b, opts.p))
+    localSocket = socket.udp()
+    assert(localSocket:setsockname(opts.b, opts.p))
     args[2] = assert(tonumber(args[2]), "Invalid port number")
-    assert(udpSocket:setpeername(args[1], args[2]))
-    --udpSocket:setCallback(listenSocket)
-    print(string.format("Listening on %s:%d", udpSocket:getsockname()))
-    listenerThread = thread.create(listenSocket, udpSocket)
-    repeat
-        local msg = term.read()
-        if (msg) then udpSocket:send(msg .. "\n") end
-    until not msg
-    udpSocket:close()
-elseif (opts.l) then
-    local tcpsocket = socket.tcp()
-    assert(tcpsocket:bind(opts.b, opts.p))
-    args[2] = assert(tonumber(args[2]), "Invalid port number")
-    assert(tcpsocket:bind(args[1], args[2]))
-    print(string.format("Listening on %s:%d", tcpsocket:getsockname()))
-    tcpsocket:listen(1)
-    local client = tcpsocket:accept()
-    if (client) then
-        listenerThread = thread.create(listenSocket, client)
-        repeat
-            local msg = term.read()
-            if (msg) then client:send(msg .. "\n") end
-        until not msg
-        client:close()
+    assert(localSocket:setpeername(args[1], args[2]))
+    print(string.format("Listening on %s:%d", localSocket:getsockname()))
+    listenerThread = thread.create(listenSocket, localSocket)
+    readUserInput(localSocket)
+    localSocket:close()
+elseif (opts.l) then --listen tcp
+    localSocket = socket.tcp()
+    args[1] = args[1] or opts.b
+    args[2] = assert(tonumber(args[2] or opts.p), "Invalid port number")
+    assert(localSocket:bind(args[1], args[2]))
+    print(string.format("Listening on %s:%d", localSocket:getsockname()))
+    localSocket:listen(1)
+    local reason
+    clientSocket = localSocket:accept()
+    localSocket:close() --client connected, we don't need the listening socket anymore
+    if (clientSocket) then
+        print(string.format("Connected to : %s:%d", clientSocket:getpeername()))
+        listenerThread = thread.create(listenSocket, clientSocket)
+        readUserInput(clientSocket)
+        clientSocket:close()
+    else
+        print(reason)
     end
-    tcpsocket:close()
 else --connect TCP
     args[2] = assert(tonumber(args[2]), "Invalid port number")
-    local tcpsocket = socket.tcp()
-    tcpsocket:settimeout(5)
-    local s = tcpsocket:connect(args[1], args[2])
+    localSocket = socket.tcp()
+    localSocket:settimeout(5)
+    local s = localSocket:connect(args[1], args[2])
     if (s ~= 1) then
         print("Timeout")
-        os.exit(1)
+    else
+        print(string.format("Connected to %s:%d", localSocket:getpeername()))
+        listenerThread = thread.create(listenSocket, localSocket)
+        readUserInput(localSocket)
     end
-    print(string.format("Connected to %s:%d", tcpsocket:getpeername()))
-    listenerThread = thread.create(listenSocket, tcpsocket)
-    repeat
-        local msg = term.read()
-        if (msg) then tcpsocket:send(msg .. "\n") end
-    until not msg
-    tcpsocket:close()
+    localSocket:close()
 end
+
+exit()

network/lib/network/tcp/TCPLayer.lua

@@ -4,7 +4,7 @@ local TCPSegment   = require("network.tcp.TCPSegment")
 local ipv4Address  = require("network.ipv4.address")
 local NetworkLayer = require('network.abstract.NetworkLayer')
 local network      = require("network")
-local utils = require("network.utils")
+local utils        = require("network.utils")
 local class        = require("libClass2")
 
 
@@ -30,7 +30,7 @@ end
 
 ---@package
 function TCPLayer:tick()
-    for _,s in pairs(utils.getTreeBottomValues(self._sockets)) do
+    for _, s in pairs(utils.getTreeBottomValues(self._sockets)) do
         ---@cast s TCPSocket
         s:_tick()
     end
@@ -42,7 +42,7 @@ function TCPLayer:payloadHandler(from, to, payload)
     if (seg:flag(TCPSegment.Flags.SYN) and not seg:flag(TCPSegment.Flags.ACK)) then
         --initial sync packet
         socket = self:getSocket(to, seg:dstPort(), 0, 0)
-        if (not socket) then
+        if (not socket) then --Send RST
             local rstseg = TCPSegment(seg:dstPort(), seg:srcPort(), "")
             rstseg:flags(TCPSegment.Flags.RST|TCPSegment.Flags.ACK)
             rstseg:seq(seg:flag(TCPSegment.Flags.ACK) and seg:ack() or 0)

network/lib/network/utils.lua

@@ -28,7 +28,7 @@ local class = require("libClass2")
 ---@class Buffer:Object
 ---@field private _data string
 ---@operator call:Buffer
-local Buffer = require('libClass2')()
+local Buffer = class()
 
 ---Comment
 ---@return Buffer

network/lib/socket/init.lua

@@ -1,12 +1,9 @@
 local TCPSocket = require("socket.tcp")
+local UDPSocket = require("socket.udp")
 
 return {
-    udp = require("socket.udp"),
-    ---Creates and returns an TCP master object. A master object can be transformed into a server object with the method listen (after a call to bind) or into a client object with the method connect. The only other method supported by a master object is the close method.
-    ---
-    ---In case of success, a new master object is returned. In case of error, nil is returned, followed by an error message.
-    ---@return TCPSocket
-    tcp = function() return TCPSocket() end,
+    udp = UDPSocket,
+    tcp = TCPSocket,
     ---@type socketdns
     dns = select(2, pcall(require, "socket.dns"))
 }

network/lib/socket/tcp.lua

@@ -155,6 +155,7 @@ function TCPSocket:_makeDataSegment(data)
     return seg
 end
 
+--set the segment `seg` acknowledgment number to acknowledge the `recived` segment
 ---@protected
 ---@param seg TCPSegment
 ---@param received TCPSegment
@@ -165,6 +166,7 @@ function TCPSocket:_setAck(seg, received)
     return seg
 end
 
+---make a new acknowledgment segment for the `recived` segment
 ---@protected
 ---@param received TCPSegment
 ---@return TCPSegment
@@ -243,7 +245,7 @@ function TCPSocket:bind(address, port)
 end
 
 function TCPSocket:close()
-    if (self._kind == "client") then
+    if (self._kind == "client" and self._state ~= "CLOSED") then
         local seg = self:_makeSegment()
         seg:flag(f.FIN, true)
         seg:flag(f.ACK, true)
@@ -338,7 +340,7 @@ end
 ---Prefix is an optional string to be concatenated to the beginning of any received data before return.\
 ---
 ---If successful, the method returns the received pattern. In case of error, the method returns nil followed by an error message, followed by a (possibly empty) string containing the partial that was received. The error message can be the string 'closed' in case the connection was closed before the transmission was completed or the string 'timeout' in case there was a timeout during the operation.
----@param pattern? string
+---@param pattern? "*a"|"*l"|number
 ---@param prefix? string
 ---@return string? data, string?reason
 function TCPSocket:recieve(pattern, prefix)
@@ -349,6 +351,7 @@ function TCPSocket:recieve(pattern, prefix)
     local data
     repeat
         data = self._inBuffer:read(pattern)
+        os.sleep()
     until data ~= nil or self:_hasTimedOut(t1)
     if (not data) then
         return nil, self:_hasTimedOut(t1) and "timeout" or ""
@@ -473,6 +476,7 @@ function TCPSocket:payloadHandler(from, to, tcpSegment)
     if (tcpSegment:offset() > 5) then
         self:handleOptions(tcpSegment)
     end
+
     if (self._state == "LISTEN") then
         if (tcpSegment:flag(f.SYN)) then
             if (#(self._backlog) < self._backlogLen) then

network/lib/socket/udp.lua

@@ -173,10 +173,4 @@ function UDPSocket:payloadHandler(from, to, udpPacket)
     table.insert(self._buffer, {udpPacket:payload(), ipv4Address.tostring(from), udpPacket:srcPort()})
 end
 
----Create and returns an unconnected UDP socket.
----@return UDPSocket
-local function udp()
-    return UDPSocket()
-end
-
-return udp
+return UDPSocket