; =============================================================================
; socket.asm - the kernel network stack + socket layer (SYS_NET).
;
; The kernel owns the link port: SLIP framing, IPv4, checksums, and the per-
; protocol logic all live here. Userland just does socket()/connect()/send()/
; recv()/close() via one syscall (DE = &netreq). No program touches SLIP or IP.
;
; Our address is 10.0.0.2. Incoming ICMP echo requests are auto-answered while
; any process pumps RX (recv/poll), so the Game Boy responds to pings for real.
;
; Phase 1: ICMP (echo).  UDP and TCP build on the same frame/checksum core.
; =============================================================================
INCLUDE "include/gbos.inc"

; ---- socket entry layout ----------------------------------------------------
DEF SK_TYPE    EQU 0     ; 0 = free
DEF SK_RIP     EQU 1     ; 4  remote IP
DEF SK_RPORT   EQU 5     ; 2  remote port
DEF SK_LPORT   EQU 7     ; 2  local port
DEF SK_HASRX   EQU 9     ; 1  rxbuf holds a datagram
DEF SK_RXLEN   EQU 10    ; 1  payload length
DEF SK_SRCIP   EQU 11    ; 4  sender IP
DEF SK_SRCPORT EQU 15    ; 2  sender port
DEF SK_RXBUF   EQU 17    ; 288 payload (holds a DHCP reply, or a max segment)
DEF SK_STATE   EQU 305   ; 1  TCP state (see TCP_* below)
DEF SK_SND     EQU 306   ; 4  snd_nxt (our next sequence number)
DEF SK_RCV     EQU 310   ; 4  rcv_nxt (next sequence we expect from peer)
DEF SK_SIZE    EQU 314
DEF TCP_MSS    EQU 200

; TCP connection states
DEF TCP_CLOSED  EQU 0
DEF TCP_SYNSENT EQU 1
DEF TCP_ESTAB   EQU 2
DEF TCP_FINWAIT EQU 3
DEF TCP_DONE    EQU 4
; TCP flags
DEF TF_FIN EQU $01
DEF TF_SYN EQU $02
DEF TF_RST EQU $04
DEF TF_PSH EQU $08
DEF TF_ACK EQU $10

SECTION "net_bss", WRAM0
wSockTab::     DS MAX_SOCKS * SK_SIZE
wNetOurIP::    DS 4
wNetReq::      DS NR_SIZE
wNetSockPtr::  DS 2
wNetIhl::      DS 1
wNetSeq::      DS 1
wNetTx::       DS 320      ; tx build buffer
wNetRxBuf::    DS 320      ; rx frame assembly / parse
wNetRxLen::    DS 2        ; 16-bit: frames can exceed 255 bytes (DHCP/BOOTP)
wNetRxEsc::    DS 1
wNetUdpPort::  DS 2      ; scratch: UDP dst port being demuxed
wNetTO::       DS 3      ; recv timeout counter (net_pump clobbers registers)
wNetTcpFlags:: DS 1      ; TCP flags for the segment being built
wNetTcpDlen::  DS 1      ; TCP payload length for the segment being built
wNetTcpIn::    DS 1      ; incoming TCP flags
wNetTcpDoff::  DS 1      ; incoming data offset (from frame start)
wNetTcpDlen2:: DS 1      ; incoming payload length
wNetTcpTh::    DS 1      ; TCP header length of the segment being built (20/24)

SECTION "socket", ROM0

; net_init - clear the socket table and set our IP (called at boot).
net_init::
    ld hl, wSockTab
    ld bc, MAX_SOCKS * SK_SIZE
.clr
    xor a
    ld [hl+], a
    dec bc
    ld a, b
    or c
    jr nz, .clr
    xor a
    ld [wNetOurIP+0], a         ; 0.0.0.0 until DHCP assigns an address
    ld [wNetOurIP+1], a
    ld [wNetOurIP+2], a
    ld [wNetOurIP+3], a
    ld [wNetRxLen], a
    ld [wNetRxLen+1], a
    ld [wNetRxEsc], a
    ld [wNetSeq], a
    ret

; net_setip - store our IPv4 address (from wNetReq.NR_IP). Used by the DHCP
; client once it has a lease.
net_op_setip:
    ld a, [wNetReq+NR_IP+0]
    ld [wNetOurIP+0], a
    ld a, [wNetReq+NR_IP+1]
    ld [wNetOurIP+1], a
    ld a, [wNetReq+NR_IP+2]
    ld [wNetOurIP+2], a
    ld a, [wNetReq+NR_IP+3]
    ld [wNetOurIP+3], a
    xor a
    ret

; -----------------------------------------------------------------------------
; low-level link port (byte at a time). Preserve BC/DE/HL; use A only.
net_putbyte:                    ; A = byte to transmit
    ld [rSB], a
    ld a, $81
    ld [rSC], a
.w
    ld a, [rSC]
    bit 7, a
    jr nz, .w
    ret

net_getbyte_nb:                 ; -> A = byte, CF set if none
    ld a, $80
    ld [rSC], a
    ld a, [rSC]
    bit 7, a
    jr z, .got
    scf
    ret
.got
    ld a, [rSB]
    and a                       ; CF = 0
    ret

; -----------------------------------------------------------------------------
; net_slip_send(HL = buf, DE = len) - transmit a SLIP frame (len may exceed 255).
net_slip_send:
    ld a, SLIP_END
    call net_putbyte
.loop
    ld a, d
    or e
    jr z, .end
    ld a, [hl+]
    cp SLIP_END
    jr z, .esc_end
    cp SLIP_ESC
    jr z, .esc_esc
    call net_putbyte
.cont
    dec de
    jr .loop
.esc_end
    ld a, SLIP_ESC
    call net_putbyte
    ld a, SLIP_ESC_END
    call net_putbyte
    jr .cont
