#lang racket/base

(require racket/pretty)

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

(provide (struct-out resolver-network-query)
	 resolve-from-zone)

(struct resolver-network-query (q zone ns-rr addresses k) #:transparent)

;; Rules:
;;
;;  - If the DB already has an answer, return it.
;;
;;  - Otherwise, find the leafmost NS record in the DB for the
;;    requested name.
;;
;;  - Query that service. Augment the DB with the answers received, if
;;    any. Loop back to the beginning, remembering that we've tried
;;    the specific service we just interacted with so we don't try it
;;    again.
;;
;;  - Eventually, the DB will have either been augmented with an
;;    answer, or we will have run out of untried nameservers to ask.
;;
;;  - Authoritative NXDOMAINs ('name-error) mean we get to stop
;;    looking.
;;
;;  - Resolve CNAMEs on the way. Remember which names we've been
;;    resolving in response to any given query, to avoid
;;    loops. Perhaps limit the length of the chain to avoid
;;    DoS. (TODO)
;;
;;  - Only performs recursive service if so requested.
;;
;;  - See RFC 1035 section 7.1.

;; 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? ks kf)
  (match-define (question name qtype qclass) q)
  (define rrset (hash-ref zone name set))
  (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)))
  (define (k qr)
    (if (set-empty? (question-result-answers qr))
	(kf)
	(ks qr)))
  ;; Algorithm for CNAME expansion from RFC 1034 4.3.2 step 3a.
  (if (and (not (set-empty? cnames)) (not (eqv? qtype 'cname)))
      (let loop ((cnames (set->list cnames))
		 (reply base-reply))
	(if (null? cnames)
	    (k reply)
	    (let ((cname-rr (car cnames)))
	      (resolve-from-zone
	       (question (rr-rdata cname-rr) qtype qclass)
	       zone
	       start-of-authority
	       recursion-desired?
	       (set)
	       (lambda (qr) (loop (cdr cnames) (merge-replies reply qr)))))))
      (k 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 (network-query q zone ns-rr k)
  (define ns-name (rr-rdata ns-rr))
  ;; ^ the rr-name is the subzone origin; the rr-rdata is the
  ;; nameserver for the subzone
  (resolve-from-zone
   (question ns-name 'a 'in) ;; TODO: 'aaaa ?
   zone
   #f
   #t
   (set)
   (lambda (qr)
     (match qr
       [#f (k 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)
	    (k zone) ;; Again, no addresses for the nameserver!
	    (resolver-network-query q
				    enhanced-zone
				    ns-rr
				    (map rr-rdata (set->list address-rrs))
				    k))]))))

;; 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 k)
  (answer-from-zone q zone start-of-authority recursion-desired?
   k
   (lambda ()
     (let ((best-nameservers (closest-untried-nameservers q zone nameservers-tried)))
       (if (null? best-nameservers)
	   (k (empty-answer q zone start-of-authority))
	   (if recursion-desired?
	       (let ((best-nameserver (random-element best-nameservers)))
		 (network-query q zone best-nameserver
				(lambda (enhanced-zone)
				  (write `(BACK-FROM-NETWORK-QUERY (original-question ,q) (best-nameserver ,best-nameserver) (qr ,enhanced-zone))) (newline)
				  (if (eq? enhanced-zone #f)
				      ;; name-error received!
				      (k #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)
							 k)))))
	       (k (build-referral q zone start-of-authority (list->set best-nameservers)))))))))