needs lib/str lib/struct lib/psrs lib/func io/stream mem/dict \
  drv/timer com/ip4 com/tcb
unit com/tcp

\ Packet section
6 const TCPPROTO
16 5 - 4* const MAXOPTSZ

\ Points to byte following the last option byte. Used for opts writing only.
\ For reading, Options+DataOffset gives our range.
0 value optptr

struct TCPMessage {
  beint Ports SeqNum AckNum ;
  beshort Flags Window Checksum UrgentPtr ;
  [void,0] Options ;
}
ip4payload TCPMessage bindstruct

: SourcePort Ports 16 rshift ;
: DestPort Ports $ffff and ;
: SwappedPorts Ports 16 lrot ;
: DataOffset Flags 12 rshift 4* ;
: Flags! doto Flags $fff invand or | ;
: FlagsOnly Flags $fff and ;
: opts[ ( -- ) Options to optptr ;
: ]opts ( -- )
  optptr Options - dup MAXOPTSZ >= if drop 0 else ( optlen )
    dup 4 mod if optptr 4 zfill then 4 /+ then ( optoff )
  5 + 12 lshift doto Flags $fff and or | ;

: tcp[] ( -- a u ) ip4payload @ DataLength DataOffset consume[] ;
\ Can only be called after optsdone has been called.
: tcp[]! ( a u -- ) DataOffset + DataLength! tcp[] cmove ;

consts 1 FIN 2 SYN 4 RST 8 PSH $10 ACK $20 URG
: FIN? Flags FIN and ;
: SYN? Flags SYN and bool ;
: RST? Flags RST and bool ;
: PSH? Flags PSH and bool ;
: ACK? Flags ACK and bool ;
: URG? Flags URG and bool ;

: tcpcksum ( -- n )
  0 addrof SourceAddr 8 ck[] ( sum )
  Protocol ck1 DataLength ck1
  doto Checksum 0 | >r ip4payload @ DataLength ck[] ckinv r> to Checksum ;

: wrap ( -- )
  TCPPROTO to Protocol
  com/ip4 wrap
  tcpcksum to Checksum ;

: .opt ( a n -- a+ )
  2- max0 case
    1 = of c@+ .x1 endof
    2 = of dup wbe@ .x2 2+ endof
    4 = of dup be@ .x 4+ endof
    2dup cspit[] + endcase ;

: addopt ( len code -- a ) optptr c!+ dipdup c! doto optptr tuck + | 2+ ;

: .opts ( -- )
  Options begin dup tcp[] drop < while c@+ ?dup while ( a c )
    dup 1 = if drop else
      . spc> c@+ .opt ." | " then
    repeat then ( a ) drop ;

: .tcp ( -- )
  ."Port   " SourcePort . ." --> " DestPort . nl>
  ."Seq    " SeqNum .x nl>
  ."Ack    " AckNum .x nl>
  ."Flags  " Flags .x2 nl>
  ."Window " Window .x2 nl>
  ."CkSum  " Checksum .x2 spc> tcpcksum .x2 nl>
  ."Urgent " UrgentPtr .x2 nl>
  ."Opts   " .opts nl>
  tcp[] dumpn ;

\ TCPCon section
536 const DEFAULTMSS
400 const MAXINACKMS
4000 const MAXOUTACKMS
32 const TCPCONCNT

enum Closed SynSent SynReceived SynAckSent AckSent Established FinReceived
enumrev statusrev

extends Stream struct TCPCon {
  uint  addr ports status ints outts tag ;
  TCBIn in ;
  TCBOut out ;
}
TCPCon typesz const TCPSZ

: contcb$ ( inseq con -- ) tuck in tcbin$ 1 swap out tcbout$ ;

: :readbuf ( n con -- a? read-n ) in tcbread dup if rot drop then ;
: :writebuf ( a n con -- written-n )
  dup status Established = if out tcbwrite else 2drop drop 0 then ;
: :close ( con -- ) >A 0 0 Closed to{ addr ports status } 0 A> contcb$ ;

