(use bokbok)
(use matchable)

(define *server* #f)

(define (open-handler con)
  (printf "OPEN from ~a @ ~a:~a\n"
          (connection-user con)
          (connection-host con)
          (connection-port con)))

(define (close-handler con)
  (printf "CLOSE from ~a @ ~a:~a\n"
          (connection-user con)
          (connection-host con)
          (connection-port con)))

(define (request-handler con request)
  (printf "REQUEST from ~a @ ~a:~a: ~s\n"
          (connection-user con)
          (connection-host con)
          (connection-port con)
          request)

  (match request
         (("error")
          (error "Error raised"))
         (("callback")
          (let ((response
                 (request! con '("called you back"))))
            (list "did it" (sprintf "~s" response))))
         (("kill")
          (if *server*
              (begin
                (stop-server! *server*)
                (list "ok"))
              (error "I'm not a server")))
         (else
          (cons "echo" request))))

(define (run-client! con)
  (printf "> ")
  (let ((req (read)))
    (if (eof-object? req)
        (close-connection! con)
        (begin
         (printf
          "Reply: ~s\n"
          (request! con req))
         (run-client! con)))))

(define (run-server! server)
  (set! *server* server)
  (printf "Running server...\n")
  (wait-until-server-stopped server))

(match
 (command-line-arguments)
 (("client" "unix" path)
  (run-client! (open-connection path #f #f #f request-handler close-handler)))
 (("client" "unix" path user pass)
  (run-client! (open-connection path #f user (passphrase->key pass) request-handler close-handler)))

 (("server" "unix" path)
  (run-server! (start-server path #f 10 #f open-handler request-handler close-handler)))
 (("server" "unix" path pass)
  (run-server! (start-server path #f 10 (lambda _ (passphrase->key pass)) open-handler request-handler close-handler)))

 (("client" "tcp" addr port)
  (run-client! (open-connection addr (string->number port) #f #f request-handler close-handler)))
 (("client" "tcp" addr port user pass)
  (run-client! (open-connection addr (string->number port) user (passphrase->key pass) request-handler close-handler)))

 (("server" "tcp" port)
  (run-server! (start-server #f (string->number port) 10 #f open-handler request-handler close-handler)))
 (("server" "tcp" port pass)
  (run-server! (start-server #f (string->number port) 10 (lambda _ (passphrase->key pass)) open-handler request-handler close-handler)))

 (else
  (printf "Imagine a comprehensive usage screen here!\n")))

((license "BSD")
 (category other)
 (needs tweetnacl socket srfi-27 matchable)
 (author "Alaric Snell-Pym")
 (synposis "A network programming toolkit"))

(module
 bokbok
 (passphrase->key

  open-connection
  close-connection!
  request!

  connection?
  connection-user
  connection-host
  connection-port

  start-server
  stop-server!
  wait-until-server-stopped
  )

 (import chicken scheme)

 (use extras)
 (use data-structures)
 (use ports)
 (use srfi-1)
 (use srfi-4)
 (use srfi-13)
 (use srfi-18)
 (use srfi-69)

 (use matchable)
 (use tweetnacl)
 (use socket)
 (use srfi-27)
 (use moa) ;; from srfi-27
 (use entropy-unix) ;; from srfi-27

 ;; FIXME: Accept a hostname for open-connection, and look it up

 ;; FIXME: User data field in a connection

 ;; FIXME: Reduce timeouts to detect gone servers quickly

 ;; FIXME: Auto reconnect on request!

 ;; FIXME: Configurable maximum message size (and use 4 bytes, so it can be >64KiB)

 (define (debug . args)
  (let ((str (apply sprintf args)))
    (printf "DEBUG: ~a\n" str)))

#; (define (debug . args)
   (void))

 (define *header* "BOKBOKBOK")

 (define (read-packet! i)
   (let ((l1 (read-byte i)))
     (if (eof-object? l1)
         l1
         (let ((l2 (read-byte i)))
           (if (eof-object? l2)
               l2
               (let ((l3 (read-byte i)))
                 (if (eof-object? l3)
                     l3
                     (let ((len (+ (* 65536 l1) (* 256 l2) l3)))
                       (let ((p (read-string len i)))
                         (if (= (string-length p) len)
                             p
                             (with-input-from-string "" read)))))))))))

 (define (write-packet-no-flush! o p)
   (let* ((len (string-length p))
          (l1 (remainder (quotient len 65536) 256))
          (l2 (remainder (quotient len 256) 256))
          (l3 (remainder len 256)))
     (when (> len 16777216)
           (error "Cannot send oversized packet, maximum is 16777216 bytes" len))

     (write-byte l1 o)
     (write-byte l2 o)
     (write-byte l3 o)
     (write-string p #f o)))

 (define (write-packet! o p)
   (write-packet-no-flush! o p)
   (flush-output o))

 ;; Join a list of strings into a packet
 (define (join-packet ps)
   (let ((plen (length ps)))
     (when (> plen 255)
           (error "Can't fit more than 255 fields into a packet"))
     (with-output-to-string
       (lambda ()
        (write-byte plen)
        (for-each
         (lambda (p)
           (write-packet-no-flush! (current-output-port) p))
         ps)))))

 ;; Split a packet into a list of strings
 (define (split-packet p)
   (if (eof-object? p)
       p
       (with-input-from-string p
         (lambda ()
           (let ((plen (read-byte)))
             (if (eof-object? plen)
                 plen)
             (let loop
                 ((todo plen)
                  (result '()))
               (if (zero? todo)
                   (reverse! result)
                   (loop
                    (- todo 1)
                    (let ((p (read-packet! (current-input-port))))
                      (if (eof-object? p)
                          p
                          (cons p
                                result)))))))))))

 (define (passphrase->key key-string)
   (let* ((key-hash (hash key-string))
          (key (string->blob (substring/shared key-hash 0 symmetric-box-keybytes))))
     (cons
      (symmetric-box key)
      (symmetric-unbox key))))

 (define (bytes->key key-bytes)
   (let ((key (string->blob key-bytes)))
    (cons
     (symmetric-box key)
     (symmetric-unbox key))))

 (define get-random-bytes
   (let* ((entropy (make-entropy-source-urandom-device))
          (random (make-random-source-moa))
          (make-random-u8vector (random-source-make-u8vectors random))
          (mutex (make-mutex 'rpc-random-source)))
     (random-source-entropy-source-set! random entropy)
     (random-source-randomize! random entropy)
     (lambda (length)
       (mutex-lock! mutex)
       (let ((rb (make-random-u8vector length)))
         (mutex-unlock! mutex)
         rb))))

 (define (generate-session-key)
   (let* ((key-u8vector (get-random-bytes symmetric-box-keybytes))
          (key (u8vector->blob/shared key-u8vector)))
     (values
      (cons
       (symmetric-box key)
       (symmetric-unbox key))
      (blob->string key))))

 (define *nonce-counter* 0)
 (define *nonce-seed* (blob->string
                       (u8vector->blob/shared
                        (get-random-bytes symmetric-box-noncebytes))))

 (define (encrypt key plaintext)
   (set! *nonce-counter* (+ *nonce-counter* 1))
   (let* ((nonce-hash (hash (string-append (number->string *nonce-counter*)
                                           *nonce-seed*)))
          (nonce (substring/shared nonce-hash 0 symmetric-box-noncebytes))
          (nonce-u8vector (blob->u8vector/shared (string->blob nonce))))
     (string-append
      nonce
      ((car key) plaintext nonce-u8vector))))

 (define (decrypt key c)
   (assert (>= (string-length c) symmetric-box-noncebytes))
   (let* ((nonce (substring/shared c 0 symmetric-box-noncebytes))
          (c-wo-n (substring/shared c symmetric-box-noncebytes))
          (nonce-u8vector (blob->u8vector/shared (string->blob nonce)))
          (plaintext ((cdr key) c-wo-n nonce-u8vector)))
     (if plaintext
         plaintext
         (error "Invalid cyphertext"))))

 (define-record-type connection
   (make-connection* host port user
                    socket input output
                    key
                    mutex
                    thread
                    waiters
                    request-handler
                    close-handler
                    counter)
   connection?
   (host connection-host)
   (port connection-port)
   (user connection-user)
   (socket connection-socket)
   (input connection-input)
   (output connection-output)
   (key connection-key)
   (mutex connection-mutex)
   (thread connection-thread)
   (waiters connection-waiters)
   (request-handler connection-request-handler)
   (close-handler connection-close-handler)
   (counter connection-counter (setter connection-counter)))

 (define (connection-send! con packet-parts)
   (let ((packet (if (connection-key con)
                     (encrypt (connection-key con) (join-packet packet-parts))
                     (join-packet packet-parts))))
     (mutex-lock! (connection-mutex con))
     (write-packet! (connection-output con) packet)
     (mutex-unlock! (connection-mutex con))))

 (define (handle-request! con id body)
   (debug "Handling request id:~a ~s" id body)
   (let ((thread
          (make-thread
           (lambda ()
             (let ((response
                    (handle-exceptions
                     error
                     (list "err"
                           id
                           (cond
                            ((condition? error)
                             (sprintf
                              "~s in ~s"
                              ((condition-property-accessor 'exn 'message "Unknown error") error)
                              (cons ((condition-property-accessor 'exn 'location (void)) error)
                                    ((condition-property-accessor 'exn 'arguments '()) error))))
                            (else (->string error))))
                     (append
                      (list "ok" id)
                      ((connection-request-handler con)
                       con
                       body)))))
               (debug "Sending response id:~a ~s" id response)
               (connection-send! con response)))
           `(bokbok-request-thread ,(connection-host con) ,(connection-port con) ,id))))
     (thread-start! thread)))

 (define (handle-response! con id body)
   (debug "Handling response id:~a ~s" id body)
   (mutex-lock! (connection-mutex con))
   (let ((waiter (hash-table-ref/default (connection-waiters con)
                                         id #f)))
     (if waiter
         (begin
           (set-cdr! waiter body)
           (mutex-unlock! (car waiter)))
         (debug "Discarding response to unknown request id:~a" id)))
   (mutex-unlock! (connection-mutex con)))

 (define (handle-connection-thread!)
   (let* ((con (thread-specific (current-thread)))
          (session-key (connection-key con)))
     (debug "Session thread starting")
    (let loop ()
      (debug "Session thread waiting for packet")
      ;; We are the only thing that reads from connection-input, so need no mutex!
      (let ((raw-request (read-packet! (connection-input con))))
        (if (eof-object? raw-request)
            ;; Terminate loop
            ((connection-close-handler con) con)
            ;; Handle request and loop
            (let* ((request-bytes (if session-key
                                      (decrypt session-key raw-request)
                                      raw-request))
                   (request (split-packet request-bytes)))
              (match
               request
               (("req" id . body)
                (handle-request! con id body))
               (("ok" id . body)
                (handle-response! con id (list 'ok body)))
               (("err" id . body)
                (handle-response! con id (list 'error body))))
              ;; Loop for next request
              (loop)))))))

 (define (make-connection host port user socket input output key request-handler close-handler)
   (let ((con
          (make-connection*
           host port user
           socket input output
           key
           (make-mutex `(bokbok-connection-mutex ,host ,port))
           (make-thread handle-connection-thread!
                        `(bokbok-connection-thread ,host ,port))
           (make-hash-table)
           request-handler
           close-handler
           0)))

     (thread-specific-set!
      (connection-thread con)
      con)
     (thread-start!
      (connection-thread con))
     con))

 (define (open-connection host port username key request-handler close-handler)
   (let ((s
          (if (not port)
              ;; UNIX domain
              (let ((s (socket af/unix sock/stream)))
                (socket-connect s (unix-address host))
                s)
              ;; TCP domain
              (let ((s (socket af/inet sock/stream)))
                (socket-connect s (inet-address host port))
                (set! (tcp-no-delay? s) #t)
                s))))
     (receive
      (input output)
      (parameterize
       ((socket-send-buffer-size 4096)
        (socket-send-size 16384))
       (socket-i/o-ports s))
      (if (and username key)
          (begin ;; Encrypted connection
            (receive
             (session-key session-key-bytes)
             (generate-session-key)
             (debug "Client writing username and session key")
             (write-packet-no-flush! output username)
             (write-packet! output (encrypt key session-key-bytes))
             (debug "Client waiting for header")
             (let ((header (read-string (string-length *header*) input)))
               (debug "Got header bytes: ~s" header)
               (if (string=? header *header*)
                   (make-connection host port username s input output session-key request-handler close-handler)
                   (error "Invalid hello from server" header)))))
          (begin ;; Plaintext connection
            (debug "Client waiting for header")
            (let ((header (read-string (string-length *header*) input)))
              (debug "Got header bytes: ~s" header)
              (if (string=? header *header*)
                  (make-connection host port #f s input output #f request-handler close-handler)
                  (error "Invalid hello from server" header))))))))

 (define (close-connection! con)
   ;; FIXME: Set flags to suppress any further responses from still-running handlers
   ;; FIXME: Return an error to all pending waiters
   (thread-terminate! (connection-thread con))
   (close-output-port (connection-output con))
   (close-input-port (connection-input con)))

 (define (request! con packet-parts)
   (mutex-lock! (connection-mutex con))
   (let* ((id (number->string (connection-counter con)))
          (waiter (cons (make-mutex `(bokbok-request-mutex ,(connection-host con) ,(connection-port con) id))
                        #f)))

     ;; Mutex starts life locked by the connection thread
     (mutex-lock! (car waiter)
                  #f
                  (connection-thread con))

     (hash-table-set! (connection-waiters con) id waiter)
     (set! (connection-counter con)
           (+ (connection-counter con) 1))
     (mutex-unlock! (connection-mutex con))

     (connection-send! con (cons "req" (cons id packet-parts)))

     ;; Wait for response, when connection thread unlocks the mutex
     (mutex-lock! (car waiter))

     ;; Return response
     (cdr waiter)))

 (define (start-server bind-addr bind-port backlog user->key open-handler request-handler close-handler)
   (let* ((s (if (not bind-port)
                 ;; UNIX domain
                 (let ((s (socket af/unix sock/stream)))
                   (socket-bind s (unix-address bind-addr))
                   (socket-listen s backlog)
                   s)
                 ;; TCP domain
                 (let ((s (socket af/inet sock/stream)))
                   (socket-bind s (inet-address bind-addr bind-port))
                   (set! (so-reuse-address? s) #t)
                   (socket-listen s backlog)
                   s)))
          (thread
           (make-thread
            (lambda ()
              (let loop ()
                (debug "Listener thread calling accept")
                (let* ((cs
                        (socket-accept s))
                       (handler-thread
                        (make-thread
                         (lambda ()
                           ;; Connection setup handler thread
                           ;; Hands over to connection thread when make-connection is called
                           (when bind-port
                                 (set! (tcp-no-delay? cs) #t))
                           (receive
                            (input output)
                            (parameterize
                             ((socket-send-buffer-size 4096)
                              (socket-send-size 16384))
                             (socket-i/o-ports cs))
                            (let* ((peer (socket-peer-name cs))
                                   (host (if bind-port
                                             (sockaddr-address peer)
                                             #f))
                                   (port (if bind-port
                                             (sockaddr-port peer)
                                             #f)))
                              (debug "Handshake thread started for ~a:~a" host port)
                              (if user->key
                                  ;; Encrypted connection
                                  (let* ((user-name (read-packet! input))
                                         (encrypted-session-key (read-packet! input)))
                                    (if (and (not (eof-object? user-name))
                                             (not (eof-object? encrypted-session-key)))
                                        (let ((user-key (user->key user-name)))
                                          (if user-key
                                              (let* ((session-key-bytes (decrypt user-key encrypted-session-key))
                                                     (session-key (bytes->key session-key-bytes)))
                                                (if session-key
                                                    (begin
                                                      (debug "Writing header")
                                                      (write-string *header* #f output)
                                                      (flush-output output)
                                                      (open-handler
                                                       (make-connection
                                                        host port user-name
                                                        cs input output
                                                        session-key
                                                        request-handler close-handler)))
                                                    (begin
                                                      (debug "Rejecting connection due to invalid session key")
                                                      (close-output-port output)
                                                      (close-input-port input))))
                                              (begin
                                                (debug "Rejecting connection due to unknown user")
                                                (close-output-port output)
                                                (close-input-port input))))
                                        (begin
                                          (debug "Rejecting connection due to EOF in session setup")
                                          (close-output-port output)
                                          (close-input-port input))))
                                  ;; Plaintext connection
                                  (begin
                                    (debug "Writing header")
                                    (write-string *header* #f output)
                                    (flush-output output)
                                    (open-handler
                                     (make-connection
                                      host port #f
                                      cs input output
                                      #f
                                      request-handler close-handler))))))
                           cs)
                         `(bokbok-connection-handler ,bind-addr ,bind-port))))
                  (debug "Listener thread starting handshake thread")
                  (thread-start! handler-thread)
                  (loop))))
            `(bokbok-listen-handler ,bind-addr ,bind-port))))
     (thread-specific-set! thread s)
     (thread-start! thread)
     thread))

 (define (stop-server! server)
   (thread-terminate! server))

 (define (wait-until-server-stopped server)
   (let ((s (thread-specific server)))
    (handle-exceptions
     exc
     (if (and
          (uncaught-exception? exc)
          (terminated-thread-exception? (uncaught-exception-reason exc)))
         (void)
         (abort exc))             ; unexpected exceptions are rethrown
     (debug "Waiting for listener thread death...")
     (thread-join! server)
     (void))
    (socket-close s))))

(use posix)

(define *version* "0.1")

(define (newer file1 file2)
  (or
   (not (get-environment-variable "FAST_BUILD"))
   (not (file-exists? file2))
   (> (file-modification-time file1)
      (file-modification-time file2))))

(define (build-module name)
  (let ((source-file (string-append name ".scm"))
        (so-file (string-append name ".so"))
        (import-file (string-append name ".import.scm"))
        (import-so-file (string-append name ".import.so"))
        (o-file (string-append name ".o")))

    (when (newer source-file so-file)
          (compile -s -optimize-level 3 -debug-level 2 ,(string->symbol source-file) -j ,(string->symbol name))
          (compile -s -optimize-level 3 -debug-level 2 ,(string->symbol import-file))
          (compile -c -optimize-level 3 -debug-level 2 ,(string->symbol source-file) -unit ,(string->symbol name)))

    (install-extension (string->symbol name) `(,so-file ,o-file ,import-so-file)
                       `((version ,*version*)
                         (static o-file)))))

(define (build-program name)
  (let ((source-file (string-append name ".scm"))
        (exec-file name))
   (when (newer source-file exec-file)
         (compile -optimize-level 3 -debug-level 2 ,(string->symbol source-file)))
   (install-program (string->symbol name) exec-file
                    `((version ,*version*)))))


(build-module "bokbok")
(build-program "bokbok-demo")