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 }