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