#lang racket/base

;; Simple imperative DNS proxy.

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

;; Instantiated with a collection of trusted roots to begin its
;; searches from. Performs recursive queries. Doesn't yet cache
;; responses, but will do so in future.

;; Rules:

;;  - Finds the leafmost NS record it can find in its rootset for the
;;    requested name.
;;  - Queries that service. If the answer is a referral, follows the
;;    chain. Remembers which servers it has seen before to avoid
;;    loops.
;;  - If it resolves CNAMEs on the way (should it?), remembers which
;;    names it has been resolving in response to any given query, to
;;    avoid loops. Perhaps limit the length of the chain to avoid DoS.
;;  - Only performs recursive service if so requested! (TODO)
;;  - Never put CNAME records anywhere in an answer section other than
;;    at the top (TODO; also check the server)
;;
;;  - See RFC 1035 section 7.1.

;; DJB's rules for handling DNS responses:

;; When a cache receives a normal DNS response, it learns exactly one
;; of the following five pieces of information:
;;
;; 1. ``The query was not answered because the query name is an
;; alias. I need to change the query name and try again.'' This
;; applies if the answer section of the response contains a CNAME
;; record for the query name and CNAME does not match the query type.
;;
;; 2. ``The query name has no records answering the query, and is also
;; guaranteed to have no records of any other type.'' This applies if
;; the response code is NXDOMAIN and #1 doesn't apply. The amount of
;; time that this information can be cached depends on the contents of
;; the SOA record in the authority section of the response, if there
;; is one.
;;
;; 3. ``The query name has one or more records answering the query.''
;; This applies if the answer section of the response contains one or
;; more records under the query name matching the query type, and #1
;; doesn't apply, and #2 doesn't apply.
;;
;; 4. ``The query was not answered because the server does not have
;; the answer. I need to contact other servers.'' This applies if the
;; authority section of the response contains NS records, and the
;; authority section of the response does not contain SOA records, and
;; #1 doesn't apply, and #2 doesn't apply, and #3 doesn't apply. The
;; ``other servers'' are named in the NS records in the authority
;; section.
;;
;; 5. ``The query name has no records answering the query, but it may
;; have records of another type.'' This applies if #1 doesn't apply,
;; and #2 doesn't apply, and #3 doesn't apply, and #4 doesn't
;; apply. The amount of time that this information can be cached
;; depends on the contents of the SOA record in the authority section,
;; if there is one.
;;
;; This procedure requires an incredible amount of bug-prone parsing
;; for a very small amount of information. The underlying problem is
;; that DNS was designed to declare information in a human-oriented
;; format, rather than to support crucial operations in the simplest
;; possible way.

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

;; start-proxy : UInt16 ListOf<RR> -> Void
;; Starts a proxy service that will answer questions received on the
;; given UDP port based on the NS RRs it is given.
(require racket/pretty)
(define (start-proxy port-number raw-roots)
  ;; Compile the table of roots
  (define roots (compile-zone-db raw-roots))
  (pretty-print roots)

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

  (define (service-loop)
    (with-handlers ((exn:break? (lambda (e) (raise e)))
		    (exn? (lambda (e)
			    (display "Error in DNS proxy handler:") (newline)
			    (write e)
			    (newline)
			    (newline))))
      (read-and-process-request))
    (service-loop))

  (define (read-and-process-request)
    (define buffer (make-bytes 512))
    (define-values (packet-length source-hostname source-port)
      (udp-receive! s buffer))

    (define (send-error error-response-code)
      (bit-string-case buffer
	([ (id :: bits 16) (:: binary) ]
	 (udp-send-to s source-hostname source-port
		      (dns-message->packet
		       (dns-message id 'response 'query
			     'non-authoritative 'not-truncated
			     'no-recursion-desired 'recursion-available
			     error-response-code '() '() '() '()))))
	(else
	 ;; We don't even have enough information in the packet to reply.
	 (void))))

    (display "----------------------------------------") (newline)
    (write (subbytes buffer 0 packet-length)) (newline)
    (dump-bytes! buffer packet-length)
    (flush-output)

    (define request-message
      (with-handlers ((exn? (lambda (e)
			      (send-error 'format-error)
			      (raise e))))
	(packet->dns-message (subbytes buffer 0 packet-length))))

    ;;(write request-message) (newline)

    (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)
		   '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 reply-packet
      (with-handlers ((exn? (lambda (e)
			      (send-error 'server-failure)
			      (raise e))))
	;; TODO: check opcode and direction in request
	(define questions (dns-message-questions request-message))
	(if (null? questions)
	    #f ;; No questions -> don't send any replies
	    (begin
	      ;; TODO: what if there are multiple questions in one
	      ;; request packet? Single reply, or multiple replies?
	      ;; Process the additional questions, or ignore them?
	      ;; djbdns looks like it handles exactly one question per
	      ;; request, ignoring any excess...
	      (dns-message->packet (answer-question (car questions) make-reply))))))

    ;; TODO: Truncation
    (when reply-packet
      (udp-send-to s source-hostname source-port reply-packet)))

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

  (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 '(#"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))))