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