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

  (ground-vm
   (generic-spy 'UDP)
   (udp-driver)
   (timer-driver)
   (spawn-vm #:debug-name 'dns-vm
	     (name-process 'dns-spy (spawn (dns-spy)))
	     (timer-relay 'timer-relay:dns)
	     (name-process 'query-id-allocator (spawn (query-id-allocator)))
	     (name-process 'server-dns-reader (spawn (dns-read-driver server-addr)))
	     (name-process 'server-dns-writer (spawn (dns-write-driver server-addr)))
	     (name-process 'client-dns-reader (spawn (dns-read-driver client-addr)))
	     (name-process 'client-dns-writer (spawn (dns-write-driver client-addr)))
	     (name-process 'packet-dispatcher (spawn (packet-dispatcher server-addr)))
	     (name-process 'question-dispatcher (spawn (question-dispatcher zone
									    roots-only
									    client-addr))))))

;; (: query-id-allocator : -> (Transition (Setof Natural)))
(define (query-id-allocator)
  ;; TODO: track how many are allocated and throttle requests if too
  ;; many are in flight
  (transition (set) ;; all active query IDs
    (subscriber `(request ,(wild) allocate-query-id)
      (match-state allocated
	(on-message
	 [`(request ,reply-addr allocate-query-id)
	  (let recheck ()
	    (define n (random 65536))
	    (if (set-member? allocated n)
		(recheck)
		(transition (set-add allocated n)
		  (send-message `(reply ,reply-addr ,n)))))])))
    (subscriber `(release-query-id ,(wild))
      (match-state allocated
	(on-message
	 [`(release-query-id ,(? exact-nonnegative-integer? n))
	  (transition (set-remove allocated n))])))))

;; (: packet-dispatcher : UdpAddress -> (Transition (Setof ActiveRequest)))
(define (packet-dispatcher s)
  (transition (set)
    (subscriber (bad-dns-packet (wild) (wild) (wild) (wild))
      (on-message [p (begin (log-error (pretty-format p)) '())]))
    (subscriber (dns-request (wild) (wild) s)
      (match-state old-active-requests
	(on-message
	 [(and r (dns-request m source (== s)))
	  ;; ^ We only listen for requests on our server socket
	  (let ((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.
	    (if (set-member? old-active-requests req-id)
		(transition old-active-requests)
		;; ^ ignore retransmitted duplicates
		(transition (set-add old-active-requests req-id)
		  (name-process (list 'packet-relay req-id)
		    (spawn (packet-relay req-id r))))))])))
    (subscriber (dns-reply (wild) s (wild))
      (match-state old-active-requests
	(on-message
	 [(and r (dns-reply m (== s) sink))
	  (let ((req-id (active-request sink (dns-message-id m))))
	    (transition (set-remove old-active-requests req-id)))])))))

;; (: packet-relay : ActiveRequest DNSRequest -> (Transition 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))
  ;; TODO: pay attention to recursion-desired flag
  (match (dns-message-questions request-message)
    ['()
     ;; No questions!
     (transition/no-state
       (send-message (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)))
     (transition/no-state
       (send-message original-question)
       (let-fresh (wait-id)
	 (name-endpoint wait-id
	   (subscriber (answered-question original-question (wild))
	     (on-message
	      [(answered-question (== original-question) answer)
	       (begin (log-debug (format "Final answer to ~v with query id ~v is ~v"
					 original-question
					 (dns-message-id request-message)
					 answer))
		      (list (delete-endpoint wait-id)
			    (send-message (answer->reply original-question answer))))])))))]))

;; (: glueless-question-handler : CompiledZone Question UdpAddress -> (Transition 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))
  (transition/no-state
    (let-fresh (relay)
      (name-endpoint relay
	(subscriber (answered-question restarted-question (wild))
	  (on-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.
	    (list (delete-endpoint relay)
		  (send-message (answered-question q ans)))]))))
    (name-process (list 'glueless-question-handler-inner restarted-question)
      (spawn (question-handler roots-only-zone restarted-question client-sock)))))

;; (: question-dispatcher : CompiledZone CompiledZone UdpAddress -> (Transition CompiledZone))
(define (question-dispatcher seed-zone roots-only client-sock)
  ;; (: transition-and-set-timers : CompiledZone (Setof (Pairof DomainName Real))
  ;;    -> (Transition CompiledZone))
  (define (transition-and-set-timers new-zone timers)
    (transition new-zone
      (for/list ([timerspec timers])
	(match-define (cons name ttl) timerspec)
	(send-message (set-timer (list 'check-dns-expiry name) (* ttl 1000) 'relative)))))
  (define-values (cleaned-seed-zone initial-timers) (zone-expire seed-zone))
  (sequence-actions (transition-and-set-timers cleaned-seed-zone initial-timers)
    ;; TODO: consider deduping questions here too?
    (subscriber `(debug-dump)
      (match-state zone
	(on-message
	 [`(debug-dump)
	  (begin
	    (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)
	    (transition zone))])))
    (subscriber (question (wild) (wild) (wild) (wild))
      (match-state zone
	(on-message
	 [(? question? q)
	  (transition zone
	    (cond
	     [(question-cyclic? q)
	      (log-warning (format "Cyclic question ~v" q))
	      (send-message (answered-question q (empty-complete-answer)))]
	     [(question-too-glueless? q)
	      (log-warning (format "Overly-glueless question ~v" q))
	      (name-process (list 'glueless-question-handler-outer q)
		(spawn (glueless-question-handler roots-only q client-sock)))]
	     [else
	      (name-process (list 'question-handler q)
		(spawn (question-handler zone q client-sock)))]))])))
    (subscriber (network-reply (wild) (wild))
      (match-state zone
	(on-message
	 [(network-reply _ answer)
	  (let-values (((new-zone timers) (incorporate-complete-answer answer zone #t)))
	    (transition-and-set-timers new-zone timers))])))
    (subscriber (timer-expired (list 'check-dns-expiry (wild)) (wild))
      (match-state zone
	(on-message
	 [(timer-expired (list 'check-dns-expiry (? domain? name)) (? number? now-msec))
	  (transition (zone-expire-name zone name (/ now-msec 1000.0)))])))))

(struct question-state (zone q client-sock nameservers-tried retry-count) #:transparent)

(struct expanding-cnames (q accumulator remaining-count) #:transparent)

;; (define-type QHState (U QuestionState ExpandingCNAMEs))

;; (: question-handler : CompiledZone Question UdpAddress -> (Transition QHState))
(define (question-handler zone q client-sock)
  (retry-question (question-state zone q client-sock (set) 0)))

;; (: send-empty-reply : QHState Question -> (Transition QHState))
(define (send-empty-reply w q)
  (transition w (send-message (answered-question q (empty-complete-answer)))))

;; (: retry-question : QHState -> (Transition QHState))
(define (retry-question w)
  (match w
    [(question-state _ q _ _ 20) ;; TODO: is this a sensible limit?
     ;; Too many retries, i.e. too many referrals.
     (log-error (format "Too many retries: ~v" w))
     (send-empty-reply w q)]
    [(question-state zone q client-sock nameservers-tried old-retry-count)
     ;; Credit remaining. Try once more (perhaps for the first time, in fact).
     (define 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 w 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))))
	(transition w
	  (network-query client-sock
			 q
			 zone-origin
			 (map rr-rdata-domain-name (set->list nameserver-rrs))
			 referral-id)
	  (name-endpoint referral-id
	    (subscriber (network-reply referral-id (wild))
	      (match-state w
		(on-message
		 [(network-reply (== referral-id) #f) ;; name-error/NXDOMAIN
		  (transition w
		    (delete-endpoint referral-id)
		    (send-message (answered-question q #f)))]
		 [(network-reply (== referral-id) ans)
		  (let-values (((new-zone ignored-timers)
				(incorporate-complete-answer ans zone #f)))
		    (when (log-level? (current-logger) 'debug)
		      (log-debug (format "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 (format "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
		      (list->set
		       (for/list ([rr nameserver-rrs])
			 (rr-rdata-domain-name rr))))
		    (sequence-actions
			(retry-question (struct-copy question-state w
					  [nameservers-tried (set-union nameservers-tried
									nameserver-names)]
					  [zone new-zone]
					  [retry-count (+ old-retry-count 1)]))
		      (delete-endpoint referral-id)))])))))]
       [(? complete-answer? ans)
	(transition w (send-message (answered-question q ans)))]
       [(partial-answer base cnames)
	(transition (expanding-cnames q base (length cnames))
	  (map (lambda (cname)
		 ;; TODO: record chains of CNAMEs to avoid pathologically-long chains
		 (define cname-q (cname-question cname q))
		 (list (send-message cname-q)
		       (let-fresh (subscription-id)
			 (name-endpoint subscription-id
			   (subscriber (answered-question cname-q (wild))
			     (match-state (expanding-cnames q acc remaining)
			       (on-message
				[(answered-question (== cname-q) ans)
				 (let ()
				   (define new-acc (if ans (merge-answers acc ans) acc))
				   (define new-remaining (- remaining 1))
				   (define new-w (expanding-cnames q new-acc new-remaining))
				   (transition new-w
				     (delete-endpoint subscription-id)
				     (if (zero? new-remaining)
					 (send-message (answered-question q new-acc))
					 '())))])))))))
	       cnames))])]))

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