#lang racket/base

;; Simple imperative DNS server.

(require racket/udp)
(require racket/set)
(require racket/bool)
(require (planet tonyg/bitsyntax))
(require "api.rkt")
(require "codec.rkt")

;; Instantiated with a SOA record for the zone it is serving as well
;; as a zone's worth of DNS data which is used to answer queries
;; authoritatively. Never caches information, never performs recursive
;; queries.

;; Rules:

;;  - Answers authoritative NXDOMAIN answers for queries falling within
;;    its zone. (This is the only responder entitled to answer NXDOMAIN!)
;;  - Answers with referrals for queries falling in subzones. It
;;    determines subzones based on the RRs it is configured with at
;;    startup.

;; compile-db : ListOf<RR> -> Hash<DomainName,ListSetOf<RR>>
;; Builds an immutable hash table from the given RRs, suitable for
;; quickly looking up answers to queries.
(define (compile-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))

(define (in-bailiwick? dn root)
  (cond
   ((equal? dn root) #t)
   ((null? dn) #f)
   (else (in-bailiwick? (cdr dn) root))))

(define (authoritativeness-for dn soa-rr)
  (if (in-bailiwick? dn (rr-name soa-rr))
      'authoritative
      'non-authoritative))

(define (filter-by-type rrset type)
  (set-filter (lambda (rr) (eqv? (rr-type rr) type)) rrset))

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

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

;; ASSUMPTION: r1 and r2 are both DNS replies to the same query.
;; ASSUMPTION: no response-codes other than no-error or name-error are in use.
(define (merge-replies r1 r2)
  (dns-message (dns-message-id r1)
	       'response
	       'query
	       (if (and (eqv? (dns-message-authoritative r1) 'authoritative)
			(eqv? (dns-message-authoritative r2) 'authoritative))
		   'authoritative
		   'non-authoritative)
	       'not-truncated
	       (dns-message-recursion-desired r1)
	       'no-recursion-available
	       (if (and (eqv? (dns-message-authoritative r1) 'name-error)
			(eqv? (dns-message-authoritative r2) 'name-error))
		   'name-error
		   'no-error)
	       (dns-message-questions r1)
	       (listset-union (dns-message-answers r1) (dns-message-answers r2))
	       (listset-union (dns-message-authorities r1) (dns-message-authorities r2))
	       (listset-union (dns-message-additional r1) (dns-message-additional r2))))

(define (listset-union xs1 xs2)
  (set->list (set-union (list->set xs1) (list->set xs2))))

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

;; start-server : UInt16 RR ListOf<RR> -> Void
;; Starts a server that will answer questions received on the given
;; UDP port based on the RRs it is given and the zone origin specified
;; in the soa-rr given.
(require racket/pretty)
(define (start-server port-number soa-rr rrs)
  ;; Compile the zone hash table
  (define zone (compile-db (cons soa-rr rrs)))

  (pretty-print zone)

  ;; Set up the socket
  (define s (udp-open-socket #f #f))
  (udp-bind! s #f port-number)

  (let service-loop ()

    (define buffer
      (make-bytes 512))
    (define-values (packet-length source-hostname source-port)
      (udp-receive! s buffer))
    (define request-message
      (packet->dns-message (sub-bit-string buffer 0 (* 8 packet-length))))

    ;; TODO: check opcode in request

    (define (make-reply name send-name-error? answers authorities additional)
      (dns-message (dns-message-id request-message)
		   'response
		   'query
		   (authoritativeness-for name soa-rr)
		   'not-truncated
		   (dns-message-recursion-desired request-message)
		   'no-recursion-available
		   (if send-name-error? 'name-error 'no-error)
		   (dns-message-questions request-message)
		   (set->list answers)
		   (set->list authorities)
		   (set->list additional)))

    ;; TODO: what if there are multiple questions in one request
    ;; packet?  Single reply, or multiple replies? djbdns looks like
    ;; it handles exactly one question per request...

    ;; TODO: Truncation

    ;; 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 (answer-question q)
      (let resolve ((name (question-name q)))
	;; Notice that we claim to be authoritative for our configured
	;; zone. If we ever answer name-error, that means there are no
	;; RRs *at all* for the queried name. If there are RRs for the
	;; queried name, but they happen not to be the ones asked for,
	;; name-error must *not* be returned: instead, a normal
	;; no-error reply is sent with an empty answer section.
	;;
	;; If we wanted to support caching of negative replies, we'd
	;; follow the guidelines in section 4.3.4 "Negative response
	;; caching" of RFC1034, adding our zone SOA with an
	;; appropriate TTL to the additional section of the reply.
	;;
	;; TODO: We support returning out-of-bailiwick records (glue)
	;; here. Reexamine the rules for doing so.
	(cond
	 ((hash-ref zone name #f) =>
	  ;; The full name matches in our zone database.
	  (lambda (rrset)
	    (define filtered-rrs (filter-rrs rrset (question-type q) (question-class q)))
	    (define cnames (filter-by-type rrset 'cname))
	    (define base-reply (make-reply name
					   #f
					   (set-union cnames filtered-rrs)
					   (set soa-rr)
					   (set)))
	    ;; Algorithm for CNAME expansion from RFC 1034 4.3.2 step 3a.
	    (if (and (not (set-empty? cnames))
		     (not (eqv? (question-type q) 'cname)))
	        (foldl (lambda (cname-rr current-reply)
			 (merge-replies current-reply
					(resolve (rr-rdata cname-rr))))
		       base-reply
		       (set->list cnames))
		base-reply)))
	 ((referral-for name soa-rr zone) =>
	  ;; No full name match, but a referral to somewhere beneath our
	  ;; SOA but within our zone.
	  (lambda (ns-rrset)
	    (make-reply name
			#f
			ns-rrset
			(set soa-rr)
			(additional-section/a zone (set-map ns-rrset rr-rdata)))))
	 (else
	  ;; Neither a full name match nor a referral is
	  ;; available. Answer name-error (NXDOMAIN) if the queried
	  ;; name is in-bailiwick, or a normal no-error otherwise.
	  (make-reply name
		      (in-bailiwick? name (rr-name soa-rr))
		      (set)
		      (set)
		      (set))))))

    ;;(display "----------------------------------------")
    ;;(newline)
    ;;(write request-message) (newline)

    ;; TODO: properly deal with multiple questions
    (for-each (lambda (q)
		(define reply-message (answer-question q))
		;;(write reply-message) (newline)
		(udp-send-to s source-hostname source-port (dns-message->packet reply-message)))
	      (dns-message-questions request-message))

    (service-loop)))

(start-server 5555
	      (rr '(#"example") 'soa 'in 30
		  (soa '(#"ns" #"example")
		       '(#"tonyg" #"example")
		       1
		       24
		       24
		       30
		       10))
	      (list (rr '(#"localhost" #"example") 'a 'in 30 '#(127 0 0 1))
		    (rr '(#"example") 'mx 'in 30 (mx 5 '(#"localhost" #"example")))
		    (rr '(#"example") 'mx 'in 30 (mx 10 '(#"subns" #"example")))
		    (rr '(#"google" #"example") 'cname 'in 30 '(#"www" #"google" #"com"))
		    (rr '(#"ns" #"example") 'a 'in 30 '#(127 0 0 1))
		    (rr '(#"subzone" #"example") 'ns 'in 30 '(#"subns" #"example"))
		    (rr '(#"subns" #"example") 'a 'in 30 '#(127 0 0 2))))