Golang Msg.Compress方法代码示例

// Forwards a DNS request to the specified nameservers. On success, the
// original reply packet will be returned to the caller. On failure, a
// new packet will be returned with `RCODE` set to `SERVFAIL`.
// Even though the original `ResponseWriter` object is taken as an argument,
// this function does not send a reply to the client. Instead, the
// packet is returned for further processing by the caller.
func getServerReply(w dns.ResponseWriter, req *dns.Msg) *dns.Msg {
	if *verbose {
		log.Print("Forwarding ", req.Question[0].Name, "/", dns.Type(req.Question[0].Qtype).String())
	}

	// create a new DNS client

	client := &dns.Client{Net: "udp", ReadTimeout: 4 * time.Second, WriteTimeout: 4 * time.Second, SingleInflight: true}

	if _, tcp := w.RemoteAddr().(*net.TCPAddr); tcp {
		client.Net = "tcp"
	}

	var r *dns.Msg
	var err error

	// loop through the specified nameservers and forward them the request

	// the request ID is used as a starting point in order to introduce at least
	// some element of randomness, instead of always hitting the first nameserver

	for i := 0; i < len(nameservers); i++ {
		r, _, err = client.Exchange(req, nameservers[(int(req.Id)+i)%len(nameservers)])

		if err == nil {
			r.Compress = true

			return r
		}
	}

	log.Print("Failed to forward request.", err)
	return getEmptyMsg(w, req, dns.RcodeServerFailure)
}

func (p ReverseProxy) ServeDNS(w dns.ResponseWriter, r *dns.Msg, extra []dns.RR) error {
	var (
		reply *dns.Msg
		err   error
	)

	switch {
	case middleware.Proto(w) == "tcp":
		reply, err = middleware.Exchange(p.Client.TCP, r, p.Host)
	default:
		reply, err = middleware.Exchange(p.Client.UDP, r, p.Host)
	}

	if reply != nil && reply.Truncated {
		// Suppress proxy error for truncated responses
		err = nil
	}

	if err != nil {
		return err
	}

	reply.Compress = true
	reply.Id = r.Id
	w.WriteMsg(reply)
	return nil
}

// ServeDNSForward forwards a request to a nameservers and returns the response.
func (s *server) ServeDNSForward(w dns.ResponseWriter, req *dns.Msg) *dns.Msg {

	if s.config.NoRec {
		m := s.ServerFailure(req)
		w.WriteMsg(m)
		return m
	}

	if len(s.config.Nameservers) == 0 || dns.CountLabel(req.Question[0].Name) < s.config.Ndots {
		if s.config.Verbose {
			if len(s.config.Nameservers) == 0 {
				logf("can not forward, no nameservers defined")
			} else {
				logf("can not forward, name too short (less than %d labels): `%s'", s.config.Ndots, req.Question[0].Name)
			}
		}
		m := s.ServerFailure(req)
		m.RecursionAvailable = true // this is still true
		w.WriteMsg(m)
		return m
	}

	tcp := isTCP(w)

	var (
		r   *dns.Msg
		err error
		try int
	)

	nsid := 0
	if s.config.NSRotate {
		// Use request Id for "random" nameserver selection.
		nsid = int(req.Id) % len(s.config.Nameservers)
	}
Redo:
	switch tcp {
	case false:
		r, _, err = s.dnsUDPclient.Exchange(req, s.config.Nameservers[nsid])
	case true:
		r, _, err = s.dnsTCPclient.Exchange(req, s.config.Nameservers[nsid])
	}
	if err == nil {
		r.Compress = true
		r.Id = req.Id
		w.WriteMsg(r)
		return r
	}
	// Seen an error, this can only mean, "server not reached", try again
	// but only if we have not exausted our nameservers.
	if try < len(s.config.Nameservers) {
		try++
		nsid = (nsid + 1) % len(s.config.Nameservers)
		goto Redo
	}

	logf("failure to forward request %q", err)
	m := s.ServerFailure(req)
	return m
}

// handlerFunc receives requests, looks up the result and returns what is found.
func handlerFunc(res dns.ResponseWriter, req *dns.Msg) {
	message := new(dns.Msg)
	switch req.Opcode {
	case dns.OpcodeQuery:
		message.SetReply(req)
		message.Compress = false
		message.Answer = make([]dns.RR, 0)

		for _, question := range message.Question {
			answers := answerQuestion(strings.ToLower(question.Name), question.Qtype)
			if len(answers) > 0 {
				for i := range answers {
					message.Answer = append(message.Answer, answers[i])
				}
			} else {
				// If there are no records, go back through and search for SOA records
				for _, question := range message.Question {
					answers := answerQuestion(strings.ToLower(question.Name), dns.TypeSOA)
					for i := range answers {
						message.Ns = append(message.Ns, answers[i])
					}
				}
			}
		}
		if len(message.Answer) == 0 && len(message.Ns) == 0 {
			message.Rcode = dns.RcodeNameError
		}
	default:
		message = message.SetRcode(req, dns.RcodeNotImplemented)
	}
	res.WriteMsg(message)
}

