#lang racket/base

;; DNS proxy using os-big-bang.rkt and os-udp.rkt.

(require racket/match)
(require racket/set)
(require racket/bool)
(require "../racket-bitsyntax/main.rkt")
(require "api.rkt")
(require "codec.rkt")
(require "zonedb.rkt")
(require "network-query.rkt")
(require "resolver.rkt")
(require "../racket-matrix/os.rkt")
(require "../racket-matrix/os-big-bang.rkt")
(require "../racket-matrix/os-udp.rkt")
(require "../racket-matrix/os-timer.rkt")
(require "os-dns.rkt")

;; Instantiated with a collection of trusted roots to begin its
;; searches from. Performs recursive queries.

;; For discarding retransmitted requests that we're still working on.
(struct active-request (source id) #:transparent)

;; start-proxy : UInt16 CompiledZone -> Void
(require racket/pretty)
(define (start-proxy port-number zone)

  (define boot-server
    (os-big-bang 'no-state/boot-server
      ;;(spawn dns-spy)
      (spawn (timer-relay 'timer-relay:dns))
      (spawn (query-id-allocator))
      (send-meta-message `(request create-server-socket (udp new ,port-number 512)))
      (subscribe/fresh wait-id
	(meta-message-handlers w
	  [`(reply create-server-socket ,s)
	   (transition w
	     (unsubscribe wait-id)
	     (send-meta-message
	      `(request create-client-socket (udp new 0 512)))
	     (client-socket-waiter s))]))))

  (define (client-socket-waiter s)
    (subscribe/fresh wait-id
      (meta-message-handlers w
	[`(reply create-client-socket ,c)
	 (display "Ready.") (newline)
	 (transition w
	   (unsubscribe wait-id)
	   (spawn (dns-read-driver s))
	   (spawn (dns-write-driver s))
	   (spawn (dns-read-driver c))
	   (spawn (dns-write-driver c))
	   (spawn (packet-dispatcher s))
	   (spawn (question-dispatcher zone c)))])))

  (ground-vm (os-big-bang (void)
	       ;;(spawn udp-spy)
	       (spawn udp-driver)
	       (spawn (timer-driver 'timer-driver))
	       (spawn (nested-vm boot-server)))))

(define (query-id-allocator)
  ;; TODO: track how many are allocated and throttle requests if too
  ;; many are in flight
  (os-big-bang (set) ;; SetOf<UInt16>, all active query IDs
    (subscribe 'query-id-request-handler
      (message-handlers allocated
	[`(request ,reply-addr allocate-query-id)
	 (let recheck ()
	   (define n (random 65536))
	   (if (set-member? allocated n)
	       (recheck)
	       (transition (set-add allocated n)
		 (send-message `(reply ,reply-addr ,n)))))]
	[`(release-query-id ,n)
	 (transition (set-remove allocated n))]))))

(define (packet-dispatcher s)
  (os-big-bang (set) ;; SetOf<ActiveRequest>
    (subscribe 'packet-dispatcher
      (message-handlers old-active-requests
	[(? bad-dns-packet? p)
	 (pretty-print p) ;; TODO: perhaps use metalevel events? perhaps don't bother though
	 old-active-requests]
	[(and r (dns-request m source (== s))) ;; We only listen for requests on our server socket
	 (define req-id (active-request source (dns-message-id m)))
	 ;; TODO: when we have presence/error-handling, remove req-id
	 ;; from active requests once request-handler pseudothread exits.
	 (if (set-member? old-active-requests req-id)
	     old-active-requests ;; ignore retransmitted duplicates
	     (transition (set-add old-active-requests req-id)
	       (spawn (packet-relay req-id r))))]
	[(and r (dns-reply m (== s) sink))
	 (define req-id (active-request sink (dns-message-id m)))
	 (set-remove old-active-requests req-id)]))))

(define (packet-relay req-id request)
  (match-define (dns-request request-message request-source request-sink) request)
  (define (answer->reply q a)
    (define-values (response-code ns us ds)
      (match a
	[#f
	 (values 'name-error '() '() '())]
	[(complete-answer ns us ds)
	 (values 'no-error (rr-set->list ns) (rr-set->list us) (rr-set->list ds))]))
    (dns-reply (dns-message (dns-message-id request-message)
			    'response
			    'query
			    'non-authoritative
			    'not-truncated
			    (dns-message-recursion-desired request-message)
			    'recursion-available
			    response-code
			    (if q (list q) '())
			    ns
			    us
			    ds)
	       request-sink
	       request-source))
  ;; TODO: pay attention to recursion-desired flag
  (match (dns-message-questions request-message)
    ['()
     ;; No questions!
     (os-big-bang 'no-state/packet-relay
       (send-message (answer->reply #f (empty-complete-answer))))]
    [(cons original-question _)
     ;; At least one question
     (begin (write `(Looking up ,original-question
			     with query id ,(dns-message-id request-message)))
	    (newline))
     (os-big-bang 'no-state/packet-relay
       (send-message original-question)
       (subscribe/fresh wait-id
	 (message-handlers w
	   [(answered-question (== original-question) answer)
	    (begin (write `(Final answer to ,original-question
				  with query id ,(dns-message-id request-message)
				  is ,answer))
		   (newline))
	    (transition w
	      (unsubscribe wait-id)
	      (send-message (answer->reply original-question answer)))])))]))

(define (question-dispatcher zone0 client-sock)
  (define (transition-and-set-timers new-zone timers)
    (transition new-zone
      (for/list ([timerspec timers])
	(match-define (cons name ttl) timerspec)
	(send-message (set-timer (list 'check-dns-expiry name) (* ttl 1000) #t)))))
  (os-big-bang/transition
   (extend-transition (let-values (((expired-zone timers) (zone-expire zone0)))
			(transition-and-set-timers expired-zone timers))
     ;; TODO: consider deduping questions here too?
     (subscribe 'question-handler-factory
       (message-handlers zone
	 [`(debug-dump)
	  (with-output-to-file "zone-proxy.zone"
	    (lambda ()
	      (write-bytes (bit-string->bytes (zone->bit-string zone))))
	    #:mode 'binary
	    #:exists 'replace)
	  (with-output-to-file "zone-proxy.dump"
	    (lambda ()
	      (display "----------------------------------------------------------------------\n")
	      (display (seconds->date (current-seconds)))
	      (newline)
	      (for* ([(name rrmap) zone] [(rr expiry) rrmap])
		(write (list rr expiry))
		(newline))
	      (newline))
	    #:mode 'text
	    #:exists 'append)
	  zone]
	 [(? question? q)
	  (transition zone
	    (spawn (question-handler zone q client-sock)))]
	 [(network-reply _ answer)
	  (define-values (new-zone timers) (incorporate-complete-answer answer zone))
	  (transition-and-set-timers new-zone timers)]
	 [(timer-expired (list 'check-dns-expiry name) _)
	  (zone-expire-name zone name)])))))

(struct question-state (zone q client-sock nameservers-tried retry-count) #:transparent)
(struct expanding-cnames (q accumulator remaining-count) #:transparent)

(define (question-handler zone q client-sock)
  (os-big-bang/transition
   (retry-question (question-state zone q client-sock (set) 0))))

(define (send-empty-reply w q)
  (transition w (send-message (answered-question q (empty-complete-answer)))))

(define (retry-question w)
  (match w
    [(question-state _ q _ _ 20) ;; TODO: is this a sensible limit?
     ;; Too many retries, i.e. too many referrals.
     (begin (write `(Too many retries ,w))
	    (newline))
     (send-empty-reply w q)]
    [(question-state zone q client-sock nameservers-tried old-retry-count)
     ;; Credit remaining. Try once more (perhaps for the first time, in fact).
     (match (resolve-from-zone q zone #f nameservers-tried)
       [#f ;; We're not authoritative so this is just a signal that we can't answer usefully
	(send-empty-reply w q)]
       [(referral zone-origin nameserver-rrs _)
	(begin
	  (write `(Referral for ,q
			    to ,(domain-labels zone-origin)
			    servers ,(map domain-labels (set-map nameserver-rrs rr-rdata))))
	  (newline))
	(define referral-id (gensym 'referral))
	(transition w
	  (spawn (network-query client-sock
				q
				zone-origin
				(map rr-rdata (set->list nameserver-rrs))
				referral-id))
	  (subscribe referral-id
	    (message-handlers w
	      [(network-reply (== referral-id) #f) ;; name-error/NXDOMAIN
	       (transition w (send-message (answered-question q #f)))]
	      [(network-reply (== referral-id) ans)
	       (define-values (new-zone ignored-timers) (incorporate-complete-answer ans zone))
	       (define nameserver-names (for/set ([rr nameserver-rrs]) (rr-rdata rr)))
	       (extend-transition
		   (retry-question (struct-copy question-state w
						[nameservers-tried (set-union nameservers-tried
									      nameserver-names)]
						[zone new-zone]
						[retry-count (+ old-retry-count 1)]))
		 (unsubscribe referral-id))])))]
       [(? complete-answer? ans)
	(transition w (send-message (answered-question q ans)))]
       [(partial-answer base cnames)
	(transition (expanding-cnames q base (length cnames))
	  (map (lambda (cname)
		 (define cname-q (question cname (question-type q) (question-class q)))
		 (list (send-message cname-q)
		       (subscribe/fresh subscription-id
			 (message-handlers (expanding-cnames q acc remaining)
			   [(answered-question (== cname-q) ans)
			    (define new-acc (if ans (merge-answers acc ans) acc))
			    (define new-remaining (- remaining 1))
			    (define new-w (expanding-cnames q new-acc new-remaining))
			    (transition new-w
			      (unsubscribe subscription-id)
			      (if (zero? new-remaining)
				  (send-message (answered-question q new-acc))
				  '()))]))))
	       cnames))])]))

(require "test-rrs.rkt")
(require racket/file)
(file-stream-buffer-mode (current-output-port) 'none)
(start-proxy (test-port-number)
	     (if (file-exists? "zone-proxy.zone")
		 (bit-string->zone (file->bytes "zone-proxy.zone"))
		 (compile-zone-db test-roots)))