#lang racket/base
;; (Temporary) example client and server

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

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

(define server-addr (tcp-listener 2322))

(define (main)
  (ground-vm
   (transition 'no-state
     (spawn (timer-driver 'timer-driver))
     ;; PAPER NOTE: remove #:debug-name for presentation economy
     (spawn tcp-driver #:debug-name 'tcp-driver)
     (spawn tcp-spy #:debug-name 'tcp-spy)
     (spawn listener #:debug-name 'ssh-tcp-listener))))

(define listener
  (transition 'no-state
    (role 'connection-waiter (topic-subscriber (tcp-channel (wild) server-addr (wild))
					       #:virtual? #t)
      #:state state
      #:topic t
      #:on-presence (match t
		      [(topic 'publisher (tcp-channel remote-addr (== server-addr) _) #f)
		       (transition state
			 (spawn (session-vm server-addr remote-addr)
				#:debug-name (list 'ssh-session-vm remote-addr)))]
		      ;; PAPER NOTE: This second clause can be replaced with
		      ;; [_ state] for presentation economy
		      [(topic 'publisher (tcp-channel remote-addr (== server-addr) _) #t)
		       ;; Ignore virtual flows. They just mean there's
		       ;; someone willing to supply connections to us
		       ;; at some point in the future.
		       state]))))

;;---------------------------------------------------------------------------

(define (check-remote-identification! peer-identification-string)
  (define required-peer-identification-regex #rx"^SSH-2\\.0-.*")
  ;; 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))
    (error 'ssh-session
	   "Invalid peer identification string ~v"
	   peer-identification-string)))

(define (repl-boot self-pid)
  (transition 'no-repl-state
    (role 'spy (or (topic-subscriber (wild) #:virtual? #t)
		   (topic-publisher (wild) #:virtual? #t))
      #:state state
      [message
       (write `(APP ,message))
       (newline)
       (flush-output)
       state])
    (at-meta-level
     (role 'channel-listener (topic-subscriber (channel-message (channel-stream-name #t (wild))
							       (wild)))
       #:state state
       #:topic t
       #:on-presence
       (if (topic-virtual? t)
	   state
	   (match t
	     [(topic _ (channel-message (channel-stream-name _ cname) _) _)
	      (transition state (spawn (repl-instance cname) #:debug-name cname))]))))))

(define (repl-instance cname)
  (define inbound-stream (channel-stream-name #t cname))
  (define outbound-stream (channel-stream-name #f cname))
  (define (handle-channel-message state body)
    (match body
      [(channel-stream-request #"pty-req" _)
       (transition state
	 (at-meta-level (send-message (channel-message inbound-stream (channel-stream-ok))
				      'subscriber)))]
      [m
       (write `(channel inbound ,m)) (newline)
       state]))
  (match (channel-name-type cname)
    [#"session"
     (transition 'no-instance-state
       (at-meta-level
	(role 'input (topic-subscriber (channel-message inbound-stream (wild)))
	  #:state state
	  #:on-presence (transition state (at-meta-level
					   (send-message
					    (channel-message inbound-stream
							     (channel-stream-config
							      (default-packet-limit)
							      #""))
					    'subscriber)))
	  [(channel-message _ body)
	   (handle-channel-message state body)]))
       (at-meta-level
	(role 'output (topic-publisher (channel-message outbound-stream (wild)))
	  #:state state
	  [m
	   (write `(channel outbound ,cname ,m)) (newline)
	   state])))]
    [type
     (transition 'no-instance-state
       (at-meta-level (send-message
		       (channel-message outbound-stream
					(channel-stream-open-failure
					 SSH_OPEN_UNKNOWN_CHANNEL_TYPE
					 (bytes-append #"Unknown channel type " type))))))]))

(define (spy marker)
  (role 'spy (or (topic-subscriber (wild) #:virtual? #t)
		 (topic-publisher (wild) #:virtual? #t))
    #:state state
    [message
     (write `(,marker ,message))
     (newline)
     (flush-output)
     state]))

(define (session-vm local-addr remote-addr)
  (define local-identification #"SSH-2.0-RacketSSH_0.0")

  (define (issue-identification-string)
    (at-meta-level
     (send-message (tcp-channel local-addr
				remote-addr
				(bytes-append local-identification #"\r\n")))))

  (define (read-handshake-and-become-reader)
    (transition 'handshake-is-stateless ;; but, crucially, the ssh-reader proper isn't!
      (at-meta-level
       (role 'socket-reader (topic-subscriber (tcp-channel remote-addr local-addr (wild)))
	 #:state state
	 [(tcp-channel _ _ (? eof-object?))
	  (transition state (kill))]
	 [(tcp-channel _ _ (? bytes? remote-identification))
	  (check-remote-identification! remote-identification)
	  ;; First, set the incoming mode to bytes. Then
	  ;; initialise the reader, switching to packet-reading
	  ;; mode. Finally, spawn the remaining processes and
	  ;; issue the initial credit to the reader.
	  (extend-transition
	      (prefix-transition (ssh-reader local-addr remote-addr)
		(at-meta-level (send-tcp-mode remote-addr local-addr 'bytes)))
	    (spawn (ssh-writer local-addr remote-addr)
		   #:monitor? #t
		   #:debug-name 'ssh-writer)
	    ;; Wait for a cycle to let the reader and writer get
	    ;; started, then tell the reader we are ready for a
	    ;; single packet and spawn the session manager.

	    ;; TODO: try using presence instead of the yield.
	    (yield #:state state
	      (transition state
		(send-message (inbound-credit 1))
		(spawn (ssh-session local-identification
				    remote-identification
				    repl-boot
				    'server)
		       #:monitor? #t
		       #:debug-name 'ssh-session))))]))))

  (define (exn->outbound-packet reason)
    (outbound-packet (ssh-msg-disconnect (exn:fail:contract:protocol-reason-code reason)
					 (string->bytes/utf-8 (exn-message reason))
					 #"")))

  (define (disconnect-message-required? reason)
    (and (exn:fail:contract:protocol? reason)
	 (not (exn:fail:contract:protocol-originated-at-peer? reason))))

  (define (active-exception-handler reason)
    ;; This is kind of gross: because the absence handler gets invoked
    ;; several times in a row because of multiple flows intersecting
    ;; this role, we have to be careful to make the transmission of
    ;; the disconnection packet idempotent.
    ;; TODO: this is likely no longer true now we're using monitors %%%
    (define interesting? (disconnect-message-required? reason))
    (transition inert-exception-handler
      (when interesting? (send-message (exn->outbound-packet reason)))
      (yield #:state state ;; gross
	(transition state (at-meta-level (kill #:reason (and interesting? reason)))))))

  (define (inert-exception-handler reason)
    inert-exception-handler)

  (nested-vm
   (list 'ssh-session-vm remote-addr)
   ;; TODO: use (not-yet-existing) macro variant of nested-vm to avoid
   ;; spuriously binding nested-boot-pid without blaming the wrong
   ;; process in case of error.
   (lambda (nested-boot-pid)
     (transition 'no-state
       (spawn (timer-relay 'ssh-timer-relay) #:debug-name 'ssh-timer-relay)
       (spy 'SSH)

       (issue-identification-string)

       ;; Expect identification string, then update (!) our inbound
       ;; subscription handler to switch to packet mode.
       (at-meta-level (send-tcp-mode remote-addr local-addr 'lines))
       (at-meta-level (send-tcp-credit remote-addr local-addr 1))

       (spawn (read-handshake-and-become-reader)
	      #:monitor? #t
	      #:debug-name 'ssh-reader)

       (spawn (transition active-exception-handler
		(role 'monitor-listener (topic-subscriber (monitor (wild) (wild)))
		  #:state current-handler
		  #:reason reason
		  #:on-absence (current-handler reason))))))))

;; TODO: module+
(main)