// handleDNS is a handler function to actualy perform the dns querey response
func (c *CatchAll) handleDNS(w dns.ResponseWriter, r *dns.Msg) {
	defer w.Close()
	var rr dns.RR

	domainSpoof := r.Question[0].Name

	msgResp := new(dns.Msg)
	msgResp.SetReply(r)
	msgResp.Compress = false

	rr = new(dns.A)

	if c.SpoofDomain {
		rr.(*dns.A).Hdr = dns.RR_Header{Name: domainSpoof, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0}
	} else {
		rr.(*dns.A).Hdr = dns.RR_Header{Name: c.Domain, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0}
	}

	rr.(*dns.A).A = c.IP

	switch r.Question[0].Qtype {
	case dns.TypeA:
		msgResp.Answer = append(msgResp.Answer, rr)
	default:
		log.Warnf("Unknown dns type %T", r.Question[0].Qtype)
		return
	}

	w.WriteMsg(msgResp)
}

// reply writes the given dns.Msg out to the given dns.ResponseWriter,
// compressing the message first and truncating it accordingly.
func reply(w dns.ResponseWriter, m *dns.Msg) {
	m.Compress = true // https://github.com/mesosphere/mesos-dns/issues/{170,173,174}

	if err := w.WriteMsg(truncate(m, isUDP(w))); err != nil {
		logging.Error.Println(err)
	}
}

func handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
	if *debug {
		Log.Printf("handleRequest: message: %+v\n", r)
	}

	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = false

	switch r.Opcode {
	case dns.OpcodeQuery:
		parseQuery(m)

	case dns.OpcodeUpdate:
		for _, question := range r.Question {
			for _, rr := range r.Ns {
				updateRecord(rr, &question)
			}
		}
	}

	if r.IsTsig() != nil {
		if w.TsigStatus() == nil {
			m.SetTsig(r.Extra[len(r.Extra)-1].(*dns.TSIG).Hdr.Name,
				dns.HmacMD5, 300, time.Now().Unix())
		} else {
			Log.Println("Status", w.TsigStatus().Error())
		}
	}

	w.WriteMsg(m)
}

// handleQuery is used to handle DNS queries in the configured domain
func (d *DNSServer) handleQuery(resp dns.ResponseWriter, req *dns.Msg) {
	q := req.Question[0]
	defer func(s time.Time) {
		metrics.MeasureSince([]string{"consul", "dns", "domain_query", d.agent.config.NodeName}, s)
		d.logger.Printf("[DEBUG] dns: request for %v (%v) from client %s (%s)",
			q, time.Now().Sub(s), resp.RemoteAddr().String(),
			resp.RemoteAddr().Network())
	}(time.Now())

	// Switch to TCP if the client is
	network := "udp"
	if _, ok := resp.RemoteAddr().(*net.TCPAddr); ok {
		network = "tcp"
	}

	// Setup the message response
	m := new(dns.Msg)
	m.SetReply(req)
	m.Compress = !d.config.DisableCompression
	m.Authoritative = true
	m.RecursionAvailable = (len(d.recursors) > 0)

	// Only add the SOA if requested
	if req.Question[0].Qtype == dns.TypeSOA {
		d.addSOA(d.domain, m)
	}

	// Dispatch the correct handler
	d.dispatch(network, req, m)

	// Write out the complete response
	if err := resp.WriteMsg(m); err != nil {
		d.logger.Printf("[WARN] dns: failed to respond: %v", err)
	}
}

