#lang racket/base

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

(require racket/match)
(require racket/class)
(require racket/port)

(require "safe-io.rkt")
(require "oakley-groups.rkt")

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

(require "functional-queue.rkt")
(require "conversation.rkt")
(require "standard-thread.rkt")
(require "ordered-rpc.rkt")

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

(provide required-peer-identification-regex
	 client-preamble-lines
	 client-identification-string
	 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 -> ConnectionState).
;; The raw received bytes of the packet are given because sometimes
;; cryptographic operations on the received bytes are mandated by the
;; protocol.

;; A ConnectionState is a (connection StreamState StreamState
;; PacketDispatcher ... TODO fix this) representing the complete state
;; of the SSH transport, authentication, and connection layers.
(struct connection (io-room-handle
		    session-room-handle
		    discard-next-packet?
		    dispatch-table
		    total-transferred
		    rekey-state
		    authentication-state
		    continuations
		    channel-map
		    is-server?
		    local-id
		    remote-id
		    session-id) ;; starts off #f until initial keying
	#:transparent)

;; A CloseState is one of
;; - 'neither, indicating that neither side has signalled closure
;; - 'local, only the local end has signalled closure
;; - 'remote, only the remote end has signalled closure
;; - 'both, both ends have signalled closure.
;; Represents local knowledge of the state of a shared shutdown state
;; machine.
;;
;;                'neither
;;                /       \
;;              \/         \/
;;           'local      'remote
;;               \         /
;;                \/     \/
;;                  'both

;; A ChannelState is a (ssh-channel ...) TODO
;; Named ssh-channel to avoid conflicts with Racket's built-in
;; synchronous channels.
(struct ssh-channel (room-handle ;; RoomHandle
		     my-ref ;; Uint32
		     your-ref ;; Maybe<Uint32>
		     type ;; String
		     continuations ;; TransactionManager (see ordered-rpc.rkt)
		     outbound-window ;; Maybe<Natural>
		     outbound-packet-size ;; Maybe<Natural>
		     inbound-window ;; Natural
		     eof-state ;; CloseState covering EOF signals
		     close-state ;; CloseState covering CLOSE signals
		     )
	#: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 identification-recogniser #rx"^SSH-")
(define (identification-line? str)
  (regexp-match identification-recogniser str))

(define required-peer-identification-regex (make-parameter #rx"^SSH-2\\.0-.*"))

(define client-preamble-lines (make-parameter '()))
(define client-identification-string (make-parameter "SSH-2.0-RacketSSH_0.0"))

(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)]))

;; ConnectionState 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)
	    (begin (write-message!/flush (ssh-msg-unimplemented seq) conn)
		   conn)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; 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)
  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))
  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 ((string->bytes/utf-8 (exchange-hash-info-client-id hash-info)) :: (t:string))
		      ((string->bytes/utf-8 (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)
;;   -> ConnectionState
;; 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))
     (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))
			(write-message!/flush (ssh-msg-kexdh-reply host-key-bytes
								   public-key-as-integer
								   h-signature)
					      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)
