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

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

(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 (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 (connection-handler local-addr remote-addr)
  (define local-identification #"SSH-2.0-RacketSSH_0.0")
  (nested-vm
   (list 'ssh-session-vm remote-addr)
   (lambda (nested-boot-pid)
     (transition 'running
       (spawn (timer-relay 'ssh-timer-relay) #:debug-name 'ssh-timer-relay)

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

       ;; 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/monitor
	(transition 'handshake-is-stateless
	  (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/monitor (ssh-writer local-addr remote-addr) #: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.
		(yield #:state state
		  (transition state
		    (send-message (inbound-credit 1))
		    (spawn/monitor (ssh-session local-identification
						remote-identification
						repl-boot
						'server)
				   #:debug-name 'ssh-session))))])))
	#:debug-name 'ssh-reader)

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

       (role 'monitor-listener (topic-subscriber (monitor (wild) (wild)))
	 #:state state
	 #:reason reason
	 #:on-absence
	 ;; 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 %%%
	 (if (eq? state 'running)
	     (if (and (exn:fail:contract:protocol? reason)
		      (not (exn:fail:contract:protocol-originated-at-peer? reason)))
		 (transition 'error-packet-sent
		   (send-message (outbound-packet (ssh-msg-disconnect
						   (exn:fail:contract:protocol-reason-code reason)
						   (string->bytes/utf-8 (exn-message reason))
						   #"")))
		   (yield #:state state
		     (transition state (at-meta-level (kill)))))
		 (transition state (at-meta-level (kill #:reason reason))))
	     state))))))

(ground-vm
 (transition 'no-state
   (spawn (timer-driver 'timer-driver))
   (spawn tcp-driver #:debug-name 'tcp-driver)
   (spawn tcp-spy #:debug-name 'tcp-spy)

   (spawn (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) _) #t)
			       ;; Ignore virtual flows. They just mean there's
			       ;; someone willing to supply connections to us
			       ;; at some point in the future.
			       state]
			      [(topic 'publisher (tcp-channel remote-addr (== server-addr) _) #f)
			       (transition state
				 (spawn (connection-handler server-addr remote-addr)
					#:debug-name (list 'ssh-session-vm remote-addr)))]))))))