#lang racket/base

;; DNS server using simple-udp-service.rkt.

(require racket/match)
(require racket/udp)
(require racket/set)
(require racket/bool)
(require "../racket-bitsyntax/main.rkt")
(require "api.rkt")
(require "codec.rkt")
(require "zonedb.rkt")
(require "dump-bytes.rkt")
(require "simple-udp-service.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.

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

;; 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-response-code r1) 'name-error)
			(eqv? (dns-message-response-code 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))))

;; 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-zone-db (cons soa-rr rrs)))
  (pretty-print zone)

  (start-udp-service
   port-number
   classify-dns-packet
   (lambda (packet old-state) (values (log-error packet) old-state))
   (lambda (message old-state) (values (handle-request soa-rr zone message) old-state))
   #f
   #:packet-size-limit 512))

(define (classify-dns-packet packet)
  (match-define (udp-packet body host port) packet)
  (with-handlers ((exn? (lambda (e) #f)))
    (list (packet->dns-message body) host port)))

;; UdpPacket -> ListOf<UdpPacket>
(define (log-error packet)
  (pretty-print `(bad-packet-received ,packet))
  (list))

(define (handle-request soa-rr zone request-data)
  (match-define (list request-message request-host request-port) request-data)

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

  (define (answer-question q make-reply)
    (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))))))

  ;; TODO: check opcode and direction in request
  ;; TODO: think again about multiple questions in one packet
  (map (lambda (q)
	 (udp-packet (dns-message->packet (answer-question q make-reply))
		     request-host
		     request-port))
       (dns-message-questions request-message)))

(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 '(#"roar" #"example") 'a 'in 30 '#(192 168 1 1))
		    (rr '(#"ns" #"example") 'a 'in 30 '#(127 0 0 1))
		    (rr '(#"hello" #"example") 'txt 'in 30 '(#"Hello CRASH"))
		    (rr '(#"subzone" #"example") 'ns 'in 30 '(#"subns" #"example"))
		    (rr '(#"subns" #"example") 'a 'in 30 '#(127 0 0 2))))