;;   -> ConnectionState
;; 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))
     (write-message!/flush (ssh-msg-kexdh-init public-key-as-integer) conn)
     (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)

  (when (rekey-wait? rekey)
    (write-message!/flush local-algs conn))

  (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))))))))))
    (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.
		       (write-message!/flush (ssh-msg-newkeys) conn)
		       (send (connection-io-room-handle conn) say
			     (new-keys (connection-is-server? conn)
				       derive-key
				       c2s-enc s2c-enc
				       c2s-mac s2c-mac
				       c2s-zip s2c-zip))
		       (set-handlers (struct-copy connection conn
						  [rekey-state
						   (rekey-in-seconds-or-bytes
						    (rekey-interval)
						    (rekey-volume)
						    (connection-total-transferred conn))])
				     SSH_MSG_SERVICE_REQUEST handle-msg-service-request))))

  (if should-discard-first-kex-packet
      (struct-copy connection (continue-after-discard conn) [discard-next-packet? #t])
      (continue-after-discard 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")
	 (begin
	   (write-message!/flush (ssh-msg-service-accept service) conn)
	   (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
    (write-message!/flush (ssh-msg-userauth-success) 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))))
   (else
    (write-message!/flush (ssh-msg-userauth-failure '(none) #f) conn)
    conn)))

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

(define (allocate-channel conn room type your-ref outbound-window outbound-packet-size)
  (define my-ref (hash-count (connection-channel-map conn)))
  (define ch (ssh-channel (join-room room 'session)
			  my-ref
			  your-ref
			  type
			  (make-transaction-manager)
			  outbound-window
			  outbound-packet-size
			  1048576 ;; TODO: parameterize? Make configurable by app?
			  'neither
			  'neither
			  ))
  (values ch
	  (struct-copy connection (send-initial-credit conn ch)
		       [channel-map (hash-set (connection-channel-map conn) my-ref ch)])))

(define (send-initial-credit conn ch)
  (define remaining-window (ssh-channel-outbound-window ch))
  (if (and remaining-window
	   (positive? remaining-window))
      (channel-notify conn ch (credit 'app remaining-window))
      conn))

(define (get-channel conn my-ref)
  (hash-ref (connection-channel-map conn) my-ref))

(define (update-channel conn ch)
  (struct-copy connection conn
	       [channel-map (hash-set (connection-channel-map conn) (ssh-channel-my-ref ch) ch)]))

(define (discard-channel ch conn)
  (struct-copy connection conn
	       [channel-map (hash-remove (connection-channel-map conn) (ssh-channel-my-ref ch))]))

;; 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 ch conn action)
  (define new-close-state (update-close-state (ssh-channel-close-state ch) action))
  (case action
    ((local) (write-message!/flush (ssh-msg-channel-close (ssh-channel-your-ref ch))
				   conn))
    ((remote) (send (ssh-channel-room-handle ch) depart 'remote-closed)))
  (if (eq? new-close-state 'both)
      (discard-channel ch conn)
      (update-channel conn (struct-copy ssh-channel ch
					[close-state new-close-state]))))

(define (channel-notify conn ch message)
  (send (ssh-channel-room-handle ch) say message)
  conn)

(define (channel-request conn ch message k)
  (update-channel conn
		  (struct-copy ssh-channel ch
			       [continuations (room-rpc (ssh-channel-room-handle ch)
							(ssh-channel-continuations ch)
							message
							k)])))

(define (finish-channel-request ch conn txn message)
  (define-values (worklist new-continuations)
    (room-rpc-finish (ssh-channel-continuations ch) txn message))
  (let loop ((worklist worklist)
	     (ch (struct-copy ssh-channel ch [continuations new-continuations]))
	     (conn conn))
    (if (null? worklist)
	(update-channel conn ch)
	(let ((item (car worklist)))
	  (define-values (new-ch new-conn) (item ch conn))
	  (loop (cdr worklist) new-ch new-conn)))))

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

(define (start-connection-service conn)
  (set-handlers conn
		;; TODO: SSH_MSG_GLOBAL_REQUEST		handle-msg-global-request
		SSH_MSG_CHANNEL_OPEN		handle-msg-channel-open
		SSH_MSG_CHANNEL_WINDOW_ADJUST	handle-msg-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))

(define (handle-msg-global-request packet message conn)
  (log-error "TODO: Unimplemented: handle-msg-global-request")
  conn)

(define (handle-msg-channel-open packet message conn)
  (match-define (ssh-msg-channel-open channel-type*
				      sender-channel
				      initial-window-size
				      maximum-packet-size
				      extra-request-data*)
		message)
  (define channel-type (bit-string->bytes channel-type*))
  (define extra-request-data (bit-string->bytes extra-request-data*))
  (app-request conn
	       `(open-channel ,(connection-username conn) ,channel-type ,extra-request-data)
	       (lambda (reply conn)
		 (match reply
		   (`(ok ,(? room? room) ,(? bytes? extra-reply-data))
		    (let-values (((ch conn) (allocate-channel conn
							      room
							      channel-type
							      sender-channel
							      initial-window-size
							      maximum-packet-size)))
		      (write-message!/flush (ssh-msg-channel-open-confirmation
					     sender-channel
					     (ssh-channel-my-ref ch)
					     (ssh-channel-inbound-window ch)
					     (default-packet-limit) ;; TODO get from reader
					     extra-reply-data)
					    conn)
		      conn))
		   (`(error ,reason-code ,description)
		    (write-message!/flush (ssh-msg-channel-open-failure
					   sender-channel
					   reason-code
					   (string->bytes/utf-8 description)
					   #"")
					  conn)
		    conn)))))

(define (handle-msg-window-adjust packet message conn)
  (match-define (ssh-msg-channel-window-adjust recipient-channel count) message)
  (define ch (get-channel conn recipient-channel))
  (channel-notify conn ch (credit 'app count)))

(define (handle-msg-channel-data packet message conn)
  (match-define (ssh-msg-channel-data recipient-channel data*) message)
  (define data (bit-string->bytes data*))
  (define ch (get-channel conn recipient-channel))
  (channel-notify conn ch `(data ,data)))

(define (handle-msg-channel-extended-data packet message conn)
  (match-define (ssh-msg-channel-extended-data recipient-channel type-code data*) message)
  (define data (bit-string->bytes data*))
  (define ch (get-channel conn recipient-channel))
  (channel-notify conn ch `(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)))
  (update-channel (channel-notify conn ch `(eof))
		  (struct-copy ssh-channel ch
			       [eof-state (update-close-state (ssh-channel-eof-state ch)
							      'remote)])))

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

(define (handle-msg-channel-request packet message conn)
  (match-define (ssh-msg-channel-request recipient-channel type* want-reply? data*) message)
  (define type (bit-string->bytes type*))
  (define data (bit-string->bytes data*))
  (define ch (get-channel conn recipient-channel))
  (if (not want-reply?)
      (channel-notify conn ch `(notify ,type ,data))
      (channel-request conn ch `(,type ,data)
		       (lambda (reply ch conn)
			 (define your-ref (ssh-channel-your-ref ch))
			 (write-message!/flush (match reply
						 ('ok (ssh-msg-channel-success your-ref))
						 ('error (ssh-msg-channel-failure your-ref)))
					       conn)
			 (values ch conn)))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Session main loop
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (write-message! message conn)
  (send (connection-io-room-handle conn) say message))

(define (flush-outbound-messages! conn)
  (send (connection-io-room-handle conn) say 'flush))

(define (write-message!/flush message conn)
  (write-message! message conn)
  (flush-outbound-messages! conn))

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

(define (app-notify conn message)
  (send (connection-session-room-handle conn) say message)
  conn)

(define (app-request conn message k)
  (struct-copy connection conn
	       [continuations (room-rpc (connection-session-room-handle conn)
					(connection-continuations conn)
					message
					k)]))

(define (finish-app-request conn txn message)
  (define-values (worklist new-continuations)
    (room-rpc-finish (connection-continuations conn) txn message))
  (foldl (lambda (item conn) (item conn))
	 (struct-copy connection conn [continuations new-continuations])
	 worklist))

(define (maybe-send-disconnect-message! e conn)
  (when (not (exn:fail:contract:protocol-originated-at-peer? e))
    (write-message!/flush (ssh-msg-disconnect (exn:fail:contract:protocol-reason-code e)
					      (string->bytes/utf-8 (exn-message e))
					      #"")
			  conn)))

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

(define io-room-message-handler
  (lambda (message)
    (lambda (conn)
      (match message
	((arrived 'read-thread)
	 (send (connection-io-room-handle conn) say (credit 'read-thread 1))
	 conn)
	((arrived _)
	 conn)
	((and departure (departed who why))
	 (if (zero? (hash-count (connection-channel-map conn)))
	     ;; No open or half-open channels. No point in complaining; just leave.
	     #f
	     ;; At least one channel. Make more of a noise.
	     (disconnect-with-error/local-info departure
					       SSH_DISCONNECT_CONNECTION_LOST
					       "I/O error")))
	((says _ amount 'output-byte-count)
	 ;; writer reporting bytes transferred
	 (bump-total amount conn))
	((says _ (received-packet seq packet message transferred-count) _)
	 (send (connection-io-room-handle conn) say (credit 'read-thread 1))
	 (bump-total
	  transferred-count
	  (if (connection-discard-next-packet? conn)
	      (struct-copy connection conn [discard-next-packet? #f])
	      (dispatch-packet seq packet message conn))))))))

(define session-room-message-handler
  (lambda (message)
    (lambda (conn)
      (match message
	((arrived _)
	 conn)
	((and departure (departed _ _))
	 (disconnect-with-error/local-info
	  departure
	  SSH_DISCONNECT_BY_APPLICATION
	  "Application disconnected"))
	((says _ (rpc-reply transaction message) _)
	 ;; TODO: not cap-secure. Introduce sealers, or indirect.
	 (finish-app-request conn transaction message))))))

;; (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 (channel-events conn)
  (hash-fold (lambda (my-ref ch evt)
	       (choice-evt evt
			   (handle-evt (send (ssh-channel-room-handle ch) listen-evt)
				       (channel-room-message-handler my-ref))))
	     never-evt
	     (connection-channel-map conn)))

(define (channel-room-message-handler my-ref)
  (lambda (message)
    (lambda (conn)
      (define ch (get-channel conn my-ref))
      (define your-ref (ssh-channel-your-ref ch))
      (match message
	((arrived _)
	 conn)
	((departed _ _)
	 (maybe-close-channel ch conn 'local))
	((says _ (credit _ amount) _)
	 (write-message!/flush (ssh-msg-channel-window-adjust your-ref amount) conn)
	 conn)
	((says _ `(data ,bits) _)
	 (write-message!/flush (ssh-msg-channel-data your-ref bits) conn)
	 conn)
	((says _ `(eof) _)
	 (write-message!/flush (ssh-msg-channel-eof your-ref) conn)
	 conn)
	((says _ (rpc-reply id m) _)
	 (finish-channel-request ch conn id m))))))

(define (run-ssh-session conn)
  (with-handlers
      ((exn:fail:contract:protocol? (lambda (e)
				      (maybe-send-disconnect-message! e conn)
				      (raise e))))
    (let loop ((conn conn))
      (define rekey (connection-rekey-state conn))
      (if (time-to-rekey? rekey conn)
	  (let ((algs ((local-algorithm-list))))
	    (write-message!/flush algs conn)
	    (loop (struct-copy connection conn [rekey-state (rekey-local algs)])))
	  (let ((handler (sync (if (rekey-wait? rekey)
				   (handle-evt (alarm-evt (* (rekey-wait-deadline rekey) 1000))
					       (lambda (dummy)
						 (lambda (conn)
						   conn)))
				   never-evt)
			       (handle-evt (send (connection-io-room-handle conn) listen-evt)
					   io-room-message-handler)
			       (handle-evt (send (connection-session-room-handle conn) listen-evt)
					   session-room-message-handler)
			       (channel-events conn))))
	    (define new-conn (handler conn))
	    ;; The handler is permitted to return #f to indicate that the session is to be
	    ;; gracefully shut down.
	    (when new-conn
	      (loop new-conn)))))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Session choreography
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (send-preamble-and-identification! out)
  (let ((my-id (client-identification-string)))
    (for-each (lambda (line)
		(when (identification-line? line)
		  (error 'ssh-session
			 "Client preamble includes forbidden line ~v"
			 line))
		(display line out)
		(display "\r\n" out))
	      (client-preamble-lines))
    (display my-id out)
    (display "\r\n" out)
    (flush-output out)
    my-id))

;; Port -> String
(define (read-preamble-and-identification! in)
  (let ((line (read-line-limited in 253))) ;; 255 incl CRLF
    (when (eof-object? line)
      (error 'ssh-session "EOF while reading connection preamble"))
    (if (identification-line? line)
	line
	(read-preamble-and-identification! in))))

;; 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 role in out)
  (define io-room (make-room (gensym 'ssh-io-room)))
  ;;(spy-on io-room)
  (define session-room (make-room (gensym 'ssh-session-room)))
  ;;(spy-on session-room)

  (define local-identification-string (send-preamble-and-identification! out))
  (define peer-identification-string (read-preamble-and-identification! in))

  ;; Each identification string is both a cleartext indicator that
  ;; we've reached some notion of the right place and also input to
  ;; the hash function used during D-H key exchange.
  (when (not (regexp-match (required-peer-identification-regex)
			   peer-identification-string))
    (display "Invalid identification\r\n" out)
    (flush-output out)
    (error 'ssh-session
	   "Invalid peer identification string ~v"
	   peer-identification-string))

  (standard-thread (lambda () (ssh-reader in io-room)))
  (standard-thread (lambda () (ssh-writer out io-room)))
  (wait-for-members io-room '(read-thread write-thread))
  (standard-thread
   (lambda ()
     (run-ssh-session (connection (join-room io-room 'session)
				  (join-room session-room 'session)
				  #f
				  base-packet-dispatcher
				  0
				  (rekey-in-seconds-or-bytes -1 -1 0)
				  #f
				  (make-transaction-manager)
				  (hash)
				  (case role ((client) #f) ((server) #t))
				  local-identification-string
				  peer-identification-string
				  #f))))

  (join-room session-room 'app))