#lang typed/racket/base

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

(require racket/set)
(require racket/match)
(require (only-in racket/math exact-floor exact-truncate))
(require "api.rkt")
(require "codec.rkt")
(require (planet tonyg/bitsyntax))
(require (rename-in marketplace/tr-struct-copy [tr-struct-copy struct-copy])) ;; PR13149

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

(define-type RelativeSeconds Real)
(define-type AbsoluteSeconds Real)
(define-predicate absolute-seconds? AbsoluteSeconds)

;; 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 : RelativeSeconds]) #:transparent)
(define-type InfiniteLifetime infinite-lifetime)

;; An Expiry is either an AbsoluteSeconds or an InfiniteLifetime.
(define-type Expiry (U AbsoluteSeconds InfiniteLifetime))

;; 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.
(define-type CompiledZone (HashTable DomainName (HashTable RR Expiry)))

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

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

(: current-inexact-seconds : -> AbsoluteSeconds)
(define (current-inexact-seconds)
  (/ (current-inexact-milliseconds) 1000.0))

(: still-valid? : Expiry AbsoluteSeconds -> Boolean)
(define (still-valid? expiry now)
  (or (infinite-lifetime? expiry)
      (>= expiry now)))

(: zone-ref : CompiledZone DomainName -> (Option (Setof RR)))
(define (zone-ref db name)
  (define expirymap (hash-ref db name (lambda () #f)))
  (and expirymap
       (let ((now (current-inexact-seconds)))
	 (for/fold: ([acc : (Setof RR) (set)])
	     ([resource : RR (in-hash-keys expirymap)])
	   (define expiry (hash-ref expirymap resource))
	   (if (still-valid? expiry now)
	       (let ((new-ttl (if (infinite-lifetime? expiry)
				  (infinite-lifetime-ttl expiry)
				  (- expiry now))))
		 (set-add acc
			  (struct-copy rr resource
			    [ttl (cast (exact-floor new-ttl) Nonnegative-Integer)])))
	       acc)))))

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

(: incorporate-rr : (Option 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
		     (if (zero? (rr-ttl resource0))
			 ;; We are definitely not caching this
			 ;; resource then, because we are not even
			 ;; called by incorporate-complete-answer in
			 ;; case of 0-TTL and the cache. This record
			 ;; is transient and used just for the current
			 ;; resolution. Storing it with a real 0-TTL
			 ;; would mean it immediately is ignored,
			 ;; which is silly, so store it with an
			 ;; infinite-lifetime instead.
			 (infinite-lifetime 0)
			 ;; Otherwise it has a normal TTL, which we
			 ;; honour.
			 (+ 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 (lambda () (ann #hash() (HashTable rr Expiry)))))
  (define old-expiry (hash-ref old-expirymap resource (lambda () 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]))

(: incorporate-complete-answer :
   (Option CompleteAnswer) CompiledZone Boolean -> (Values CompiledZone Timers))
(define (incorporate-complete-answer ans db is-cache?)
  (match ans
    [#f
     (values db (set))]
    [(complete-answer ns us ds)
     (define now (current-inexact-seconds))
     (for/fold ([db db] [timers ((inst set Timer))])
	 ([rr (in-list (append (set->list ns)
			       (set->list us)
			       (set->list ds)))]) ;; no in-sequences in typed racket
       (if (and is-cache? (zero? (rr-ttl rr))) ;; Do not *cache* 0-TTL RRs (RFC 1034 3.6)
	   (values db timers)
	   (values ((incorporate-rr now) rr db)
		   (set-add timers (cons (rr-name rr) (rr-ttl rr))))))]))

(: zone-expire-name : CompiledZone DomainName AbsoluteSeconds -> 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 empty-expirymap (ann #hash() (HashTable RR Expiry)))
  (define old-expirymap (hash-ref db name (lambda () empty-expirymap)))
  (define new-expirymap
    (for/fold: ([acc : (HashTable RR Expiry) empty-expirymap])
	([resource : RR (in-hash-keys old-expirymap)])
      (define expiry (hash-ref old-expirymap resource))
      (if (still-valid? expiry now-seconds)
	  (hash-set acc resource expiry)
	  acc)))
  (if (zero? (hash-count new-expirymap))
      (hash-remove db name)
      (hash-set db name new-expirymap)))

(: zone-expire : 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 : CompiledZone zone] [timers : Timers (set)])
      ([name : DomainName (in-hash-keys zone)])
    (define new-zone (zone-expire-name zone name now-seconds))
    (define expirymap (hash-ref new-zone name (lambda () #f)))
    (values new-zone
	    (if expirymap
		(set-union (list->set
			    (map (lambda: ([e : AbsoluteSeconds]) (cons name (- e now-seconds)))
				 (filter absolute-seconds? (hash-values expirymap))))
			   timers)
		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))

(: in-bailiwick? : DomainName DomainName -> Boolean)
;; Answers #t iff dn falls within the bailiwick of the zone with
;; origin o.
(define (in-bailiwick? dn o)
  (or (equal? dn o)
      (let ((p (domain-parent dn)))
	(and p (in-bailiwick? p o)))))

(: set-filter : (All (X) (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/fold: ([acc : (Setof X) (set)]) ([x : X (in-list (set->list in))])
    (if (predicate x) (set-add acc x) acc)))

(: filter-by-type : (Setof RR) RRType -> (Setof RR))
;; Selects only those members of rrset having rr-type type.
(define (filter-by-type rrset type)
  (define p? (rdata-type-pred type))
  (set-filter (lambda: ([rr : RR]) (p? (rr-rdata rr))) rrset))

(: no-rrs : (Setof RR))
(define no-rrs (set))

(: filter-rrs : (Setof RR) QueryType QueryClass -> (Setof RR))
;; 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
    (cond
     ((eq? qtype '*) rrs)
     ((eq? qtype 'axfr) no-rrs) ;; TODO: warn? error? AXFR is not currently supported.
     ((eq? qtype 'mailb) no-rrs) ;; TODO: warn? error? MAILB is not currently supported.
     ((eq? qtype 'maila) no-rrs) ;; TODO: warn? error? MAILA is not currently supported.
     (else (filter-by-type rrs qtype))))
  (define filtered-by-type-and-class
    (case qclass
      ((*) filtered-by-type)
      (else (set-filter (lambda: ([rr : 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)
  (define cnames (filter-by-type rrs 'cname))
  (append (cname-sort (set->list cnames))
	  (set->list (set-subtract rrs cnames))))

(: rr-rdata-domain-name : RR -> DomainName)
(define (rr-rdata-domain-name rr)
  (rdata-domain-name (cast (rr-rdata rr) rdata-domain)))

(: cname-sort : (Listof RR) -> (Listof RR))
;; 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-domain-name cnames)))
  (define roots (set->list (set-subtract lhss rhss))) ;; Nodes not the targets of some edge.
  (: targets-of : DomainName -> (Listof RR))
  (define (targets-of name) (filter (lambda: ([rr : RR]) (equal? (rr-name rr) name)) cnames))
  (let: iterate ((remaining : (Listof DomainName) roots)
		 (seen : (Setof DomainName) (set))
		 (acc : (Listof RR) '()))
    (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-domain-name rrs)))
		(iterate (append targets (cdr remaining))
			 (set-add seen source)
			 (append rrs acc))))))))

(: zone->bit-string : CompiledZone -> BitString)
;; Produces a serialized form of the zone suitable for saving to disk.
(define (zone->bit-string zone)
  (for/fold: ([acc : BitString (bit-string)]) ([name : DomainName (in-hash-keys zone)])
    (define rrmap (hash-ref zone name))
    (for/fold: ([acc : BitString acc]) ([rr : RR (in-hash-keys rrmap)])
      (define expiry (hash-ref rrmap rr))
      (bit-string-append
       acc
       (cond
	[(infinite-lifetime? expiry)
	 (bit-string (rr :: (t:rr)) 1 ((exact-truncate (infinite-lifetime-ttl expiry)) :: bits 32))]
	[else
	 (bit-string (rr :: (t:rr)) 0 ((exact-truncate expiry) :: bits 32))])))))

(: bit-string->zone : 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) ]
       (define new-ttl (exact-floor (- expirytime now)))
       (if (negative? new-ttl)
	   (loop db rest)
	   (loop ((incorporate-rr now) (struct-copy rr rr0 [ttl new-ttl]) db) rest))))))