#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
	 referral-for
	 additional-section/a
	 filter-rrs

	 resolve)

;; 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.

;; 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)
  ;; RR Hash -> Hash
  (define (incorporate-rr rr db)
    (hash-set db (rr-name rr) (set-add (hash-ref db (rr-name rr) (set)) rr)))
  (foldl incorporate-rr (make-immutable-hash) rrs))

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

;; 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))

;; TODO: Make limit maybe #f?? Representing no limit, for DNS roots??
(define (referral-for name soa-rr zone)
  (define limit (rr-name soa-rr))
  (let search ((name name))
    (cond
     ((or (null? name) (equal? name limit))
      ;; We've walked up the tree to the top of the zone. No referrals
      ;; are possible.
      #f)
     ((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)))
     (else
      ;; Nothing for this suffix. Keep lookup.
      (search (cdr name))))))

;; 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))

;; 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)

;; 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 (resolve q soa-rr knowledge recursive-resolver)
  ;; Extract the pieces of the question:
  (match-define (question name qtype qclass) q)
  ;; Examine knowledgebase:
  (cond
   [(hash-ref knowledge name #f) =>
    ;; The full name matches in our collection of trusted facts.
    (lambda (rrset)
      (define filtered-rrs (filter-rrs rrset qtype qclass))
      (define cnames (filter-by-type rrset 'cname))
      (define base-reply (question-result q
					  knowledge
					  (set-union cnames filtered-rrs)
					  (if (in-bailiwick? name (rr-name soa-rr))
					      (set soa-rr)
					      (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 (question (rr-rdata cname-rr) qtype qclass)
					   soa-rr
					   (question-result-knowledge current-reply)
					   recursive-resolver)))
		 base-reply
		 (set->list cnames))
	  base-reply))]
   [(referral-for name soa-rr knowledge) =>
    ;; No full name match, but a referral to somewhere beneath our SOA
    ;; but within the knowledge we have.
    (lambda (ns-rrset)
      (recursive-resolver q ns-rrset))]
   [else
    ;; Neither a full name match nor a referral is available. Answer
    ;; that we have no relevant information in the zone. It's up to
    ;; the caller to decide whether this means NXDOMAIN or simply an
    ;; empty reply.
    #f]))