// ServeDNSStubForward forwards a request to a nameservers and returns the response.
func (s *server) ServeDNSStubForward(w dns.ResponseWriter, req *dns.Msg, ns []string) *dns.Msg {
	// Check EDNS0 Stub option, if set drop the packet.
	option := req.IsEdns0()
	if option != nil {
		for _, o := range option.Option {
			if o.Option() == ednsStubCode && len(o.(*dns.EDNS0_LOCAL).Data) == 1 &&
				o.(*dns.EDNS0_LOCAL).Data[0] == 1 {
				// Maybe log source IP here?
				logf("not fowarding stub request to another stub")
				return nil
			}
		}
	}

	tcp := isTCP(w)

	// Add a custom EDNS0 option to the packet, so we can detect loops
	// when 2 stubs are forwarding to each other.
	if option != nil {
		option.Option = append(option.Option, &dns.EDNS0_LOCAL{ednsStubCode, []byte{1}})
	} else {
		req.Extra = append(req.Extra, ednsStub)
	}

	var (
		r   *dns.Msg
		err error
		try int
	)

	// Use request Id for "random" nameserver selection.
	nsid := int(req.Id) % len(ns)
Redo:
	switch tcp {
	case false:
		r, _, err = s.dnsUDPclient.Exchange(req, ns[nsid])
	case true:
		r, _, err = s.dnsTCPclient.Exchange(req, ns[nsid])
	}
	if err == nil {
		r.Compress = true
		r.Id = req.Id
		w.WriteMsg(r)
		return r
	}
	// Seen an error, this can only mean, "server not reached", try again
	// but only if we have not exausted our nameservers.
	if try < len(ns) {
		try++
		nsid = (nsid + 1) % len(ns)
		goto Redo
	}

	logf("failure to forward stub request %q", err)
	m := s.ServerFailure(req)
	w.WriteMsg(m)
	return m
}

func CompressIfLarge(m *dns.Msg) {
	bytes, err := m.Pack()
	if err != nil {
		return
	}
	if len(bytes) > 512 { // may not fit into UDP packet
		m.Compress = true // will be compressed in WriteMsg
	}
}

func (s *Server) handle(w dns.ResponseWriter, request *dns.Msg) {
	// Always close the writer
	defer w.Close()

	// Capture starting time for measuring message response time
	var start time.Time
	start = time.Now()

	// Setup the default response
	var response *dns.Msg
	response = &dns.Msg{}
	response.SetReply(request)
	response.Compress = true

	// Lookup answers to any of the questions
	for _, question := range request.Question {
		// Capture starting time for measuring lookup
		var lookupStart time.Time
		lookupStart = time.Now()

		// Perform lookup for this question
		var records []dns.RR
		records = s.registry.Lookup(question.Name, question.Qtype, question.Qclass)

		// Capture ending and elapsed time
		var lookupElapsed time.Duration
		lookupElapsed = time.Since(lookupStart)

		// Append results to the response
		response.Answer = append(response.Answer, records...)

		// If StatsD is enabled, record some metrics
		if s.statsd != nil {
			var tags []string
			tags = []string{
				fmt.Sprintf("name:%s", question.Name),
				fmt.Sprintf("qtype:%s", dns.TypeToString[question.Qtype]),
				fmt.Sprintf("qclass:%s", dns.ClassToString[question.Qclass]),
			}

			s.statsd.TimeInMilliseconds("lookup.time", lookupElapsed.Seconds()*1000.0, tags, 1)
			s.statsd.Histogram("lookup.answer", float64(len(records)), tags, 1)
			s.statsd.Count("request.question", 1, tags, 1)
		}
	}

	// Respond to the request
	w.WriteMsg(response)

	// Record any ending metrics
	if s.statsd != nil {
		var elapsed time.Duration
		elapsed = time.Since(start)
		s.statsd.TimeInMilliseconds("request.time", elapsed.Seconds()*1000.0, nil, 1)
	}
}

// HandleForwarding forwards a request to the nameservers and returns the response
func (s *Server) HandleForwarding(w mdns.ResponseWriter, r *mdns.Msg) (bool, error) {
	defer trace.End(trace.Begin(r.String()))

	var m *mdns.Msg
	var err error
	var try int

	if len(s.Nameservers) == 0 {
		log.Errorf("No nameservers defined, can not forward")
		return false, respServerFailure(w, r)
	}

	// which protocol are they talking
	tcp := false
	if _, ok := w.RemoteAddr().(*net.TCPAddr); ok {
		tcp = true
	}

	// Use request ID for "random" nameserver selection.
	nsid := int(r.Id) % len(s.Nameservers)

Redo:
	nameserver := s.Nameservers[nsid]
	if i := strings.Index(nameserver, ":"); i < 0 {
		nameserver += ":53"
	}

	if tcp {
		m, _, err = s.tcpclient.Exchange(r, nameserver)
	} else {
		m, _, err = s.udpclient.Exchange(r, nameserver)
	}
	if err != nil {
		// Seen an error, this can only mean, "server not reached", try again but only if we have not exausted our nameservers.
		if try < len(s.Nameservers) {
			try++
			nsid = (nsid + 1) % len(s.Nameservers)
			goto Redo
		}

		log.Errorf("Failure to forward request: %q", err)
		return false, respServerFailure(w, r)
	}

	// We have a response so cache it
	s.cache.Add(m)

	m.Compress = true
	if err := w.WriteMsg(m); err != nil {
		log.Errorf("Error writing response: %q", err)
		return true, err
	}
	return true, nil
}

