#lang typed/racket/base
;; DNS drivers using marketplace.

(require racket/set)
(require racket/match)
(require "codec.rkt")
(require marketplace/sugar-typed)
(require marketplace/drivers/udp)
(require marketplace/support/pseudo-substruct)

(provide (struct-out bad-dns-packet-repr)
	 BadDnsPacket bad-dns-packet bad-dns-packet?
	 BadDnsPacketPattern bad-dns-packet-pattern bad-dns-packet-pattern?

	 (struct-out dns-request-repr)
	 DNSRequest dns-request dns-request?
	 DNSRequestPattern dns-request-pattern dns-request-pattern?

	 (struct-out dns-reply-repr)
	 DNSReply dns-reply dns-reply?
	 DNSReplyPattern dns-reply-pattern dns-reply-pattern?

	 dns-read-driver
	 dns-write-driver
	 dns-spy)

(struct: (TDetail TSource TSink TReason)
	 bad-dns-packet-repr
	 ([detail : TDetail] [source : TSource] [sink : TSink] [reason : TReason]) #:transparent)
(pseudo-substruct: (bad-dns-packet-repr Any UdpAddress UdpAddress Symbol)
		   BadDnsPacket bad-dns-packet bad-dns-packet?)
(pseudo-substruct: (bad-dns-packet-repr Any
					(U Wild UdpAddressPattern)
					(U Wild UdpAddressPattern)
					(U Wild Symbol))
		   BadDnsPacketPattern bad-dns-packet-pattern bad-dns-packet-pattern?)

(struct: (TMessage TSource TSink)
	 dns-request-repr
	 ([message : TMessage] [source : TSource] [sink : TSink]) #:transparent)
(pseudo-substruct: (dns-request-repr DNSMessage UdpAddress UdpAddress)
		   DNSRequest dns-request dns-request?)
(pseudo-substruct: (dns-request-repr (U Wild DNSMessage)
				     (U Wild UdpAddressPattern)
				     (U Wild UdpAddressPattern))
		   DNSRequestPattern dns-request-pattern dns-request-pattern?)

(struct: (TMessage TSource TSink)
	 dns-reply-repr
	 ([message : TMessage] [source : TSource] [sink : TSink]) #:transparent)
(pseudo-substruct: (dns-reply-repr DNSMessage UdpAddress UdpAddress)
		   DNSReply dns-reply dns-reply?)
(pseudo-substruct: (dns-reply-repr (U Wild DNSMessage)
				   (U Wild UdpAddressPattern)
				   (U Wild UdpAddressPattern))
		   DNSReplyPattern dns-reply-pattern dns-reply-pattern?)

(: dns-read-driver : UdpAddress -> (Transition Void))
(define (dns-read-driver s)
  (transition: (void) : Void
    (at-meta-level
     (endpoint: : Void
		#:subscriber (udp-packet-pattern (wild) s (wild))
		[(udp-packet source (== s) #"")
		 (begin (log-info "Debug dump packet received")
			(send-message `(debug-dump)))]
		[(udp-packet source (== s) body)
		 (send-message
		  (with-handlers ((exn:fail? (lambda (e)
					       (bad-dns-packet body source s 'unparseable))))
		    (define message (packet->dns-message body))
		    (case (dns-message-direction message)
		      ((request) (dns-request message source s))
		      ((response) (dns-reply message source s)))))]))))

(: dns-write-driver : UdpAddress -> (Transition Void))
(define (dns-write-driver s)
  (: translate : DNSMessage UdpAddress -> (ActionTree Void))
  (define (translate message sink)
    (with-handlers ((exn:fail? (lambda (e)
				 (send-message (bad-dns-packet message s sink 'unencodable)))))
      (at-meta-level
       (send-message (udp-packet s sink (dns-message->packet message))))))
  (transition: (void) : Void
    (endpoint: : Void
	       #:subscriber (dns-request (wild) s (wild))
	       [(dns-request message (== s) sink) (translate message sink)])
    (endpoint: : Void
	       #:subscriber (dns-reply   (wild) s (wild))
	       [(dns-reply   message (== s) sink) (translate message sink)])))

(: dns-spy : -> (Transition Void))
(define (dns-spy)
  (transition: (void) : Void
    (endpoint: : Void
	       #:subscriber (wild) #:observer
	       [(dns-request message source sink)
		(begin (log-info (format "DNS: ~v asks ~v ~v~n   : ~v"
					 source sink (dns-message-id message)
					 (dns-message-questions message)))
		       (void))]
	       [(dns-reply message source sink)
		(begin (log-info (format "DNS: ~v answers ~v~n   : ~v"
					 source sink
					 message))
		       (void))]
	       [x
		(begin (log-info (format "DNS: ~v" x))
		       (void))])))