require unix/socket.fs

require libs/xstring/xstring.4th

require util.4th
require extensions.4th
require connections.4th
require commandline.4th
require motd.4th
require motd-parser.4th
require proxyline-parser.4th
require torcontrol.4th
require dos.4th
require sendbuffer.4th

AF_INET constant SERVER_SOCKET_DOMAIN
  SOCK_STREAM
  SOCK_NONBLOCK or
        constant SERVER_SOCKET_TYPE
      0 constant SERVER_SOCKET_PROTOCOL
      0 constant SERVER_ADDR
CONFIG_SERVER_PORT constant SERVER_PORT
      4 constant SERVER_LISTEN_BACKLOG

\ Listening file descriptor.
0 variable! listenfd

\ If we should accept new connections.
true variable! accept-connections

\ If we should echo back command responses.
true variable! command-echo

\ Idle detection.
false variable! idle
: server-idle? ( -- flag ) idle @ ;
: server-idle! ( flag -- ) idle ! ;

\ Temporary structs.
create con /CONNECTION allot
create saddr /sockaddr_in allot
create optval /option_value allot

: (assert-setsockopt) ( result-n -- )
  0< abort" making socket reusable failed" ;
: (optval-true!) ( -- )
  1 optval l! ;
: (make-reusable) ( socket-fd-n -- )
  (optval-true!)
  SOL_SOCKET SO_REUSEADDR optval /option_value setsockopt (assert-setsockopt) ;

: (saddr!) ( protocol-n sin_addr-n port-n -- )
  htons saddr port w!
  saddr sin_addr l!
  saddr family w! ;

: (assert-socket) ( result-n -- result-n )
  dup 0< abort" socket failed." ;
: (assert-bind) ( result-n -- )
  0< abort" bind failed." ;
: (assert-listen) ( result-n -- )
  0< abort" listen failed." ;

: (erase-saddr) ( -- )
  saddr /sockaddr_in erase ;
: (create-socket)  ( -- )
  SERVER_SOCKET_DOMAIN
  SERVER_SOCKET_TYPE
  SERVER_SOCKET_PROTOCOL
    socket (assert-socket) listenfd ! ;
: (make-socket-reusable) ( -- )
  listenfd @ (make-reusable) ;
: (set-saddr) ( -- )
  SERVER_SOCKET_DOMAIN SERVER_ADDR SERVER_PORT (saddr!) ;
: (bind-socket) ( -- )
  listenfd @ saddr /sockaddr_in bind (assert-bind) ;
: (listen-socket) ( -- )
  listenfd @ SERVER_LISTEN_BACKLOG listen() (assert-listen) ;
: (server-info) ( -- )
  cr cr ." Server listening at port: " SERVER_PORT . cr ;
: initialize-server ( -- )
  (erase-saddr)
  (create-socket)
  (make-socket-reusable)
  (set-saddr)
  (bind-socket)
  (listen-socket)
  (server-info) ;

: (perform-disconnect) ( connection-addr -- )
  dup connection.circuitid @ 0<> IF
    dup connections.indexOf dos-remove-connection
  THEN
  dup connection.connected false swap !
      connection.fd @ close throw ;

: (close-clients) ( -- )
  connections.count 0= IF
    EXIT
  THEN
  connections.count 0 DO
    I connections.at connection.connected @ true = IF
      I connections.at (perform-disconnect)
    THEN
  LOOP ;

: (assert-close) ( result-n -- )
  0<> abort" close failed" ;
: (close-server) ( -- )
  listenfd @ close (assert-close) ;

: (close-server-info) ( -- )
  cr ." Closed server connections." cr ;
: close-server ( -- )
  (close-clients) (close-server) (close-server-info) ;

: (queue-disconnect) ( connection-addr -- )
  dup (perform-disconnect) EVENT_CONNECTION_CLOSED events.enqueue ;

: (try-accept) ( -- c-result-n )
  listenfd @ 0 0 accept() ;
: (accept-error) ( accept-result-n -- )
  errno EAGAIN <> abort" accept error" drop ;
: (erase-connection) ( connection-addr -- )
  /CONNECTION erase ;
: (set-connection-number) ( -- )
  connections.last dup connections.indexOf 1+ swap connection.number ! ;
: (enqueue-new-connection) ( -- )
  connections.last EVENT_CONNECTION_NEW events.enqueue ;
: (store-connection) ( connection-addr -- )
  dup >r connections.append IF
    (set-connection-number)
    (enqueue-new-connection)
    rdrop
  ELSE
    ." Warning: failed to store connection, disconnecting client." cr
    r> (perform-disconnect)
  THEN ;
: (con!) ( fd-n connected-flag -- )
  con connection.connected !
  con connection.fd ! ;