.esc_esc
    ld a, SLIP_ESC
    call net_putbyte
    ld a, SLIP_ESC_ESC
    call net_putbyte
    jr .cont
.end
    ld a, SLIP_END
    call net_putbyte
    ret

; -----------------------------------------------------------------------------
; net_sum(HL = buf, B = len, DE = seed) -> DE = folded 16-bit one's-complement
; sum (RFC 1071), carry-folded each add. len <= 255. Used to chain a pseudo-
; header into a segment checksum.
net_sum:
.loop
    ld a, b
    cp 2
    jr c, .tail
    ld a, [hl+]                 ; high byte
    ld c, a
    ld a, [hl+]                 ; low byte
    add e
    ld e, a
    ld a, c
    adc d
    ld d, a
    jr nc, .nc1
    inc de
.nc1
    dec b
    dec b
    jr .loop
.tail
    ld a, b
    or a
    ret z
    ld a, [hl]                  ; last odd byte -> high half
    add d
    ld d, a
    ret nc
    inc e
    ret nz
    inc d
    ret

; net_cksum(HL, B, DE=seed) -> HL = inverted 16-bit checksum.
net_cksum:
    call net_sum
    ld a, d
    cpl
    ld h, a
    ld a, e
    cpl
    ld l, a
    ret

; -----------------------------------------------------------------------------
; net_sock_ptr(A = sockid) -> HL = &wSockTab[A]  (clobbers B, DE)
net_sock_ptr:
    ld hl, wSockTab
    or a
    ret z
    ld b, a
    ld de, SK_SIZE
.l
    add hl, de
    dec b
    jr nz, .l
    ret

; =============================================================================
; sys_net(DE = &netreq) -> A = result
; =============================================================================
sys_net::
    ld hl, wNetReq              ; copy the request out of process RAM
    ld b, NR_SIZE
.cp
    ld a, [de]
    ld [hl+], a
    inc de
    dec b
    jr nz, .cp
    ld a, [wNetReq + NR_OP]
    cp NET_SOCKET
    jp z, net_op_socket
    cp NET_CONNECT
    jp z, net_op_connect
    cp NET_SEND
    jp z, net_op_send
    cp NET_RECV
    jp z, net_op_recv
    cp NET_CLOSE
    jp z, net_op_close
    cp NET_BIND
    jp z, net_op_bind
    cp NET_POLL
    jp z, net_op_poll
    cp NET_SETIP
    jp z, net_op_setip
    ld a, $FF
    ret

; ---- bind(sock, port) - set local port -----------------------------------
net_op_bind:
    ld a, [wNetReq + NR_SOCK]
    call net_sock_ptr
    ld a, l
    add SK_LPORT
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [wNetReq + NR_PORT]
    ld [hl+], a
    ld a, [wNetReq + NR_PORT + 1]
    ld [hl], a
    xor a
    ret

; ---- socket(type) -> A = sockid or $FF ---------------------------------------
net_op_socket:
    ld c, 0
.scan
    ld a, c
    call net_sock_ptr           ; HL = &sock[c]
    ld a, [hl]
    or a                        ; SK_TYPE == 0 (free)?
    jr z, .free
    inc c
    ld a, c
    cp MAX_SOCKS
    jr c, .scan
    ld a, $FF                   ; table full
    ret
.free
    ld a, [wNetReq + NR_SOCK]   ; requested type
    ld [hl], a                  ; SK_TYPE
    ld a, l                     ; clear SK_HASRX
    add SK_HASRX
    ld l, a
    ld a, h
    adc 0
    ld h, a
    xor a
    ld [hl], a
    ld a, c                     ; return sockid
    ret

; ---- connect(sock, ip, port) - set remote (+ TCP handshake) ---------------
net_op_connect:
    ld a, [wNetReq + NR_SOCK]
    call net_sock_ptr           ; HL = &sock
    ld a, l
    ld [wNetSockPtr], a
    ld a, h
    ld [wNetSockPtr+1], a
    ld a, l
    add SK_RIP
    ld l, a
    ld a, h
    adc 0
    ld h, a                     ; HL = &sock.RIP
    ld de, wNetReq + NR_IP
    ld b, 6                     ; 4 IP + 2 port (contiguous: RIP then RPORT)
.cp
    ld a, [de]
    ld [hl+], a
    inc de
    dec b
    jr nz, .cp
    call sock_ptr_hl
    ld a, [hl]                  ; SK_TYPE
    cp SOCK_TCP
    jp z, tcp_connect
    xor a
    ret

; ---- close(sock) -------------------------------------------------------------
net_op_close:
    ld a, [wNetReq + NR_SOCK]
    call net_sock_ptr
    ld a, l
    ld [wNetSockPtr], a
    ld a, h
    ld [wNetSockPtr+1], a
    ld a, [hl]                  ; SK_TYPE
    cp SOCK_TCP
    jp z, tcp_close
    call sock_ptr_hl
    xor a
    ld [hl], a                  ; SK_TYPE = free
    ret

; ---- poll - pump RX once (answers pings) -------------------------------------
net_op_poll:
    call net_pump
    xor a
    ret

; ---- send(sock, buf, len) ----------------------------------------------------
net_op_send:
    ld a, [wNetReq + NR_SOCK]
    call net_sock_ptr
    ld a, l
    ld [wNetSockPtr], a
    ld a, h
    ld [wNetSockPtr+1], a
    ld a, [hl]                  ; SK_TYPE
    cp SOCK_ICMP
    jp z, icmp_send
    cp SOCK_UDP
    jp z, udp_send
    cp SOCK_TCP
    jp z, tcp_data_send
    ld a, $FF
    ret

