#lang racket/base
;;
;;; Copyright 2010, 2011, 2012, 2013 Tony Garnock-Jones <tonyg@ccs.neu.edu>
;;;
;;; This file is part of marketplace-ssh.
;;;
;;; marketplace-ssh is free software: you can redistribute it and/or
;;; modify it under the terms of the GNU General Public License as
;;; published by the Free Software Foundation, either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; marketplace-ssh is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with marketplace-ssh. If not, see
;;; <http://www.gnu.org/licenses/>.

(require (planet tonyg/bitsyntax))
(require (planet vyzo/crypto:2:3))

(require racket/set)
(require racket/match)

(require "oakley-groups.rkt")

(require "ssh-host-key.rkt")

(require "ssh-numbers.rkt")
(require "ssh-message-types.rkt")
(require "ssh-exceptions.rkt")
(require "ssh-transport.rkt")
(require "ssh-channel.rkt")

(require "marketplace-support.rkt")

(provide rekey-interval
	 rekey-volume
	 ssh-session)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Data definitions
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; A RekeyState is one of
;; - a (rekey-wait Number Number), representing a time or
;;   transfer-amount by which rekeying should be started
;; - a (rekey-local SshMsgKexinit), when we've sent our local
;;   algorithm list and are waiting for the other party to send theirs
;; - a (rekey-in-progress KeyExchangeState), when both our local
;;   algorithm list has been sent and the remote one has arrived and the
;;   actual key exchange has begun
(struct rekey-wait (deadline threshold-bytes) #:transparent)
(struct rekey-local (local-algorithms) #:transparent)
(struct rekey-in-progress (state) #:transparent)

;; An AuthenticationState is one of
;; - #f, for not-yet-authenticated
;; - an (authenticated String String), recording successful completion
;;   of the authentication protocol after a request to be identified
;;   as the given username for the given service.
;; TODO: When authentication is properly implemented, we will need
;; intermediate states here too.
(struct authenticated (username service) #:transparent)

;; A PacketDispatcher is a Hashtable mapping Byte to PacketHandler.

;; A PacketHandler is a (Bytes DecodedPacket ConnectionState -> Transition).
;; The raw received bytes of the packet are given because sometimes
;; cryptographic operations on the received bytes are mandated by the
;; protocol.

;; TODO: Remove dispatch-table in favour of using the os2 subscription
;; mechanism to dispatch packets. I could do this now, but I'd lose
;; SSH_MSG_UNIMPLEMENTED support: I would need to be able to query the
;; current routing table to see whether there was an active listener
;; ready to take a given packet.

;; A ConnectionState is a (connection ... TODO fix this) representing
;; the complete state of the SSH transport, authentication, and
;; connection layers.
(struct connection (discard-next-packet?
		    dispatch-table
		    total-transferred
		    rekey-state
		    authentication-state
		    channels ;; ListOf<ChannelState>
		    is-server?
		    local-id
		    remote-id
		    session-id ;; starts off #f until initial keying
		    application-boot) ;; used when authentication completes
	#:transparent)

;; Generic inputs into the exchange-hash part of key
;; exchange. Diffie-Hellman uses these fields along with the host key,
;; the exchange values, and the shared secret to get the final hash.
(struct exchange-hash-info (client-id
			    server-id
			    client-kexinit-bytes
			    server-kexinit-bytes)
	#:transparent)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Parameters
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define rekey-interval (make-parameter 3600))
(define rekey-volume (make-parameter 1000000000))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Packet dispatch and handling
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; Bytes -> Byte
;; Retrieves the packet type byte from a packet.
(define (encoded-packet-msg-type encoded-packet)
  (bytes-ref encoded-packet 0))

;; PacketDispatcher [ Byte Maybe<PacketHandler> ]* -> PacketDispatcher
;; Adds or removes handlers to or from the given PacketDispatcher.
(define (extend-packet-dispatcher core-dispatcher . key-value-pairs)
  (let loop ((d core-dispatcher)
	     (key-value-pairs key-value-pairs))
    (cond
     ((null? key-value-pairs)
      d)
     ((null? (cdr key-value-pairs))
      (error 'extend-packet-dispatcher
	     "Must call extend-packet-dispatcher with matched key/value pairs"))
     (else
      (loop (let ((packet-type-number (car key-value-pairs))
		  (packet-handler-or-false (cadr key-value-pairs)))
	      (if packet-handler-or-false
		  (hash-set d packet-type-number packet-handler-or-false)
		  (hash-remove d packet-type-number)))
	    (cddr key-value-pairs))))))

;; ConnectionState [ Byte Maybe<PacketHandler> ]* -> ConnectionState
;; Installs (or removes) PacketHandlers in the given connection state;
;; see extend-packet-dispatcher.
(define (set-handlers conn . key-value-pairs)
  (struct-copy connection conn
	       [dispatch-table (apply extend-packet-dispatcher
				      (connection-dispatch-table conn)
				      key-value-pairs)]))

;; Transition Byte PacketHandler -> ConnectionState
;; Installs a PacketHandler that removes the installed dispatch entry
;; and then delegates to its argument.
(define (oneshot-handler conn packet-type-number packet-handler)
  (set-handlers conn
		packet-type-number
		(lambda (packet message conn)
		  (packet-handler packet
				  message
				  (set-handlers conn packet-type-number #f)))))

(define (dispatch-packet seq packet message conn)
  (define packet-type-number (encoded-packet-msg-type packet))
  (if (and (not (rekey-wait? (connection-rekey-state conn)))
	   (or (not (ssh-msg-type-transport-layer? packet-type-number))
	       (= packet-type-number SSH_MSG_SERVICE_REQUEST)
	       (= packet-type-number SSH_MSG_SERVICE_ACCEPT)))
      ;; We're in the middle of some phase of an active key-exchange,
      ;; and received a packet that's for a higher layer than the
      ;; transport layer, or one of the forbidden types given at the
      ;; send of RFC4253 section 7.1.
      (disconnect-with-error SSH_DISCONNECT_PROTOCOL_ERROR
			     "Packets of type ~v forbidden while in key-exchange"
			     packet-type-number)
      ;; We're either idling, or it's a permitted packet type while
      ;; performing key exchange. Look it up in the dispatch table.
      (let ((handler (hash-ref (connection-dispatch-table conn)
			       packet-type-number
			       #f)))
	(if handler
	    (handler packet message conn)
	    (transition conn
	      (send-message (outbound-packet (ssh-msg-unimplemented seq))))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Handlers for core transport packet types
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; PacketHandler for handling SSH_MSG_DISCONNECT.
(define (handle-msg-disconnect packet message conn)
  (disconnect-with-error* #t
			  '()
			  (ssh-msg-disconnect-reason-code message)
			  "Received SSH_MSG_DISCONNECT with reason code ~a and message ~s"
			  (ssh-msg-disconnect-reason-code message)
			  (bytes->string/utf-8 (bit-string->bytes
						(ssh-msg-disconnect-description message)))))

;; PacketHandler for handling SSH_MSG_IGNORE.
(define (handle-msg-ignore packet message conn)
  (transition conn))

;; PacketHandler for handling SSH_MSG_UNIMPLEMENTED.
(define (handle-msg-unimplemented packet message conn)
  (disconnect-with-error/local-info
   `((offending-sequence-number ,(ssh-msg-unimplemented-sequence-number message)))
   SSH_DISCONNECT_PROTOCOL_ERROR
   "Disconnecting because of received SSH_MSG_UNIMPLEMENTED."))

;; PacketHandler for handling SSH_MSG_DEBUG.
(define (handle-msg-debug packet message conn)
  (log-debug (format "Received SSHv2 SSH_MSG_DEBUG packet ~v" message))
  (transition conn))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Key Exchange
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (rekey-in-seconds-or-bytes delta-seconds delta-bytes total-transferred)
  (rekey-wait (+ (current-seconds) delta-seconds)
	      (+ total-transferred delta-bytes)))

(define (time-to-rekey? rekey conn)
  (and (rekey-wait? rekey)
       (or (>= (current-seconds) (rekey-wait-deadline rekey))
	   (>= (connection-total-transferred conn) (rekey-wait-threshold-bytes rekey)))))

;; (SshMsgKexinit -> Symbol) SshMsgKexinit SshMsgKexinit -> Symbol
;; Computes the name of the "best" algorithm choice at the given
;; getter, using the rules from the RFC and the client and server
;; algorithm precedence lists.
(define (best-result getter client-algs server-algs)
  (define client-list0 (getter client-algs))
  (define server-list (getter server-algs))
  (let loop ((client-list client-list0))
    (cond
     ((null? client-list) (disconnect-with-error/local-info
			   `((client-list ,client-list0)
			     (server-list ,server-list))
			   SSH_DISCONNECT_KEY_EXCHANGE_FAILED
			   "Could not agree on a suitable algorithm for ~v"
			   getter))
     ((memq (car client-list) server-list) (car client-list))
     (else (loop (cdr client-list))))))

;; ExchangeHashInfo Bytes Natural Natural Natural -> Bytes
;; Computes the session ID as defined by SSH's DH key exchange method.
(define (dh-exchange-hash hash-info host-key e f k)
  (let ((block-to-hash
	 (bit-string->bytes
	  (bit-string ((exchange-hash-info-client-id hash-info) :: (t:string))
		      ((exchange-hash-info-server-id hash-info) :: (t:string))
		      ((exchange-hash-info-client-kexinit-bytes hash-info) :: (t:string))
		      ((exchange-hash-info-server-kexinit-bytes hash-info) :: (t:string))
		      (host-key :: (t:string))
		      (e :: (t:mpint))
		      (f :: (t:mpint))
		      (k :: (t:mpint))))))
    (sha1 block-to-hash)))

;; ExchangeHashInfo Symbol Symbol ConnectionState
;;   (Bytes Bytes Symbol ConnectionState -> ConnectionState)
;;   -> Transition
;; Performs the server's half of the Diffie-Hellman key exchange protocol.
(define (perform-server-key-exchange hash-info kex-alg host-key-alg conn finish)
  (case kex-alg
    [(diffie-hellman-group14-sha1 diffie-hellman-group1-sha1)
     (define group (if (eq? kex-alg 'diffie-hellman-group14-sha1)
		       dh:oakley-group-14
		       dh:oakley-group-2)) ;; yes, SSH's group1 == Oakley/RFC2409 group 2
     (define-values (private-key public-key) (generate-key group))
     (define public-key-as-integer (bit-string->integer public-key #t #f))
     (transition
	 (oneshot-handler conn
			  SSH_MSG_KEXDH_INIT
			  (lambda (packet message conn)
			    (define e (ssh-msg-kexdh-init-e message))
			    (define e-width (mpint-width e))
			    (define e-as-bytes (integer->bit-string e (* 8 e-width) #t))
			    (define shared-secret (compute-key private-key e-as-bytes))
			    (define hash-alg sha1)
			    (define-values (host-key-private host-key-public)
			      (host-key-algorithm->keys host-key-alg))
			    (define host-key-bytes
			      (pieces->ssh-host-key (public-key->pieces host-key-public)))
			    (define exchange-hash
			      (dh-exchange-hash hash-info
						host-key-bytes
						e
						public-key-as-integer
						(bit-string->integer shared-secret #t #f)))
			    (define h-signature (host-key-signature host-key-private
								    host-key-alg
								    exchange-hash))
			    (sequence-actions (transition conn)
			      (send-message (outbound-packet
					     (ssh-msg-kexdh-reply (bit-string->bytes host-key-bytes)
								  public-key-as-integer
								  (bit-string->bytes h-signature))))
			      (lambda (conn)
				(finish shared-secret exchange-hash hash-alg conn))))))]
    [else (disconnect-with-error SSH_DISCONNECT_KEY_EXCHANGE_FAILED
				 "Bad key-exchange algorithm ~v" kex-alg)]))

;; ExchangeHashInfo Symbol Symbol ConnectionState
;;   (Bytes Bytes Symbol ConnectionState -> ConnectionState)
;;   -> Transition
;; Performs the client's half of the Diffie-Hellman key exchange protocol.
(define (perform-client-key-exchange hash-info kex-alg host-key-alg conn finish)
  (case kex-alg
    [(diffie-hellman-group14-sha1 diffie-hellman-group1-sha1)
     (define group (if (eq? kex-alg 'diffie-hellman-group14-sha1)
		       dh:oakley-group-14
		       dh:oakley-group-2)) ;; yes, SSH's group1 == Oakley/RFC2409 group 2
     (define-values (private-key public-key) (generate-key group))
     (define public-key-as-integer (bit-string->integer public-key #t #f))
     (sequence-actions (transition conn)
       (send-message (outbound-packet (ssh-msg-kexdh-init public-key-as-integer)))
       (lambda (conn)
	 (transition
	     (oneshot-handler conn
			      SSH_MSG_KEXDH_REPLY
			      (lambda (packet message conn)
				(define f (ssh-msg-kexdh-reply-f message))
				(define f-width (mpint-width f))
				(define f-as-bytes (integer->bit-string f (* 8 f-width) #t))
				(define shared-secret (compute-key private-key f-as-bytes))
				(define hash-alg sha1)
				(define host-key-bytes (ssh-msg-kexdh-reply-host-key message))
				(define host-public-key
				  (pieces->public-key (ssh-host-key->pieces host-key-bytes)))
				(define exchange-hash
				  (dh-exchange-hash hash-info
						    host-key-bytes
						    public-key-as-integer
						    f
						    (bit-string->integer shared-secret #t #f)))
				(verify-host-key-signature! host-public-key
							    host-key-alg
							    exchange-hash
							    (ssh-msg-kexdh-reply-h-signature
							     message))
				(finish shared-secret exchange-hash hash-alg conn))))))]
    [else (disconnect-with-error SSH_DISCONNECT_KEY_EXCHANGE_FAILED
				 "Bad key-exchange algorithm ~v" kex-alg)]))

;; PacketHandler for handling SSH_MSG_KEXINIT.
(define (handle-msg-kexinit packet message conn)
  (define rekey (connection-rekey-state conn))
  (when (rekey-in-progress? rekey)
    (disconnect-with-error SSH_DISCONNECT_PROTOCOL_ERROR
			   "Received SSH_MSG_KEXINIT during ongoing key exchange"))
  (define local-algs (if (rekey-local? rekey)
			 (rekey-local-local-algorithms rekey)
			 ((local-algorithm-list))))
  (define encoded-local-algs (ssh-message-encode local-algs))
  (define remote-algs message)
  (define encoded-remote-algs packet)

  (define is-server? (connection-is-server? conn))
  (define c (if is-server? remote-algs local-algs))
  (define s (if is-server? local-algs remote-algs))

  (define kex-alg (best-result ssh-msg-kexinit-kex_algorithms c s))
  (define host-key-alg (best-result ssh-msg-kexinit-server_host_key_algorithms c s))
  (define c2s-enc (best-result ssh-msg-kexinit-encryption_algorithms_client_to_server c s))
  (define s2c-enc (best-result ssh-msg-kexinit-encryption_algorithms_server_to_client c s))
  (define c2s-mac (best-result ssh-msg-kexinit-mac_algorithms_client_to_server c s))
  (define s2c-mac (best-result ssh-msg-kexinit-mac_algorithms_server_to_client c s))
  (define c2s-zip (best-result ssh-msg-kexinit-compression_algorithms_client_to_server c s))
  (define s2c-zip (best-result ssh-msg-kexinit-compression_algorithms_server_to_client c s))
  ;; Ignore languages.
  ;; Don't check the reserved field here, either. TODO: should we?

  (define (guess-matches? chosen-value getter)
    (let ((remote-choices (getter remote-algs)))
      (and (pair? remote-choices) ;; not strictly necessary because of
				  ;; the error behaviour of
				  ;; best-result.
	   (eq? (car remote-choices) ;; the remote peer's guess for this parameter
		chosen-value))))

  (define should-discard-first-kex-packet
    (and (ssh-msg-kexinit-first_kex_packet_follows remote-algs)
	 ;; They've already transmitted their guess. Does their guess match
	 ;; what we've actually selected?
	 (not (and
	       (guess-matches? kex-alg ssh-msg-kexinit-kex_algorithms)
	       (guess-matches? host-key-alg ssh-msg-kexinit-server_host_key_algorithms)
	       (guess-matches? c2s-enc ssh-msg-kexinit-encryption_algorithms_client_to_server)
	       (guess-matches? s2c-enc ssh-msg-kexinit-encryption_algorithms_server_to_client)
	       (guess-matches? c2s-mac ssh-msg-kexinit-mac_algorithms_client_to_server)
	       (guess-matches? s2c-mac ssh-msg-kexinit-mac_algorithms_server_to_client)
	       (guess-matches? c2s-zip ssh-msg-kexinit-compression_algorithms_client_to_server)
	       (guess-matches? s2c-zip ssh-msg-kexinit-compression_algorithms_server_to_client)))))

  (define (continue-after-discard conn)
    ((if is-server?
	 perform-server-key-exchange
	 perform-client-key-exchange)
     (if is-server?
	 (exchange-hash-info (connection-remote-id conn)
			     (connection-local-id conn)
			     encoded-remote-algs
			     encoded-local-algs)
	 (exchange-hash-info (connection-local-id conn)
			     (connection-remote-id conn)
			     encoded-local-algs
			     encoded-remote-algs))
     kex-alg
     host-key-alg
     conn
     continue-after-key-exchange))

  (define (continue-after-key-exchange shared-secret exchange-hash hash-alg conn)
    (define session-id (if (connection-session-id conn)
			   (connection-session-id conn) ;; don't overwrite existing ID
			   exchange-hash))
    (define k-h-prefix (bit-string ((bit-string->integer shared-secret #t #f) :: (t:mpint))
				   (exchange-hash :: binary)))
    (define (derive-key kind needed-bytes-or-false)
      (let extend ((key (hash-alg (bit-string->bytes
				   (bit-string (k-h-prefix :: binary)
					       (kind :: binary)
					       (session-id :: binary))))))
	(cond
	 ((eq? #f needed-bytes-or-false)
	  key)
	 ((>= (bytes-length key) needed-bytes-or-false)
	  (subbytes key 0 needed-bytes-or-false))
	 (else
	  (extend (bytes-append key (hash-alg (bit-string->bytes
					       (bit-string (k-h-prefix :: binary)
							   (key :: binary))))))))))
    (transition
	(oneshot-handler (struct-copy connection conn
				      [session-id session-id]) ;; just in case it changed
			 SSH_MSG_NEWKEYS
			 (lambda (newkeys-packet newkeys-message conn)
			   ;; First, send our SSH_MSG_NEWKEYS,
			   ;; incrementing the various counters, and then
			   ;; apply the new algorithms. Also arm our rekey
			   ;; timer.
			   (define new-rekey-state (rekey-in-seconds-or-bytes
						    (rekey-interval)
						    (rekey-volume)
						    (connection-total-transferred conn)))
			   (transition
			       (set-handlers
				(struct-copy connection conn [rekey-state new-rekey-state])
				SSH_MSG_SERVICE_REQUEST handle-msg-service-request)
			     (send-message (outbound-packet (ssh-msg-newkeys)))
			     (send-message
			      (new-keys (connection-is-server? conn)
					derive-key
					c2s-enc s2c-enc
					c2s-mac s2c-mac
					c2s-zip s2c-zip))
			     (send-message (set-timer 'rekey-timer
						      (* (rekey-wait-deadline new-rekey-state) 1000)
						      'absolute)))))))

  (sequence-actions (continue-after-discard conn)
    (when should-discard-first-kex-packet
      (lambda (conn) (transition (struct-copy connection conn [discard-next-packet? #t]))))
    (lambda (conn)
      (if (rekey-wait? (connection-rekey-state conn))
	  (transition (struct-copy connection conn [rekey-state (rekey-local local-algs)])
	    (send-message (outbound-packet local-algs)))
	  (transition conn)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Service request manager
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (handle-msg-service-request packet message conn)
  (define service (bit-string->bytes (ssh-msg-service-request-service-name message)))
  (match service
    [#"ssh-userauth"
     (if (connection-authentication-state conn)
	 (disconnect-with-error SSH_DISCONNECT_SERVICE_NOT_AVAILABLE
				"Repeated authentication is not permitted")
	 (sequence-actions (transition conn)
	   (send-message (outbound-packet (ssh-msg-service-accept service)))
	   (lambda (conn) (transition
			      (oneshot-handler conn
					       SSH_MSG_USERAUTH_REQUEST
					       handle-msg-userauth-request)))))]
    [else
     (disconnect-with-error SSH_DISCONNECT_SERVICE_NOT_AVAILABLE
			    "Service ~v not supported"
			    service)]))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; User authentication
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (handle-msg-userauth-request packet message conn)
  (define user-name (bit-string->bytes (ssh-msg-userauth-request-user-name message)))
  (define service-name (bit-string->bytes (ssh-msg-userauth-request-service-name message)))
  (cond
   [(and (positive? (bytes-length user-name))
	 (equal? service-name #"ssh-connection"))
    ;; TODO: Actually implement client authentication
    (sequence-actions (transition conn)
      (send-message (outbound-packet (ssh-msg-userauth-success)))
      (lambda (conn)
	(start-connection-service
	 (set-handlers (struct-copy connection conn
				    [authentication-state (authenticated user-name service-name)])
		       SSH_MSG_USERAUTH_REQUEST
		       (lambda (packet message conn)
			 ;; RFC4252 section 5.1 page 6
			 conn))))
      (lambda (conn)
	(transition conn
	  ;; TODO: canary for NESTED VM!: #:exit-signal? #t
	  (spawn-vm #:debug-name 'ssh-application-vm
		    ((connection-application-boot conn) user-name)))))]
   [else
    (transition conn
      (send-message (outbound-packet (ssh-msg-userauth-failure '(none) #f))))]))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Channel management
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (unused-local-channel-ref conn)
  (define (bump-candidate candidate)
    (modulo (+ candidate 1) #x100000000))
  (define first-candidate (match (connection-channels conn)
			    ['() 0]
			    [(cons ch _) (bump-candidate (ssh-channel-local-ref ch))]))
  (let examine-candidate ((candidate first-candidate))
    (let loop ((chs (connection-channels conn)))
      (cond
       [(null? chs) candidate]
       [(= (ssh-channel-local-ref (car chs)) candidate)
	(examine-candidate (bump-candidate candidate))]
       [else (loop (cdr chs))]))))

(define (replacef proc updater creator lst)
  (let loop ((lst lst))
    (cond [(null? lst) (list (creator))]
	  [(proc (car lst)) (cons (updater (car lst)) (cdr lst))]
	  [else (cons (car lst) (loop (cdr lst)))])))

(define (remf proc lst)
  (cond [(null? lst) '()]
	[(proc (car lst)) (cdr lst)]
	[else (cons (car lst) (remf proc (cdr lst)))]))

;; ChannelName -> ChannelState -> Boolean
(define ((ssh-channel-name=? cname) c)
  (equal? (ssh-channel-name c) cname))

;; Connection Uint32 -> ChannelState
(define (get-channel conn local-ref)
  (define ch (findf (lambda (c) (equal? (ssh-channel-local-ref c) local-ref))
		    (connection-channels conn)))
  (when (not ch)
    (disconnect-with-error SSH_DISCONNECT_PROTOCOL_ERROR
			   "Attempt to use known channel local-ref ~v"
			   local-ref))
  ch)

;; ChannelName Maybe<Uint32> Connection -> Connection
(define (update-channel cname updater conn)
  (struct-copy connection conn
   [channels
    (replacef (ssh-channel-name=? cname)
	      updater
	      (lambda () (updater (ssh-channel cname
					       (unused-local-channel-ref conn)
					       #f
					       #f
					       'neither)))
	      (connection-channels conn))]))

;; ChannelName Connection -> Connection
(define (discard-channel cname conn)
  (struct-copy connection conn
   [channels
    (remf (ssh-channel-name=? cname) (connection-channels conn))]))

;; CloseState Either<'local,'remote> -> CloseState
(define (update-close-state old-state action)
  (define local? (case action ((local) #t) ((remote) #f)))
  (case old-state
    ((neither) (if local? 'local 'remote))
    ((local) (if local? 'local 'both))
    ((remote) (if local? 'both 'remote))
    ((both) 'both)))

(define (maybe-close-channel cname conn action)
  (cond
   [(findf (ssh-channel-name=? cname) (connection-channels conn)) =>
    (lambda (ch)
      (define old-close-state (ssh-channel-close-state ch))
      (define new-close-state (update-close-state old-close-state action))
      (transition (if (eq? new-close-state 'both)
		      (discard-channel ch conn)
		      (update-channel cname
				      (lambda (ch)
					(struct-copy ssh-channel ch
						     [close-state new-close-state]))
				      conn))
	(case action
	  [(local)
	   (case old-close-state
	     [(neither remote)
	      (list (send-message (outbound-packet
				   (ssh-msg-channel-close (ssh-channel-remote-ref ch)))))]
	     [else (list)])]
	  [(remote)
	   (case old-close-state
	     [(neither local)
	      (list (delete-endpoint (list cname 'outbound))
		    (delete-endpoint (list cname 'inbound)))]
	     [else (list)])])))]
   [else (transition conn)]))

(define (channel-endpoints cname initial-message-producer)
  (define inbound-stream-name (channel-stream-name #t cname))
  (define outbound-stream-name (channel-stream-name #f cname))
  (define (! conn message)
    (transition conn (send-message (outbound-packet message))))
  (list
   (name-endpoint (list cname 'outbound)
     (subscriber (channel-message outbound-stream-name (wild))
       (match-state conn
	 (on-presence (transition conn
			(initial-message-producer inbound-stream-name outbound-stream-name)))
	 (on-absence (maybe-close-channel cname conn 'local))
	 (on-message
	  [(channel-message _ body)
	   (let ()
	     (define ch (findf (ssh-channel-name=? cname) (connection-channels conn)))
	     (define remote-ref (ssh-channel-remote-ref ch))
	     (match body
	       [(channel-stream-data data-bytes)
		;; TODO: split data-bytes into packets if longer than max packet size
		(! conn (ssh-msg-channel-data remote-ref data-bytes))]
	       [(channel-stream-extended-data type data-bytes)
		(! conn (ssh-msg-channel-extended-data remote-ref type data-bytes))]
	       [(channel-stream-eof)
		(! conn (ssh-msg-channel-eof remote-ref))]
	       [(channel-stream-notify type data-bytes)
		(! conn (ssh-msg-channel-request remote-ref type #f data-bytes))]
	       [(channel-stream-request type data-bytes)
		(! conn (ssh-msg-channel-request remote-ref type #t data-bytes))]
	       [(channel-stream-open-failure reason description)
		(! (discard-channel cname conn)
		   (ssh-msg-channel-open-failure remote-ref reason description #""))]))]))))
   (name-endpoint (list cname 'inbound)
     (publisher (channel-message inbound-stream-name (wild))
       (match-state conn
	 (on-message
	  [(channel-message _ body)
	   (let ()
	     (define ch (findf (ssh-channel-name=? cname) (connection-channels conn)))
	     (define remote-ref (ssh-channel-remote-ref ch))
	     (match body
	       [(channel-stream-config maximum-packet-size extra-data)
		(if (channel-name-locally-originated? cname)
		    ;; This must be intended to form the SSH_MSG_CHANNEL_OPEN.
		    (! conn (ssh-msg-channel-open (channel-name-type cname)
						  (ssh-channel-local-ref ch)
						  0
						  maximum-packet-size
						  extra-data))
		    ;; This must be intended to form the SSH_MSG_CHANNEL_OPEN_CONFIRMATION.
		    (! conn (ssh-msg-channel-open-confirmation remote-ref
							       (ssh-channel-local-ref ch)
							       0
							       maximum-packet-size
							       extra-data)))]
	       [(channel-stream-credit count)
		(! conn (ssh-msg-channel-window-adjust remote-ref count))]
	       [(channel-stream-ok)
		(! conn (ssh-msg-channel-success remote-ref))]
	       [(channel-stream-fail)
		(! conn (ssh-msg-channel-failure remote-ref))]))]))))))

(define (channel-notify conn ch inbound? body)
  (transition conn
    (send-message (channel-message (channel-stream-name inbound? (ssh-channel-name ch))
				   body)
		  (if inbound? 'publisher 'subscriber))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Connection service
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (respond-to-opened-outbound-channel conn cname)
  (if (and (ground? cname)
	   (not (memf (ssh-channel-name=? cname) (connection-channels conn))))
      (transition (update-channel cname values conn)
	(channel-endpoints cname (lambda (inbound-stream outbound-stream)
				   '())))
      (transition conn)))

(define (start-connection-service conn)
  (sequence-actions
      (transition
	  (set-handlers conn
			;; TODO: SSH_MSG_GLOBAL_REQUEST		handle-msg-global-request
			SSH_MSG_CHANNEL_OPEN			handle-msg-channel-open
			SSH_MSG_CHANNEL_OPEN_CONFIRMATION	handle-msg-channel-open-confirmation
			SSH_MSG_CHANNEL_OPEN_FAILURE		handle-msg-channel-open-failure
			SSH_MSG_CHANNEL_WINDOW_ADJUST		handle-msg-channel-window-adjust
			SSH_MSG_CHANNEL_DATA			handle-msg-channel-data
			SSH_MSG_CHANNEL_EXTENDED_DATA		handle-msg-channel-extended-data
			SSH_MSG_CHANNEL_EOF			handle-msg-channel-eof
			SSH_MSG_CHANNEL_CLOSE			handle-msg-channel-close
			SSH_MSG_CHANNEL_REQUEST			handle-msg-channel-request
			SSH_MSG_CHANNEL_SUCCESS			handle-msg-channel-success
			SSH_MSG_CHANNEL_FAILURE			handle-msg-channel-failure))
    ;; Start responding to channel interest coming from the
    ;; application. We are responding to channels appearing from the
    ;; remote peer by virtue of our installation of the handler for
    ;; SSH_MSG_CHANNEL_OPEN above.
    (observe-subscribers (channel-message (channel-stream-name ? (channel-name #t ? ?)) ?)
      (match-state conn
	(match-conversation (channel-message (channel-stream-name #t cname) _)
	  (on-presence (respond-to-opened-outbound-channel conn cname)))))
    (observe-publishers (channel-message (channel-stream-name ? (channel-name #t ? ?)) ?)
      (match-state conn
	(match-conversation (channel-message (channel-stream-name #f cname) _)
	  (on-presence (respond-to-opened-outbound-channel conn cname)))))))

(define (handle-msg-channel-open packet message conn)
  (match-define (ssh-msg-channel-open channel-type*
				      remote-ref
				      initial-window-size
				      maximum-packet-size
				      extra-request-data*)
    message)

  (when (memf (lambda (c) (equal? (ssh-channel-remote-ref c) remote-ref))
	      (connection-channels conn))
    (disconnect-with-error SSH_DISCONNECT_PROTOCOL_ERROR
			   "Attempt to open already-open channel ~v"
			   remote-ref))

  (define channel-type (bit-string->bytes channel-type*))
  (define extra-request-data (bit-string->bytes extra-request-data*))
  (define cname (channel-name #f channel-type remote-ref))

  (transition (update-channel cname
			      (lambda (e) (struct-copy ssh-channel e [remote-ref remote-ref]))
			      conn)
    (channel-endpoints cname
		       (lambda (inbound-stream outbound-stream)
			 (list (send-feedback
				(channel-message outbound-stream
						 (channel-stream-config maximum-packet-size
									extra-request-data)))
			       (send-feedback
				(channel-message outbound-stream
						 (channel-stream-credit initial-window-size))))))))

(define (handle-msg-channel-open-confirmation packet message conn)
  (match-define (ssh-msg-channel-open-confirmation local-ref
						   remote-ref
						   initial-window-size
						   maximum-packet-size
						   extra-request-data*)
    message)
  (define ch (get-channel conn local-ref))
  (define extra-request-data (bit-string->bytes extra-request-data*))
  (define outbound-stream (channel-stream-name #f (ssh-channel-name ch)))
  (transition (update-channel (ssh-channel-name ch)
			      (lambda (c)
				(struct-copy ssh-channel c
					     [remote-ref remote-ref]
					     [outbound-packet-size maximum-packet-size]))
			      conn)
    (send-feedback (channel-message outbound-stream
				    (channel-stream-config maximum-packet-size
							   extra-request-data)))
    (send-feedback (channel-message outbound-stream
				    (channel-stream-credit initial-window-size)))))

(define (handle-msg-channel-open-failure packet message conn)
  (match-define (ssh-msg-channel-open-failure local-ref
					      reason
					      description*
					      _)
    message)
  (define ch (get-channel conn local-ref))
  (define description (bit-string->bytes description*))
  (define inbound-stream (channel-stream-name #t (ssh-channel-name ch)))
  (sequence-actions (transition conn)
    (send-message (channel-message inbound-stream
				   (channel-stream-open-failure reason description)))
    (lambda (conn) (maybe-close-channel (ssh-channel-name ch) conn 'remote))))

(define (handle-msg-channel-window-adjust packet message conn)
  (match-define (ssh-msg-channel-window-adjust local-ref count) message)
  (define ch (get-channel conn local-ref))
  (channel-notify conn ch #f (channel-stream-credit count)))

(define (handle-msg-channel-data packet message conn)
  (match-define (ssh-msg-channel-data local-ref data*) message)
  (define data (bit-string->bytes data*))
  (define ch (get-channel conn local-ref))
  (channel-notify conn ch #t (channel-stream-data data)))

(define (handle-msg-channel-extended-data packet message conn)
  (match-define (ssh-msg-channel-extended-data local-ref type-code data*) message)
  (define data (bit-string->bytes data*))
  (define ch (get-channel conn local-ref))
  (channel-notify conn ch #t (channel-stream-extended-data type-code data)))

(define (handle-msg-channel-eof packet message conn)
  (define ch (get-channel conn (ssh-msg-channel-eof-recipient-channel message)))
  (channel-notify conn ch #t (channel-stream-eof)))

(define (handle-msg-channel-close packet message conn)
  (define ch (get-channel conn (ssh-msg-channel-close-recipient-channel message)))
  (maybe-close-channel (ssh-channel-name ch) conn 'remote))

(define (handle-msg-channel-request packet message conn)
  (match-define (ssh-msg-channel-request local-ref type* want-reply? data*) message)
  (define type (bit-string->bytes type*))
  (define data (bit-string->bytes data*))
  (define ch (get-channel conn local-ref))
  (channel-notify conn ch #t
		  (if want-reply?
		      (channel-stream-request type data)
		      (channel-stream-notify type data))))

(define (handle-msg-channel-success packet message conn)
  (match-define (ssh-msg-channel-success local-ref) message)
  (define ch (get-channel conn local-ref))
  (channel-notify conn ch #f (channel-stream-ok)))

(define (handle-msg-channel-failure packet message conn)
  (match-define (ssh-msg-channel-failure local-ref) message)
  (define ch (get-channel conn local-ref))
  (channel-notify conn ch #f (channel-stream-fail)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Session main process
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (connection-username conn)
  (match (connection-authentication-state conn)
    ((authenticated username servicename)
     username)
    (else (disconnect-with-error SSH_DISCONNECT_PROTOCOL_ERROR
				 "Not authenticated"))))

(define ((bump-total amount) conn)
  (transition
      (struct-copy connection conn
		   [total-transferred (+ (connection-total-transferred conn) amount)])))

;; (K V A -> A) A Hash<K,V> -> A
(define (hash-fold fn seed hash)
  (do ((pos (hash-iterate-first hash) (hash-iterate-next hash pos))
       (seed seed (fn (hash-iterate-key hash pos) (hash-iterate-value hash pos) seed)))
      ((not pos) seed)))

(define (maybe-rekey conn)
  (define rekey (connection-rekey-state conn))
  (if (time-to-rekey? rekey conn)
      (let ((algs ((local-algorithm-list))))
	(transition (struct-copy connection conn [rekey-state (rekey-local algs)])
	  (send-message (outbound-packet algs))))
      (transition conn)))

;; PacketDispatcher. Handles the core transport message types.
(define base-packet-dispatcher
  (hasheq SSH_MSG_DISCONNECT handle-msg-disconnect
	  SSH_MSG_IGNORE handle-msg-ignore
	  SSH_MSG_UNIMPLEMENTED handle-msg-unimplemented
	  SSH_MSG_DEBUG handle-msg-debug
	  SSH_MSG_KEXINIT handle-msg-kexinit))

(define (ssh-session self-pid
		     local-identification-string
		     peer-identification-string
		     application-boot
		     session-role)
  (transition (connection #f
			  base-packet-dispatcher
			  0
			  (rekey-in-seconds-or-bytes -1 -1 0)
			  #f
			  '()
			  (case session-role ((client) #f) ((server) #t))
			  local-identification-string
			  peer-identification-string
			  #f
			  application-boot)

    (subscriber (timer-expired 'rekey-timer (wild))
      (match-state conn
	(on-message [(timer-expired 'rekey-timer now)
		     (sequence-actions (transition conn)
		       maybe-rekey)])))

    (subscriber (outbound-byte-credit (wild))
      (match-state conn
	(on-message [(outbound-byte-credit amount)
		     (sequence-actions (transition conn)
		       (bump-total amount)
		       maybe-rekey)])))

    (subscriber (inbound-packet (wild) (wild) (wild) (wild))
      (match-state conn
	(on-message
	 [(inbound-packet sequence-number payload message transfer-size)
	  (sequence-actions (transition conn)
	    (lambda (conn)
	      (if (connection-discard-next-packet? conn)
		  (transition
		      (struct-copy connection conn [discard-next-packet? #f]))
		  (dispatch-packet sequence-number payload message conn)))
	    (bump-total transfer-size)
	    (send-message (inbound-credit 1))
	    maybe-rekey)])))))