IN: udp
USING: kernel alien io-internals sequences strings generic namespaces threads io-internals arrays errors ;
USE: unix-internals

LIBRARY: libc
FUNCTION: ssize_t recv ( int s, void* buf, size_t nbytes, int flags ) ;
FUNCTION: ssize_t recvfrom ( int s, void* buf, size_t nbytes, int flags, sockaddr-in* from, socklen_t* fromlen ) ;
FUNCTION: ssize_t sendto ( int s, void* buf, size_t len, int flags, sockaddr-in* to, socklen_t tolen ) ;

: SOCK_DGRAM 2 ;

: udp-socket-fd ( -- socket )
    PF_INET SOCK_DGRAM 0 socket dup io-error dup init-handle ;

: do-recv ( socket -- data )
    1500 "char" <c-array> dup >r 1500 0 recv r> swap >string [ swap char-nth ] map-with ;

: do-recvfrom ( socket -- data addr port )
    "sockaddr-in" <c-object> dup >r 
    1500 "char" <c-array> dup >r
    rot
    r> 1500 0 r> "sockaddr-in" c-size <int> recvfrom
    dup -1 = [
	drop 2drop f
    ] [
        >string [ swap char-nth ] map-with
        swap dup sockaddr-in-addr inet-ntoa
        swap sockaddr-in-port ntohs
    ] if ;

: do-sendto ( socket data addr port -- ? )
    client-sockaddr ! socket data sockaddr
    >r dup length swap string>char-alien swap
    0
    r> "sockaddr-in" c-size sendto -1 = not ;
    

! ----------------------------------------------------------------

TUPLE: datagram packet addr port ;

: datagram-socket ( port -- fd )
    server-sockaddr udp-socket-fd
    dup SOL_SOCKET SO_REUSEADDR sockopt
    swap dupd "sockaddr-in" c-size bind drop ;

C: datagram ( port -- datagram )
    [ >r datagram-socket f <port> r> set-delegate ] keep
    datagram over set-port-type ;

! ----------------------------------------------------------------

TUPLE: recv-task ;

C: recv-task ( port -- task )
    [ >r <io-task> r> set-delegate ] keep ;

M: recv-task do-io-task
    io-task-port dup port-handle do-recv
    swap set-datagram-packet t ;

M: recv-task task-container drop read-tasks get ;

: wait-receive ( stream -- )
    [ swap <recv-task> add-io-task stop ] callcc0 drop ;

: receive ( stream -- packet )
    dup wait-receive dup pending-error datagram-packet ;

! ----------------------------------------------------------------

TUPLE: recvfrom-task ;

C: recvfrom-task ( port -- task )
    [ >r <io-task> r> set-delegate ] keep ;

M: recvfrom-task do-io-task
    io-task-port dup >r port-handle do-recvfrom
    r> 
    swap over set-datagram-port
    swap over set-datagram-addr
    set-datagram-packet    
    t ;

M: recvfrom-task task-container drop read-tasks get ;

: wait-receive-from ( stream -- )
    [ swap <recvfrom-task> add-io-task stop ] callcc0 drop ;

: receive-from ( stream -- packet addr port )
    dup wait-receive-from dup pending-error
    dup datagram-packet
    over datagram-addr
    rot datagram-port ;

! ----------------------------------------------------------------

TUPLE: sendto-task packet addr dport ;

C: sendto-task ( port -- task )
    [ >r <io-task> r> set-delegate ] keep ;

M: sendto-task do-io-task
    dup dup io-task-port port-handle
    over sendto-task-packet
    rot dup sendto-task-addr
    swap sendto-task-dport
    [ do-sendto drop ] [ >r 2drop 2drop io-task-port r> swap set-port-error ] recover t ;

M: sendto-task task-container drop write-tasks get ;

: wait-send-to ( stream packet addr port -- )
    [ >r >r rot <sendto-task>
      r> over set-sendto-task-dport
      swap over set-sendto-task-addr
      swap over set-sendto-task-packet
      r> swap add-io-task stop ] callcc0 2drop 2drop ;

: send-to ( stream packet addr port -- )
    2swap over >r 2swap wait-send-to r> pending-error ;

! ----------------------------------------------------------------

: >2array< ( a -- x y )
  dup first swap second ;

: (udp-server) ( quot daragram -- )
	[
		tuck
		receive-from 2array
		rot call
		>2array< send-to
	] 2keep
	(udp-server) ;

: udp-server ( port quot -- )
	swap <datagram> (udp-server) ;

: (test-udp-server) ( -- )
	[ 9999 [ swap "-ok" append swap ] udp-server ] in-thread
	9998 <datagram>
	[ "ay" "localhost" 9999 send-to ] keep
	[ receive-from ] keep
	[ "bi" "localhost" 9999 send-to ] keep
	[ receive-from ] keep
	drop ;

USE: test

: test-udp-server ( -- )
	{ t }
	[ 0 <datagram> <sendto-task>
	  "world" over set-sendto-task-dport
	  "localhost" over set-sendto-task-addr
	  "packet" over set-sendto-task-packet
	  do-io-task
	] unit-test
	

	{ "ay-ok" "127.0.0.1" 9999 "bi-ok" "127.0.0.1" 9999 }
	[ (test-udp-server) ] unit-test

	{ { kernel-error 3 0 "world" } }
	[ [ 0 <datagram> "hello" "localhost" "world" send-to ] catch ] unit-test ;

implementation of udp for factor language