#lang racket/base
;; DNS wire-protocol codec.

(provide value->query-opcode query-opcode->value
	 value->query-response-code query-response-code->value

	 (struct-out dns-message)

	 packet->dns-message
	 dns-message->packet

	 max-ttl

	 ;; For the use of zonedb's save/load routines, etc.
	 t:rr)

(require "api.rkt")
(require "mapping.rkt")

(require racket/match)
(require "../racket-bitsyntax/main.rkt")

;; An Opcode is a Symbol or a Number, one of the possibilities given
;; in the following define-mapping. It represents a DNS message
;; operation; see the RFC for details.
(define-mapping value->query-opcode query-opcode->value
  #:forward-default values
  #:backward-default values
  (0 query)
  (1 iquery)
  (2 status))

;; A ResponseCode is a Symbol or a Number, one of the possibilities
;; given in the following define-mapping. It represents the outcome of
;; a DNS query.
(define-mapping value->query-response-code query-response-code->value
  (0 no-error)
  (1 format-error)
  (2 server-failure)
  (3 name-error) ;; most frequently known on the internet as NXDOMAIN.
  (4 not-implemented)
  (5 refused))

;; A DNSMessage is a
;; (dns-message Uint16 Direction Opcode Authoritativeness
;; Truncatedness RecursionDesired RecursionAvailable ResponseCode
;; ListOf<Question> ListOf<RR> ListOf<RR> ListOf<RR>).
;;
;; Interpreted as either a DNS request or reply, depending on the
;; Direction.
(struct dns-message (id
		     direction
		     opcode
		     authoritative
		     truncated
		     recursion-desired
		     recursion-available
		     response-code
		     questions
		     answers
		     authorities
		     additional)
	#:prefab)

;; Bit-syntax type for counted repeats of a value.
;; Example: Length-prefixed list of 32-bit unsigned words:
;;   (bit-string-case input ([ len (vals :: (t:ntimes len bits 32)) ] vals))
;;   (bit-string (vals :: (t:ntimes bits 32)))
(define-syntax t:ntimes
  (syntax-rules ()
    ((_ #t input ks kf times-to-repeat option ...)
     (let loop ((count times-to-repeat)
		(acc '())
		(input input))
       (cond
	((positive? count) (bit-string-case input
			     ([ (v :: option ...) (rest :: binary) ]
			      (loop (- count 1) (cons v acc) rest))
			     (else (kf))))
	(else (ks (reverse acc) input)))))
    ((_ #f vs option ...)
     (t:listof #f vs option ...))))

;; Bit-syntax type for repeats of a value until no more input available.
;; Example: List of 32-bit unsigned words:
;;   (bit-string-case input ([ (vals :: (t:listof bits 32)) ] vals))
;;   (bit-string (vals :: (t:listof bits 32)))
(define-syntax t:listof
  (syntax-rules ()
    ((_ #t input ks kf option ...)
     (let loop ((acc '())
		(input input))
       (bit-string-case input
	 ([ (v :: option ...) (rest :: binary) ]
	  (loop (cons v acc) rest))
	 ([]
	  (ks (reverse acc) #""))
	 (else
	  (kf)))))
    ((_ #f vs option ...)
     (let loop ((vs vs))
       (cond
	((pair? vs) (bit-string ((car vs) :: option ...)
				((loop (cdr vs)) :: binary)))
	(else (bit-string)))))))

;; <rfc1035>
;; All communications inside of the domain protocol are carried in a single
;; format called a message.  The top level format of message is divided
;; into 5 sections (some of which are empty in certain cases) shown below:
;;
;;     +---------------------+
;;     |        Header       |
;;     +---------------------+
;;     |       Question      | the question for the name server
;;     +---------------------+
;;     |        Answer       | RRs answering the question
;;     +---------------------+
;;     |      Authority      | RRs pointing toward an authority
;;     +---------------------+
;;     |      Additional     | RRs holding additional information
;;     +---------------------+
;; </rfc1035>

;; <rfc1035>
;; The header contains the following fields:
;;
;;                                     1  1  1  1  1  1
;;       0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                      ID                       |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |QR|   Opcode  |AA|TC|RD|RA|   Z    |   RCODE   |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                    QDCOUNT                    |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                    ANCOUNT                    |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                    NSCOUNT                    |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                    ARCOUNT                    |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;; </rfc1035>

;; Bytes -> DNSMessage
;; Parse an encoded DNS message packet into the corresponding Racket
;; structure. Raises an exception on failure.
(define (packet->dns-message packet)
  (bit-string-case packet
    ([ (id :: bits 16)

       ;; 16 bits of flags, opcode, and response-code:
       (qr :: (t:named-bit 'request 'response))
       (opcode :: bits 4)
       (aa :: (t:named-bit 'non-authoritative 'authoritative))
       (tc :: (t:named-bit 'not-truncated 'truncated))
       (rd :: (t:named-bit 'no-recursion-desired 'recursion-desired))
       (ra :: (t:named-bit 'no-recursion-available 'recursion-available))
       (= 0 :: bits 3)
       (rcode :: bits 4)

       (qdcount :: bits 16)
       (ancount :: bits 16)
       (nscount :: bits 16)
       (arcount :: bits 16)
       (q-section :: (t:ntimes qdcount (t:question packet)))
       (a-section :: (t:ntimes ancount (t:rr packet)))
       (auth-section :: (t:ntimes nscount (t:rr packet)))
       (additional-section :: (t:ntimes arcount (t:rr packet))) ]
     (dns-message id qr (value->query-opcode opcode)
		  aa tc rd ra
		  (value->query-response-code rcode)
		  q-section a-section auth-section additional-section))))

;; DNSMessage -> Bytes
;; Render a Racket structured DNS message using the DNS binary encoding.
(define (dns-message->packet m)
  (bit-string->bytes
   (bit-string ((dns-message-id m) :: bits 16)
	       ((dns-message-direction m) :: (t:named-bit 'request 'response))
	       ((query-opcode->value (dns-message-opcode m)) :: bits 4)
	       ((dns-message-authoritative m) :: (t:named-bit 'non-authoritative 'authoritative))
	       ((dns-message-truncated m) :: (t:named-bit 'not-truncated 'truncated))
	       ((dns-message-recursion-desired m)
		:: (t:named-bit 'no-recursion-desired 'recursion-desired))
	       ((dns-message-recursion-available m)
		:: (t:named-bit 'no-recursion-available 'recursion-available))
	       (0 :: bits 3)
	       ((query-response-code->value (dns-message-response-code m)) :: bits 4)
	       ((length (dns-message-questions m)) :: bits 16)
	       ((length (dns-message-answers m)) :: bits 16)
	       ((length (dns-message-authorities m)) :: bits 16)
	       ((length (dns-message-additional m)) :: bits 16)
	       ((dns-message-questions m) :: (t:ntimes (t:question)))
	       ((dns-message-answers m) :: (t:ntimes (t:rr)))
	       ((dns-message-authorities m) :: (t:ntimes (t:rr)))
	       ((dns-message-additional m) :: (t:ntimes (t:rr))))))

;; Bit-syntax type for a single bit, represented in Racket as one of
;; two possible symbolic values.
;; Example: a bit represented by 'zero when it is zero, and 'one when it is one.
;;   (bit-string-case input ([ (v :: (t:named-bit 'zero 'one)) ] v))
;;   (bit-string (v :: (t:named-bit 'zero 'one)))
(define-syntax t:named-bit
  (syntax-rules ()
    ((_ #t input ks kf name0 name1)
     (bit-string-case input
       ([ (v :: bits 1) (rest :: binary) ]
	(ks (if (zero? v) name0 name1) rest))
       (else (kf))))
    ((_ #f v name0 name1)
     (cond
      ((eq? v name1) (bit-string (1 :: bits 1)))
      ((eq? v name0) (bit-string (0 :: bits 1)))
      (else (error 't:named-bit
		   "Value supplied is neither ~v nor ~v: ~v"
		   name0 name1 v))))))

;; Bit-syntax type for a DomainName. When decoding (but not when
;; encoding!), we support DNS's weird compressed domain-name syntax;
;; this requires us to pass in the *whole packet* to the decoder to
;; let it refer to random substrings within the packet.
(define-syntax t:domain-name
  (syntax-rules ()
    ((_ #t input ks kf whole-packet)
     (let-values (((name rest) (parse-domain-name whole-packet input '())))
       (ks (domain name) rest)))
    ((_ #f val)
     (encode-domain-name val))))

;; DomainName -> Bitstring
(define (encode-domain-name name)
  (define labels (domain-labels name))
  (bit-string (labels :: (t:listof (t:pascal-string "Label" 64)))
	      (0 :: integer bytes 1))) ;; end of list of labels!

;; Bytes Bytes ListOf<Natural> -> ListOf<Bytes>
;; PRECONDITION: input never empty
;; INVARIANT: pointers-followed contains every "jump target" we have
;; jumped to so far during decoding of this domain-name, in order to
;; prevent us from getting stuck in a pointer loop. It should start as
;; the empty list.
(define (parse-domain-name whole-packet input pointers-followed)
  (bit-string-case input

    ([(= 3 :: bits 2) (offset :: bits 14) (rest :: binary)]
     (if (member offset pointers-followed)
	 (error 'parse-domain-name "DNS compressed-pointer loop detected")
	 (let-values (((lhs rhs) (bit-string-split-at whole-packet (* 8 offset))))
	   (let-values (((labels ignored-tail)
			 (parse-domain-name whole-packet rhs (cons offset pointers-followed))))
	     (values labels rest)))))

    ([(= 0 :: bits 8) (rest :: binary)]
     (values '() rest))

    ([(= 0 :: bits 2) (len :: bits 6) (label :: binary bytes len) (rest :: binary)]
     ;; TODO: validate labels: make sure they conform to the prescribed syntax
     (let-values (((labels leftover)
		   (parse-domain-name whole-packet rest pointers-followed)))
       (values (cons (bit-string->bytes label) labels) leftover)))))

;; Bit-syntax type for single-byte-length-prefixed strings of
;; bytes. No character codec is applied to the bytes. During encoding,
;; expects two extra arguments: the name of the kind of value, for use
;; in error reports, and the maximum permissible length (plus one). If
;; the encoder is given a string of length greater than or equal to
;; the given maximum, an error is signalled.
(define-syntax t:pascal-string
  (syntax-rules ()
    ((_ #t input ks kf)
     (bit-string-case input
       ([ len (body :: binary bytes len) (rest :: binary) ]
	(ks (bit-string->bytes body) rest))
       (else (kf))))
    ((_ #f s)
     (t:pascal-string #f s "Character-string" 256))
    ((_ #f s string-kind length-limit)
     (let ((len (bytes-length s)))
       (when (>= len length-limit)
	 (error 't:pascal-string "~s too long: ~v" string-kind s))
       (bit-string len (s :: binary))))))

;; <rfc1035>
;; The question section is used to carry the "question" in most queries,
;; i.e., the parameters that define what is being asked.  The section
;; contains QDCOUNT (usually 1) entries, each of the following format:
;;
;;                                     1  1  1  1  1  1
;;       0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                                               |
;;     /                     QNAME                     /
;;     /                                               /
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                     QTYPE                     |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                     QCLASS                    |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;; </rfc1035>

;; Bit-syntax type for Questions. The decoder needs to be given the
;; whole packet because the question may contain nested domain-names.
(define-syntax t:question
  (syntax-rules ()
    ((_ #t input ks kf whole-packet)
     (bit-string-case input
       ([ (qname :: (t:domain-name whole-packet))
	  (qtype :: bits 16)
	  (qclass :: bits 16)
	  (tail :: binary) ]
	(ks (question qname
		      (value->qtype qtype)
		      (value->qclass qclass)
		      #f)
	    tail))))
    ((_ #f q)
     (bit-string ((question-name q) :: (t:domain-name))
		 ((qtype->value (question-type q)) :: bits 16)
		 ((qclass->value (question-class q)) :: bits 16)))))

;; <rfc1035>
;; All RRs have the same top level format shown below:
;;
;;                                     1  1  1  1  1  1
;;       0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                                               |
;;     /                                               /
;;     /                      NAME                     /
;;     |                                               |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                      TYPE                     |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                     CLASS                     |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                      TTL                      |
;;     |                                               |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;;     |                   RDLENGTH                    |
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--|
;;     /                     RDATA                     /
;;     /                                               /
;;     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
;; </rfc1035>

;; Bit-syntax type for RRs. The decoder needs to be given the whole
;; packet because the RR may contain nested domain-names.
(define-syntax t:rr
  (syntax-rules ()
    ((_ #t input ks kf whole-packet)
     (decode-rr whole-packet input ks kf))
    ((_ #f rr)
     (encode-rr rr))))

;; Bytes Bytes (RR Bytes -> A) ( -> A) -> A
;; Helper for t:rr.
(define (decode-rr whole-packet input ks kf)
  (bit-string-case input
    ([ (name :: (t:domain-name whole-packet))
       (type-number :: bits 16)
       (class :: bits 16)
       (ttl :: bits 32)
       (rdlength :: bits 16)
       (rdata :: binary bytes rdlength)
       (tail :: binary) ]
     (let ((type (value->type type-number)))
       (ks (rr name
	       type
	       (value->class class)
	       ttl
	       (decode-rdata whole-packet type rdata))
	   tail)))
    (else (kf))))

;; RR -> Bitstring
;; Helper for t:rr.
(define (encode-rr rr)
  (let ((encoded-rdata (encode-rdata (rr-type rr) (rr-rdata rr))))
    (bit-string ((rr-name rr) :: (t:domain-name))
		((type->value (rr-type rr)) :: bits 16)
		((class->value (rr-class rr)) :: bits 16)
		((rr-ttl rr) :: bits 32)
		((/ (bit-string-length encoded-rdata) 8) :: bits 16)
		(encoded-rdata :: binary))))

;; Bytes RRType Bytes -> RData
;; Decode RData according to the RRType. Takes the whole packet for
;; the same reason as t:rr does.
(define (decode-rdata whole-packet type rdata)
  (case type
    ((cname mb md mf mg mr ns ptr) (bit-string-case rdata
				     ([ (name :: (t:domain-name whole-packet)) ] name)))
    ((hinfo) (bit-string-case rdata
	       ([ (cpu :: (t:pascal-string)) (os :: (t:pascal-string)) ]
		(hinfo cpu os))))
    ((minfo) (bit-string-case rdata
	       ([ (rmailbx :: (t:domain-name whole-packet))
		  (emailbx :: (t:domain-name whole-packet)) ]
		(minfo rmailbx emailbx))))
    ((mx) (bit-string-case rdata
	    ([ (preference :: bits 16)
	       (exchange :: (t:domain-name whole-packet)) ]
	     (mx preference exchange))))
    ((null) (bit-string->bytes rdata))
    ((soa) (bit-string-case rdata
	     ([ (mname :: (t:domain-name whole-packet))
		(rname :: (t:domain-name whole-packet))
		(serial :: bits 32)
		(refresh :: bits 32)
		(retry :: bits 32)
		(expire :: bits 32)
		(minimum :: bits 32) ]
	      (soa mname rname serial refresh retry expire minimum))))
    ((txt) (bit-string-case rdata
	     ([ (strs :: (t:listof (t:pascal-string))) ]
	      strs)))
    ((a) (bit-string-case rdata
	   ([ a b c d ]
	    (vector a b c d))))
    ((aaaa) (bit-string-case rdata
	      ([ (ipv6-addr :: binary bits 128) ]
	       (list->vector (bytes->list (bit-string->bytes ipv6-addr))))))
    ((wks) (bit-string-case rdata
	     ([ a b c d protocol (bitmap :: binary) ]
	      (wks (vector a b c d) protocol bitmap))))
    ((srv) (bit-string-case rdata
	     ([ (priority :: bits 16)
		(weight :: bits 16)
		(port :: bits 16)
		(target :: (t:domain-name whole-packet)) ]
	      (srv priority weight port target))))
    (else (bit-string->bytes rdata))))

;; RRType RData -> Bitstring
;; Encode RData according to the RRType.
(define (encode-rdata type rdata)
  (case type
    ((cname mb md mf mg mr ns ptr) (encode-domain-name rdata))
    ((hinfo) (bit-string ((hinfo-cpu rdata) :: (t:pascal-string))
			 ((hinfo-os rdata) :: (t:pascal-string))))
    ((minfo) (bit-string ((minfo-rmailbx rdata) :: (t:domain-name))
			 ((minfo-emailbx rdata) :: (t:domain-name))))
    ((mx) (bit-string ((mx-preference rdata) :: bits 16)
		      ((mx-exchange rdata) :: (t:domain-name))))
    ((null) rdata)
    ((soa) (bit-string ((soa-mname rdata) :: (t:domain-name))
		       ((soa-rname rdata) :: (t:domain-name))
		       ((soa-serial rdata) :: bits 32)
		       ((soa-refresh rdata) :: bits 32)
		       ((soa-retry rdata) :: bits 32)
		       ((soa-expire rdata) :: bits 32)
		       ((soa-minimum rdata) :: bits 32)))
    ((txt) (bit-string (rdata :: (t:listof (t:pascal-string)))))
    ((a) (match rdata ((vector a b c d) (bit-string a b c d))))
    ((aaaa) (bit-string ((list->bytes (vector->list rdata)) :: binary bits 128)))
    ((wks) (match (wks-address rdata)
	     ((vector a b c d)
	      (bit-string a b c d (wks-protocol rdata) ((wks-bitmap rdata) :: binary)))))
    ((srv) (bit-string ((srv-priority rdata) :: bits 16)
		       ((srv-weight rdata) :: bits 16)
		       ((srv-port rdata) :: bits 16)
		       ((srv-target rdata) :: (t:domain-name))))
    (else rdata)))

;; UInt32
(define max-ttl #xffffffff)