luajitos

TCP.lua (13703B)

---

 -- TCP Protocol Implementation
 -- Implements a basic TCP state machine and connection handling
 
 local TCP = {}
 
 -- TCP Flags
 TCP.FLAGS = {
     FIN = 0x01,
     SYN = 0x02,
     RST = 0x04,
     PSH = 0x08,
     ACK = 0x10,
     URG = 0x20,
 }
 
 -- TCP States
 TCP.STATE = {
     CLOSED = 0,
     LISTEN = 1,
     SYN_SENT = 2,
     SYN_RECEIVED = 3,
     ESTABLISHED = 4,
     FIN_WAIT_1 = 5,
     FIN_WAIT_2 = 6,
     CLOSE_WAIT = 7,
     CLOSING = 8,
     LAST_ACK = 9,
     TIME_WAIT = 10,
 }
 
 -- TCP connections
 TCP.connections = {}
 TCP.next_port = 49152
 
 ---Calculate TCP checksum (with pseudo-header)
 ---@param src_ip table Source IP (4 bytes)
 ---@param dst_ip table Destination IP (4 bytes)
 ---@param tcp_packet string TCP packet
 ---@return number checksum 16-bit checksum
 function TCP.calculate_checksum(src_ip, dst_ip, tcp_packet)
     local pseudo = {}
 
     -- Source IP
     for i = 1, 4 do
         pseudo[#pseudo + 1] = string.char(src_ip[i])
     end
 
     -- Destination IP
     for i = 1, 4 do
         pseudo[#pseudo + 1] = string.char(dst_ip[i])
     end
 
     -- Zero + Protocol (6 = TCP)
     pseudo[#pseudo + 1] = string.char(0x00, 0x06)
 
     -- TCP length
     local len = #tcp_packet
     pseudo[#pseudo + 1] = string.char(math.floor(len / 256), len % 256)
 
     -- TCP packet
     pseudo[#pseudo + 1] = tcp_packet
 
     -- Calculate checksum on pseudo-header + TCP packet
     local data = table.concat(pseudo)
     local sum = 0
     local i = 1
 
     while i < #data do
         local word = string.byte(data, i) * 256 + string.byte(data, i + 1)
         sum = sum + word
         i = i + 2
     end
 
     if i == #data then
         sum = sum + string.byte(data, i) * 256
     end
 
     -- Fold to 16 bits
     while sum > 0xFFFF do
         sum = (sum & 0xFFFF) + (sum >> 16)
     end
 
     return ~sum & 0xFFFF
 end
 
 ---Build TCP packet
 ---@param src_port number Source port
 ---@param dst_port number Destination port
 ---@param seq number Sequence number
 ---@param ack number Acknowledgment number
 ---@param flags number TCP flags
 ---@param window number Window size
 ---@param payload string Payload data
 ---@param src_ip table Source IP
 ---@param dst_ip table Destination IP
 ---@return string packet TCP packet
 function TCP.build_packet(src_port, dst_port, seq, ack, flags, window, payload, src_ip, dst_ip)
     local tcp = {}
 
     -- Source port
     tcp[#tcp + 1] = string.char(math.floor(src_port / 256), src_port % 256)
 
     -- Destination port
     tcp[#tcp + 1] = string.char(math.floor(dst_port / 256), dst_port % 256)
 
     -- Sequence number
     tcp[#tcp + 1] = string.char(
         (seq >> 24) & 0xFF,
         (seq >> 16) & 0xFF,
         (seq >> 8) & 0xFF,
         seq & 0xFF
     )
 
     -- Acknowledgment number
     tcp[#tcp + 1] = string.char(
         (ack >> 24) & 0xFF,
         (ack >> 16) & 0xFF,
         (ack >> 8) & 0xFF,
         ack & 0xFF
     )
 
     -- Data offset (5 = 20 bytes, no options) + reserved + flags
     local data_offset = 5 << 4  -- 20 bytes header
     tcp[#tcp + 1] = string.char(data_offset, flags)
 
     -- Window size
     tcp[#tcp + 1] = string.char(math.floor(window / 256), window % 256)
 
     -- Checksum placeholder
     local checksum_pos = #tcp + 1
     tcp[#tcp + 1] = string.char(0x00, 0x00)
 
     -- Urgent pointer
     tcp[#tcp + 1] = string.char(0x00, 0x00)
 
     -- Payload
     tcp[#tcp + 1] = payload
 
     -- Calculate checksum
     local packet = table.concat(tcp)
     local checksum = TCP.calculate_checksum(src_ip, dst_ip, packet)
 
     -- Insert checksum
     tcp[checksum_pos] = string.char(math.floor(checksum / 256), checksum % 256)
 
     return table.concat(tcp)
 end
 
 ---Parse TCP packet
 ---@param payload string TCP payload
 ---@return table|nil tcp Parsed TCP packet
 function TCP.parse_packet(payload)
     if #payload < 20 then
         return nil
     end
 
     local tcp = {}
 
     -- Source port
     tcp.src_port = string.byte(payload, 1) * 256 + string.byte(payload, 2)
 
     -- Destination port
     tcp.dst_port = string.byte(payload, 3) * 256 + string.byte(payload, 4)
 
     -- Sequence number
     tcp.seq = (string.byte(payload, 5) << 24) |
               (string.byte(payload, 6) << 16) |
               (string.byte(payload, 7) << 8) |
               string.byte(payload, 8)
 
     -- Acknowledgment number
     tcp.ack = (string.byte(payload, 9) << 24) |
               (string.byte(payload, 10) << 16) |
               (string.byte(payload, 11) << 8) |
               string.byte(payload, 12)
 
     -- Data offset and flags
     local data_offset_flags = string.byte(payload, 13)
     tcp.data_offset = (data_offset_flags >> 4) * 4  -- Convert to bytes
     tcp.flags = string.byte(payload, 14)
 
     -- Window size
     tcp.window = string.byte(payload, 15) * 256 + string.byte(payload, 16)
 
     -- Checksum
     tcp.checksum = string.byte(payload, 17) * 256 + string.byte(payload, 18)
 
     -- Payload (after header)
     tcp.payload = string.sub(payload, tcp.data_offset + 1)
 
     return tcp
 end
 
 ---Create a new TCP connection
 ---@param NetworkStack table Network stack instance
 ---@param dst_ip table Destination IP
 ---@param dst_port number Destination port
 ---@param src_port number|nil Source port (auto-assigned if nil)
 ---@return table|nil connection Connection object
 function TCP.connect(NetworkStack, dst_ip, dst_port, src_port)
     src_port = src_port or TCP.next_port
     TCP.next_port = TCP.next_port + 1
     if TCP.next_port > 65535 then
         TCP.next_port = 49152
     end
 
     local conn_id = src_port .. "_" .. table.concat(dst_ip, ".") .. "_" .. dst_port
 
     local connection = {
         id = conn_id,
         src_port = src_port,
         dst_ip = dst_ip,
         dst_port = dst_port,
         state = TCP.STATE.SYN_SENT,
         seq = math.random(1, 0x7FFFFFFF),  -- Initial sequence number
         ack = 0,
         window = 8192,
         receive_buffer = {},
         send_buffer = {},
         callbacks = {},
         NetworkStack = NetworkStack,
     }
 
     TCP.connections[conn_id] = connection
 
     -- Send SYN
     local syn_packet = TCP.build_packet(
         src_port, dst_port,
         connection.seq, 0,
         TCP.FLAGS.SYN,
         connection.window,
         "",
         NetworkStack.config.ip, dst_ip
     )
 
     local ip_packet = NetworkStack.build_ipv4(6, NetworkStack.config.ip, dst_ip, syn_packet)
 
     -- Try to send (may need ARP resolution)
     local dst_mac = NetworkStack.arp_resolve(NetworkStack.RTL8139, dst_ip, 5)
     if dst_mac then
         local our_mac = NetworkStack.RTL8139.getMACAddress()
         local frame = NetworkStack.RTL8139.buildEthernetFrame(dst_mac, our_mac, 0x0800, ip_packet)
         NetworkStack.RTL8139.send(frame)
         connection.seq = connection.seq + 1  -- SYN consumes one sequence number
     end
 
     return connection
 end
 
 ---Send data on TCP connection
 ---@param connection table Connection object
 ---@param data string Data to send
 ---@return boolean success True if queued successfully
 function TCP.send(connection, data)
     if connection.state ~= TCP.STATE.ESTABLISHED then
         return false
     end
 
     local tcp_packet = TCP.build_packet(
         connection.src_port, connection.dst_port,
         connection.seq, connection.ack,
         TCP.FLAGS.PSH | TCP.FLAGS.ACK,
         connection.window,
         data,
         connection.NetworkStack.config.ip, connection.dst_ip
     )
 
     local ip_packet = connection.NetworkStack.build_ipv4(
         6, connection.NetworkStack.config.ip, connection.dst_ip, tcp_packet
     )
 
     local dst_mac = connection.NetworkStack.arp_resolve(
         connection.NetworkStack.RTL8139, connection.dst_ip, 5
     )
     if dst_mac then
         local our_mac = connection.NetworkStack.RTL8139.getMACAddress()
         local frame = connection.NetworkStack.RTL8139.buildEthernetFrame(
             dst_mac, our_mac, 0x0800, ip_packet
         )
         connection.NetworkStack.RTL8139.send(frame)
         connection.seq = connection.seq + #data
         return true
     end
 
     return false
 end
 
 ---Close TCP connection
 ---@param connection table Connection object
 function TCP.close(connection)
     if connection.state == TCP.STATE.ESTABLISHED then
         local fin_packet = TCP.build_packet(
             connection.src_port, connection.dst_port,
             connection.seq, connection.ack,
             TCP.FLAGS.FIN | TCP.FLAGS.ACK,
             connection.window,
             "",
             connection.NetworkStack.config.ip, connection.dst_ip
         )
 
         local ip_packet = connection.NetworkStack.build_ipv4(
             6, connection.NetworkStack.config.ip, connection.dst_ip, fin_packet
         )
 
         local dst_mac = connection.NetworkStack.arp_resolve(
             connection.NetworkStack.RTL8139, connection.dst_ip, 5
         )
         if dst_mac then
             local our_mac = connection.NetworkStack.RTL8139.getMACAddress()
             local frame = connection.NetworkStack.RTL8139.buildEthernetFrame(
                 dst_mac, our_mac, 0x0800, ip_packet
             )
             connection.NetworkStack.RTL8139.send(frame)
             connection.seq = connection.seq + 1  -- FIN consumes one sequence number
             connection.state = TCP.STATE.FIN_WAIT_1
         end
     end
 end
 
 ---Handle incoming TCP packet
 ---@param NetworkStack table Network stack instance
 ---@param src_ip table Source IP
 ---@param dst_ip table Destination IP
 ---@param tcp table Parsed TCP packet
 function TCP.handle_packet(NetworkStack, src_ip, dst_ip, tcp)
     -- Find matching connection
     local conn_id = tcp.dst_port .. "_" .. table.concat(src_ip, ".") .. "_" .. tcp.src_port
     local connection = TCP.connections[conn_id]
 
     if not connection then
         -- No connection found - send RST
         return
     end
 
     -- Handle based on state
     if connection.state == TCP.STATE.SYN_SENT then
         if (tcp.flags & TCP.FLAGS.SYN) ~= 0 and (tcp.flags & TCP.FLAGS.ACK) ~= 0 then
             -- SYN-ACK received
             connection.ack = tcp.seq + 1
             connection.state = TCP.STATE.ESTABLISHED
 
             -- Send ACK
             local ack_packet = TCP.build_packet(
                 connection.src_port, connection.dst_port,
                 connection.seq, connection.ack,
                 TCP.FLAGS.ACK,
                 connection.window,
                 "",
                 NetworkStack.config.ip, connection.dst_ip
             )
 
             local ip_packet = NetworkStack.build_ipv4(6, NetworkStack.config.ip, connection.dst_ip, ack_packet)
             local dst_mac = NetworkStack.arp_resolve(NetworkStack.RTL8139, connection.dst_ip, 5)
             if dst_mac then
                 local our_mac = NetworkStack.RTL8139.getMACAddress()
                 local frame = NetworkStack.RTL8139.buildEthernetFrame(dst_mac, our_mac, 0x0800, ip_packet)
                 NetworkStack.RTL8139.send(frame)
             end
 
             -- Call connected callback
             if connection.callbacks.on_connected then
                 connection.callbacks.on_connected()
             end
         end
     elseif connection.state == TCP.STATE.ESTABLISHED then
         if #tcp.payload > 0 then
             -- Data received
             connection.ack = tcp.seq + #tcp.payload
 
             -- Store in receive buffer
             table.insert(connection.receive_buffer, tcp.payload)
 
             -- Send ACK
             local ack_packet = TCP.build_packet(
                 connection.src_port, connection.dst_port,
                 connection.seq, connection.ack,
                 TCP.FLAGS.ACK,
                 connection.window,
                 "",
                 NetworkStack.config.ip, connection.dst_ip
             )
 
             local ip_packet = NetworkStack.build_ipv4(6, NetworkStack.config.ip, connection.dst_ip, ack_packet)
             local dst_mac = NetworkStack.arp_resolve(NetworkStack.RTL8139, connection.dst_ip, 5)
             if dst_mac then
                 local our_mac = NetworkStack.RTL8139.getMACAddress()
                 local frame = NetworkStack.RTL8139.buildEthernetFrame(dst_mac, our_mac, 0x0800, ip_packet)
                 NetworkStack.RTL8139.send(frame)
             end
 
             -- Call data callback
             if connection.callbacks.on_data then
                 connection.callbacks.on_data(tcp.payload)
             end
         end
 
         if (tcp.flags & TCP.FLAGS.FIN) ~= 0 then
             -- FIN received
             connection.ack = tcp.seq + 1
             connection.state = TCP.STATE.CLOSE_WAIT
 
             -- Send ACK
             local ack_packet = TCP.build_packet(
                 connection.src_port, connection.dst_port,
                 connection.seq, connection.ack,
                 TCP.FLAGS.ACK,
                 connection.window,
                 "",
                 NetworkStack.config.ip, connection.dst_ip
             )
 
             local ip_packet = NetworkStack.build_ipv4(6, NetworkStack.config.ip, connection.dst_ip, ack_packet)
             local dst_mac = NetworkStack.arp_resolve(NetworkStack.RTL8139, connection.dst_ip, 5)
             if dst_mac then
                 local our_mac = NetworkStack.RTL8139.getMACAddress()
                 local frame = NetworkStack.RTL8139.buildEthernetFrame(dst_mac, our_mac, 0x0800, ip_packet)
                 NetworkStack.RTL8139.send(frame)
             end
 
             -- Call closed callback
             if connection.callbacks.on_closed then
                 connection.callbacks.on_closed()
             end
         end
     end
 end
 
 ---Read data from connection
 ---@param connection table Connection object
 ---@return string|nil data Received data
 function TCP.read(connection)
     if #connection.receive_buffer > 0 then
         local data = table.concat(connection.receive_buffer)
         connection.receive_buffer = {}
         return data
     end
     return nil
 end
 
 return TCP