#lang racket/unit

(require racket/pretty)

(require racket/match)
(require "api.rkt")
(require "codec.rkt")
(require "zonedb.rkt")
(require "udp-operations-sig.rkt")
(require "network-query-sig.rkt")

(import udp-operations^)
(export network-query^)

;; DJB's rules for handling DNS responses. Some of these are handled
;; here (specifically, rules 2 through 5, in the action of
;; incorporate-dns-reply), some are handled in resolver-unit.rkt (rule
;; 1, in the action of answer-from-zone):

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

(define first-timeout 3) ;; seconds

;; seconds -> Maybe<seconds>
(define (next-timeout timeout)
  (case timeout
    ((3) 11)
    ((11) 45)
    ((45) #f)))

;; IPv4 -> String
(define (ip->host-name ip-address)
  (match-define (vector a b c d) ip-address)
  (format "~a.~a.~a.~a" a b c d))

(define (make-network-query-packet q)
  (dns-message->packet
   (dns-message (random 65536)
		'request
		'query
		'non-authoritative
		'not-truncated
		'no-recursion-desired
		'no-recursion-available
		'no-error
		(list q)
		'()
		'()
		'())))

;; incorporate-dns-reply :
;;   DNSMessage CompiledZone RR<NS> ( -> Maybe<CompiledZone> )
;;   -> Maybe<CompiledZone>
;;
;; Incorporates RRs from the answer, authorities, and additional
;; sections of the passed-in `message` to the passed-in `zone`,
;; returning the augmented zone. RRs are only incorporated if their
;; `rr-name` falls in the bailiwick of the given `ns-rr`. All of this
;; only happens if the passed-in message's `dns-message-response-code`
;; is `'no-error`: if it's `'name-error`, then `#f` is returned, and
;; if it's any other code,the `keep-trying` thunk is invoked. (If the
;; caller is `network-query/addresses`, then `keep-trying` will try
;; other servers from the list of IPs available.)
(define (incorporate-dns-reply message zone ns-rr keep-trying)
  (case (dns-message-response-code message)
    [(no-error) 
     (foldl (lambda (claim-rr zone)
	      (if (in-bailiwick? (rr-name claim-rr) ns-rr)
		  (incorporate-rr claim-rr zone)
		  zone))
	    zone
	    (append (dns-message-answers message)
		    (dns-message-authorities message)
		    (dns-message-additional message)))]
    [(name-error) #f]
    [else (keep-trying)]))

;; network-query/addresses :
;;   Question CompiledZone RR<NS> ListOf<IPv4> -> Maybe<CompiledZone>
;;
;; Repeatedly uses `network-query/addresses/timeout` to try asking the
;; whole of `server-ips` the question `q`, starting with a timeout of
;; `first-timeout` seconds and increasing each time
;; `network-query/addresses/timeout` returns `'no-answer` up to a
;; give-up timeout limit.
(define (network-query/addresses q zone ns-rr server-ips)
  (let ((s (udp-open-socket #f #f)))
    (let try-with-timeout ((timeout first-timeout))
      (match (network-query/addresses/timeout s q zone ns-rr server-ips timeout)
	['no-answer
	 (define new-timeout (next-timeout timeout))
	 (if new-timeout
	     (try-with-timeout new-timeout)
	     zone)]
	[result result]))))

;; network-query/addresses/timeout :
;;   UdpSocket Question CompiledZone RR<NS> ListOf<IPv4> Seconds
;;   -> (or Maybe<CompiledZone> 'no-answer)
;;
;; Sends the question to each of the servers whose addresses are given
;; in `server-ips` using `network-query/address/timeout`, one at a
;; time, in order, trying the next in the list only if `'no-answer`
;; results from the most recent communication attempt. If and when the
;; list is exhausted, `'no-answer` is returned.
(define (network-query/addresses/timeout s q zone ns-rr server-ips timeout)
  ;; TODO: randomize ordering of servers in list.
  (let search ((remaining-ips server-ips))
    (if (null? remaining-ips)
	'no-answer
	(match (network-query/address/timeout s q zone ns-rr (car remaining-ips) timeout)
	  ['no-answer (search (cdr remaining-ips))]
	  [result result]))))

;; network-query/address/timeout :
;;   UdpSocket Question CompiledZone RR<NS> IPv4 Seconds
;;   -> (or Maybe<CompiledZone> 'no-answer)
;;
;; Sends the question to the server address `server-ip` given. Waits
;; `timeout` seconds for an answer: if one arrives, it is incorporated
;; into the passed-in `zone` (using `incorporate-dns-reply`), and the
;; result is returned to the caller. If the timeout expires before a
;; reply is received, or some error result is received from the
;; server, `'no-answer` is returned to the caller.
(define (network-query/address/timeout s q zone ns-rr server-ip timeout)
  (define server-host-name (ip->host-name server-ip))
  (define server-port 53)
  (write `(querying ,server-host-name ,server-port with timeout ,timeout)) (newline)
  (udp-send-to s server-host-name server-port (make-network-query-packet q))
  (define buffer (make-bytes 512)) ;; maximum DNS reply length
  (define result (udp-receive/timeout s buffer timeout))
  ;; TODO: correlate query-ID
  ;; TODO: maybe receive only specifically from the queried IP address?
  (if result
      (let* ((reply-length (car result))
	     (packet (subbytes buffer 0 reply-length))
	     (reply-message (packet->dns-message packet)))
	(pretty-print `(response ,result ,reply-message))
	(incorporate-dns-reply reply-message
			       zone
			       ns-rr
			       (lambda () 'no-answer)))
      'no-answer))