racket-dns-2012/proxy.rkt

308 lines
12 KiB
Racket

#lang racket/base
;; DNS proxy using os-big-bang.rkt and os-udp.rkt.
(require racket/match)
(require racket/set)
(require racket/bool)
(require racket-bitsyntax)
(require "api.rkt")
(require "codec.rkt")
(require "zonedb.rkt")
(require "network-query.rkt")
(require "resolver.rkt")
(require racket-typed-matrix/sugar-untyped)
(require racket-typed-matrix/support/spy)
(require racket-typed-matrix/drivers/timer-untyped)
(require racket-typed-matrix/drivers/udp-untyped)
(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) #:prefab)
;; start-proxy : UInt16 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)
(nested-vm #:debug-name 'dns-vm
(dns-spy)
(timer-relay 'timer-relay:dns)
(spawn #:debug-name 'query-id-allocator #:child (query-id-allocator))
(spawn #:debug-name 'server-dns-reader #:child (dns-read-driver server-addr))
(spawn #:debug-name 'server-dns-writer #:child (dns-write-driver server-addr))
(spawn #:debug-name 'client-dns-reader #:child (dns-read-driver client-addr))
(spawn #:debug-name 'client-dns-writer #:child (dns-write-driver client-addr))
(spawn #:debug-name 'packet-dispatcher #:child (packet-dispatcher server-addr))
(spawn #:debug-name 'question-dispatcher
#:child (question-dispatcher zone roots-only client-addr)))))
(define (query-id-allocator)
;; TODO: track how many are allocated and throttle requests if too
;; many are in flight
(transition (set) ;; SetOf<UInt16>, all active query IDs
(endpoint #:subscriber `(request ,(wild) allocate-query-id)
#:state allocated
[`(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)))))])
(endpoint #:subscriber `(release-query-id ,(wild))
#:state allocated
[`(release-query-id ,n)
(transition (set-remove allocated n))])))
(define (packet-dispatcher s)
(transition (set) ;; SetOf<ActiveRequest>
(endpoint #:subscriber (bad-dns-packet (wild) (wild) (wild) (wild))
[p (begin (log-error (pretty-format p)) '())])
(endpoint #:subscriber (dns-request (wild) (wild) s)
#:state old-active-requests
[(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)
(spawn #:debug-name (list 'packet-relay req-id)
#:child (packet-relay req-id r)))))])
(endpoint #:subscriber (dns-reply (wild) s (wild))
#:state old-active-requests
[(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)))])))
(define (packet-relay req-id request)
(match-define (dns-request request-message request-source request-sink) request)
(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/packet-relay
(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/packet-relay
(send-message original-question)
(endpoint #:subscriber (answered-question original-question (wild))
#:state w
#:let-name wait-id
[(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))
(transition w
(delete-endpoint wait-id)
(send-message (answer->reply original-question answer))))]))]))
(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
(endpoint #:subscriber (answered-question restarted-question (wild))
#:state w
#:let-name relay
[(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.
(transition w
(delete-endpoint relay)
(send-message (answered-question q ans)))])
(spawn #:debug-name (list 'glueless-question-handler-inner restarted-question)
#:child (question-handler roots-only-zone restarted-question client-sock))))
(define (question-dispatcher seed-zone roots-only client-sock)
(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?
(endpoint #:subscriber `(debug-dump)
#:state zone
[`(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 rrmap) zone] [(rr expiry) rrmap])
(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)
'())])
(endpoint #:subscriber (question (wild) (wild) (wild) (wild))
#:state zone
[(? 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))
(spawn #:debug-name (list 'glueless-question-handler-outer q)
#:child (glueless-question-handler roots-only q client-sock))]
[else
(spawn #:debug-name (list 'question-handler q)
#:child (question-handler zone q client-sock))]))])
(endpoint #:subscriber (network-reply (wild) (wild))
#:state zone
[(network-reply _ answer)
(let-values (((new-zone timers) (incorporate-complete-answer answer zone #t)))
(transition-and-set-timers new-zone timers))])
(endpoint #:subscriber (timer-expired (list 'check-dns-expiry (wild)) (wild))
#:state zone
[(timer-expired (list 'check-dns-expiry name) now-msec)
(transition (zone-expire-name zone name (/ now-msec 1000.0)))])))
(struct question-state (zone q client-sock nameservers-tried retry-count) #:prefab)
(struct expanding-cnames (q accumulator remaining-count) #:prefab)
(define (question-handler zone q client-sock)
(retry-question (question-state zone q client-sock (set) 0)))
(define (send-empty-reply w q)
(transition w (send-message (answered-question q (empty-complete-answer)))))
(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))))
(transition w
(network-query client-sock
q
zone-origin
(map rr-rdata (set->list nameserver-rrs))
referral-id)
(endpoint #:subscriber (network-reply referral-id (wild))
#:name referral-id
#:state w
[(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 'missing)
k (hash-ref new-zone k 'missing))))
(log-debug "=-=-=-=-=-="))
(define nameserver-names (for/set ([rr nameserver-rrs]) (rr-rdata 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)
(endpoint #:subscriber (answered-question cname-q (wild))
#:state (expanding-cnames q acc remaining)
#:let-name subscription-id
[(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))