#lang racket/base

;; Noddy representation of a zone, and various zone and RRSet utilities.

(require racket/set)
(require racket/match)
(require "api.rkt")
(require "codec.rkt")

(provide compile-zone-db
	 in-bailiwick?
	 set-filter
	 filter-by-type
	 filter-rrs
	 rr-set->list

	 resolve-from-zone)

;; A CompiledZone is a Hash<DomainName,SetOf<RR>>, representing a
;; collection of DNS RRSets indexed by DomainName.

;; TODO: maybe store domain names big-end first? It'd make bailiwick
;; and subzone checks into prefix rather than suffix checks. It makes
;; domain names into paths through the DNS DB tree.

;; RR Hash -> Hash
(define (incorporate-rr rr db)
  (hash-set db (rr-name rr) (set-add (hash-ref db (rr-name rr) (set)) rr)))

;; compile-zone-db : ListOf<RR> -> CompiledZone
;; Builds an immutable hash table from the given RRs, suitable for
;; quickly looking up answers to queries.
(define (compile-zone-db rrs)
  (foldl incorporate-rr (make-immutable-hash) rrs))

;; in-bailiwick? : DomainName RR -> Boolean
;; Answers #t iff dn falls within the bailiwick of the zone with
;; origin rr.
(define (in-bailiwick? dn rr)
  (cond
   ((equal? dn (rr-name rr)) #t)
   ((null? dn) #f)
   (else (in-bailiwick? (cdr dn) rr))))

;; set-filter : (X -> Boolean) SetOf<X> -> SetOf<X>
;; Retains only those elements of its argument for which the predicate
;; answers #t.
(define (set-filter predicate in)
  (for/set ([x (in-set in)]
	    #:when (predicate x))
	   x))

;; filter-by-type : SetOf<RR> RRType -> SetOf<RR>
;; Selects only those members of rrset having rr-type type.
(define (filter-by-type rrset type)
  (set-filter (lambda (rr) (eqv? (rr-type rr) type)) rrset))

;; filter-rrs : SetOf<RR> QueryType QueryClass
;; Returns a set like its argument with RRs not matching the given
;; type and class removed.
(define (filter-rrs rrs qtype qclass)
  (define filtered-by-type
    (case qtype
      ((*) rrs)
      (else (filter-by-type rrs qtype))))
  (define filtered-by-type-and-class
    (case qclass
      ((*) filtered-by-type)
      (else (set-filter (lambda (rr) (eqv? (rr-class rr) qclass)) filtered-by-type))))
  filtered-by-type-and-class)

;; rr-set->list : SetOf<RR> -> ListOf<RR>
;; Like set->list, but places all CNAME records first.
;; This is apparently to work around bugs in old versions of BIND?
(define (rr-set->list rrs)
  (append (set->list (filter-by-type rrs 'cname))
	  (set->list (set-filter (lambda (rr) (not (eqv? (rr-type rr) 'cname))) rrs))))

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

(define (answer-available? q zone)
  (hash-has-key? zone (question-name q)))

;; QuestionResult Maybe<QuestionResult> -> QuestionResult
;; Add the supporting facts from r2 into r1, keeping r1's
;; question. Replaces the knowledge from r1 with the knowledge from
;; r2.  Suitable for use when r2 is answering some sub-question of
;; r1's question.
(define (merge-replies r1 r2)
  (match r2
    [#f r1]
    [(question-result _ k2 n2 u2 d2) ;; a normal result
     (match-define (question-result q1 k1 n1 u1 d1) r1)
     (question-result q1
		      k2
		      (set-union n1 n2)
		      (set-union u1 u2)
		      (set-union d1 d2))]))

(define (answer-from-zone q zone start-of-authority recursion-desired?)
  (match-define (question name qtype qclass) q)
  (define rrset (hash-ref zone name))
  (define filtered-rrs (filter-rrs rrset qtype qclass))
  (define cnames (filter-by-type rrset 'cname)) ;; TODO: filter by class too??
  (define base-reply (question-result q
				      zone
				      (set-union cnames filtered-rrs)
				      (if (and start-of-authority
					       (in-bailiwick? name start-of-authority))
					  (set start-of-authority)
					  (set))
				      (set)))
  ;; Algorithm for CNAME expansion from RFC 1034 4.3.2 step 3a.
  (if (and (not (set-empty? cnames)) (not (eqv? qtype 'cname)))
      (foldl (lambda (cname-rr current-reply)
	       (merge-replies current-reply
			      (resolve-from-zone
			       (question (rr-rdata cname-rr) qtype qclass)
			       zone
			       start-of-authority
			       recursion-desired?
			       (set))))
	     base-reply
	     (set->list cnames))
      base-reply))

(define (closest-nameservers name zone)
  (let search ((name name))
    (cond
     ((hash-ref zone name #f) =>
      ;; There's an entry for this suffix of the original name. Check
      ;; to see if it has an NS record indicating a subzone.
      (lambda (rrset)
	(define ns-rrset (filter-by-type rrset 'ns))
	(if (set-empty? ns-rrset)
	    (search (cdr name)) ;; no NS records for this suffix. Keep looking.
	    ns-rrset)))
     ((null? name)
      ;; The root, and we don't have an RRSet for it. Give up.
      (set))
     (else
      ;; Remove a label and keep looking.
      (search (cdr name))))))

;; Returns a list of NS RRs in some priority order: records for which
;; we know the associated A record are listed before records for which
;; we don't.
(define (closest-untried-nameservers q zone nameservers-tried)
  (define name (question-name q))
  (define ns-rrset (closest-nameservers name zone))
  (let loop ((untried (set->list (set-subtract ns-rrset nameservers-tried)))
	     (with-address '())
	     (without-address '()))
    (if (null? untried)
	(append with-address without-address)
	(let ((ns-rr (car untried)))
	  (define rrs (hash-ref zone (rr-rdata ns-rr) (set)))
	  (define a-rrs (filter-by-type rrs 'a))
	  (define has-address? (not (set-empty? a-rrs)))
	  (loop (cdr untried)
		(if has-address? (cons ns-rr with-address) with-address)
		(if has-address? without-address (cons ns-rr without-address)))))))

(define (empty-answer q zone start-of-authority)
  (if (and start-of-authority
	   (in-bailiwick? (question-name q) start-of-authority))
      ;; NXDOMAIN/name-error if the question is something we're qualified to answer
      #f
      ;; A normal no-answers packet otherwise.
      (question-result q
		       zone
		       (set)
		       (set)
		       (set))))

(define (random-element a-nonempty-list)
  (car a-nonempty-list))

(define first-timeout 3) ;; seconds

;; seconds -> Maybe<seconds>
(define (next-timeout timeout)
  (case timeout
    ((3) 11)
    ((11) 45)
    ((45) #f)))

;; IPv4 -> String
(define (ip->host-name ip-address)
  (match-define (vector a b c d) ip-address)
  (format "~a.~a.~a.~a" a b c d))

(define (negative-network-query-result zone)
  zone)

(define (make-network-query-packet q)
  (dns-message->packet
   (dns-message (random 65536)
		'request
		'query
		'non-authoritative
		'not-truncated
		'no-recursion-desired
		'no-recursion-available
		'no-error
		(list q)
		'()
		'()
		'())))

(define (incorporate-claims claim-rrset ns-rr zone)
  (foldl (lambda (claim-rr zone)
	   (if (in-bailiwick? (rr-name claim-rr) ns-rr)
	       (incorporate-rr claim-rr zone)
	       zone))
	 zone
	 claim-rrset))

(define (incorporate-dns-reply m zone ns-rr keep-trying)
  (case (dns-message-response-code m)
    [(no-error) 
     (foldl (lambda (claim-rr zone)
	      (if (in-bailiwick? (rr-name claim-rr) ns-rr)
		  (incorporate-rr claim-rr zone)
		  zone))
	    zone
	    (append (dns-message-answers m)
		    (dns-message-authorities m)
		    (dns-message-additional m)))]
    [(name-error) #f]
    [else (keep-trying)]))

(require racket/udp)
(require racket/pretty)
(define (network-query/addresses q zone ns-rr server-ips)
  (let ((s (udp-open-socket #f #f)))
    ;; TODO: randomize ordering of servers in list.
    (let search ((timeout 3)
		 (remaining-ips server-ips))
      (if (null? remaining-ips)
	  (let ((new-timeout (next-timeout timeout)))
	    (if new-timeout
		(search new-timeout server-ips)
		(negative-network-query-result zone)))
	  (let ((ip (car remaining-ips)))
	    (define server-host-name (ip->host-name ip))
	    (define server-port 53)
	    (write `(querying ,server-host-name ,server-port with timeout ,timeout)) (newline)
	    (udp-send-to s server-host-name server-port (make-network-query-packet q))
	    (define buffer (make-bytes 512)) ;; maximum DNS reply length
	    (define result (sync/timeout timeout (udp-receive!-evt s buffer)))
	    ;; TODO: correlate query-ID
	    ;; TODO: maybe receive only specifically from the queried IP address?
	    (if result
		(let* ((reply-length (car result))
		       (packet (subbytes buffer 0 reply-length))
		       (reply-message (packet->dns-message packet)))
		  (pretty-print `(response ,result ,reply-message))
		  (incorporate-dns-reply reply-message
					 zone
					 ns-rr
					 (lambda ()
					   (search timeout (cdr remaining-ips)))))
		(search timeout (cdr remaining-ips))))))))

(define (network-query q zone ns-rr)
  (define ns-name (rr-rdata ns-rr))
  ;; ^ the rr-name is the subzone origin; the rr-rdata is the
  ;; nameserver for the subzone
  (match (resolve-from-zone (question ns-name 'a 'in) ;; TODO: 'aaaa ?
			    zone
			    #f
			    #t
			    (set))
    [#f (negative-network-query-result zone)] ;; Can't find the address of the nameserver!
    [(question-result _ enhanced-zone answers _ _)
     (define address-rrs (filter-by-type answers 'a))
     (if (set-empty? address-rrs)
	 (negative-network-query-result zone) ;; Again, no addresses for the nameserver!
	 (network-query/addresses q
				  enhanced-zone
				  ns-rr
				  (map rr-rdata (set->list address-rrs))))]))

;; additional-section/a : CompiledZone ListOf<DomainName>
;; Implements the "additional section" rules from RFC 1035 (and the
;; rules for IPv6 from RFC 3596). Provides A and AAAA records for
;; names mentioned in the "names" list that have entries in "zone".
(define (additional-section/a zone names)
  ;; RFC 3596 (section 3) requires that we process AAAA here as well
  ;; as A.
  (foldl (lambda (name section)
	   (set-union section
		      (set-filter (lambda (rr) (and (memv (rr-type rr) '(a aaaa))
						    (eqv? (rr-class rr) 'in)))
				  (hash-ref zone name))))
	 (set)
	 names))

;; build-referral : Question CompiledZone RR SetOf<RR> -> QuestionResult
;; Used when servers choose iterative referral over recursive
;; resolution. The RRs in ns-rrset must be NS RRs.
(define (build-referral q zone start-of-authority ns-rrset)
  (question-result q
		   zone
		   ns-rrset
		   (and start-of-authority (set start-of-authority))
		   (additional-section/a zone (set-map ns-rrset rr-rdata))))

(define (resolve-from-zone q zone start-of-authority recursion-desired? nameservers-tried)
  (if (answer-available? q zone)
      (answer-from-zone q zone start-of-authority recursion-desired?)
      (let ((best-nameservers (closest-untried-nameservers q zone nameservers-tried)))
	(if (null? best-nameservers)
	    (empty-answer q zone start-of-authority)
	    (if recursion-desired?
		(let ((best-nameserver (random-element best-nameservers)))
		  (define enhanced-zone (network-query q zone best-nameserver))
		  (if (eq? enhanced-zone #f)
		      ;; name-error received!
		      #f
		      ;; we presumably learned something that might help us
		      (resolve-from-zone q
					 enhanced-zone
					 start-of-authority
					 recursion-desired?
					 (set-add nameservers-tried best-nameserver))))
		(build-referral q zone start-of-authority (list->set best-nameservers)))))))

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

;; (require racket/trace)
;; (trace ;;resolve-from-zone
;;        ;;build-referral
;;        ;;incorporate-claims
;;        ;;additional-section/a
;;        ;;network-query
;;        ;;network-query/addresses
;;        ;;dns-message->claims
;;        ;;negative-network-query-result
;;        ;;closest-untried-nameservers
;;        ;;answer-from-zone
;;        ;;merge-replies
;;        ;;in-bailiwick?
;;        )

;; (pretty-print
;;  (resolve-from-zone (question
;; 		     ;;'(#"www" #"google" #"com")
;; 		     ;;'(#"foo" #"rallyx" #"ccs" #"neu" #"edu")
;; 		     '(#"rallyx" #"ccs" #"neu" #"edu")
;; 		     'a
;; 		     'in)
;; 		    (compile-zone-db
;; 		     ;; (list (rr '() 'ns 'in 30 '(#"google-public-dns-a" #"google" #"com"))
;; 		     ;; 	  (rr '(#"google-public-dns-a" #"google" #"com") 'a 'in 30 '#(8 8 8 8)))
;; 		     (list (rr '() 'ns 'in 30 '(#"f" #"root-servers" #"net"))
;; 			   (rr '(#"f" #"root-servers" #"net") 'a 'in 30 '#(198 41 0 4))
;; 			   (rr '(#"f" #"root-servers" #"net") 'a 'in 30 '#(192 228 79 201))
;; 			   (rr '(#"f" #"root-servers" #"net") 'a 'in 30 '#(192 33 4 12))
;; 			   (rr '(#"f" #"root-servers" #"net") 'a 'in 30 '#(192 203 230 10))
;; 			   (rr '(#"f" #"root-servers" #"net") 'a 'in 30 '#(192 112 36 4))
;; 			   (rr '(#"f" #"root-servers" #"net") 'a 'in 30 '#(128 63 2 53))
;; 			   (rr '(#"f" #"root-servers" #"net") 'a 'in 30 '#(192 58 128 30))
;; 			   (rr '(#"f" #"root-servers" #"net") 'a 'in 30 '#(193 0 14 129)))
;; 		     )
;; 		    #f
;; 		    #t
;; 		    (set)))