luajitos

NetworkStack.lua (19556B)

---

 #!/usr/bin/env luajit
 -- Network Stack Library
 -- Implements ARP, IPv4, ICMP protocols on top of RTL8139
 
 local NetworkStack = {}
 
 -- IP configuration
 NetworkStack.config = {
     ip = {10, 0, 2, 15},      -- Our IP (QEMU user network default guest IP)
     netmask = {255, 255, 255, 0},
     gateway = {10, 0, 2, 2},   -- QEMU user network gateway
     dns = {10, 0, 2, 3}        -- QEMU user network DNS
 }
 
 -- ARP cache
 NetworkStack.arp_cache = {}
 
 -- Pending ICMP requests
 NetworkStack.icmp_requests = {}
 
 -- TCP instance (loaded on demand)
 NetworkStack.TCP = nil
 
 ---Convert IP address array to string
 ---@param ip table Array of 4 bytes
 ---@return string ip "A.B.C.D"
 function NetworkStack.ip_to_string(ip)
     return string.format("%d.%d.%d.%d", ip[1], ip[2], ip[3], ip[4])
 end
 
 ---Parse IP address string
 ---@param str string "A.B.C.D" format
 ---@return table|nil ip Array of 4 bytes
 function NetworkStack.parse_ip(str)
     local parts = {}
     for part in string.gmatch(str, "[^%.]+") do
         local num = tonumber(part)
         if not num or num < 0 or num > 255 then
             return nil
         end
         parts[#parts + 1] = num
     end
     if #parts ~= 4 then
         return nil
     end
     return parts
 end
 
 ---Calculate IP checksum
 ---@param data string Data to checksum
 ---@return number checksum 16-bit checksum
 function NetworkStack.ip_checksum(data)
     local sum = 0
     local i = 1
 
     -- Sum 16-bit words
     while i < #data do
         local word = string.byte(data, i) * 256 + string.byte(data, i + 1)
         sum = sum + word
         i = i + 2
     end
 
     -- Add remaining byte if odd length
     if i == #data then
         sum = sum + string.byte(data, i) * 256
     end
 
     -- Fold 32-bit sum to 16 bits
     while sum > 0xFFFF do
         sum = (sum & 0xFFFF) + (sum >> 16)
     end
 
     return ~sum & 0xFFFF
 end
 
 ---Build ARP request packet
 ---@param target_ip table Target IP address (4 bytes)
 ---@param our_mac table Our MAC address (6 bytes)
 ---@param our_ip table Our IP address (4 bytes)
 ---@return string packet ARP request packet
 function NetworkStack.build_arp_request(target_ip, our_mac, our_ip)
     local arp = {}
 
     -- Hardware type (Ethernet = 1)
     arp[#arp + 1] = string.char(0x00, 0x01)
 
     -- Protocol type (IPv4 = 0x0800)
     arp[#arp + 1] = string.char(0x08, 0x00)
 
     -- Hardware size (6 bytes for MAC)
     arp[#arp + 1] = string.char(0x06)
 
     -- Protocol size (4 bytes for IPv4)
     arp[#arp + 1] = string.char(0x04)
 
     -- Opcode (Request = 1)
     arp[#arp + 1] = string.char(0x00, 0x01)
 
     -- Sender MAC
     for i = 1, 6 do
         arp[#arp + 1] = string.char(our_mac[i])
     end
 
     -- Sender IP
     for i = 1, 4 do
         arp[#arp + 1] = string.char(our_ip[i])
     end
 
     -- Target MAC (unknown, all zeros)
     arp[#arp + 1] = string.char(0x00, 0x00, 0x00, 0x00, 0x00, 0x00)
 
     -- Target IP
     for i = 1, 4 do
         arp[#arp + 1] = string.char(target_ip[i])
     end
 
     return table.concat(arp)
 end
 
 ---Parse ARP packet
 ---@param payload string ARP payload
 ---@return table|nil arp Parsed ARP packet
 function NetworkStack.parse_arp(payload)
     if #payload < 28 then
         return nil
     end
 
     local arp = {}
 
     -- Hardware type
     arp.hw_type = string.byte(payload, 1) * 256 + string.byte(payload, 2)
 
     -- Protocol type
     arp.proto_type = string.byte(payload, 3) * 256 + string.byte(payload, 4)
 
     -- Opcode
     arp.opcode = string.byte(payload, 7) * 256 + string.byte(payload, 8)
 
     -- Sender MAC
     arp.sender_mac = {}
     for i = 1, 6 do
         arp.sender_mac[i] = string.byte(payload, 8 + i)
     end
 
     -- Sender IP
     arp.sender_ip = {}
     for i = 1, 4 do
         arp.sender_ip[i] = string.byte(payload, 14 + i)
     end
 
     -- Target MAC
     arp.target_mac = {}
     for i = 1, 6 do
         arp.target_mac[i] = string.byte(payload, 18 + i)
     end
 
     -- Target IP
     arp.target_ip = {}
     for i = 1, 4 do
         arp.target_ip[i] = string.byte(payload, 24 + i)
     end
 
     return arp
 end
 
 ---Build IPv4 packet
 ---@param protocol number Protocol (1=ICMP, 6=TCP, 17=UDP)
 ---@param src_ip table Source IP (4 bytes)
 ---@param dst_ip table Destination IP (4 bytes)
 ---@param payload string IP payload
 ---@return string packet IPv4 packet
 function NetworkStack.build_ipv4(protocol, src_ip, dst_ip, payload)
     local ip = {}
 
     -- Version (4) and IHL (5 = 20 bytes)
     ip[#ip + 1] = string.char(0x45)
 
     -- DSCP and ECN
     ip[#ip + 1] = string.char(0x00)
 
     -- Total length
     local total_len = 20 + #payload
     ip[#ip + 1] = string.char(math.floor(total_len / 256), total_len % 256)
 
     -- Identification
     ip[#ip + 1] = string.char(0x00, 0x01)
 
     -- Flags and fragment offset
     ip[#ip + 1] = string.char(0x00, 0x00)
 
     -- TTL
     ip[#ip + 1] = string.char(64)
 
     -- Protocol
     ip[#ip + 1] = string.char(protocol)
 
     -- Checksum (will calculate later)
     local checksum_pos = #ip + 1
     ip[#ip + 1] = string.char(0x00, 0x00)
 
     -- Source IP
     for i = 1, 4 do
         ip[#ip + 1] = string.char(src_ip[i])
     end
 
     -- Destination IP
     for i = 1, 4 do
         ip[#ip + 1] = string.char(dst_ip[i])
     end
 
     -- Calculate checksum
     local header = table.concat(ip)
     local checksum = NetworkStack.ip_checksum(header)
 
     -- Insert checksum
     ip[checksum_pos] = string.char(math.floor(checksum / 256), checksum % 256)
 
     -- Rebuild header with correct checksum and add payload
     return table.concat(ip) .. payload
 end
 
 ---Parse IPv4 packet
 ---@param payload string IPv4 payload
 ---@return table|nil ip Parsed IPv4 packet
 function NetworkStack.parse_ipv4(payload)
     if #payload < 20 then
         return nil
     end
 
     local ip = {}
 
     -- Version and IHL
     local ver_ihl = string.byte(payload, 1)
     ip.version = ver_ihl >> 4
     ip.ihl = ver_ihl & 0x0F
 
     if ip.version ~= 4 then
         return nil
     end
 
     -- Header length in bytes
     local header_len = ip.ihl * 4
 
     -- Protocol
     ip.protocol = string.byte(payload, 10)
 
     -- Source IP
     ip.src_ip = {}
     for i = 1, 4 do
         ip.src_ip[i] = string.byte(payload, 12 + i)
     end
 
     -- Destination IP
     ip.dst_ip = {}
     for i = 1, 4 do
         ip.dst_ip[i] = string.byte(payload, 16 + i)
     end
 
     -- Payload
     ip.payload = string.sub(payload, header_len + 1)
 
     return ip
 end
 
 ---Build ICMP Echo Request
 ---@param id number Identifier
 ---@param seq number Sequence number
 ---@param data string Payload data
 ---@return string packet ICMP packet
 function NetworkStack.build_icmp_echo(id, seq, data)
     local icmp = {}
 
     -- Type (8 = Echo Request)
     icmp[#icmp + 1] = string.char(0x08)
 
     -- Code (0)
     icmp[#icmp + 1] = string.char(0x00)
 
     -- Checksum (will calculate later)
     local checksum_pos = #icmp + 1
     icmp[#icmp + 1] = string.char(0x00, 0x00)
 
     -- Identifier
     icmp[#icmp + 1] = string.char(math.floor(id / 256), id % 256)
 
     -- Sequence number
     icmp[#icmp + 1] = string.char(math.floor(seq / 256), seq % 256)
 
     -- Data
     icmp[#icmp + 1] = data
 
     -- Calculate checksum
     local packet = table.concat(icmp)
     local checksum = NetworkStack.ip_checksum(packet)
 
     -- Insert checksum
     icmp[checksum_pos] = string.char(math.floor(checksum / 256), checksum % 256)
 
     return table.concat(icmp)
 end
 
 ---Parse ICMP packet
 ---@param payload string ICMP payload
 ---@return table|nil icmp Parsed ICMP packet
 function NetworkStack.parse_icmp(payload)
     if #payload < 8 then
         return nil
     end
 
     local icmp = {}
 
     -- Type
     icmp.type = string.byte(payload, 1)
 
     -- Code
     icmp.code = string.byte(payload, 2)
 
     -- Identifier
     icmp.id = string.byte(payload, 5) * 256 + string.byte(payload, 6)
 
     -- Sequence
     icmp.seq = string.byte(payload, 7) * 256 + string.byte(payload, 8)
 
     -- Data
     icmp.data = string.sub(payload, 9)
 
     return icmp
 end
 
 ---Send ARP request and wait for reply
 ---@param RTL8139 table RTL8139 driver
 ---@param target_ip table Target IP address
 ---@param timeout number Timeout in seconds
 ---@return table|nil mac Target MAC address
 function NetworkStack.arp_resolve(RTL8139, target_ip, timeout)
     local ip_str = NetworkStack.ip_to_string(target_ip)
 
     -- Check cache first
     if NetworkStack.arp_cache[ip_str] then
         return NetworkStack.arp_cache[ip_str]
     end
 
     local our_mac = RTL8139.getMACAddress()
     local arp_request = NetworkStack.build_arp_request(target_ip, our_mac, NetworkStack.config.ip)
 
     -- Send ARP request (broadcast)
     RTL8139.sendBroadcast(0x0806, arp_request)
 
     -- Wait for reply
     local start_time = os.time or function() return 0 end
     local start = start_time()
 
     while (start_time() - start) < timeout do
         local packet = RTL8139.receive()
         if packet then
             local eth = RTL8139.parseEthernetHeader(packet)
             if eth and eth.ethertype == 0x0806 then
                 local arp = NetworkStack.parse_arp(eth.payload)
                 if arp and arp.opcode == 2 then -- ARP Reply
                     -- Cache the result
                     local sender_ip_str = NetworkStack.ip_to_string(arp.sender_ip)
                     NetworkStack.arp_cache[sender_ip_str] = arp.sender_mac
 
                     -- Check if this is the reply we're waiting for
                     if sender_ip_str == ip_str then
                         return arp.sender_mac
                     end
                 end
             end
         end
     end
 
     return nil
 end
 
 ---Build UDP packet
 ---@param src_port number Source port
 ---@param dst_port number Destination port
 ---@param payload string UDP payload
 ---@return string packet UDP packet
 function NetworkStack.build_udp(src_port, dst_port, payload)
     local udp = {}
 
     -- Source port
     udp[#udp + 1] = string.char(math.floor(src_port / 256), src_port % 256)
 
     -- Destination port
     udp[#udp + 1] = string.char(math.floor(dst_port / 256), dst_port % 256)
 
     -- Length (header + payload)
     local length = 8 + #payload
     udp[#udp + 1] = string.char(math.floor(length / 256), length % 256)
 
     -- Checksum (optional for IPv4, set to 0)
     udp[#udp + 1] = string.char(0x00, 0x00)
 
     -- Payload
     udp[#udp + 1] = payload
 
     return table.concat(udp)
 end
 
 ---Parse UDP packet
 ---@param payload string UDP payload
 ---@return table|nil udp Parsed UDP packet
 function NetworkStack.parse_udp(payload)
     if #payload < 8 then
         return nil
     end
 
     local udp = {}
 
     -- Source port
     udp.src_port = string.byte(payload, 1) * 256 + string.byte(payload, 2)
 
     -- Destination port
     udp.dst_port = string.byte(payload, 3) * 256 + string.byte(payload, 4)
 
     -- Length
     udp.length = string.byte(payload, 5) * 256 + string.byte(payload, 6)
 
     -- Checksum
     udp.checksum = string.byte(payload, 7) * 256 + string.byte(payload, 8)
 
     -- Payload
     udp.payload = string.sub(payload, 9)
 
     return udp
 end
 
 ---Send UDP packet
 ---@param dst_ip table Destination IP address
 ---@param src_port number Source port
 ---@param dst_port number Destination port
 ---@param payload string Payload data
 ---@return boolean success True if sent successfully
 function NetworkStack.send_udp(dst_ip, src_port, dst_port, payload)
     if not NetworkStack.RTL8139 then
         return false
     end
 
     -- Build UDP packet
     local udp_packet = NetworkStack.build_udp(src_port, dst_port, payload)
 
     -- Build IP packet (protocol 17 = UDP)
     local ip_packet = NetworkStack.build_ipv4(17, NetworkStack.config.ip, dst_ip, udp_packet)
 
     -- Resolve MAC address
     local dst_mac = NetworkStack.arp_resolve(NetworkStack.RTL8139, dst_ip, 5)
     if not dst_mac then
         return false
     end
 
     -- Build and send Ethernet frame
     local our_mac = NetworkStack.RTL8139.getMACAddress()
     local frame = NetworkStack.RTL8139.buildEthernetFrame(dst_mac, our_mac, 0x0800, ip_packet)
 
     local result = NetworkStack.RTL8139.send(frame)
     return result > 0
 end
 
 -- UDP socket tracking
 NetworkStack.udp_sockets = {}
 NetworkStack.next_ephemeral_port = 49152  -- Start of ephemeral port range
 
 ---Bind a UDP socket to a port
 ---@param port number Port to bind to
 ---@param callback function Callback(src_ip, src_port, data) when packet received
 ---@return number|nil socket_id Socket ID, or nil on error
 function NetworkStack.udp_bind(port, callback)
     if NetworkStack.udp_sockets[port] then
         return nil  -- Port already in use
     end
 
     NetworkStack.udp_sockets[port] = {
         port = port,
         callback = callback
     }
 
     return port
 end
 
 ---Unbind a UDP socket
 ---@param port number Port to unbind
 function NetworkStack.udp_unbind(port)
     NetworkStack.udp_sockets[port] = nil
 end
 
 ---Get next available ephemeral port
 ---@return number port Next ephemeral port
 function NetworkStack.get_ephemeral_port()
     local port = NetworkStack.next_ephemeral_port
     NetworkStack.next_ephemeral_port = NetworkStack.next_ephemeral_port + 1
     if NetworkStack.next_ephemeral_port > 65535 then
         NetworkStack.next_ephemeral_port = 49152
     end
     return port
 end
 
 ---Initialize network stack with RTL8139 driver
 ---@param RTL8139 table RTL8139 driver instance
 function NetworkStack.init(RTL8139)
     NetworkStack.RTL8139 = RTL8139
 
     -- Register packet handler for IP packets
     RTL8139.onReceive(function(packet, eth)
         if eth.ethertype == 0x0800 then
             -- IPv4 packet
             local ip = NetworkStack.parse_ipv4(eth.payload)
             if not ip then
                 return
             end
 
             -- Check if packet is for us
             local is_for_us = true
             for i = 1, 4 do
                 if ip.dst_ip[i] ~= NetworkStack.config.ip[i] then
                     is_for_us = false
                     break
                 end
             end
 
             if not is_for_us then
                 return  -- Not for us
             end
 
             if ip.protocol == 1 then
                 -- ICMP packet
                 local icmp = NetworkStack.parse_icmp(ip.payload)
                 if icmp then
                     if icmp.type == 0 then
                         -- Echo Reply
                         local key = icmp.id .. "_" .. icmp.seq
                         if NetworkStack.icmp_requests[key] then
                             NetworkStack.icmp_requests[key].reply = {
                                 src_ip = ip.src_ip,
                                 icmp = icmp,
                                 time = os.time and os.time() or 0
                             }
                         end
                     elseif icmp.type == 8 then
                         -- Echo Request - respond to ping
                         NetworkStack.handle_ping_request(ip.src_ip, icmp)
                     end
                 end
             elseif ip.protocol == 6 then
                 -- TCP packet
                 if NetworkStack.TCP then
                     local tcp = NetworkStack.TCP.parse_packet(ip.payload)
                     if tcp then
                         NetworkStack.TCP.handle_packet(NetworkStack, ip.src_ip, ip.dst_ip, tcp)
                     end
                 end
             elseif ip.protocol == 17 then
                 -- UDP packet
                 local udp = NetworkStack.parse_udp(ip.payload)
                 if udp then
                     -- Check if we have a socket bound to this port
                     local socket = NetworkStack.udp_sockets[udp.dst_port]
                     if socket and socket.callback then
                         socket.callback(ip.src_ip, udp.src_port, udp.payload)
                     end
                 end
             end
         end
     end, {ethertype = 0x0800})
 end
 
 ---Handle ping request (respond to ICMP echo)
 ---@param src_ip table Source IP address
 ---@param icmp table Parsed ICMP packet
 function NetworkStack.handle_ping_request(src_ip, icmp)
     if not NetworkStack.RTL8139 then
         return
     end
 
     -- Build ICMP Echo Reply
     local reply = {}
 
     -- Type (0 = Echo Reply)
     reply[#reply + 1] = string.char(0x00)
 
     -- Code (0)
     reply[#reply + 1] = string.char(0x00)
 
     -- Checksum (will calculate later)
     local checksum_pos = #reply + 1
     reply[#reply + 1] = string.char(0x00, 0x00)
 
     -- Identifier
     reply[#reply + 1] = string.char(math.floor(icmp.id / 256), icmp.id % 256)
 
     -- Sequence number
     reply[#reply + 1] = string.char(math.floor(icmp.seq / 256), icmp.seq % 256)
 
     -- Data (echo back the same data)
     reply[#reply + 1] = icmp.data
 
     -- Calculate checksum
     local packet = table.concat(reply)
     local checksum = NetworkStack.ip_checksum(packet)
 
     -- Insert checksum
     reply[checksum_pos] = string.char(math.floor(checksum / 256), checksum % 256)
 
     -- Build IP packet
     local icmp_packet = table.concat(reply)
     local ip_packet = NetworkStack.build_ipv4(1, NetworkStack.config.ip, src_ip, icmp_packet)
 
     -- Resolve MAC address
     local dst_mac = NetworkStack.arp_resolve(NetworkStack.RTL8139, src_ip, 5)
     if not dst_mac then
         return
     end
 
     -- Send response
     local our_mac = NetworkStack.RTL8139.getMACAddress()
     local frame = NetworkStack.RTL8139.buildEthernetFrame(dst_mac, our_mac, 0x0800, ip_packet)
     NetworkStack.RTL8139.send(frame)
 end
 
 ---Load TCP module
 ---@return boolean success True if loaded successfully
 function NetworkStack.load_tcp()
     if NetworkStack.TCP then
         return true  -- Already loaded
     end
 
     -- Try to load TCP module
     if CRamdiskOpen then
         local tcp_handle = CRamdiskOpen("/os/libs/TCP.lua", "r")
         if tcp_handle then
             local tcp_code = CRamdiskRead(tcp_handle)
             CRamdiskClose(tcp_handle)
 
             if tcp_code then
                 local tcp_func, err = load(tcp_code, "/os/libs/TCP.lua", "t")
                 if tcp_func then
                     NetworkStack.TCP = tcp_func()
                     return true
                 end
             end
         end
     end
 
     return false
 end
 
 ---Send ping (ICMP Echo Request)
 ---@param dst_ip table Destination IP address
 ---@param timeout number Timeout in seconds
 ---@return table|nil reply Reply data or nil if timeout
 function NetworkStack.ping(dst_ip, timeout)
     if not NetworkStack.RTL8139 then
         return nil
     end
 
     local id = math.random(1, 65535)
     local seq = 1
     local data = "LuajitOS ping"
 
     -- Build ICMP echo request
     local icmp_packet = NetworkStack.build_icmp_echo(id, seq, data)
 
     -- Build IP packet
     local ip_packet = NetworkStack.build_ipv4(1, NetworkStack.config.ip, dst_ip, icmp_packet)
 
     -- Resolve MAC address
     local dst_mac = NetworkStack.arp_resolve(NetworkStack.RTL8139, dst_ip, 5)
     if not dst_mac then
         return nil
     end
 
     -- Register request
     local key = id .. "_" .. seq
     NetworkStack.icmp_requests[key] = {
         sent_time = os.time and os.time() or 0,
         reply = nil
     }
 
     -- Send packet
     local our_mac = NetworkStack.RTL8139.getMACAddress()
     local frame = NetworkStack.RTL8139.buildEthernetFrame(dst_mac, our_mac, 0x0800, ip_packet)
     NetworkStack.RTL8139.send(frame)
 
     -- Wait for reply
     local start_time = os.time and os.time() or 0
     while (os.time and os.time() or 0) - start_time < timeout do
         NetworkStack.RTL8139.poll()
         if NetworkStack.icmp_requests[key].reply then
             local reply = NetworkStack.icmp_requests[key].reply
             NetworkStack.icmp_requests[key] = nil
             return reply
         end
     end
 
     NetworkStack.icmp_requests[key] = nil
     return nil
 end
 
 return NetworkStack