func (r *resolver) ServeDNS(w dns.ResponseWriter, query *dns.Msg) {
	var (
		resp *dns.Msg
		err  error
	)

	if query == nil || len(query.Question) == 0 {
		return
	}
	name := query.Question[0].Name
	if query.Question[0].Qtype == dns.TypeA {
		resp, err = r.handleIPv4Query(name, query)
	} else if query.Question[0].Qtype == dns.TypePTR {
		resp, err = r.handlePTRQuery(name, query)
	}

	if err != nil {
		log.Error(err)
		return
	}

	if resp == nil {
		if len(r.extDNS) == 0 {
			return
		}

		num := maxExtDNS
		if len(r.extDNS) < maxExtDNS {
			num = len(r.extDNS)
		}
		for i := 0; i < num; i++ {
			log.Debugf("Querying ext dns %s:%s for %s[%d]", w.LocalAddr().Network(), r.extDNS[i], name, query.Question[0].Qtype)

			c := &dns.Client{Net: w.LocalAddr().Network()}
			addr := fmt.Sprintf("%s:%d", r.extDNS[i], 53)

			resp, _, err = c.Exchange(query, addr)
			if err == nil {
				resp.Compress = true
				break
			}
			log.Errorf("external resolution failed, %s", err)
		}
		if resp == nil {
			return
		}
	}

	err = w.WriteMsg(resp)
	if err != nil {
		log.Errorf("error writing resolver resp, %s", err)
	}
}

// ServeDNSForward forwards a request to a nameservers and returns the response.
func (s *server) ServeDNSForward(w dns.ResponseWriter, req *dns.Msg) *dns.Msg {
	if s.config.NoRec {
		m := s.ServerFailure(req)
		w.WriteMsg(m)
		return m
	}

	if len(s.config.Nameservers) == 0 || dns.CountLabel(req.Question[0].Name) < s.config.Ndots {
		if s.config.Verbose {
			if len(s.config.Nameservers) == 0 {
				logf("can not forward, no nameservers defined")
			} else {
				logf("can not forward, name too short (less than %d labels): `%s'", s.config.Ndots, req.Question[0].Name)
			}
		}
		m := s.ServerFailure(req)
		m.RecursionAvailable = true // this is still true
		w.WriteMsg(m)
		return m
	}

	var (
		r   *dns.Msg
		err error
	)

	nsid := s.randomNameserverID(req.Id)
	try := 0
Redo:
	if isTCP(w) {
		r, err = exchangeWithRetry(s.dnsTCPclient, req, s.config.Nameservers[nsid])
	} else {
		r, err = exchangeWithRetry(s.dnsUDPclient, req, s.config.Nameservers[nsid])
	}
	if err == nil {
		r.Compress = true
		r.Id = req.Id
		w.WriteMsg(r)
		return r
	}
	// Seen an error, this can only mean, "server not reached", try again
	// but only if we have not exausted our nameservers.
	if try < len(s.config.Nameservers) {
		try++
		nsid = (nsid + 1) % len(s.config.Nameservers)
		goto Redo
	}

	logf("failure to forward request %q", err)
	m := s.ServerFailure(req)
	return m
}

func handleForwardingRaw(nameservers []string, req *dns.Msg, remote net.Addr) *dns.Msg {
	if len(nameservers) == 0 {
		log.Printf("no nameservers defined, can not forward\n")
		m := new(dns.Msg)
		m.SetReply(req)
		m.SetRcode(req, dns.RcodeServerFailure)
		m.Authoritative = false     // no matter what set to false
		m.RecursionAvailable = true // and this is still true
		return m
	}

	tcp := false
	if _, ok := remote.(*net.TCPAddr); ok {
		tcp = true
	}

	var (
		r   *dns.Msg
		err error
		try int
	)
	// Use request Id for "random" nameserver selection.
	nsid := int(req.Id) % len(nameservers)
	dnsClient := &dns.Client{Net: "udp", ReadTimeout: 4 * time.Second, WriteTimeout: 4 * time.Second, SingleInflight: true}
	if tcp {
		dnsClient.Net = "tcp"
	}
Redo:
	nameserver := nameservers[nsid]
	if i := strings.Index(nameserver, ":"); i < 0 {
		nameserver += ":53"
	}
	r, _, err = dnsClient.Exchange(req, nameserver)
	if err == nil {
		r.Compress = true
		return r
	}
	// Seen an error, this can only mean, "server not reached", try again
	// but only if we have not exausted our nameservers.
	if try < len(nameservers) {
		try++
		nsid = (nsid + 1) % len(nameservers)
		goto Redo
	}

	log.Printf("failure to forward request %q\n", err)
	m := new(dns.Msg)
	m.SetReply(req)
	m.SetRcode(req, dns.RcodeServerFailure)
	return m
}