dns

DNS client and server implementations using the Go project's dnsmessage package
Log | Files | Refs | README | LICENSE

example_test.go (3098B)


      1 package dns_test
      2 
      3 import (
      4 	"fmt"
      5 	"io"
      6 	"os"
      7 	"strings"
      8 
      9 	"golang.org/x/net/dns/dnsmessage"
     10 	"olowe.co/dns"
     11 )
     12 
     13 // reflectA returns the same A records for every name.
     14 func reflectA(name dnsmessage.Name) (dnsmessage.ResourceHeader, []dnsmessage.AResource) {
     15 	h := dnsmessage.ResourceHeader{
     16 		Name:  name,
     17 		Type:  dnsmessage.TypeA,
     18 		Class: dnsmessage.ClassINET,
     19 		TTL:   uint32(3600),
     20 	}
     21 	r1 := dnsmessage.AResource{
     22 		A: [4]byte{192, 0, 2, 1},
     23 	}
     24 	r2 := dnsmessage.AResource{
     25 		A: [4]byte{192, 0, 2, 2},
     26 	}
     27 	return h, []dnsmessage.AResource{r1, r2}
     28 }
     29 
     30 func lookup(name dnsmessage.Name, t dnsmessage.Type) (dnsmessage.ResourceHeader, []dnsmessage.AResource) {
     31 	if t != dnsmessage.TypeA {
     32 		// we only have A records
     33 		return dnsmessage.ResourceHeader{}, nil
     34 	}
     35 	return reflectA(name)
     36 }
     37 
     38 // authoritativeHandler answers questions for the zone ".test.".
     39 func authoritativeHandler(w dns.ResponseWriter, qmsg *dnsmessage.Message) {
     40 	var rmsg dnsmessage.Message
     41 	rmsg.Header.ID = rmsg.Header.ID
     42 	rmsg.Questions = qmsg.Questions
     43 
     44 	// reject empty questions, and any messages with more than 1 question;
     45 	// even BIND doesn't support more than 1 question per message.
     46 	if len(qmsg.Questions) != 1 {
     47 		rmsg.Header.RCode = dnsmessage.RCodeNotImplemented
     48 		w.WriteMsg(rmsg)
     49 		return
     50 	}
     51 
     52 	// reject questions for anything other than our our test zone ".test."
     53 	q := qmsg.Questions[0]
     54 	if !strings.HasSuffix(q.Name.String(), ".test.") {
     55 		rmsg.Header.RCode = dnsmessage.RCodeRefused
     56 		w.WriteMsg(rmsg)
     57 		return
     58 	}
     59 
     60 	header := dnsmessage.Header{
     61 		ID:            qmsg.Header.ID,
     62 		Response:      true,
     63 		Authoritative: true,
     64 	}
     65 	buf := make([]byte, 2, 512+2)
     66 	builder := dnsmessage.NewBuilder(buf, header)
     67 	builder.EnableCompression()
     68 	rmsg.Header.RCode = dnsmessage.RCodeServerFailure
     69 	if err := builder.StartQuestions(); err != nil {
     70 		w.WriteMsg(rmsg)
     71 		return
     72 	}
     73 	if err := builder.Question(q); err != nil {
     74 		w.WriteMsg(rmsg)
     75 		return
     76 	}
     77 	if err := builder.StartAnswers(); err != nil {
     78 		w.WriteMsg(rmsg)
     79 		return
     80 	}
     81 
     82 	resourceHeader, records := lookup(q.Name, q.Type)
     83 	for _, r := range records {
     84 		if err := builder.AResource(resourceHeader, r); err != nil {
     85 			w.WriteMsg(rmsg)
     86 			return
     87 		}
     88 	}
     89 	buf, err := builder.Finish()
     90 	if err != nil {
     91 		w.WriteMsg(rmsg)
     92 		return
     93 	}
     94 	// finished message starts at a 2-byte offset for some reason
     95 	w.Write(buf[2:])
     96 }
     97 
     98 func ExampleHandler() {
     99 	qmsg := dnsmessage.Message{
    100 		Header: dnsmessage.Header{ID: uint16(69)},
    101 		Questions: []dnsmessage.Question{
    102 			{
    103 				Name:  dnsmessage.MustNewName("www.example.test."),
    104 				Type:  dnsmessage.TypeA,
    105 				Class: dnsmessage.ClassINET,
    106 			},
    107 		},
    108 	}
    109 
    110 	pipe := pipe{w: os.Stdout}
    111 	authoritativeHandler(pipe, &qmsg)
    112 	// Output: [192.0.2.1 192.0.2.2]
    113 }
    114 
    115 type pipe struct {
    116 	w io.Writer
    117 }
    118 
    119 func (p pipe) Write(b []byte) (n int, err error) {
    120 	var m dnsmessage.Message
    121 	if err := m.Unpack(b); err != nil {
    122 		return 0, err
    123 	}
    124 	ips := dns.ExtractIPs(m.Answers)
    125 	fmt.Fprintln(p.w, ips)
    126 	return len(b), nil
    127 }
    128 
    129 func (p pipe) WriteMsg(m dnsmessage.Message) error {
    130 	ips := dns.ExtractIPs(m.Answers)
    131 	fmt.Fprintln(p.w, ips)
    132 	return nil
    133 }