; ---- recv(sock, buf, maxlen) -> A = len or $FF (timeout) ---------------------
net_op_recv:
    ld a, [wNetReq + NR_SOCK]
    call net_sock_ptr
    ld a, l
    ld [wNetSockPtr], a
    ld a, h
    ld [wNetSockPtr+1], a
    xor a                       ; 24-bit timeout counter in WRAM (net_pump
    ld [wNetTO], a              ; clobbers BC/DE/HL, so we can't count in regs)
    ld [wNetTO+1], a
    ld [wNetTO+2], a
.wait
    call net_pump
    ld a, [wNetReq+NR_SOCK]     ; tcp_in may have moved wNetSockPtr; restore ours
    call net_sock_ptr
    ld a, l
    ld [wNetSockPtr], a
    ld a, h
    ld [wNetSockPtr+1], a
    call sock_has_rx            ; nz if a datagram is buffered
    jr nz, .got
    call sock_ptr_hl            ; TCP peer closed + no data -> EOF (return 0)
    ld a, l
    add SK_STATE
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [hl]
    cp TCP_DONE
    jr z, .eof
    ld hl, wNetTO
    inc [hl]
    jr nz, .wait                ; low byte didn't wrap -> spin (no yield)
    call SchedYield             ; every 256 pumps, yield to the scheduler
    inc hl
    inc [hl]
    jr nz, .wait
    inc hl
    inc [hl]
    ld a, [hl]
    cp 250                      ; timeout after ~250*65536 pumps
    jr c, .wait
    ld a, $FF
    ret
.eof
    xor a                       ; connection closed, no data -> 0 bytes (EOF)
    ret
.got
    ; copy sock.RXBUF -> process NR_BUF, len = sock.RXLEN, NR_IP = sock.SRCIP
    call sock_ptr_hl            ; HL = &sock
    push hl
    ld a, l
    add SK_RXLEN
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [hl]                  ; RXLEN
    ld c, a                     ; C = length
    pop hl
    ; source IP -> wNetReq.NR_IP (returned to caller via buffer? no - copy to caller req is skipped; give src via NR_IP in kernel copy only)
    push bc
    push hl
    ; dst = process buf pointer
    ld a, [wNetReq + NR_BUF]
    ld e, a
    ld a, [wNetReq + NR_BUF + 1]
    ld d, a                     ; DE = process buf
    ld a, l
    add SK_RXBUF
    ld l, a
    ld a, h
    adc 0
    ld h, a                     ; HL = &sock.RXBUF
    ld a, c
    or a
    jr z, .nocopy
.cpl
    ld a, [hl+]
    ld [de], a
    inc de
    dec c
    jr nz, .cpl
.nocopy
    pop hl                      ; &sock
    pop bc                      ; C = length
    ; clear HASRX
    ld a, l
    add SK_HASRX
    ld l, a
    ld a, h
    adc 0
    ld h, a
    xor a
    ld [hl], a
    ld a, c                     ; return length
    ret

; sock_ptr_hl -> HL = current socket pointer (from wNetSockPtr)
sock_ptr_hl:
    ld a, [wNetSockPtr]
    ld l, a
    ld a, [wNetSockPtr+1]
    ld h, a
    ret

; sock_has_rx -> Z clear (nz) if the current socket has a buffered datagram
sock_has_rx:
    call sock_ptr_hl
    ld a, l
    add SK_HASRX
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [hl]
    or a
    ret

; =============================================================================
; icmp_send - build + transmit an ICMP echo request from the current socket.
; =============================================================================
icmp_send:
    ; IP header (20 bytes)
    ld a, $45
    ld [wNetTx+0], a
    xor a
    ld [wNetTx+1], a
    ld [wNetTx+2], a            ; total len hi
    ld a, [wNetReq+NR_LEN]
    add 28
    ld [wNetTx+3], a            ; total len lo = 28 + payload
    xor a
    ld [wNetTx+4], a
    ld [wNetTx+5], a            ; id
    ld [wNetTx+6], a
    ld [wNetTx+7], a            ; flags/frag
    ld a, 64
    ld [wNetTx+8], a           ; ttl
    ld a, 1
    ld [wNetTx+9], a           ; proto ICMP
    xor a
    ld [wNetTx+10], a
    ld [wNetTx+11], a          ; ip cksum
    ld a, [wNetOurIP+0]
    ld [wNetTx+12], a
    ld a, [wNetOurIP+1]
    ld [wNetTx+13], a
    ld a, [wNetOurIP+2]
    ld [wNetTx+14], a
    ld a, [wNetOurIP+3]
    ld [wNetTx+15], a
    ; dst IP = sock.RIP
    call sock_ptr_hl
    ld a, l
    add SK_RIP
    ld l, a
    ld a, h
    adc 0
    ld h, a                     ; HL = &sock.RIP
    ld de, wNetTx+16
    ld b, 4
.dip
    ld a, [hl+]
    ld [de], a
    inc de
    dec b
    jr nz, .dip
    ; ICMP header (8 bytes) at wNetTx+20
    ld a, 8
    ld [wNetTx+20], a          ; type = echo request
    xor a
    ld [wNetTx+21], a          ; code
    ld [wNetTx+22], a
    ld [wNetTx+23], a          ; cksum
    ld a, $12
    ld [wNetTx+24], a          ; id
    ld a, $34
    ld [wNetTx+25], a
    xor a
    ld [wNetTx+26], a          ; seq hi
    ld a, [wNetSeq]
    inc a
    ld [wNetSeq], a
    ld [wNetTx+27], a          ; seq lo
    ; payload from process buf
    ld a, [wNetReq+NR_BUF]
    ld l, a
    ld a, [wNetReq+NR_BUF+1]
    ld h, a
    ld de, wNetTx+28
    ld a, [wNetReq+NR_LEN]
    ld b, a
    or a
    jr z, .nopl
.cppl
    ld a, [hl+]
    ld [de], a
    inc de
    dec b
    jr nz, .cppl
.nopl
    ; ICMP checksum over 8 + payload
    ld a, [wNetReq+NR_LEN]
    add 8
    ld b, a
    ld hl, wNetTx+20
    ld de, 0
    call net_cksum
    ld a, h
    ld [wNetTx+22], a
    ld a, l
    ld [wNetTx+23], a
    ; IP checksum over 20
    ld hl, wNetTx
    ld b, 20
    ld de, 0
    call net_cksum
    ld a, h
    ld [wNetTx+10], a
    ld a, l
    ld [wNetTx+11], a
    ; transmit: 28 + payload
    ld a, [wNetReq+NR_LEN]
    add 28
    ld e, a
    ld d, 0
    ld hl, wNetTx
    call net_slip_send
    xor a
    ret

; =============================================================================
; udp_send - build + transmit a UDP datagram from the current socket.
; =============================================================================
udp_send:
    ld a, $45
    ld [wNetTx+0], a
    xor a
    ld [wNetTx+1], a
    ld a, [wNetReq+NR_LEN]
    add 28                      ; 20 IP + 8 UDP + payload
    ld [wNetTx+3], a            ; IP total length low
    ld a, 0
    adc 0
    ld [wNetTx+2], a            ; IP total length high (carry, for DHCP)
    xor a
    ld [wNetTx+4], a
    ld [wNetTx+5], a
    ld [wNetTx+6], a
    ld [wNetTx+7], a
    ld a, 64
    ld [wNetTx+8], a
    ld a, 17
    ld [wNetTx+9], a           ; proto UDP
    xor a
    ld [wNetTx+10], a
    ld [wNetTx+11], a
    ld a, [wNetOurIP+0]
    ld [wNetTx+12], a
    ld a, [wNetOurIP+1]
    ld [wNetTx+13], a
    ld a, [wNetOurIP+2]
    ld [wNetTx+14], a
    ld a, [wNetOurIP+3]
    ld [wNetTx+15], a
    ; dst IP = sock.RIP
    call sock_ptr_hl
    ld a, l
    add SK_RIP
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld de, wNetTx+16
    ld b, 4
.dip
    ld a, [hl+]
    ld [de], a
    inc de
    dec b
    jr nz, .dip
    ; UDP header at wNetTx+20: src port = SK_LPORT
    call sock_ptr_hl
    ld a, l
    add SK_LPORT
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [hl+]
    ld [wNetTx+20], a
    ld a, [hl]
    ld [wNetTx+21], a
    ; dst port = SK_RPORT
    call sock_ptr_hl
    ld a, l
    add SK_RPORT
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [hl+]
    ld [wNetTx+22], a
    ld a, [hl]
    ld [wNetTx+23], a
    ; length = 8 + payload
    xor a
    ld [wNetTx+24], a
    ld a, [wNetReq+NR_LEN]
    add 8
    ld [wNetTx+25], a
    xor a
    ld [wNetTx+26], a          ; cksum = 0
    ld [wNetTx+27], a
    ; payload from process buf
    ld a, [wNetReq+NR_BUF]
    ld l, a
    ld a, [wNetReq+NR_BUF+1]
    ld h, a
    ld de, wNetTx+28
    ld a, [wNetReq+NR_LEN]
    ld b, a
    or a
    jr z, .nopl
.cppl
    ld a, [hl+]
    ld [de], a
    inc de
    dec b
    jr nz, .cppl
.nopl
    ; UDP checksum with pseudo-header: sum(src+dst IP) + proto + udplen, then seg
    ld hl, wNetTx+12
    ld b, 8
    ld de, 0
    call net_sum               ; DE = IP-address sum
    ld a, e
    add 17                     ; + protocol
    ld e, a
    ld a, d
    adc 0
    ld d, a
    ld a, [wNetReq+NR_LEN]
    add 8                      ; + udp length
    add e
    ld e, a
    ld a, d
    adc 0
    ld d, a                    ; DE = pseudo-header seed
    ld a, [wNetReq+NR_LEN]
    add 8
    ld b, a
    ld hl, wNetTx+20
    call net_cksum             ; HL = inverted checksum
    ld a, h
    or l
    jr nz, .cknz
    ld hl, $FFFF               ; 0 means "no checksum" -> send all-ones
.cknz
    ld a, h
    ld [wNetTx+26], a
    ld a, l
    ld [wNetTx+27], a
    ; IP checksum
    ld hl, wNetTx
    ld b, 20
    ld de, 0
    call net_cksum
    ld a, h
    ld [wNetTx+10], a
    ld a, l
    ld [wNetTx+11], a
    ld a, [wNetReq+NR_LEN]
    add 28
    ld e, a
    ld a, 0
    adc 0
    ld d, a
    ld hl, wNetTx
    call net_slip_send
    xor a
    ret

; =============================================================================
; net_pump - read available serial bytes, assemble SLIP frames, process each.
; Non-blocking: returns when the link has no more bytes ready.
; =============================================================================
net_pump:
.next
    call net_getbyte_nb
    ret c                       ; no byte ready
    ; SLIP decode using wNetRxLen / wNetRxEsc
    cp SLIP_END
    jr z, .frame
    ld c, a                     ; save byte
    ld a, [wNetRxEsc]
    or a
    jr nz, .escaped
    ld a, c
    cp SLIP_ESC
    jr z, .setesc
    jr .store
.escaped
    ld a, c
    cp SLIP_ESC_END
    jr nz, .notend
    ld c, SLIP_END
.notend
    ld a, c
    cp SLIP_ESC_ESC
    jr nz, .noteesc
    ld c, SLIP_ESC
.noteesc
    xor a
    ld [wNetRxEsc], a
    ld a, c
    jr .store2
.setesc
    ld a, 1
    ld [wNetRxEsc], a
    jr .next
.store
    ld a, c
.store2
    ld c, a                     ; C = byte to store
    ; index = wNetRxLen (16-bit); drop if >= 320 (buffer size)
    ld a, [wNetRxLen+1]
    or a
    jr z, .instore              ; high 0 -> index < 256, room
    cp 1
    jr nz, .next                ; high >= 2 -> overflow, drop
    ld a, [wNetRxLen]
    cp 64                       ; high 1 -> room iff low < 64 (index < 320)
    jr nc, .next
.instore
    ld a, [wNetRxLen]
    ld e, a
    ld a, [wNetRxLen+1]
    ld d, a
    ld hl, wNetRxBuf
    add hl, de
    ld a, c
    ld [hl], a
    ld a, [wNetRxLen]           ; index++
    inc a
    ld [wNetRxLen], a
    jr nz, .next
    ld a, [wNetRxLen+1]
    inc a
    ld [wNetRxLen+1], a
    jr .next
.frame
    ; END: if we have a frame, process it and return (one frame per pump, so a
    ; recv drains each segment before the next is delivered - no rx overwrite).
    xor a
    ld [wNetRxEsc], a
    ld a, [wNetRxLen]
    ld hl, wNetRxLen+1
    or [hl]
    jr z, .next                 ; empty frame, keep reading
    call process_frame
    xor a
    ld [wNetRxLen], a
    ld [wNetRxLen+1], a
    ret

; -----------------------------------------------------------------------------
; process_frame - IPv4 demux of the frame in wNetRxBuf (wNetRxLen bytes).
process_frame:
    ld a, [wNetRxBuf]
    and $F0
    cp $40
    ret nz                      ; not IPv4
    ld a, [wNetRxBuf]
    and $0F
    add a, a
    add a, a
    ld [wNetIhl], a             ; ihl in bytes
    ld a, [wNetRxBuf+9]         ; protocol
    cp 1
    jp z, icmp_in
    cp 17
    jp z, udp_in
    cp 6
    jp z, tcp_in
    ret

; -----------------------------------------------------------------------------
; icmp_in - handle an inbound ICMP message.
icmp_in:
    ld a, [wNetIhl]
    ld e, a
    ld d, 0
    ld hl, wNetRxBuf
    add hl, de                  ; HL = &icmp
    ld a, [hl]                  ; type
    cp 8
    jr z, icmp_request
    or a                        ; type 0 = echo reply
    jr z, icmp_reply
    ret

; incoming echo request -> answer it in place
icmp_request:
    ld a, [wNetIhl]
    ld e, a
    ld d, 0
    ld hl, wNetRxBuf
    add hl, de                  ; HL = &icmp
    xor a
    ld [hl+], a                 ; type = 0
    inc hl                      ; skip code
    xor a
    ld [hl+], a
    ld [hl], a                  ; cksum = 0
    call net_swap_ip
    ; icmp cksum over (rxlen - ihl)
    ld a, [wNetIhl]
    ld e, a
    ld d, 0
    ld hl, wNetRxBuf
    add hl, de                  ; HL = &icmp
    ld a, [wNetRxLen]
    sub e
    ld b, a
    ld de, 0
    push hl
    call net_cksum              ; HL = cksum
    pop de                      ; DE = &icmp
    ld a, e
    add 2
    ld e, a
    ld a, d
    adc 0
    ld d, a                     ; DE = &icmp+2
    ld a, h
    ld [de], a
    inc de
    ld a, l
    ld [de], a
    ; IP cksum
    xor a
    ld [wNetRxBuf+10], a
    ld [wNetRxBuf+11], a
    ld a, [wNetIhl]
    ld b, a
    ld hl, wNetRxBuf
    ld de, 0
    call net_cksum
    ld a, h
    ld [wNetRxBuf+10], a
    ld a, l
    ld [wNetRxBuf+11], a
    ; send it back
    ld a, [wNetRxLen]
    ld e, a
    ld a, [wNetRxLen+1]
    ld d, a
    ld hl, wNetRxBuf
    call net_slip_send
    ret

; incoming echo reply -> deliver to the first ICMP socket
icmp_reply:
    ld c, 0
.scan
    ld a, c
    call net_sock_ptr           ; HL = &sock
    ld a, [hl]
    cp SOCK_ICMP
    jr z, .deliver
    inc c
    ld a, c
    cp MAX_SOCKS
    jr c, .scan
    ret                         ; no ICMP socket
.deliver
    ; HL = &sock. payload = wNetRxBuf + ihl + 8, len = rxlen - ihl - 8
    push hl
    ; set HASRX = 1
    ld a, l
    add SK_HASRX
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, 1
    ld [hl], a
    pop hl
    ; SRCIP = wNetRxBuf[12..15]
    push hl
    ld a, l
    add SK_SRCIP
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld de, wNetRxBuf+12
    ld b, 4
.sip
    ld a, [de]
    ld [hl+], a
    inc de
    dec b
    jr nz, .sip
    pop hl
    ; RXLEN + RXBUF
    ld a, [wNetIhl]
    add 8
    ld e, a
    ld d, 0                     ; DE = ihl+8 (payload offset)
    ld a, [wNetRxLen]
    sub e
    ld c, a                     ; C = payload length
    push hl
    ld a, l
    add SK_RXLEN
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, c
    ld [hl], a                  ; RXLEN
    pop hl
    ; copy payload
    ld a, l
    add SK_RXBUF
    ld l, a
    ld a, h
    adc 0
    ld h, a                     ; HL = &sock.RXBUF (dest)
    push hl
    ld hl, wNetRxBuf
    add hl, de                  ; HL = &payload (src)
    pop de                      ; DE = dest
    ld a, c
    or a
    ret z
.cp
    ld a, [hl+]
    ld [de], a
    inc de
    dec c
    jr nz, .cp
    ret

; -----------------------------------------------------------------------------
; udp_in - deliver an inbound UDP datagram to the socket bound to its dst port.
udp_in:
    ld a, [wNetIhl]
    add 2
    ld e, a
    ld d, 0
    ld hl, wNetRxBuf
    add hl, de                  ; HL = &udp.dstport
    ld a, [hl+]
    ld [wNetUdpPort], a
    ld a, [hl]
    ld [wNetUdpPort+1], a
    call net_find_udp           ; HL = &sock, CF if none
    ret c
    ; deliver: HASRX=1, SRCIP, RXLEN + RXBUF (payload after ihl+8)
    push hl
    ld a, l
    add SK_HASRX
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, 1
    ld [hl], a
    pop hl
    push hl
    ld a, l
    add SK_SRCIP
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld de, wNetRxBuf+12
    ld b, 4
.sip
    ld a, [de]
    ld [hl+], a
    inc de
    dec b
    jr nz, .sip
    pop hl
    ; payload offset = ihl + 8, length = rxlen - ihl - 8
    ld a, [wNetIhl]
    add 8
    ld e, a
    ld d, 0
    ld a, [wNetRxLen]
    sub e
    ld c, a                     ; C = payload length
    push hl
    ld a, l
    add SK_RXLEN
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, c
    ld [hl], a
    pop hl
    ld a, l
    add SK_RXBUF
    ld l, a
    ld a, h
    adc 0
    ld h, a                     ; HL = &sock.RXBUF (dest)
    push hl
    ld hl, wNetRxBuf
    add hl, de                  ; HL = &payload (src)
    pop de                      ; DE = dest
    ld a, c
    or a
    ret z
.cp
    ld a, [hl+]
    ld [de], a
    inc de
    dec c
    jr nz, .cp
    ret

; net_find_udp -> HL = &UDP socket bound to wNetUdpPort, CF set if none
net_find_udp:
    ld c, 0
.l
    ld a, c
    call net_sock_ptr           ; HL = &sock (clobbers B, DE)
    ld a, [hl]
    cp SOCK_UDP
    jr nz, .nx
    push hl
    ld a, l
    add SK_LPORT
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [wNetUdpPort]
    cp [hl]
    jr nz, .nomatch
    inc hl
    ld a, [wNetUdpPort+1]
    cp [hl]
    jr nz, .nomatch
    pop hl                      ; match
    or a                        ; CF = 0
    ret
.nomatch
    pop hl
.nx
    inc c
    ld a, c
    cp MAX_SOCKS
    jr c, .l
    scf
    ret

; net_swap_ip - swap the src/dst IP fields of the frame in wNetRxBuf
net_swap_ip:
    ld hl, wNetRxBuf+12
    ld de, wNetRxBuf+16
    ld b, 4
.l
    ld a, [hl]
    ld c, a
    ld a, [de]
    ld [hl+], a
    ld a, c
    ld [de], a
    inc de
    dec b
    jr nz, .l
    ret

; =============================================================================
; TCP - a minimal client: 3-way handshake, in-order data, FIN close. No
; retransmission (the GB<->host link is lossless; the host's real TCP handles
; the internet side). One rx segment buffered per socket at a time.
; =============================================================================

; sk_field_to(A = socket field offset, DE = dest, B = count) - copy from socket
sk_field_to:
    push bc
    push de
    ld c, a
    call sock_ptr_hl
    ld a, l
    add c
    ld l, a
    ld a, h
    adc 0
    ld h, a
    pop de
    pop bc
.l
    ld a, [hl+]
    ld [de], a
    inc de
    dec b
    jr nz, .l
    ret

; sk_add_snd/sk_add_rcv(A = amount) - add to the 32-bit big-endian seq field
sk_add_snd:
    ld c, SK_SND
    jr sk_add_seq
sk_add_rcv:
    ld c, SK_RCV
sk_add_seq:
    ld b, a
    call sock_ptr_hl
    ld a, l
    add c
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, l
    add 3
    ld l, a
    ld a, h
    adc 0
    ld h, a                     ; HL = &field[3] (LSB)
    ld a, [hl]
    add b
    ld [hl], a
    ret nc
    dec hl
    inc [hl]
    ret nz
    dec hl
    inc [hl]
    ret nz
    dec hl
    inc [hl]
    ret

; sk_set_rcv_from_seq(A = addend) - rcv_nxt = their_seq(from packet) + addend
sk_set_rcv_from_seq:
    push af
    ld a, [wNetIhl]
    add 4
    ld e, a
    ld d, 0
    ld hl, wNetRxBuf
    add hl, de                  ; HL = &their_seq
    push hl
    call sock_ptr_hl
    ld a, l
    add SK_RCV
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld d, h
    ld e, l                     ; DE = &SK_RCV
    pop hl                      ; HL = &their_seq
    ld a, [hl+]
    ld [de], a
    inc de
    ld a, [hl+]
    ld [de], a
    inc de
    ld a, [hl+]
    ld [de], a
    inc de
    ld a, [hl]
    ld [de], a
    pop af
    call sk_add_rcv
    ret

; net_find_tcp -> HL = &TCP socket bound to wNetUdpPort, CF set if none
net_find_tcp:
    ld c, 0
.l
    ld a, c
    call net_sock_ptr
    ld a, [hl]
    cp SOCK_TCP
    jr nz, .nx
    push hl
    ld a, l
    add SK_LPORT
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [wNetUdpPort]
    cp [hl]
    jr nz, .nomatch
    inc hl
    ld a, [wNetUdpPort+1]
    cp [hl]
    jr nz, .nomatch
    pop hl
    or a
    ret
.nomatch
    pop hl
.nx
    inc c
    ld a, c
    cp MAX_SOCKS
    jr c, .l
    scf
    ret

; tcp_send_seg - build+send a segment from the current socket. wNetTcpFlags =
; flags, wNetTcpDlen = payload length (payload read from wNetReq.NR_BUF). A SYN
; also carries an MSS option so the peer never sends us an oversized segment.
tcp_send_seg:
    ld a, [wNetTcpFlags]
    and TF_SYN
    ld a, 20
    jr z, .th
    ld a, 24
.th
    ld [wNetTcpTh], a
    ld a, $45
    ld [wNetTx+0], a
    xor a
    ld [wNetTx+1], a
    ld [wNetTx+2], a
    ld a, [wNetTcpTh]
    add 20
    ld b, a
    ld a, [wNetTcpDlen]
    add b
    ld [wNetTx+3], a            ; total = 20 + thlen + dlen
    xor a
    ld [wNetTx+4], a
    ld [wNetTx+5], a
    ld [wNetTx+6], a
    ld [wNetTx+7], a
    ld a, 64
    ld [wNetTx+8], a
    ld a, 6
    ld [wNetTx+9], a
    xor a
    ld [wNetTx+10], a
    ld [wNetTx+11], a
    ld a, [wNetOurIP+0]
    ld [wNetTx+12], a
    ld a, [wNetOurIP+1]
    ld [wNetTx+13], a
    ld a, [wNetOurIP+2]
    ld [wNetTx+14], a
    ld a, [wNetOurIP+3]
    ld [wNetTx+15], a
    ld a, SK_RIP
    ld de, wNetTx+16
    ld b, 4
    call sk_field_to
    ld a, SK_LPORT
    ld de, wNetTx+20
    ld b, 2
    call sk_field_to
    ld a, SK_RPORT
    ld de, wNetTx+22
    ld b, 2
    call sk_field_to
    ld a, SK_SND
    ld de, wNetTx+24
    ld b, 4
    call sk_field_to
    ld a, SK_RCV
    ld de, wNetTx+28
    ld b, 4
    call sk_field_to
    ld a, [wNetTcpTh]
    cp 24
    ld a, $50
    jr nz, .doff
    ld a, $60
.doff
    ld [wNetTx+32], a
    ld a, [wNetTcpFlags]
    ld [wNetTx+33], a
    ld a, 2
    ld [wNetTx+34], a           ; window = 512
    xor a
    ld [wNetTx+35], a
    ld [wNetTx+36], a
    ld [wNetTx+37], a
    ld [wNetTx+38], a
    ld [wNetTx+39], a
    ld a, [wNetTcpTh]
    cp 24
    jr nz, .nomss
    ld a, 2
    ld [wNetTx+40], a           ; MSS option kind
    ld a, 4
    ld [wNetTx+41], a           ; length
    xor a
    ld [wNetTx+42], a
    ld a, TCP_MSS
    ld [wNetTx+43], a
.nomss
    ld a, [wNetTcpDlen]
    or a
    jr z, .nodata
    ld a, [wNetTcpTh]
    add 20
    ld e, a
    ld d, 0
    ld hl, wNetTx
    add hl, de
    ld d, h
    ld e, l                     ; DE = payload dest
    ld a, [wNetReq+NR_BUF]
    ld l, a
    ld a, [wNetReq+NR_BUF+1]
    ld h, a
    ld a, [wNetTcpDlen]
    ld b, a
.cpd
    ld a, [hl+]
    ld [de], a
    inc de
    dec b
    jr nz, .cpd
.nodata
    ld hl, wNetTx+12            ; TCP checksum: pseudo-header + segment
    ld b, 8
    ld de, 0
    call net_sum
    ld a, e
    add 6
    ld e, a
    ld a, d
    adc 0
    ld d, a
    ld a, [wNetTcpTh]
    ld b, a
    ld a, [wNetTcpDlen]
    add b                       ; tcplen = thlen + dlen
    add e
    ld e, a
    ld a, d
    adc 0
    ld d, a
    ld a, [wNetTcpTh]
    ld b, a
    ld a, [wNetTcpDlen]
    add b
    ld b, a                     ; B = tcplen
    ld hl, wNetTx+20
    call net_cksum
    ld a, h
    ld [wNetTx+36], a
    ld a, l
    ld [wNetTx+37], a
    ld hl, wNetTx
    ld b, 20
    ld de, 0
    call net_cksum
    ld a, h
    ld [wNetTx+10], a
    ld a, l
    ld [wNetTx+11], a
    ld a, [wNetTcpTh]
    add 20
    ld b, a
    ld a, [wNetTcpDlen]
    add b
    ld e, a
    ld a, 0
    adc 0
    ld d, a
    ld hl, wNetTx
    call net_slip_send
    ret

; sock_state_set(A = state) - store SK_STATE for the current socket
sock_state_set:
    ld b, a
    call sock_ptr_hl
    ld a, l
    add SK_STATE
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, b
    ld [hl], a
    ret

; tcp_connect - handshake for the current socket (SYN -> SYN-ACK -> ACK)
tcp_connect:
    call sock_ptr_hl            ; SK_SND = 1000
    ld a, l
    add SK_SND
    ld l, a
    ld a, h
    adc 0
    ld h, a
    xor a
    ld [hl+], a
    ld [hl+], a
    ld a, 3
    ld [hl+], a
    ld a, $E8
    ld [hl], a
    call sock_ptr_hl            ; SK_RCV = 0
    ld a, l
    add SK_RCV
    ld l, a
    ld a, h
    adc 0
    ld h, a
    xor a
    ld [hl+], a
    ld [hl+], a
    ld [hl+], a
    ld [hl], a
    call sock_ptr_hl            ; HASRX = 0
    ld a, l
    add SK_HASRX
    ld l, a
    ld a, h
    adc 0
    ld h, a
    xor a
    ld [hl], a
    ld a, TCP_SYNSENT
    call sock_state_set
    ld a, TF_SYN               ; send SYN (with MSS)
    ld [wNetTcpFlags], a
    xor a
    ld [wNetTcpDlen], a
    call tcp_send_seg
    ld a, 1
    call sk_add_snd
    xor a                      ; wait for ESTABLISHED
    ld [wNetTO], a
    ld [wNetTO+1], a
    ld [wNetTO+2], a
.wait
    call net_pump
    ld a, [wNetReq+NR_SOCK]
    call net_sock_ptr
    ld a, l
    ld [wNetSockPtr], a
    ld a, h
    ld [wNetSockPtr+1], a
    call sock_ptr_hl
    ld a, l
    add SK_STATE
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [hl]
    cp TCP_ESTAB
    jr z, .ok
    ld hl, wNetTO
    inc [hl]
    jr nz, .wait
    call SchedYield
    inc hl
    inc [hl]
    jr nz, .wait
    inc hl
    inc [hl]
    ld a, [hl]
    cp 250
    jr c, .wait
    ld a, $FF
    ret
.ok
    xor a
    ret

; tcp_data_send - net_send for a TCP socket (PSH|ACK + advance snd_nxt)
tcp_data_send:
    ld a, TF_PSH | TF_ACK
    ld [wNetTcpFlags], a
    ld a, [wNetReq+NR_LEN]
    ld [wNetTcpDlen], a
    call tcp_send_seg
    ld a, [wNetReq+NR_LEN]
    call sk_add_snd
    xor a
    ret

; tcp_close - net_close for a TCP socket (FIN, brief settle, free)
tcp_close:
    call sock_ptr_hl
    ld a, l
    add SK_STATE
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [hl]
    cp TCP_ESTAB
    jr nz, .free                ; only FIN if we were established
    ld a, TF_FIN | TF_ACK
    ld [wNetTcpFlags], a
    xor a
    ld [wNetTcpDlen], a
    call tcp_send_seg
    ld a, 1
    call sk_add_snd
    ld a, TCP_FINWAIT
    call sock_state_set
    ld b, 8                     ; brief settle: pump for the peer's FIN/ACK
.settle
    ld de, 8000
.s2
    call net_pump
    dec de
    ld a, d
    or e
    jr nz, .s2
    dec b
    jr nz, .settle
.free
    call sock_ptr_hl
    xor a
    ld [hl], a                  ; SK_TYPE = free
    ret

; tcp_buffer_data - store the current segment's payload into the socket rx slot
tcp_buffer_data:
    call sock_ptr_hl
    ld a, l
    add SK_HASRX
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, 1
    ld [hl], a
    ld a, [wNetTcpDlen2]
    cp 209
    jr c, .oklen
    ld a, 208
.oklen
    ld c, a
    call sock_ptr_hl
    ld a, l
    add SK_RXLEN
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, c
    ld [hl], a
    call sock_ptr_hl
    ld a, l
    add SK_RXBUF
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld d, h
    ld e, l                     ; DE = dest
    ld a, [wNetTcpDoff]
    ld l, a
    ld h, 0
    ld bc, wNetRxBuf
    add hl, bc                  ; HL = src
    ld a, c
    ; (C from wNetRxBuf low byte is irrelevant; use fresh length)
    ld a, [wNetTcpDlen2]
    cp 209
    jr c, .cplen
    ld a, 208
.cplen
    ld c, a
    or a
    ret z
.cp
    ld a, [hl+]
    ld [de], a
    inc de
    dec c
    jr nz, .cp
    ret

; tcp_in - the receive state machine for the matching TCP socket
tcp_in:
    ld a, [wNetIhl]
    add 2
    ld e, a
    ld d, 0
    ld hl, wNetRxBuf
    add hl, de
    ld a, [hl+]
    ld [wNetUdpPort], a
    ld a, [hl]
    ld [wNetUdpPort+1], a
    call net_find_tcp
    ret c
    ld a, l
    ld [wNetSockPtr], a
    ld a, h
    ld [wNetSockPtr+1], a
    ld a, [wNetIhl]             ; incoming flags
    add 13
    ld e, a
    ld d, 0
    ld hl, wNetRxBuf
    add hl, de
    ld a, [hl]
    ld [wNetTcpIn], a
    call sock_ptr_hl           ; current state
    ld a, l
    add SK_STATE
    ld l, a
    ld a, h
    adc 0
    ld h, a
    ld a, [hl]
    cp TCP_SYNSENT
    jr z, .synsent
    cp TCP_ESTAB
    jp z, .estab
    cp TCP_FINWAIT
    jp z, .estab                ; still ACK data/FIN while closing
    ret
.synsent
    ld a, [wNetTcpIn]
    and TF_SYN
    ret z
    ld a, 1                     ; rcv_nxt = their_seq + 1
    call sk_set_rcv_from_seq
    ld a, TF_ACK
    ld [wNetTcpFlags], a
    xor a
    ld [wNetTcpDlen], a
    call tcp_send_seg
    ld a, TCP_ESTAB
    call sock_state_set
    ret
.estab
    ; data offset (frame-relative) = ihl + (tcp[12]>>4)*4
    ld a, [wNetIhl]
    add 12
    ld e, a
    ld d, 0
    ld hl, wNetRxBuf
    add hl, de
    ld a, [hl]
    and $F0
    swap a
    add a, a
    add a, a                    ; tcphdrlen
    ld c, a
    ld a, [wNetIhl]
    add c
    ld [wNetTcpDoff], a
    ld a, [wNetRxLen]           ; datalen = rxlen - dataoff
    ld b, a
    ld a, [wNetTcpDoff]
    ld c, a
    ld a, b
    sub c
    ld [wNetTcpDlen2], a
    or a
    jr z, .checkfin
    call tcp_buffer_data
    ld a, [wNetTcpDlen2]
    call sk_add_rcv
    ld a, TF_ACK
    ld [wNetTcpFlags], a
    xor a
    ld [wNetTcpDlen], a
    call tcp_send_seg
.checkfin
    ld a, [wNetTcpIn]
    and TF_FIN
    ret z
    ld a, 1
    call sk_add_rcv
    ld a, TF_ACK
    ld [wNetTcpFlags], a
    xor a
    ld [wNetTcpDlen], a
    call tcp_send_seg
    ld a, TCP_DONE
    call sock_state_set
    ret