serve_mux.go 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148
  1. package dns
  2. import (
  3. "strings"
  4. "sync"
  5. )
  6. // ServeMux is an DNS request multiplexer. It matches the zone name of
  7. // each incoming request against a list of registered patterns add calls
  8. // the handler for the pattern that most closely matches the zone name.
  9. //
  10. // ServeMux is DNSSEC aware, meaning that queries for the DS record are
  11. // redirected to the parent zone (if that is also registered), otherwise
  12. // the child gets the query.
  13. //
  14. // ServeMux is also safe for concurrent access from multiple goroutines.
  15. //
  16. // The zero ServeMux is empty and ready for use.
  17. type ServeMux struct {
  18. z map[string]Handler
  19. m sync.RWMutex
  20. }
  21. // NewServeMux allocates and returns a new ServeMux.
  22. func NewServeMux() *ServeMux {
  23. return new(ServeMux)
  24. }
  25. // DefaultServeMux is the default ServeMux used by Serve.
  26. var DefaultServeMux = NewServeMux()
  27. func (mux *ServeMux) match(q string, t uint16) Handler {
  28. mux.m.RLock()
  29. defer mux.m.RUnlock()
  30. if mux.z == nil {
  31. return nil
  32. }
  33. var handler Handler
  34. // TODO(tmthrgd): Once https://go-review.googlesource.com/c/go/+/137575
  35. // lands in a go release, replace the following with strings.ToLower.
  36. var sb strings.Builder
  37. for i := 0; i < len(q); i++ {
  38. c := q[i]
  39. if !(c >= 'A' && c <= 'Z') {
  40. continue
  41. }
  42. sb.Grow(len(q))
  43. sb.WriteString(q[:i])
  44. for ; i < len(q); i++ {
  45. c := q[i]
  46. if c >= 'A' && c <= 'Z' {
  47. c += 'a' - 'A'
  48. }
  49. sb.WriteByte(c)
  50. }
  51. q = sb.String()
  52. break
  53. }
  54. for off, end := 0, false; !end; off, end = NextLabel(q, off) {
  55. if h, ok := mux.z[q[off:]]; ok {
  56. if t != TypeDS {
  57. return h
  58. }
  59. // Continue for DS to see if we have a parent too, if so delegate to the parent
  60. handler = h
  61. }
  62. }
  63. // Wildcard match, if we have found nothing try the root zone as a last resort.
  64. if h, ok := mux.z["."]; ok {
  65. return h
  66. }
  67. return handler
  68. }
  69. // Handle adds a handler to the ServeMux for pattern.
  70. func (mux *ServeMux) Handle(pattern string, handler Handler) {
  71. if pattern == "" {
  72. panic("dns: invalid pattern " + pattern)
  73. }
  74. mux.m.Lock()
  75. if mux.z == nil {
  76. mux.z = make(map[string]Handler)
  77. }
  78. mux.z[Fqdn(pattern)] = handler
  79. mux.m.Unlock()
  80. }
  81. // HandleFunc adds a handler function to the ServeMux for pattern.
  82. func (mux *ServeMux) HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
  83. mux.Handle(pattern, HandlerFunc(handler))
  84. }
  85. // HandleRemove deregisters the handler specific for pattern from the ServeMux.
  86. func (mux *ServeMux) HandleRemove(pattern string) {
  87. if pattern == "" {
  88. panic("dns: invalid pattern " + pattern)
  89. }
  90. mux.m.Lock()
  91. delete(mux.z, Fqdn(pattern))
  92. mux.m.Unlock()
  93. }
  94. // ServeDNS dispatches the request to the handler whose pattern most
  95. // closely matches the request message.
  96. //
  97. // ServeDNS is DNSSEC aware, meaning that queries for the DS record
  98. // are redirected to the parent zone (if that is also registered),
  99. // otherwise the child gets the query.
  100. //
  101. // If no handler is found, or there is no question, a standard SERVFAIL
  102. // message is returned
  103. func (mux *ServeMux) ServeDNS(w ResponseWriter, req *Msg) {
  104. var h Handler
  105. if len(req.Question) >= 1 { // allow more than one question
  106. h = mux.match(req.Question[0].Name, req.Question[0].Qtype)
  107. }
  108. if h != nil {
  109. h.ServeDNS(w, req)
  110. } else {
  111. HandleFailed(w, req)
  112. }
  113. }
  114. // Handle registers the handler with the given pattern
  115. // in the DefaultServeMux. The documentation for
  116. // ServeMux explains how patterns are matched.
  117. func Handle(pattern string, handler Handler) { DefaultServeMux.Handle(pattern, handler) }
  118. // HandleRemove deregisters the handle with the given pattern
  119. // in the DefaultServeMux.
  120. func HandleRemove(pattern string) { DefaultServeMux.HandleRemove(pattern) }
  121. // HandleFunc registers the handler function with the given pattern
  122. // in the DefaultServeMux.
  123. func HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
  124. DefaultServeMux.HandleFunc(pattern, handler)
  125. }