TCP can now have multiple remotes on the same listening Port

2 files changed, 126 insertions, 61 deletions

diff --git a/templates/genericdev/device.lua b/templates/genericdev/device.lua
index 1b23af0..5435172 100644
--- a/templates/genericdev/device.lua
+++ b/templates/genericdev/device.lua
@@ -285,12 +285,15 @@ function cmdline(cmdline)
 	  addr = parse_ip(argt[3])
 	  sock_number = tonumber(argt[4])
 	  local open_sock = socket_tcp(l_sock_number)
-	  connect_tcp(open_sock, sock_number, addr, tcpopen_success, {open_sock})
+	  connect_tcp(open_sock, l_sock_number, sock_number, addr, tcpopen_success, {open_sock})
+	  nsprint("Socket no: "..open_sock)
    elseif(argt[1] == "tcpsend") then
 	  sock_number = tonumber(argt[2])
 	  send_tcp(sock_number, argt[3], "")
    elseif(argt[1] == "tcpclose") then
 	  sock_number = tonumber(argt[2])
+	  table_uninsert(communicating_sockets, sock_number)
+	  table_uninsert(listening_sockets, sock_number)
 	  close_tcp(sock_number)
    end
 end
@@ -371,9 +374,10 @@ function timestamp()
    end
 
    for k,v in pairs(listening_sockets) do
-	  if(accept_tcp(v) == 1) then
-		 listening_sockets[k] = nil
-		 table.insert(communicating_sockets, v)
+	  accept_sock = accept_tcp(v)
+	  if(accept_sock > 0) then
+		 nsprint("Socket no: "..accept_sock)
+		 table.insert(communicating_sockets, accept_sock)
 	  end
 	  
    end
diff --git a/templates/tcplib/tcplib.lua b/templates/tcplib/tcplib.lua
index 5199c8d..5c85c5f 100644
--- a/templates/tcplib/tcplib.lua
+++ b/templates/tcplib/tcplib.lua
@@ -13,9 +13,8 @@
 
 ns_openlib("devlib")
 
--- @todo add a send buffer
--- @todo allow a socket to have multiple remotes
 -- @todo Currently a port number can only be used with one address as the tcp_sockets table isn't bound to any specific address, this should be changed
+-- @todo allow for a closing function that can be set by the device to handle closing sockets from the remote
 
 ------- TCP --------
 tcp_header = ">I2>I2>I4>I4>c2>I2>I2>I2"
