#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")
(require "simple-udp-service.rkt")

(require racket/pretty)

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

;; An Address can be an (address String Uint16) or #f, where an
;; address struct represents nonlocal UDP sockets, and #f represents
;; the local socket. This way, we don't need to know the IP or port of
;; the local socket, and we can be "multihomed".
(struct address (host port) #:prefab) ;; a UDP IP/port-number combination

(struct bad-dns-packet (detail source target reason) #:prefab)
(struct world-message (body source target) #:prefab)

;; ServerState
(struct world (roots continuations) #:prefab)

(define action-prompt (make-continuation-prompt-tag 'world-action))

;; TODO: Avoid attack amplification by not starting work on questions
;; that are already underway

;; TODO: Timeouts!!

(define (send/suspend outbound-messages awaken-key)
  (call-with-composable-continuation
   (lambda (k)
     (abort-current-continuation action-prompt
       (lambda () (values (lambda (w k)
			    (values outbound-messages
				    (struct-copy world w
						 [continuations (hash-set (world-continuations w)
									  awaken-key
									  k)])))
			  k))))
   action-prompt))

;; ( -> X) ServerState -> X ServerState
;; In this specific instance, X is likely to be ListOf<WorldMessage>.
(define (run-inferior boot world)
  (call-with-continuation-barrier ;; TODO: ???
   (lambda ()
     (define-values (computation-step-result computation-step-continuation)
       (call-with-continuation-prompt (lambda () (values (boot) #f)) action-prompt))
     (cond
      ((eq? computation-step-continuation #f)
       ;; The computation is finished, and has yielded a result.
       (values computation-step-result world))
      (else
       ;; The computation is not finished, but is waiting for an
       ;; action to complete.
       (computation-step-result world computation-step-continuation))))))

;; 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.
(define (start-proxy port-number raw-roots)
  ;; Compile the table of roots
  (define roots (compile-zone-db raw-roots))
  (pretty-print roots)

  (define initial-world (world roots (make-immutable-hash)))

  (start-udp-service
   port-number
   udp-packet->message
   outbound-message?
   message->udp-packet
   (message-handlers old-world
		     [(? bad-dns-packet? p)
		      (pretty-print p)
		      (values '() old-world)]
		     [(? request-from-downstream? r)
		      (handle-request r old-world)]
		     [(? reply-from-upstream? r)
		      (handle-reply r old-world)])
   (lambda (unhandled state)
     (error 'dns-server "Unhandled packet ~v" unhandled))
   initial-world
   #:packet-size-limit 512))

(define (udp-packet->message packet)
  (match-define (udp-packet body host port) packet)
  (define a (address host port))
  (with-handlers ((exn? (lambda (e) (bad-dns-packet body a #f 'unparseable))))
    (define message (packet->dns-message body))
    (world-message message a #f)))

(define (message->udp-packet m)
  (match-define (world-message body _ (address host port)) m)
  (udp-packet (dns-message->packet body) host port))

(define (local-address? a)
  (eq? a #f))

(define (remote-address? a)
  (address? a))

(define (outbound-message? m)
  (and (world-message? m)
       (local-address? (world-message-source m))
       (remote-address? (world-message-target m))))

(define (inbound-message? m)
  (and (world-message? m)
       (remote-address? (world-message-source m))
       (local-address? (world-message-target m))))

(define (request-from-downstream? m)
  (and (inbound-message? m)
       (eq? (dns-message-direction (world-message-body m)) 'request)
       (eq? (dns-message-opcode (world-message-body m)) 'query)))

(define (reply-from-upstream? m)
  (and (inbound-message? m)
       (eq? (dns-message-direction (world-message-body m)) 'response)
       (eq? (dns-message-opcode (world-message-body m)) 'query)))

(define (handle-request r old-world)
  (match-define (world-message (struct* dns-message
					([id query-id]
					 [recursion-desired recursion-desired]
					 [questions questions]))
			       request-source
			       request-target)
		r)
  (if (null? questions)
      (values '() old-world)
      ;; TODO: ignoring all but the car - good? bad? hmm?
      (answer-question (car questions) old-world (world-roots old-world)
		       query-id recursion-desired request-source)))

;; resolve-iteratively : Question SetOf<RR> -> QuestionResult
;; Follows a chain of referrals until it finds an answer to its
;; question.
(define (resolve-iteratively q ns-rrset)
  (let search ((seen (set))
	       (remaining (set->list ns-rrset)))
    (cond
     [(null? remaining) #f] ;; no answer available
     [(set-member? (car remaining) seen) (search seen (cdr remaining))]
     [else
      (define first-ns-rr (car remaining))
      (define ns-name (rr-name first-ns-rr))
      (define ns-addr 
	.......
	Should the main algorithm iterate to solution/fixpoint instead of recursing?
	If so, how should it treat cnames?
  (pretty-print 'resolve-iteratively)
  (define sub-query-id (random 65536)
  (define sub-query (dns-message sub-query-id
				 'request
				 'query
				 'non-authoritative
				 'not-truncated
				 #f
				 'no-recursion-available
				 'no-error
				 (list q)
				 (list)
				 (list)
				 (list)))
  (pretty-print `(back with ,(send/suspend 
			      
  (error 'resolve-iteratively "Gah!"))

;; TODO: Make sure we follow the guidelines and rules for implementing
;; DNS proxies more strictly.
(define (answer-question q w cache query-id recursion-desired request-source)
  (define (make-answer ns us ds)
    (list (world-message (dns-message query-id
				      'response
				      'query
				      'non-authoritative
				      'not-truncated
				      recursion-desired
				      'recursion-available
				      'no-error
				      (list q)
				      ns
				      us
				      ds)
			 #f
			 request-source)))
  (run-inferior (lambda ()
		  (match (resolve-from-zone q #f cache resolve-iteratively)
		    [#f
		     (make-answer '() '() '())]
		    [(question-result _ new-cache answers authorities additional)
		     (make-answer answers authorities additional)]))
		w))

(define (handle-reply r old-world)
  (error 'handle-reply "Unimplemented"))

(start-proxy 5555
	     (list (rr '() 'ns 'in 30 '(#"google-public-dns-a" #"google" #"com"))
		   (rr '(#"google-public-dns-a" #"google" #"com") 'a 'in 30 '#(8 8 8 8))))