#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")
(require "../racket-bitsyntax/main.rkt")

(provide zone-ref
	 zone-includes-name?
	 incorporate-complete-answer
	 zone-expire-name
	 zone-expire
	 empty-zone-db
	 compile-zone-db
	 compiled-zone?
	 in-bailiwick?
	 set-filter
	 filter-by-type
	 filter-rrs
	 rr-set->list
	 cname-sort ;; provided for unit tests
	 zone->bit-string
	 bit-string->zone)

;; An InfiniteLifetime is an (infinite-lifetime RelativeSeconds), a
;; specification of the TTL to use when sending a non-expiring RR to a
;; peer.
(struct infinite-lifetime (ttl) #:transparent)

;; A CompiledZone is a Hash<DomainName,Hash<RR,(or AbsoluteSeconds
;; InfiniteLifetime)>>, representing a collection of DNS RRSets
;; indexed by DomainName. Each RR in an RRSet either has an expiry
;; time associated with it or has an InfiniteLifetime associated with
;; it, in which case it should not expire.

;; A Timers is a SetOf<(cons DomainName AbsoluteSeconds)>,
;; representing a collection of timeouts that should be set against
;; names to to see if their associated RRs have expired.

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

(define (current-inexact-seconds)
  (/ (current-inexact-milliseconds) 1000.0))

(define (still-valid? expiry now)
  (or (infinite-lifetime? expiry)
      (>= expiry now)))

;; CompiledZone DomainName -> Maybe<Set<RR>>
(define (zone-ref db name)
  (cond
   [(hash-ref db name #f) =>
    (lambda (expirymap)
      (define now (current-inexact-seconds))
      (for/set ([(resource expiry) expirymap] #:when (still-valid? expiry now))
	(struct-copy rr resource [ttl (if (infinite-lifetime? expiry)
					  (infinite-lifetime-ttl expiry)
					  (inexact->exact (floor (- expiry now))))])))]
   [else #f]))

;; CompiledZone DomainName -> Boolean
(define (zone-includes-name? db name)
  (hash-has-key? db name))

;; incorporate-rr : Maybe<AbsoluteSeconds> -> (RR CompiledZone -> CompiledZone)
;;
;; Incorporates the given RR into our database. If base-time is a
;; number of seconds, we treat the RR as having a TTL that decreases
;; as time goes by; otherwise base-time is #f, and we treat the RR as
;; being non-expiring with an InfiniteLifetime.
(define ((incorporate-rr base-time) resource0 db)
  (define expiry (if base-time
		     (+ base-time (rr-ttl resource0))
		     (infinite-lifetime (rr-ttl resource0))))
  (define resource (struct-copy rr resource0 [ttl 0]))
  (define name (rr-name resource))
  (define old-expirymap (hash-ref db name hash))
  (define old-expiry (hash-ref old-expirymap resource 0))
  (cond
   [(infinite-lifetime? old-expiry) ;; don't update TTL: the existing record should live forever
    db]
   [(or (infinite-lifetime? expiry) (> expiry old-expiry)) ;; update TTL
    (hash-set db name (hash-set old-expirymap resource expiry))]
   [else ;; old record finite-lifetime but expiring after the new expiry: leave it alone
    db]))

;; Maybe<CompleteAnswer> CompiledZone -> (values CompiledZone Timers)
(define (incorporate-complete-answer ans db)
  (match ans
    [#f
     (values db (set))]
    [(complete-answer ns us ds)
     (define now (current-inexact-seconds))
     (for/fold ([db db] [timers (set)])
	 ([rr (in-sequences ns us ds)])
       (values ((incorporate-rr now) rr db)
	       (set-add timers (cons (rr-name rr) (rr-ttl rr)))))]))

;; CompiledZone DomainName -> CompiledZone
;; Checks the given name to see if there are any expiring records, and
;; if so, removes them.
(define (zone-expire-name db name now-seconds)
  (define new-expirymap
    (if (zone-includes-name? db name)
	(for/hash ([(resource expiry) (hash-ref db name)] #:when (still-valid? expiry now-seconds))
	  (values resource expiry))
	(hash)))
  (if (zero? (hash-count new-expirymap))
      (hash-remove db name)
      (hash-set db name new-expirymap)))

;; CompiledZone -> (values CompiledZone Timers)
;; Used to freshen a saved zone when it is loaded from disk.
(define (zone-expire zone)
  (define now-seconds (current-inexact-seconds))
  (for/fold ([zone zone] [timers (set)])
      ([name (hash-keys zone)])
    (define new-zone (zone-expire-name zone name now-seconds))
    (cond
     [(hash-ref new-zone name #f) =>
      (lambda (expirymap)
	(values new-zone
		(set-union (list->set
			    (map (lambda (e) (cons name (- e now-seconds)))
				 (filter (lambda (e) (not (infinite-lifetime? e)))
					 (hash-values expirymap))))
			   timers)))]
     [else
      (values new-zone timers)])))

;; empty-zone-db : -> CompiledZone
(define (empty-zone-db)
  (make-immutable-hash))

;; 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 #f) (empty-zone-db) rrs))

(define (compiled-zone? z)
  (hash? z)) ;; hm

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

;; 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?
;;
;; The CNAMEs even need to be in topologically-sorted order.
;; http://homepage.ntlworld.com./jonathan.deboynepollard/FGA/dns-response-taxonomy.html
;; has this to say on this topic:
;;   "A content DNS server following the algorithm in § 4.3.2 of RFC
;;    1034 will insert this chain in first-to-last order in the
;;    response. The response parsing code in most resolving proxy DNS
;;    servers and DNS client libraries expects this order. However,
;;    the actual text of RFC 1034 itself does not guarantee it."
;; Sure enough, the resolver in Firefox seems not to be able to handle
;; CNAMEs in any order other than strictly causal. While we could be
;; more careful about retaining the ordering of RRs all the way
;; through the resolution and CNAME expansion processes, that would
;; pollute the logic with a bunch of noise about RR order which isn't
;; even supposed to be relevant. So we *recover* the order here, which
;; is a bit expensive.
(define (rr-set->list rrs)
  (append (cname-sort (set->list (filter-by-type rrs 'cname)))
	  (set->list (set-filter (lambda (rr) (not (eqv? (rr-type rr) 'cname))) rrs))))

;; cname-sort : ListOf<RR<CNAME>> -> ListOf<RR<CNAME>>
;; Produce an ordering of the CNAMEs given that respects their
;; "causality". For example, if a CNAME b and b CNAME c, then the RRs
;; will be presented in that order (and not the other order, with b
;; CNAME c first).
(define (cname-sort cnames)
  (define lhss (list->set (map rr-name cnames)))
  (define rhss (list->set (map rr-rdata cnames)))
  (define roots (set->list (set-subtract lhss rhss))) ;; Nodes not the targets of some edge.
  (define (targets-of name) (for/list [(rr cnames) #:when (equal? (rr-name rr) name)] rr))
  (let iterate ((remaining roots)
		(seen (set))
		(acc '()))
    (if (null? remaining)
	(reverse acc)
	(let ((source (car remaining)))
	  (if (set-member? seen source)
	      (iterate (cdr remaining) seen acc)
	      (let* ((rrs (targets-of source))
		     (targets (map rr-rdata rrs)))
		(iterate (append targets (cdr remaining))
			 (set-add seen source)
			 (append rrs acc))))))))

;; CompiledZone -> Bitstring
;; Produces a serialized form of the zone suitable for saving to disk.
(define (zone->bit-string zone)
  (for*/fold ([acc (bit-string)])
      ([(name rrmap) zone] [(rr expiry) rrmap])
    (bit-string-append
     acc
     (match expiry
       [(infinite-lifetime ttl)
	(bit-string (rr :: (t:rr)) 1 (ttl :: bits 32))]
       [expirytime
	(bit-string (rr :: (t:rr)) 0 ((truncate (inexact->exact expirytime)) :: bits 32))]))))

;; Bitstring -> CompiledZone
;; Produces a deserialized form of the zone. Suitable for use in loading from disk.
(define (bit-string->zone bs)
  (define now (current-inexact-seconds))
  (define empty-packet (bytes))
  (let loop ((db (empty-zone-db))
	     (bs bs))
    (bit-string-case bs
      ([ ]
       db)
      ([ (rr0 :: (t:rr empty-packet)) (= 1) (ttl :: bits 32) (rest :: binary) ]
       (loop ((incorporate-rr #f) (struct-copy rr rr0 [ttl ttl]) db) rest))
      ([ (rr0 :: (t:rr empty-packet)) (= 0) (expirytime :: bits 32) (rest :: binary) ]
       (loop ((incorporate-rr now) (struct-copy rr rr0 [ttl (- expirytime now)]) db) rest)))))