: (accept-connection) ( fd-n -- )
  con (erase-connection)
  true (con!)
  con (store-connection) ;
: (server-idle-accept) ( -- )
  accept-connections @ invert IF
    EXIT
  THEN
  (try-accept) dup 0< IF
    (accept-error)
  ELSE
    (accept-connection)
  THEN ;

: (connected?) ( connection-addr -- flag )
  connection.connected @ ;
: (try-recv) ( connection-addr -- recv-result-n )
  dup connection.fd @
  swap connection.buffer
  CONNECTION_BUFFER_SIZE
  MSG_DONTWAIT
    recv ;
: (recv-error?) ( c-result-n -- flag ) 0< ;
: (recv) ( recv-result-n connection-addr -- )
  2dup connection.bufferlen ! swap
  0> IF
    EVENT_CONNECTION_RECV events.enqueue
  ELSE \ disconnected
    (queue-disconnect)
  THEN ;
: (recv-warning) ( -- ) ." Warning: recv failed, disconnecting client." cr ;
: (recv-error) ( recv-result-n connection-addr -- )
  errno EAGAIN <> IF
    (queue-disconnect) drop
    (recv-warning)
  ELSE
    2drop
  THEN ;
: (connection-recv) ( connection-addr -- )
  dup (try-recv) tuck (recv-error?) IF
    (recv-error)
  ELSE
    (recv)
  THEN ;

: (server-idle-recv) ( -- )
  true server-idle!
  connections.count 0= IF
    EXIT
  THEN
  connections.count 0 DO
    I connections.at dup (connected?) IF
      (connection-recv)
    ELSE
      drop
    THEN
  LOOP ;

: (connection.number>string) ( connection-addr -- c-addr u )
  connection.number @ to-string ;
: (connection.buffer>string) ( connection-addr -- c-addr u )
  dup connection.buffer swap connection.bufferlen @  ;
: (connection>name) ( connection-addr -- c-addr u )
  s" Anon " pad place
  (connection.number>string) pad +place
  pad count ;
: (expect-proxyline?) ( connection-addr -- flag )
  connection.circuitid @ 0= ;
: (parse-proxyline) ( connection-addr -- )
  dup >r (connection.buffer>string) proxyline>circuitid
  dup r@ connection.bufferlen ! r@ connection.buffer swap move
  r> connection.circuitid ! ;
: (last-sendbuffer-char) ( -- c )
  sendbuffer@ + 1- c@ ;
: (maybe-append-newline) ( -- )
  (last-sendbuffer-char) 10 <> IF
    s\" \n" sendbuffer-append
  THEN ;
: (format-sendbuffer) ( msg-str from-str -- )
  sendbuffer-reset
  sendbuffer-append
  s" : " sendbuffer-append
         sendbuffer-append
  (maybe-append-newline)
  sendbuffer-sanitize ;
: (connected?) ( connection-addr -- )
  connection.connected @ ;
: (different-connection?) ( from-connection-addr to-connection-addr -- )
  <> ;
: (send?) ( from-connection-addr to-connection-addr -- )
  tuck (different-connection?) swap (connected?) and ;
: (check-send) ( result-n -- )
  0< IF ." Warning: send failed." cr THEN ;
: (send-sendbuffer) ( to-connection-addr -- )
  connection.fd @ sendbuffer@ MSG_NOSIGNAL send (check-send) ;
: (try-send) ( from-connection-addr to-connection-addr -- )
  2dup (send?) IF
    nip (send-sendbuffer)
  ELSE
    2drop
  THEN ;
: (dos-update-stats) ( from-connection-addr -- )
  dup connections.indexOf
  swap (connection.buffer>string) nip over dos-add-bytes
  1 swap dos-add-lines ;
: (dos-protect?) ( connection-addr -- flag )
  connections.indexOf dos? ;
: (dos-protect) ( connection-addr -- )
  ." DOS protection enabled for circuit:" cr
  dup connections.indexOf .dos-info
  dup connections.indexOf true swap dos-handled!
  connection.circuitid @ torcontrol-close-circuit ;
: (is-command?) ( str -- flag )
  1 min s" /" compare 0= ;
create command-parser PARSER_SIZE allot
: (extract-command) ( str -- str )
  command-parser new-parser 1 parser>> parser-remaining ;
: (parse-command) ( str -- str flag )
  2dup (is-command?) IF
    (extract-command) true
  ELSE
    false
  THEN ;
4096 constant REDIRECT_BUFFER_SIZE
create server-redirect-buffer REDIRECT_BUFFER_SIZE allot
create server-emit-buffer 1 chars allot
variable redirect-broadcast-xt
: (server-redirect-reset) ( -- )
  s" " server-redirect-buffer xplace ;
