racket-dns-2012/proxy.rkt

222 lines
7.9 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/main.rkt")
(require "api.rkt")
(require "codec.rkt")
(require "zonedb.rkt")
(require "network-query.rkt")
(require "resolver.rkt")
(require "../racket-matrix/os.rkt")
(require "../racket-matrix/os-big-bang.rkt")
(require "../racket-matrix/os-udp.rkt")
(require "../racket-matrix/os-timer.rkt")
(require "os-dns.rkt")
;; Instantiated with a collection of trusted roots to begin its
;; searches from. Performs recursive queries. Caches in the stupidest
;; possible way.
;; For discarding retransmitted requests that we're still working on.
(struct active-request (source id) #:transparent)
;; start-proxy : UInt16 ListOf<RR> -> Void
(require racket/pretty)
(define (start-proxy port-number rrs)
(define boot-server
(os-big-bang 'no-state/boot-server
;;(spawn dns-spy)
(spawn (timer-relay))
(spawn (query-id-allocator))
(send-meta-message `(request create-server-socket (udp new ,port-number 512)))
(subscribe/fresh wait-id
(meta-message-handlers w
[`(reply create-server-socket ,s)
(transition w
(unsubscribe wait-id)
(send-meta-message
`(request create-client-socket (udp new 0 512)))
(client-socket-waiter s))]))))
(define (client-socket-waiter s)
(subscribe/fresh wait-id
(meta-message-handlers w
[`(reply create-client-socket ,c)
(display "Ready.") (newline)
(transition w
(unsubscribe wait-id)
(spawn (dns-read-driver s))
(spawn (dns-write-driver s))
(spawn (dns-read-driver c))
(spawn (dns-write-driver c))
(spawn (packet-dispatcher s))
(spawn (question-dispatcher (compile-zone-db rrs) c)))])))
(ground-vm (os-big-bang (void)
;;(spawn udp-spy)
(spawn udp-driver)
(spawn (timer-driver))
(spawn (nested-vm boot-server)))))
(define (query-id-allocator)
;; TODO: track how many are allocated and throttle requests if too
;; many are in flight
(os-big-bang (set) ;; SetOf<UInt16>, all active query IDs
(subscribe 'query-id-request-handler
(message-handlers 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)))))]
[`(release-query-id ,n)
(transition (set-remove allocated n))]))))
(define (packet-dispatcher s)
(os-big-bang (set) ;; SetOf<ActiveRequest>
(subscribe 'packet-dispatcher
(message-handlers old-active-requests
[(? bad-dns-packet? p)
(pretty-print p) ;; TODO: perhaps use metalevel events? perhaps don't bother though
old-active-requests]
[(and 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.
(if (set-member? old-active-requests req-id)
old-active-requests ;; ignore retransmitted duplicates
(transition (set-add old-active-requests req-id)
(spawn (packet-relay req-id r))))]
[(and r (dns-reply m (== s) sink))
(define req-id (active-request sink (dns-message-id m)))
(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!
(os-big-bang 'no-state/packet-relay
(send-message (answer->reply #f (empty-complete-answer))))]
[(cons original-question _)
;; At least one question
(write `(Looking up ,original-question
with query id ,(dns-message-id request-message))) (newline)
(os-big-bang 'no-state/packet-relay
(send-message original-question)
(subscribe/fresh wait-id
(message-handlers w
[(answered-question (== original-question) answer)
(transition w
(unsubscribe wait-id)
(send-message (answer->reply original-question answer)))])))]))
(define (question-dispatcher zone0 client-sock)
(os-big-bang zone0
;; TODO: consider deduping questions here too?
(subscribe 'question-handler-factory
(message-handlers zone
[(? question? q)
(transition zone
(spawn (question-handler zone q client-sock)))]
[(network-reply _ answer)
(define-values (new-zone timers) (incorporate-complete-answer answer zone))
(transition new-zone
(for/list ([(name ttl) timers])
(send-message (set-timer (list 'check-dns-expiry name) (* ttl 1000) #t))))]
[(timer-expired (list 'check-dns-expiry name) _)
(zone-expire-name zone name)]))))
(struct question-state (zone q client-sock nameservers-tried retry-count) #:transparent)
(struct expanding-cnames (q accumulator remaining-count) #:transparent)
(define (question-handler zone q client-sock)
(os-big-bang/transition
(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.
(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).
(match (resolve-from-zone q zone #f nameservers-tried)
[#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))
(transition w
(spawn (network-query client-sock
q
zone-origin
(map rr-rdata (set->list nameserver-rrs))
referral-id))
(subscribe referral-id
(message-handlers w
[(network-reply (== referral-id) #f) ;; name-error/NXDOMAIN
(transition w (send-message (answered-question q #f)))]
[(network-reply (== referral-id) ans)
(define-values (new-zone ignored-timers) (incorporate-complete-answer ans zone))
(extend-transition
(retry-question (struct-copy question-state w
[nameservers-tried (set-union nameservers-tried
nameserver-rrs)]
[zone new-zone]
[retry-count (+ old-retry-count 1)]))
(unsubscribe 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)
(define cname-q (question cname (question-type q) (question-class q)))
(list (send-message cname-q)
(subscribe/fresh subscription-id
(message-handlers (expanding-cnames q acc remaining)
[(answered-question (== cname-q) ans)
(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
(unsubscribe subscription-id)
(if (zero? new-remaining)
(send-message (answered-question q new-acc))
'()))]))))
cnames))])]))
(require "test-rrs.rkt")
(start-proxy (test-port-number) test-roots)