#lang syndicate
;;; SPDX-License-Identifier: LGPL-3.0-or-later
;;; SPDX-FileCopyrightText: Copyright © 2023-2024 Tony Garnock-Jones <tonyg@leastfixedpoint.com>

(provide noise-initiator
         noise-responder)

(require syndicate/relay)
(require syndicate/schemas/noise)
(require syndicate/schemas/gatekeeper)
(require noise-protocol)

(define-logger syndicate/distributed/noise)

(define (css->procedures css)
  (match-define (list S R) css)
  (values (lambda (bs) (S 'encrypt #"" bs))
          (lambda (bs) (R 'decrypt #"" bs))))

(define (extract-packets m)
  (match (parse-Packet! m)
    [(Packet-complete bs) (list bs)]
    [(Packet-fragmented bss) bss]))

(define PACKET-LIMIT (- 65535 16)) ;; 16 bytes for authentication tag

(define (chunk-bytes bs chunk-size)
  (define len (bytes-length bs))
  (for/list [(i (in-range 0 len chunk-size))]
    (subbytes bs i (min (+ i chunk-size) len))))

(define (fragment bs [f values])
  (if (> (bytes-length bs) PACKET-LIMIT)
      (Packet-fragmented (map f (chunk-bytes bs PACKET-LIMIT)))
      (Packet-complete (f bs))))

(define (noise-initiator #:service-selector service-selector
                         #:remote-static-pk remote-static-pk
                         #:acceptor-ref acceptor-ref
                         #:import import-handler
                         #:preshared-keys [psks #f]
                         #:pattern [pattern #f])
  (noise* #:role 'initiator
          #:service-selector service-selector
          #:remote-static-pk remote-static-pk
          #:import import-handler
          #:preshared-keys psks
          #:pattern pattern
          (lambda (set-peer-session! handle-message)
            (at acceptor-ref
              (assert (Resolve (Step (NoiseStepType) service-selector)
                               (object #:name 'noise-observer
                                       [#:asserted (Resolved-accepted responder-session)
                                        (at responder-session
                                          (assert (Initiator
                                                   (object #:name 'noise-initiator
                                                           [#:message m
                                                            (handle-message m)]))))
                                        (set-peer-session! responder-session)
                                        #:retracted
                                        (stop-current-facet)])))))))

(define (noise-responder #:service-selector service-selector
                         #:static-keypair static-keypair
                         #:export initial-ref
                         #:observer observer
                         #:preshared-keys [psks #f]
                         #:pattern [pattern #f])
  (noise* #:role 'responder
          #:service-selector service-selector
          #:static-keypair static-keypair
          #:export initial-ref
          #:preshared-keys psks
          #:pattern pattern
          (lambda (set-peer-session! handle-message)
            (at observer
              (assert (Resolved-accepted
                       (object #:name (list 'noise-responder initial-ref observer)
                               [#:asserted (Initiator s)
                                (set-peer-session! s)
                                #:retracted
                                (stop-current-facet)]
                               [#:message m
                                (handle-message m)])))))))

(define (noise* #:role role
                #:service-selector service-selector
                #:static-keypair [static-keypair #f]
                #:remote-static-pk [remote-static-pk #f]
                #:export [initial-ref #f]
                #:import [import-handler #f]
                #:preshared-keys [psks #f]
                #:pattern [pattern #f]
                f)
  (define H (Noise-*-25519_ChaChaPoly_BLAKE2s
             (or pattern "NK")
             #:role role
             #:prologue (encode/canonicalization service-selector)
             #:static-keypair static-keypair
             #:remote-static-pk remote-static-pk
             #:preshared-keys psks))

  (define encrypt! #f)
  (define decrypt! #f)
  (define relay #f)
  (define peer-session #f)
  (define buffered-inputs-rev '())

  (define (start-relay css)
    (set!-values (encrypt! decrypt!) (css->procedures css))
    (run-relay #:packet-writer (lambda (bs) (send! peer-session (fragment bs encrypt!)))
               #:setup-inputs (lambda (tr)
                                (set! relay tr)
                                (for [(m (in-list (reverse buffered-inputs-rev)))]
                                  (accept-bytes relay m))
                                (set! buffered-inputs-rev #f))
               #:then (and import-handler (object [(embedded a) (import-handler a)]))
               #:initial-oid (and import-handler 0)
               #:initial-ref initial-ref))

  (define (handshake-step)
    (define-values (packet css) (H 'write-message #""))
    (send! peer-session (Packet-complete packet))
    (when css (start-relay css)))

  (define (handle-message m)
    (log-syndicate/distributed/noise-debug "~v ~a got: ~v" service-selector role m)
    (for [(p (in-list (extract-packets m)))]
      (if relay
          (accept-bytes relay (decrypt! p))
          (let-values (((message css) (H 'read-message p)))
            (set! buffered-inputs-rev (cons message buffered-inputs-rev))
            (if css
                (start-relay css)
                (handshake-step))))))

  (define (set-peer-session! session)
    (when peer-session (error 'noise* "Double-setting of peer-session"))
    (set! peer-session session)
    (when (eq? role 'initiator) (handshake-step)))

  (f set-peer-session! handle-message))

(module+ test
  (require libsodium)

  (when (file-stream-buffer-mode (current-output-port))
    ;; ^ Only try to set it if we get a non-#f answer when we ask for it
    (file-stream-buffer-mode (current-output-port) 'none))

  (standard-actor-system (ds)
    (define server-keys (make-crypto-box-keypair))

    (spawn #:name 'test-responder
           (define service-object
             (object [a
                      (printf "service+: ~v\n" a)
                      (on-stop (printf "service-: ~v\n" a))]
                     [#:message m
                      (printf "service!: ~v\n" m)
                      (stop-current-facet)]))

           (at ds
             (during (Resolve (Step (NoiseStepType) 'test-service) $observer)
               (noise-responder #:service-selector 'test-service
                                #:static-keypair server-keys
                                #:observer observer
                                #:export service-object))))

    (spawn #:name 'test-initiator
           (noise-initiator #:service-selector 'test-service
                            #:remote-static-pk (crypto-box-keypair-pk server-keys)
                            #:acceptor-ref ds
                            #:import (lambda (r)
                                       (on-stop (printf "BYEEEE\n"))
                                       (at r (assert "HELLO!"))
                                       (send! r "HIIIII!!!!!"))))))