: (server-redirect-flush) ( -- )
  server-redirect-buffer xcount redirect-broadcast-xt @ execute
  (server-redirect-reset) ;
: (server-type) ( str -- )
\ overflow check 
  dup cell + server-redirect-buffer xcount nip + REDIRECT_BUFFER_SIZE <= IF
    server-redirect-buffer +xplace
  ELSE
    2drop
  THEN ;
: (server-emit) ( c -- )
  server-emit-buffer c!
  server-emit-buffer 1 chars (server-type) ;
: (enable-redirect) ( -- )
  ['] (server-emit) stdout-hook-emit
  ['] (server-type) stdout-hook-type
  (server-redirect-reset) ;
: (disable-redirect) ( -- )
  (server-redirect-flush)
  stdout-hook-reset ;
: (depth-evaluate) ( command-str -- )
  depth 2 - >r
  ['] evaluate catch IF
    2drop ." An error has occured." cr
  THEN
  depth r> <> abort" aborting to fix stack." ;
: (dispatch-admin-command) ( connection-addr command-str -- flag )
  rot connection.admin @ IF
    ['] (depth-evaluate) catch IF 2drop THEN true
  ELSE
    2drop false
  THEN ;
\ TODO: user command dispatching is very basic for now
\ TODO: maybe make commands extendible at runtime?
defer user-command-help   ( -- )
defer user-command-users  ( -- )
defer user-command-whoami ( connection-addr -- )
' noop is user-command-help
' noop is user-command-users
' drop is user-command-whoami
: (dispatch-user-command) ( connection-addr command-str -- )
  2dup s" help" startswith IF
    3drop user-command-help
  ELSE 2dup s" users" startswith IF
    3drop user-command-users
  ELSE 2dup s" whoami" startswith IF
    2drop user-command-whoami
  ELSE
    3drop ." Unknown user command." cr
  THEN THEN THEN ;
: (handle-command) ( connection-addr -- )
  dup (connection.buffer>string) (parse-command) IF
    (enable-redirect)
    3dup (dispatch-admin-command) IF
      3drop
    ELSE
      (dispatch-user-command)
    THEN
    (disable-redirect)
  ELSE
    2drop drop
  THEN ;
: (handle-broadcast) ( connection-addr -- )
  dup >r (connection.buffer>string) r@ (connection>name) (format-sendbuffer)
  r> (dos-update-stats)
  sendbuffer@ type
  connections.count 0 DO
    dup I connections.at (try-send)
  LOOP ;
: server-recv ( from-connection-addr eventid-n )
  drop
  dup (expect-proxyline?) IF
    dup (parse-proxyline)
    dup connection.circuitid @ over connections.indexOf dos-add-connection
    dup (connection.buffer>string) nip 0= IF
        drop EXIT
    THEN
  THEN
  dup connections.indexOf dos-handled? IF
    drop EXIT
  THEN
  dup (dos-protect?) IF
    (dos-protect)
  ELSE
    dup (handle-broadcast)
    (handle-command)
  THEN ;

: server-idle-accept ( eventdata-n eventid-n -- )
  2drop (server-idle-accept) ;
: server-idle-recv ( eventdata-n eventid-n -- )
  2drop (server-idle-recv) ;

false variable! motd-cached
create motd-cache SENDBUFFER_SIZE allot
0 variable! motd-cache-length
: (sendbuffer-motd-line-append) ( str -- )
  s" Server: " sendbuffer-append
               sendbuffer-append
  s\" \n"      sendbuffer-append ;
: (prepare-motd) ( -- )
  sendbuffer-reset
  motd-cached @ IF
    motd-cache motd-cache-length @ sendbuffer-append
    EXIT
  THEN
  motd@ ['] (sendbuffer-motd-line-append) parse-motd
  sendbuffer@ dup motd-cache-length ! motd-cache swap move ;
: (prepare-empty-line) ( -- )
  sendbuffer-reset s\" Server: \n" sendbuffer-append ;
: (prepare-identity) ( connection-addr -- )
  sendbuffer-reset
  s\" Server: You are now known as \"" sendbuffer-append
  (connection>name)                    sendbuffer-append
  s\" \".\n"                           sendbuffer-append ;
: server-connection-new ( connection-addr eventid-n -- )
  drop ." New client connected!" cr
  dup (prepare-motd) (send-sendbuffer)
  dup (prepare-empty-line) (send-sendbuffer)
  dup (prepare-identity)
      (send-sendbuffer) ;

: server-connection-closed ( connection-addr eventid-n -- )
  2drop ." Client disconnected." cr ;

: server-commandline ( eventdata-n eventid-n -- )
  2drop commandline-ready? IF
    commandline-getline 2dup logger.log cr
    ['] evaluate catch dup 0= IF
      drop
    ELSE
      ." error code: " . 2drop
    THEN
    cr
    commandline-reset 
  ELSE
    commandline-key? IF
      commandline-key commandline-handlekey
      commandline-redraw
    THEN
  THEN ;

: server-motd-changed ( eventdata-n eventid-n -- )
  2drop (prepare-motd)
  connections.count 0 DO
    I connections.at (send-sendbuffer)
  LOOP ;

: user-help ( -- )
  ." User commands: " cr
  ." help   ( -- ) \ this help command" cr
  ." users  ( -- ) \ display the connected users" cr
  ." whoami ( -- ) \ display your name" cr ;

: user-users ( -- )
  connections.count 0= IF
    EXIT
  THEN
  connections.count 0 DO
    I connections.at connection.connected @ IF
      ." Anon " I connections.at (connection.number>string) type cr
    THEN
  LOOP ." TODO: implement last active time." cr ;

: user-whoami ( connection-addr -- )
  ." You are Anon " (connection.number>string) type ." ." cr ;

' user-help IS user-command-help
' user-users IS user-command-users
' user-whoami IS user-command-whoami

: server-commands ( -- )
\ List server commands.
  ." Server commands: " cr cr
  ." You may enter any valid forth expression" cr cr
  ." server-commands    ( -- )         \ this help command" cr
  ." server-admin       ( user-n -- )  \ make a user admin" cr
  ." server-users       ( -- )         \ list connected users" cr
  ." server-accept      ( flag -- )    \ accept new connections" cr
  ." server-accepting?  ( -- )         \ check if the server is" cr
  ."                                   \ accepting connections" cr
  ." server-disconnect  ( user-n -- )  \ disconnect a user by closing the circuit" cr
  ." server-broadcast   ( msg-str -- ) \ broadcast a server message to" cr
  ."                                   \ all users" cr
  ." server-message     ( msg-str user-n -- ) \ send a server message to" cr
  ."                                          \ a specific user" cr
;

: help ( -- ) server-commands ;

: (userid>connection) ( user-n -- connection-addr )
  1- connections.at ;

: server-admin ( user-n -- )
  (userid>connection) connection.admin true swap ! ;

: server-users ( -- )
  connections.count 0= IF
    ." No connected users." cr
    EXIT
  THEN
  connections.count 0 DO
    I connections.at dup connection.connected @ IF
      dup ." Anon " (connection.number>string) type
          ."  CircuitID " connection.circuitid @ . cr
    ELSE
      drop
    THEN
  LOOP ;

: server-accept ( flag -- )
  dup accept-connections ! IF
    ." Server is set to accept new connections." cr
  ELSE
    ." Server is set to not accept new connections." cr
  THEN ;

: server-accepting? ( -- )
  accept-connections @ IF
    ." Server is currently accepting new connnections." cr
  ELSE
    ." Server is currently not accepting new connections." cr
  THEN ;

: server-disconnect ( user-n -- )
  (userid>connection) dup connection.connected @ IF
    connection.circuitid @ torcontrol-close-circuit
    ." Tor circuit closed." cr
  ELSE
    drop ." User not connected." cr
  THEN ;

create broadcast-parser PARSER_SIZE allot
: (nextline) ( -- line-str flag )
  s\" \n" parser>>string IF
    parser-extract 1 parser>>
    parser-mark true
  ELSE
    parser-remaining 2dup nip 0> IF
      parser>>| true
    ELSE
      false
    THEN
  THEN ;
: server-broadcast ( msg-str -- )
  connections.count 0= IF
    EXIT
  THEN
  broadcast-parser new-parser
  BEGIN
    (nextline)
  WHILE
    s" Server" (format-sendbuffer)
    connections.count 0 DO
      I connections.at dup connection.connected @ IF
        (send-sendbuffer)
      ELSE
        drop
      THEN
    LOOP
  REPEAT 2drop ;
: (assert-connected) ( connection-addr -- )
  connection.connected @ invert abort" Not connected" ;
: server-message ( msg-str user-n -- )
  >r 2dup type
  s" Server" (format-sendbuffer)
  r> (userid>connection) dup (assert-connected)
  (send-sendbuffer) ;

' server-broadcast redirect-broadcast-xt !

' server-idle-accept       EVENT_IDLE              eventhandlers.append
' server-idle-recv         EVENT_IDLE              eventhandlers.append
' server-connection-new    EVENT_CONNECTION_NEW    eventhandlers.append
' server-connection-closed EVENT_CONNECTION_CLOSED eventhandlers.append
' server-recv              EVENT_CONNECTION_RECV   eventhandlers.append
' server-commandline       EVENT_COMMANDLINE       eventhandlers.append
' server-motd-changed      EVENT_MOTD_CHANGED      eventhandlers.append

initialize-server