#lang racket/base
;; DNS proxy using os-big-bang.rkt and os-udp.rkt.
;;
;;; Copyright 2010, 2011, 2012, 2013 Tony Garnock-Jones <tonyg@ccs.neu.edu>
;;;
;;; This file is part of marketplace-dns.
;;;
;;; marketplace-dns is free software: you can redistribute it and/or
;;; modify it under the terms of the GNU General Public License as
;;; published by the Free Software Foundation, either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; marketplace-dns is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with marketplace-dns. If not, see
;;; <http://www.gnu.org/licenses/>.

(require racket/match)
(require racket/set)
(require racket/bool)
(require bitsyntax)
(require "api.rkt")
(require "codec.rkt")
(require "zonedb.rkt")
(require "network-query.rkt")
(require "resolver.rkt")
(require (except-in syndicate dataspace assert))
(require syndicate/actor)
(require syndicate/drivers/timer)
(require syndicate/drivers/udp)
(require "tk-dns.rkt")

(require racket/pretty)

;; Instantiated with a collection of trusted roots to begin its
;; searches from. Performs recursive queries.

;; For discarding retransmitted requests that we're still working on.
(struct active-request (source id) #:transparent)

;; (: start-proxy : Natural CompiledZone CompiledZone -> Void)
(define (start-proxy port-number zone roots-only)
  (define server-addr (udp-listener port-number))
  (define client-addr (udp-handle 'dns-client))

  (log-info "Ready.")

  (run-ground
   (spawn-timer-driver)
   (spawn-udp-driver)
   (dataspace #:name 'dns-vm
              (dns-spy)
              (spawn #:name 'timer-relay:dns
                     (on (message (inbound ($ m (timer-expired _ _)))) (send! m))
                     (on (message ($ m (set-timer _ _ _))) (send! (outbound m))))
              (query-id-allocator)
              (dns-read-driver server-addr)
              (dns-write-driver server-addr)
              (dns-read-driver client-addr)
              (dns-write-driver client-addr)
              (packet-dispatcher server-addr)
              (question-dispatcher zone roots-only client-addr)
              (forever))))

;; (: query-id-allocator : -> Void)
(define (query-id-allocator)
  ;; TODO: track how many are allocated and throttle requests if too
  ;; many are in flight
  (spawn #:name 'query-id-allocator
         (field [allocated (set)])
         (on (message `(request ,$reply-addr allocate-query-id))
             (let recheck ()
               (define n (random 65536))
               (if (set-member? (allocated) n)
                   (recheck)
                   (begin (allocated (set-add (allocated) n))
                          (send! `(reply ,reply-addr ,n))))))
         (on (message `(release-query-id ,(? exact-nonnegative-integer? $n)))
             (allocated (set-remove (allocated) n)))))

;; (: packet-dispatcher : UdpAddress -> Void)
(define (packet-dispatcher s)
  (spawn #:name 'packet-dispatcher
         (field [old-active-requests (set)])
         (on (message ($ p (bad-dns-packet _ _ _ _)))
             (log-error "~a" (pretty-format p)))
         (on (message ($ r (dns-request $m $source s)))
             ;; ^ We only listen for requests on our server socket
             (define req-id (active-request source (dns-message-id m)))
             ;; TODO: when we have presence/error-handling, remove req-id
             ;; from active requests once request-handler pseudothread exits.
             (when (not (set-member? (old-active-requests) req-id))
               ;; ^ ignore retransmitted duplicates
               (old-active-requests (set-add (old-active-requests) req-id))
               (packet-relay req-id r)))
         (on (message ($ r (dns-reply $m s $sink)))
             (define req-id (active-request sink (dns-message-id m)))
             (old-active-requests (set-remove (old-active-requests) req-id)))))

;; (: packet-relay : ActiveRequest DNSRequest -> Void)
(define (packet-relay req-id request)
  (match-define (dns-request request-message request-source request-sink) request)

  ;; (: answer->reply : (Option Question) (Option CompleteAnswer) -> DNSReply)
  (define (answer->reply q a)
    (define-values (response-code ns us ds)
      (match a
        [#f
         (values 'name-error '() '() '())]
        [(complete-answer ns us ds)
         (values 'no-error (rr-set->list ns) (rr-set->list us) (rr-set->list ds))]))
    (dns-reply
     (dns-message (dns-message-id request-message)
                  'response
                  'query
                  'non-authoritative
                  'not-truncated
                  (dns-message-recursion-desired request-message)
                  'recursion-available
                  response-code
                  (if q (list q) '())
                  ns
                  us
                  ds)
     request-sink
     request-source))

  (spawn*
   #:name (list 'packet-relay req-id)

   ;; TODO: pay attention to recursion-desired flag
   (match (dns-message-questions request-message)
     ['()
      ;; No questions!
      (send! (answer->reply #f (empty-complete-answer)))]
     [(cons original-question _)
      ;; At least one question
      (log-debug (format "Looking up ~v with query id ~v"
                         original-question (dns-message-id request-message)))
      (send! original-question)
      (react (stop-when (message (answered-question original-question $answer))
                        (log-debug "Final answer to ~v with query id ~v is ~v"
                                   original-question
                                   (dns-message-id request-message)
                                   answer)
                        (send! (answer->reply original-question answer))))])))

;; (: glueless-question-handler : CompiledZone Question UdpAddress -> Void)
(define (glueless-question-handler roots-only-zone q client-sock)
  ;; Restart q, an overly-glueless question, from the roots.
  (define restarted-question (restart-question q))
  (spawn #:name (list 'glueless-question-handler q)
         (stop-when (message (answered-question restarted-question $ans))
                    ;; We got the answer to our restarted question; now transform
                    ;; it into an answer to the original question, to unblock the
                    ;; original questioner.
                    (send! (answered-question q ans)))
         (on-start (question-handler roots-only-zone restarted-question client-sock))))

;; (: question-dispatcher : CompiledZone CompiledZone UdpAddress -> Void)
(define (question-dispatcher seed-zone roots-only client-sock)
  (define-values (cleaned-seed-zone initial-timers) (zone-expire seed-zone))
  (spawn #:name 'question-dispatcher
         (field [zone cleaned-seed-zone])
         (on-start (set-timers! initial-timers))

         (define (set-timers! timers)
           (for/list ([timerspec timers])
             (match-define (cons name ttl) timerspec)
             (send! (set-timer (list 'check-dns-expiry name) (* ttl 1000) 'relative))))

         ;; TODO: consider deduping questions here too?

         (on (message `(debug-dump))
             (with-output-to-file "zone-proxy.zone"
               (lambda ()
                 (write-bytes (bit-string->bytes (zone->bit-string (zone)))))
               #:mode 'binary
               #:exists 'replace)
             (with-output-to-file "zone-proxy.dump"
               (lambda ()
                 (display "----------------------------------------------------------------------\n")
                 (display (seconds->date (current-seconds)))
                 (newline)
                 (for ([name (in-hash-keys (zone))])
                   (define rrmap (hash-ref (zone) name))
                   (for ([rr (in-hash-keys rrmap)])
                     (define expiry (hash-ref rrmap rr))
                     (write (list rr expiry))
                     (newline)))
                 (newline))
               #:mode 'text
               #:exists 'append)
             ;; (with-output-to-file "zone-proxy.debug"
             ;;   (lambda ()
             ;;     (display "----------------------------------------------------------------------\n")
             ;;     (display (seconds->date (current-seconds)))
             ;;     (newline)
             ;;     (pretty-write current-ground-transition))
             ;;   #:mode 'text
             ;;   #:exists 'append)
             )

         (on (message ($ q (question _ _ _ _)))
             (cond
               [(question-cyclic? q)
                (log-warning (format "Cyclic question ~v" q))
                (send! (answered-question q (empty-complete-answer)))]
               [(question-too-glueless? q)
                (log-warning (format "Overly-glueless question ~v" q))
                (glueless-question-handler roots-only q client-sock)]
               [else
		(question-handler (zone) q client-sock)]))

         (on (message (network-reply _ $answer))
             (define-values (new-zone timers) (incorporate-complete-answer answer (zone) #t))
             (zone new-zone)
             (set-timers! timers))

         (on (message (timer-expired (list 'check-dns-expiry (? domain? $name))
                                     (? number? $now-msec)))
             (zone (zone-expire-name (zone) name (/ now-msec 1000.0))))))

;; (: send-empty-reply! : Question -> Void)
(define (send-empty-reply! q)
  (send! (answered-question q (empty-complete-answer))))

;; (: question-handler : CompiledZone Question UdpAddress -> Void)
(define (question-handler zone0 q client-sock)
  (spawn*
   #:name (list 'question-handler q)
   (let retry-question ((zone zone0)
                        (nameservers-tried (set))
                        (retry-count 0))
     (if (= retry-count 20) ;; TODO: is this a sensible limit?
         ;; Too many retries, i.e. too many referrals.
         (begin (log-error (format "Too many retries: ~v" q))
                (send-empty-reply! q))
         ;; Credit remaining. Try once more (perhaps for the first time, in fact).
         (let ((resolution-result (resolve-from-zone q zone #f nameservers-tried)))
           (log-debug (format "Resolution result: ~v" resolution-result))

           (match resolution-result
             [#f ;; We're not authoritative so this is just a signal that we can't answer usefully
              (send-empty-reply! q)]

             [(referral zone-origin nameserver-rrs _)
              (define referral-id (gensym 'referral))
              (log-debug (format "Referral for ~v id ~v to ~v servers ~v"
                                 q referral-id (domain-labels zone-origin)
                                 (map domain-labels (set-map nameserver-rrs rr-rdata-domain-name))))
              (network-query client-sock
                             q
                             zone-origin
                             (map rr-rdata-domain-name (set->list nameserver-rrs))
                             referral-id)
              (react (stop-when
                      (message (network-reply referral-id $ans))
                      (cond [(not ans) ;; name-error/NXDOMAIN
                             (send! (answered-question q #f))]
                            [else
                             (define-values (new-zone _ignored-timers)
                               (incorporate-complete-answer ans zone #f))
                             (when (log-level? (current-logger) 'debug)
                               (log-debug "Referral ~v results in origin ~v:~n"
                                          referral-id zone-origin)
                               (for ([k (set-union (list->set (hash-keys zone))
                                                   (list->set (hash-keys new-zone)))]
                                     #:when (in-bailiwick? k zone-origin))
                                 (log-debug "Old ~v ~v~nNew ~v ~v"
                                            k (hash-ref zone k (lambda () 'missing))
                                            k (hash-ref new-zone k (lambda () 'missing))))
                               (log-debug "=-=-=-=-=-="))
                             (define nameserver-names
                               (for/set ([rr nameserver-rrs]) (rr-rdata-domain-name rr)))
                             (retry-question new-zone
                                             (set-union nameservers-tried nameserver-names)
                                             (+ retry-count 1))])))]

             [(? complete-answer? ans)
              (send!  (answered-question q ans))]

             [(partial-answer base cnames)
              ;; TODO: record chains of CNAMEs to avoid pathologically-long chains
              (react (field [acc base]
                            [remaining (length cnames)])
                     (on-start (for [(cname cnames)]
                                 (define cname-q (cname-question cname q))
                                 (react (on-start (send! cname-q))
                                        (stop-when (message (answered-question cname-q $ans))
                                                   (acc (if ans (merge-answers (acc) ans) (acc)))
                                                   (remaining (- (remaining) 1))))))
                     (stop-when (rising-edge (zero? (remaining)))
                                (send! (answered-question q (acc)))))]))))))

(require "test-rrs.rkt")
(require racket/file)
(file-stream-buffer-mode (current-output-port) 'none)
(start-proxy (test-port-number)
	     (if (file-exists? "zone-proxy.zone")
		 (bit-string->zone (file->bytes "zone-proxy.zone"))
		 (compile-zone-db test-roots))
	     (compile-zone-db test-roots))