#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 "dump-bytes.rkt")
(require "os.rkt")
(require "os-big-bang.rkt")
(require "os-udp.rkt")
(require "os-dns.rkt")
(require "os-timer.rkt")

;; Instantiated with a collection of trusted roots to begin its
;; searches from. Performs recursive queries. Doesn't yet cache
;; responses, but will do so in future.

;; start-proxy : UInt16 ListOf<RR> -> Void
(require racket/pretty)
(define (start-proxy port-number rrs)
  ;; Compile the zone hash table
  (define zone (compile-zone-db rrs))
  (pretty-print zone)

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

  (define (client-socket-waiter s)
    (subscribe 'wait-for-client-socket
	       (meta-message-handlers w
		[`(reply create-client-socket ,c)
		 (transition w
			     (unsubscribe 'wait-for-client-socket)
			     (spawn (dns-read-driver c))
			     (spawn (dns-write-driver c))
			     (subscribe 'packet-handler (packet-handler s c)))])))

  (define (packet-handler s c)
    (message-handlers old-state
     [(? bad-dns-packet? p)
      (pretty-print p) ;; TODO: perhaps use metalevel events? perhaps don't bother though
      old-state]
     [(and r (dns-request _ _ (== s))) ;; We only listen for requests on our server socket
      (transition old-state (spawn (request-handler zone r c)))]))

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

(define (request-handler zone request client-sock)
  ;; 1. try resolving locally
  ;; 2. if it answers, send that out. otherwise, it needs to request something recursively.
  ;; 3. if the socket doesn't exist, request it, and wait for the reply.
  ;; 4. start processing the network query: it will result in a send/receive/timeout combo
  ;; 5. on timeout, try a different server, or if there aren't any
  ;;    more, report failure to the resolver
  ;; 6. on packet, report success to the resolver and goto 2.
  ;; -. remember to release the socket when we're done!

  (match-define (dns-request request-message request-source request-sink) request)
  (define question (and (pair? (dns-message-questions request-message))
			(car (dns-message-questions request-message))))

  (write `(Looking up ,question with query id ,(dns-message-id request-message))) (newline)

  (define (make-reply answers authorities additional)
    (dns-message (dns-message-id request-message)
		 'response
		 'query
		 'non-authoritative
		 'not-truncated
		 (dns-message-recursion-desired request-message)
		 'recursion-available
		 'no-error
		 (list question)
		 (rr-set->list answers)
		 (rr-set->list authorities)
		 (rr-set->list additional)))

  (define (resolver-actions qr)
    (match qr
      [(resolver-network-query q zone ns-rr addresses k) ;; need subquestion answered
       ;;(write `(INTERMEDIATE ,q ,ns-rr (,(length addresses) addresses))) (newline)
       (network-query/addresses client-sock q zone ns-rr addresses
				(lambda (qr) (resolver-actions (k qr))))]
      [#f ;; got a name-error/NXDOMAIN from some nameserver
       ;; TODO: re-examine my reasoning for not sending name-error/NXDOMAIN here
       (send-message (dns-reply (make-reply (set) (set) (set)) request-sink request-source))]
      [(question-result _ _ anss auths adds)
       (send-message (dns-reply (make-reply anss auths adds) request-sink request-source))]))

  (if (eq? question #f)
      (os-big-bang 'no-questions-no-processing-no-answers)
      (os-big-bang '???
		   (resolver-actions
		    (resolve-from-zone question
				       zone
				       #f ;; no SOA, since we're not authoritative for anything
				       #t ;; we *are* however recursive.
				       (set) ;; haven't tried any nameservers yet
				       values)))))

(require "test-rrs.rkt")
(start-proxy 5555 test-roots)