From 087d28c9ec16f69faa2572ca8523f91228364ec3 Mon Sep 17 00:00:00 2001 From: Tony Garnock-Jones Date: Wed, 20 Mar 2013 17:10:51 -0400 Subject: [PATCH] TR proxy. --- proxy.rkt | 284 +++++++++++++++++++++++++++++++----------------------- 1 file changed, 164 insertions(+), 120 deletions(-) diff --git a/proxy.rkt b/proxy.rkt index ac30421..6767b3d 100644 --- a/proxy.rkt +++ b/proxy.rkt @@ -1,4 +1,4 @@ -#lang racket/base +#lang typed/racket/base ;; DNS proxy using os-big-bang.rkt and os-udp.rkt. @@ -11,10 +11,10 @@ (require "zonedb.rkt") (require "network-query.rkt") (require "resolver.rkt") -(require racket-typed-matrix/sugar-untyped) +(require racket-typed-matrix/sugar-typed) (require racket-typed-matrix/support/spy) -(require racket-typed-matrix/drivers/timer-untyped) -(require racket-typed-matrix/drivers/udp-untyped) +(require racket-typed-matrix/drivers/timer) +(require racket-typed-matrix/drivers/udp) (require "tk-dns.rkt") (require racket/pretty) @@ -23,72 +23,89 @@ ;; searches from. Performs recursive queries. ;; For discarding retransmitted requests that we're still working on. -(struct active-request (source id) #:prefab) +(struct: active-request ([source : UdpAddress] [id : Natural]) #:prefab) +(define-type ActiveRequest active-request) -;; start-proxy : UInt16 CompiledZone CompiledZone -> Void +(: 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) - (nested-vm #:debug-name 'dns-vm - (spawn #:debug-name 'dns-spy #:child (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))))) + (ground-vm: + ((inst generic-spy Void) 'UDP) + ((inst udp-driver Void)) + ((inst timer-driver Void)) + (nested-vm: : Void + #:debug-name 'dns-vm + (spawn: #:debug-name 'dns-spy #:parent : Void #:child : Void (dns-spy)) + ((inst timer-relay Void) 'timer-relay:dns) + (spawn: #:debug-name 'query-id-allocator #:parent : Void + #:child : (Setof Natural) + (query-id-allocator)) + (spawn: #:debug-name 'server-dns-reader #:parent : Void + #:child : Void (dns-read-driver server-addr)) + (spawn: #:debug-name 'server-dns-writer #:parent : Void + #:child : Void (dns-write-driver server-addr)) + (spawn: #:debug-name 'client-dns-reader #:parent : Void + #:child : Void (dns-read-driver client-addr)) + (spawn: #:debug-name 'client-dns-writer #:parent : Void + #:child : Void (dns-write-driver client-addr)) + (spawn: #:debug-name 'packet-dispatcher #:parent : Void + #:child : (Setof ActiveRequest) (packet-dispatcher server-addr)) + (spawn: #:debug-name 'question-dispatcher #:parent : Void + #:child : CompiledZone (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) ;; SetOf, all active query IDs - (endpoint #:subscriber `(request ,(wild) allocate-query-id) - #:state allocated + (transition: ((inst set Natural)) : (Setof Natural) ;; all active query IDs + (endpoint: allocated : (Setof Natural) + #:subscriber `(request ,(wild) allocate-query-id) [`(request ,reply-addr allocate-query-id) - (let recheck () + (let: recheck : (Transition (Setof Natural)) () (define n (random 65536)) (if (set-member? allocated n) (recheck) - (transition (set-add allocated n) + (transition: (set-add allocated n) : (Setof Natural) (send-message `(reply ,reply-addr ,n)))))]) - (endpoint #:subscriber `(release-query-id ,(wild)) - #:state allocated - [`(release-query-id ,n) - (transition (set-remove allocated n))]))) + (endpoint: allocated : (Setof Natural) + #:subscriber `(release-query-id ,(wild)) + [`(release-query-id ,(? exact-nonnegative-integer? n)) + (transition: (set-remove allocated n) : (Setof Natural))]))) +(: packet-dispatcher : UdpAddress -> (Transition (Setof ActiveRequest))) (define (packet-dispatcher s) - (transition (set) ;; SetOf - (endpoint #:subscriber (bad-dns-packet-repr (wild) (wild) (wild) (wild)) - [p (begin (log-error (pretty-format p)) '())]) - (endpoint #:subscriber (dns-request-repr (wild) (wild) s) - #:state old-active-requests - [(and r (dns-request-repr m source (== s))) ;; We only listen for requests on our server socket + (transition: ((inst set ActiveRequest)) : (Setof ActiveRequest) + (endpoint: : (Setof ActiveRequest) + #:subscriber (bad-dns-packet-pattern (wild) (wild) (wild) (wild)) + [p (begin (log-error (pretty-format p)) '())]) + (endpoint: old-active-requests : (Setof ActiveRequest) + #:subscriber (dns-request-pattern (wild) (wild) s) + [(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-repr (wild) s (wild)) - #:state old-active-requests - [(and r (dns-reply-repr m (== s) sink)) + (transition: old-active-requests : (Setof ActiveRequest)) + ;; ^ ignore retransmitted duplicates + (transition: (set-add old-active-requests req-id) : (Setof ActiveRequest) + (spawn: #:debug-name (list 'packet-relay req-id) + #:parent : (Setof ActiveRequest) + #:child : Void (packet-relay req-id r)))))]) + (endpoint: old-active-requests : (Setof ActiveRequest) + #:subscriber (dns-reply-pattern (wild) s (wild)) + [(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)))]))) + (transition: (set-remove old-active-requests req-id) : (Setof ActiveRequest)))]))) +(: packet-relay : ActiveRequest DNSRequest -> (Transition Void)) (define (packet-relay req-id request) - (match-define (dns-request-repr request-message request-source request-sink) 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 @@ -96,7 +113,7 @@ (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-repr + (dns-reply (dns-message (dns-message-id request-message) 'response 'query @@ -115,54 +132,58 @@ (match (dns-message-questions request-message) ['() ;; No questions! - (transition 'no-state/packet-relay + (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/packet-relay + (transition/no-state (send-message original-question) - (endpoint #:subscriber (answered-question-repr original-question (wild)) - #:state w - #:let-name wait-id - [(answered-question-repr (== original-question) answer) + (endpoint: : Void + #:subscriber (answered-question-pattern original-question (wild)) + #: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))))]))])) + (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 - (endpoint #:subscriber (answered-question-repr restarted-question (wild)) - #:state w - #:let-name relay - [(answered-question-repr (== 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-repr q ans)))]) - (spawn #:debug-name (list 'glueless-question-handler-inner restarted-question) - #:child (question-handler roots-only-zone restarted-question client-sock)))) + (transition/no-state + (endpoint: : Void + #:subscriber (answered-question-pattern restarted-question (wild)) + #: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. + (list (delete-endpoint relay) + (send-message (answered-question q ans)))]) + (spawn: #:debug-name (list 'glueless-question-handler-inner restarted-question) + #:parent : Void + #:child : QHState + (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]) + (transition: new-zone : CompiledZone + (for/list: : (Listof (Action CompiledZone)) ([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 + (endpoint: zone : CompiledZone + #:subscriber `(debug-dump) [`(debug-dump) (begin (with-output-to-file "zone-proxy.zone" @@ -175,9 +196,12 @@ (display "----------------------------------------------------------------------\n") (display (seconds->date (current-seconds))) (newline) - (for* ([(name rrmap) zone] [(rr expiry) rrmap]) - (write (list rr expiry)) - (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) @@ -189,41 +213,59 @@ ;; (pretty-write current-ground-transition)) ;; #:mode 'text ;; #:exists 'append) - '())]) - (endpoint #:subscriber (question-repr (wild) (wild) (wild) (wild)) - #:state zone + (transition: zone : CompiledZone))]) + (endpoint: zone : CompiledZone + #:subscriber (question-pattern (wild) (wild) (wild) (wild)) [(? question? q) - (transition zone + (transition: zone : CompiledZone (cond [(question-cyclic? q) (log-warning (format "Cyclic question ~v" q)) - (send-message (answered-question-repr q (empty-complete-answer)))] + (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))] + (spawn: #:debug-name (list 'glueless-question-handler-outer q) + #:parent : CompiledZone + #:child : Void + (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-repr (wild) (wild)) - #:state zone - [(network-reply-repr _ answer) + (spawn: #:debug-name (list 'question-handler q) + #:parent : CompiledZone + #:child : QHState + (question-handler zone q client-sock))]))]) + (endpoint: zone : CompiledZone + #:subscriber (network-reply-pattern (wild) (wild)) + [(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)))]))) + (endpoint: zone : CompiledZone + #:subscriber (timer-expired (list 'check-dns-expiry (wild)) (wild)) + [(timer-expired (list 'check-dns-expiry (? domain? name)) (? number? now-msec)) + (transition: (zone-expire-name zone name (/ now-msec 1000.0)) : CompiledZone)]))) -(struct question-state (zone q client-sock nameservers-tried retry-count) #:prefab) -(struct expanding-cnames (q accumulator remaining-count) #:prefab) +(struct: question-state ([zone : CompiledZone] + [q : Question] + [client-sock : UdpAddress] + [nameservers-tried : (Setof DomainName)] + [retry-count : Natural]) #:prefab) +(define-type QuestionState question-state) +(struct: expanding-cnames ([q : Question] + [accumulator : CompleteAnswer] + [remaining-count : Integer]) #:prefab) +(define-type ExpandingCNAMEs expanding-cnames) + +(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))) + (retry-question (question-state zone q client-sock ((inst set DomainName)) 0))) +(: send-empty-reply : QHState Question -> (Transition QHState)) (define (send-empty-reply w q) - (transition w (send-message (answered-question-repr q (empty-complete-answer))))) + (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? @@ -242,20 +284,20 @@ (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) - (endpoint #:subscriber (network-reply-repr referral-id (wild)) - #:name referral-id - #:state w - [(network-reply-repr (== referral-id) #f) ;; name-error/NXDOMAIN - (transition w + (transition: w : QHState + ((inst network-query QHState) client-sock + q + zone-origin + (map rr-rdata-domain-name (set->list nameserver-rrs)) + referral-id) + (endpoint: w : QHState + #:subscriber (network-reply-pattern referral-id (wild)) + #:name referral-id + [(network-reply (== referral-id) #f) ;; name-error/NXDOMAIN + (transition: w : QHState (delete-endpoint referral-id) - (send-message (answered-question-repr q #f)))] - [(network-reply-repr (== referral-id) ans) + (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" @@ -264,10 +306,12 @@ (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)))) + 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))) + (define nameserver-names + (list->set + (for/list: : (Listof DomainName) ([rr nameserver-rrs]) (rr-rdata-domain-name rr)))) (sequence-actions (retry-question (struct-copy question-state w [nameservers-tried (set-union nameservers-tried @@ -276,25 +320,25 @@ [retry-count (+ old-retry-count 1)])) (delete-endpoint referral-id)))]))] [(? complete-answer? ans) - (transition w (send-message (answered-question-repr q ans)))] + (transition: w : QHState (send-message (answered-question q ans)))] [(partial-answer base cnames) - (transition (expanding-cnames q base (length cnames)) - (map (lambda (cname) + (transition: (expanding-cnames q base (length cnames)) : QHState + (map (lambda: ([cname : DomainName]) ;; 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-repr cname-q (wild)) - #:state (expanding-cnames q acc remaining) - #:let-name subscription-id - [(answered-question-repr (== cname-q) ans) + (endpoint: (expanding-cnames q acc remaining) : QHState + #:subscriber (answered-question-pattern cname-q (wild)) + #: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 + (transition: new-w : QHState (delete-endpoint subscription-id) (if (zero? new-remaining) - (send-message (answered-question-repr q new-acc)) + (send-message (answered-question q new-acc)) '())))]))) cnames))])]))