create tcpcons TCPSZ TCPCONCNT * allot0
: con@ ( idx -- con ) TCPSZ * tcpcons + ;
:> ( i -- )
  ['] :readbuf ['] :writebuf rot con@ r! stream$
  ['] :close r> to close ;
( xt ) TCPCONCNT timesi

: findcon ( addr ports -- con-or-0 )
  2>r TCPCONCNT 0 do ( V1=addr V2=ports )
    i con@ bi addr V1 = | ports V2 = and if
      i con@ unloop 2rdrop exit then loop 0 2rdrop ;
: findfreecon ( -- con-or-0 ) 0 0 findcon ;
: opencon ( addr ports con -- con )
  >A 0 Closed to{ addr ports tag status } SeqNum 1+ A> tuck contcb$ ;
: ?opencon ( addr ports -- con-or-0 )
  2dup findcon ?dup if nip nip else ( addr ports )
    findfreecon dup if opencon else nip nip then then ;

: acceptcon ( -- con-or-0 )
  SourceAddr SwappedPorts ?opencon dup if ( con )
    A! SynReceived 0 to{ status outts } then ;

: doTCPin ( -- ) \ evarg1=con
  ACK? if
    AckNum evarg1 out tcback
    evarg1 status SynAckSent = if Established evarg1 to status then then
  FIN? if FinReceived evarg1 to status then
  tcp[] ?dup not if drop else ( a u )
    SeqNum evarg1 in tcbin
    ticks evarg1 to ints then ;

addeventtype const IP4TCPRECV
: handletcp ( -- )
  IP4TCPRECV to evtype
  SourceAddr SwappedPorts findcon dup to evarg1 ( con-or-0 )
  if doTCPin rootdispatch else SYN? if rootdispatch then then ;
TCPPROTO current registerproto

: preparemsg ( con -- )
  0 Flags!
  dup addr to DestAddr
  dup ports to Ports
  MAXWINDOWSZ to Window
  in TCB.orig to AckNum
  opts[ 1460 4 2 addopt wbe! ]opts ;

: doSynReceived ( con -- )
  dup outts if drop exit then ( con )
  SynAckSent over to status
  curlink newdgram
  dup preparemsg ( con )
  dup out TCB.orig 1- to SeqNum
  ticks swap to outts ( )
  SYN ACK or Flags!
  0 0 tcp[]!
  wrap ?log curlink sendframe ;

: doFinReceived ( con -- )
  curlink newdgram
  dup preparemsg ( con )
  dup out TCB.orig to SeqNum ( con )
  doto AckNum 1+ | \ FIN's phantom byte
  FIN ACK or Flags!
  0 0 tcp[]!
  wrap ?log curlink sendframe ( con )
  close ;

: doSend ( con seqn a u -- )
  curlink newdgram
  rot to SeqNum ( con a u )
  oover preparemsg
  tcp[]! ( con )
  dup doto ints 0 | if ACK Flags! then
  ticks swap to outts ( )
  wrap ?log curlink sendframe ;

: ?doSend ( con -- )
  dup outts ?dup if MAXOUTACKMS swap elapsedms? if
    0 over to outts dup out tcbrewind then then
  DEFAULTMSS over out tcbout ?dup if doSend else drop then ;

: sendACK ( con -- )
  0 over to ints
  curlink newdgram
  dup preparemsg ( con )
  out TCB.orig to SeqNum
  ACK Flags!
  0 0 tcp[]!
  wrap ?log curlink sendframe ;

: ?doACK ( con -- )
  dup ints ?dup not if drop exit then ( con ts )
  MAXINACKMS swap elapsedms? if sendACK else drop then ;

: tcpone ( con -- )
  dup addr not if drop exit then ( con )
  dup status case ( con )
    SynReceived = of doSynReceived endof
    FinReceived = of doFinReceived endof
    Established = of dup ?doSend ?doACK endof
    2drop endcase ;
: tcpwork ( -- ) TCPCONCNT 0 do i con@ tcpone loop ;