@@ -75,7 +74,7 @@ function make_tcpheader(flags, source_port, dest_port, seq, ack, window, cksum,
 end
 
 tcp_sockets = {}
-udp_sockets = {}
+tcp_used_ports = {}
 
 port_max = 65535
 
@@ -84,11 +83,11 @@ function socket_tcp(port, addr)
    -- Select an unprivileged port
    if(not port) then
 	  port = 1025
-	  while(tcp_sockets[port]) do
+	  while(tcp_used_ports[port]) do
 		 port = port + 1
 	  end
    end
-   if(tcp_sockets[port]) then
+   if(tcp_used_ports[port]) then
 	  nsprint("Port in use")
 	  return -1
    end
@@ -96,8 +95,24 @@ function socket_tcp(port, addr)
 	  nsprint("Port number to high")
 	  return -1
    end
-   tcp_sockets[port] = {["addr"] = addr, ["passive"] = false, ["queue"] = {}, ["r_addr"] = nil, ["r_port"] = nil, ["seq"] = nil, ["ack"] = nil, ["buffer"] = {}, ["send_buffer"] = "", ["recv_func"] = nil, ["recv_func_vals"] = nil}
-   return port
+   
+   table.insert(tcp_sockets, {
+				   ["port"] = port,
+				   ["addr"] = addr,
+				   ["passive"] = false,
+				   ["child_socks"] = {},
+				   ["queue"] = {},
+				   ["r_addr"] = nil,
+				   ["r_port"] = nil,
+				   ["seq"] = nil,
+				   ["ack"] = nil,
+				   ["buffer"] = {},
+				   ["send_buffer"] = "",
+				   ["recv_func"] = nil,
+				   ["recv_func_vals"] = nil
+   })
+   tcp_used_ports[port] = 1
+   return #tcp_sockets
 end
 
 -- req_func will get called with req_func_vals once a remote machine tries to establish a connection
@@ -118,26 +133,18 @@ function new_queue(r_addr, header_data, sock)
 end
 
 -- Returns -1 on error, 0 if there is nothing to accept and 1 if a connection has successfully been accepted
--- @todo: For now a per port there can only be one peer connected, this is not how it should be in the future
 function accept_tcp(sock)
    if(not tcp_sockets[sock]) then
-	  nsprint("Is not a port")
+	  nsprint("Is not a socket")
 	  return -1
    end
    if(not tcp_sockets[sock]["passive"]) then
-	  nsprint("Is not a passive port")
+	  nsprint("Is not a passive socket")
 	  return -1
    end
    if(not tcp_sockets[sock]["queue"][1]) then
 	  return 0
    end
-   -- We are already connected
-   -- @todo allow for multiple remotes on the same port
-   if(tcp_sockets[sock]["r_addr"]) then
-	  return 1
-   end
-
-   nsprint("ACCEPTING CONNECTION")
    
    -- Do all the accepting here
    flags = new_tcp_flagtable()
@@ -145,19 +152,35 @@ function accept_tcp(sock)
    flags["syn"] = 1
    flags["ack"] = 1
 
-   header = make_tcpheader(flags, sock, tcp_sockets[sock]["queue"][1]["port"], 0, tcp_sockets[sock]["queue"][1]["ack"] + 1, 65535, 0, 0, "")
+   header = make_tcpheader(flags, tcp_sockets[sock]["port"], tcp_sockets[sock]["queue"][1]["port"], 0, tcp_sockets[sock]["queue"][1]["ack"] + 1, 65535, 0, 0, "")
 
    send_ip(header, tcp_sockets[sock]["queue"][1]["addr"], tcp_sockets[sock]["addr"], iptypes["tcp"], 0, "tcp")
-
-   tcp_sockets[sock]["r_addr"] = tcp_sockets[sock]["queue"][1]["addr"]
-   tcp_sockets[sock]["r_port"] = tcp_sockets[sock]["queue"][1]["port"]
-   tcp_sockets[sock]["ack"] = tcp_sockets[sock]["queue"][1]["ack"] + 1
-   tcp_sockets[sock]["seq"] = 1
+   
+   table.insert(tcp_sockets, {
+				   ["port"] = tcp_sockets[sock]["port"],
+				   ["addr"] = tcp_sockets[sock]["addr"],
+				   ["child_socks"] = {},
+				   ["passive"] = false,
+				   ["queue"] = {},
+				   ["r_addr"] = tcp_sockets[sock]["queue"][1]["addr"],
+				   ["r_port"] = tcp_sockets[sock]["queue"][1]["port"],
+				   ["seq"] = 1,
+				   ["ack"] = tcp_sockets[sock]["queue"][1]["ack"] + 1,
+				   ["buffer"] = {},
+				   ["send_buffer"] = "",
+				   ["recv_func"] = nil,
+				   ["recv_func_vals"] = nil
+   })
+
+   table.insert(tcp_sockets[sock]["child_socks"], #tcp_sockets)
+   nsprint("ACCEPTING CONNECTION "..ip_to_string(tcp_sockets[sock]["queue"][1]["addr"]).." "..tcp_sockets[sock]["queue"][1]["port"])
    table.remove(tcp_sockets[sock]["queue"], 1)
-   return 1
+
+   
+   return #tcp_sockets
 end
 
-function connect_tcp(sock, port, addr, after_connect, after_connect_vals)
+function connect_tcp(sock, port, r_port, addr, after_connect, after_connect_vals)
    if(not tcp_sockets[sock]) then
 	  nsprint("Is not a port")
 	  return -1
@@ -167,7 +190,7 @@ function connect_tcp(sock, port, addr, after_connect, after_connect_vals)
 
    flags["syn"] = 1
 
-   header = make_tcpheader(flags, sock, port, 0, 0, 65535, 0, 0, "")
+   header = make_tcpheader(flags, tcp_sockets[sock]["port"], r_port, 0, 0, 65535, 0, 0, "")
 --   print("TCP HEADER: "..bin_to_hex(header))
 
    tcp_sockets[sock]["seq"] = 1
@@ -177,7 +200,9 @@ function connect_tcp(sock, port, addr, after_connect, after_connect_vals)
    send_ip(header, addr, tcp_sockets[sock]["addr"], iptypes["tcp"], 0, "tcp")
 
    tcp_sockets[sock]["r_addr"] = addr
-   tcp_sockets[sock]["r_port"] = port
+   tcp_sockets[sock]["r_port"] = r_port
+
+   tcp_sockets[sock]["port"] = port
 
    tcp_sockets[sock]["recv_func"] = connect_waiting
    tcp_sockets[sock]["recv_func_vals"] = {sock, after_connect, after_connect_vals}
@@ -191,7 +216,7 @@ function connect_waiting(r_addr, header_fields, sock, after_connect, after_conne
    tcp_sockets[sock]["ack"] = header_fields["seq"] + 1
 
    -- Acknowledge to complete handshake
-   header = make_tcpheader(flags, sock, tcp_sockets[sock]["r_port"], tcp_sockets[sock]["seq"], tcp_sockets[sock]["ack"], 65535, 0, 0, "")
+   header = make_tcpheader(flags, tcp_sockets[sock]["port"], tcp_sockets[sock]["r_port"], tcp_sockets[sock]["seq"], tcp_sockets[sock]["ack"], 65535, 0, 0, "")
    send_ip(header, tcp_sockets[sock]["r_addr"], tcp_sockets[sock]["addr"], iptypes["tcp"], 0, "tcp")
    
    after_connect(table.unpack(after_connect_vals))
@@ -240,24 +265,35 @@ function close_tcp(sock)
 	  return -1
    end
 
+
+   for k,v in pairs(tcp_sockets[sock]["child_socks"]) do
+	  nsprint("Closing child "..v)
+	  close_tcp(v)
+   end
+   if(tcp_sockets[sock]["passive"]) then
+	  tcp_sockets[sock] = nil
+	  return
+   end
+
    local flags = new_tcp_flagtable()
    flags["fin"] = 1
    flags["ack"] = 1
-   header = make_tcpheader(flags, sock, tcp_sockets[sock]["r_port"], tcp_sockets[port_d]["seq"], tcp_sockets[port_d]["ack"], 65535, 0, 0, "")
+   header = make_tcpheader(flags, tcp_sockets[sock]["port"], tcp_sockets[sock]["r_port"], tcp_sockets[sock]["seq"], tcp_sockets[sock]["ack"], 65535, 0, 0, "")
 
    send_ip(header, tcp_sockets[sock]["r_addr"], tcp_sockets[sock]["addr"], iptypes["tcp"], 0, "tcp")
    tcp_sockets[sock]["closed"] = "fin_wait"
 end
 
 function reset_tcp_sock(sock)
-   if(tcp_sockets[sock]["passive"]) then
-	  tcp_sockets[sock]["closed"] = nil
-	  tcp_sockets[sock]["r_addr"] = nil
-	  tcp_sockets[sock]["r_port"] = nil
-	  listen_tcp(sock)
-   else
-	  tcp_sockets[sock] = nil
-   end
+--   if(tcp_sockets[sock]["passive"]) then
+--	  tcp_sockets[sock]["closed"] = nil
+--	  tcp_sockets[sock]["r_addr"] = nil
+--	  tcp_sockets[sock]["r_port"] = nil
+--	  listen_tcp(sock)
+--   else
+--	  tcp_sockets[sock] = nil
+   --   end
+   tcp_sockets[sock] = nil
 end
 
 function handle_tcp(interface_index, data, source, dest)
@@ -265,12 +301,37 @@ function handle_tcp(interface_index, data, source, dest)
 
    header_size, flags = unpack_dorsvflags(dorsvflags)
 
-   if(not tcp_sockets[port_d]) then
+   if(not tcp_used_ports[port_d]) then
 	  nsprint("handle_tcp: Is not a socket")
 	  return -1
    end
+
+   -- Look for a matching port that matches the remote and local sockets, alternatively look for the 
+   port_match = nil
+   full_match = nil
+   for k,v in pairs(tcp_sockets) do
+	  if(v["port"] == port_d) then
+		 if(v["passive"]) then
+			port_match = k
+		 end
+		 if(v["r_port"] == port_s and v["r_addr"] == source) then
+			full_match = k
+		 end
+	  end
+   end
+   
+   sock = port_match
+   if(full_match) then
+	  sock = full_match
+   end
+
+   -- No matching socket found
+   if(not sock) then
+	  nsprint("No matching sock found")
+	  return
+   end
    
-   if(not tcp_sockets[port_d]["r_addr"] == source or not tcp_sockets[port_d]["s_addr"] == dest) then
+   if(not tcp_sockets[sock]["r_addr"] == source or not tcp_sockets[sock]["s_addr"] == dest) then
 	  nsprint("Someone wrong tried to connect to my socket")
 	  return
    end
@@ -278,45 +339,45 @@ function handle_tcp(interface_index, data, source, dest)
    proto_data = string.sub(data, header_size * 4 + 1, #data)
    
    -- Handle final ack
-   if(flags["ack"] == 1 and tcp_sockets[port_d]["closed"] == "ack_wait") then
-	  reset_tcp_sock(port_d)
+   if(flags["ack"] == 1 and tcp_sockets[sock]["closed"] == "ack_wait") then
+	  reset_tcp_sock(sock)
 	  return
    end
    
    if(#proto_data > 0 or flags["fin"] == 1) then
 	  -- Record data to the buffer
-	  tcp_sockets[port_d]["buffer"][seq] = proto_data
+	  tcp_sockets[sock]["buffer"][seq] = proto_data
 --	  print(tcp_sockets[port_d]["buffer"][seq])
-	  table.sort(tcp_sockets[port_d]["buffer"])
+	  table.sort(tcp_sockets[sock]["buffer"])
 
 	  -- Packets without data should increase ack as well
 	  if(#proto_data == 0) then
 		 ack_increment = 1
 	  else
 		 ack_increment = #proto_data
-	  end
-	  tcp_sockets[port_d]["ack"] = tcp_sockets[port_d]["ack"] + ack_increment
+	  end  
+	  tcp_sockets[sock]["ack"] = tcp_sockets[sock]["ack"] + ack_increment
 	  
 	  if(flags["fin"] == 1) then
-		 if(tcp_sockets[port_d]["closed"] == "fin_wait") then
+		 if(tcp_sockets[sock]["closed"] == "fin_wait") then
 			flags["fin"] = 0
 		 else
-			tcp_sockets[port_d]["closed"] = "ack_wait"
+			tcp_sockets[sock]["closed"] = "ack_wait"
 		 end
 	  end
 	  -- Ack the frame
-	  header = make_tcpheader(flags, port_d, tcp_sockets[port_d]["r_port"], tcp_sockets[port_d]["seq"], tcp_sockets[port_d]["ack"], 65535, 0, 0, "")
+	  header = make_tcpheader(flags, port_d, tcp_sockets[sock]["r_port"], tcp_sockets[sock]["seq"], tcp_sockets[sock]["ack"], 65535, 0, 0, "")
 	  -- @todo unify this with the timestep function so data is sent along with the acknowledge
-	  send_ip(header, tcp_sockets[port_d]["r_addr"], tcp_sockets[port_d]["addr"], iptypes["tcp"], 0, "tcp")
-	  if(tcp_sockets[port_d]["closed"] == "fin_wait") then
-		 reset_tcp_sock(port_d)
+	  send_ip(header, tcp_sockets[sock]["r_addr"], tcp_sockets[sock]["addr"], iptypes["tcp"], 0, "tcp")
+	  if(tcp_sockets[sock]["closed"] == "fin_wait") then
+		 reset_tcp_sock(sock)
 		 return
 	  end
    end
 
    -- Call the handler function
-   if(tcp_sockets[port_d]["recv_func"]) then
-	  tcp_sockets[port_d]["recv_func"](source, {["port_s"] = port_s,
+   if(tcp_sockets[sock]["recv_func"]) then
+	  tcp_sockets[sock]["recv_func"](source, {["port_s"] = port_s,
 										  ["r_port"] = port_s,
 										  ["seq"] = seq,
 										  ["ack"] = ack,
@@ -324,10 +385,10 @@ function handle_tcp(interface_index, data, source, dest)
 										  ["flags"] = flags,
 										  ["window"] = window,
 										  ["cksum"] = cksum,
-										  ["urg"] = urg}, table.unpack(tcp_sockets[port_d]["recv_func_vals"]))
+										  ["urg"] = urg}, table.unpack(tcp_sockets[sock]["recv_func_vals"]))
    end
 
-   tcp_sockets[port_d]["recv_func"] = nil   
+   --tcp_sockets[sock]["recv_func"] = nil   
 end
 
 function tcplib_timestep()
@@ -335,14 +396,14 @@ function tcplib_timestep()
 	  if(#v["send_buffer"] > 0) then		 
 		 flags = new_tcp_flagtable()
 		 flags["ack"] = 1
-		 header = make_tcpheader(flags, k, v["r_port"], v["seq"], v["ack"], 65535, 0, 0, "")
+		 header = make_tcpheader(flags, v["port"], v["r_port"], v["seq"], v["ack"], 65535, 0, 0, "")
 
 		 -- @todo use window size to work off buffer slowly
 		 data = v["send_buffer"]
 		 full_data = header..data
 		 v["send_buffer"] = ""
 
-		 send_ip(full_data, v["r_addr"], 0, iptypes["tcp"], 0, "tcp,".."RAW" )
+		 send_ip(full_data, v["r_addr"], v["addr"], iptypes["tcp"], 0, "tcp,".."RAW" )
 		 
 		 v["seq"] = v["seq"] + #data
 	  end