1package dns
2
3// Truncate ensures the reply message will fit into the requested buffer
4// size by removing records that exceed the requested size.
5//
6// It will first check if the reply fits without compression and then with
7// compression. If it won't fit with compression, Truncate then walks the
8// record adding as many records as possible without exceeding the
9// requested buffer size.
10//
11// The TC bit will be set if any records were excluded from the message.
12// This indicates to that the client should retry over TCP.
13//
14// According to RFC 2181, the TC bit should only be set if not all of the
15// "required" RRs can be included in the response. Unfortunately, we have
16// no way of knowing which RRs are required so we set the TC bit if any RR
17// had to be omitted from the response.
18//
19// The appropriate buffer size can be retrieved from the requests OPT
20// record, if present, and is transport specific otherwise. dns.MinMsgSize
21// should be used for UDP requests without an OPT record, and
22// dns.MaxMsgSize for TCP requests without an OPT record.
23func (dns *Msg) Truncate(size int) {
24	if dns.IsTsig() != nil {
25		// To simplify this implementation, we don't perform
26		// truncation on responses with a TSIG record.
27		return
28	}
29
30	// RFC 6891 mandates that the payload size in an OPT record
31	// less than 512 bytes must be treated as equal to 512 bytes.
32	//
33	// For ease of use, we impose that restriction here.
34	if size < 512 {
35		size = 512
36	}
37
38	l := msgLenWithCompressionMap(dns, nil) // uncompressed length
39	if l <= size {
40		// Don't waste effort compressing this message.
41		dns.Compress = false
42		return
43	}
44
45	dns.Compress = true
46
47	edns0 := dns.popEdns0()
48	if edns0 != nil {
49		// Account for the OPT record that gets added at the end,
50		// by subtracting that length from our budget.
51		//
52		// The EDNS(0) OPT record must have the root domain and
53		// it's length is thus unaffected by compression.
54		size -= Len(edns0)
55	}
56
57	compression := make(map[string]struct{})
58
59	l = headerSize
60	for _, r := range dns.Question {
61		l += r.len(l, compression)
62	}
63
64	var numAnswer int
65	if l < size {
66		l, numAnswer = truncateLoop(dns.Answer, size, l, compression)
67	}
68
69	var numNS int
70	if l < size {
71		l, numNS = truncateLoop(dns.Ns, size, l, compression)
72	}
73
74	var numExtra int
75	if l < size {
76		l, numExtra = truncateLoop(dns.Extra, size, l, compression)
77	}
78
79	// See the function documentation for when we set this.
80	dns.Truncated = len(dns.Answer) > numAnswer ||
81		len(dns.Ns) > numNS || len(dns.Extra) > numExtra
82
83	dns.Answer = dns.Answer[:numAnswer]
84	dns.Ns = dns.Ns[:numNS]
85	dns.Extra = dns.Extra[:numExtra]
86
87	if edns0 != nil {
88		// Add the OPT record back onto the additional section.
89		dns.Extra = append(dns.Extra, edns0)
90	}
91}
92
93func truncateLoop(rrs []RR, size, l int, compression map[string]struct{}) (int, int) {
94	for i, r := range rrs {
95		if r == nil {
96			continue
97		}
98
99		l += r.len(l, compression)
100		if l > size {
101			// Return size, rather than l prior to this record,
102			// to prevent any further records being added.
103			return size, i
104		}
105		if l == size {
106			return l, i + 1
107		}
108	}
109
110	return l, len(rrs)
111}
112