1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40
41. 41
42. 42
43. 43
44. 44
45. 45
46. 46
47. 47
48. 48
49. 49
50. 50
51. 51
52. 52
53. 53
54. 54
55. 55
56. 56
57. 57
58. 58
59. 59
60. 60
61. 61
62. 62
63. 63
64. 64
65. 65
66. 66
67. 67
68. 68
69. 69
70. 70
71. 71
72. 72
73. 73
74. 74
75. 75
76. 76
77. 77
78. 78
79. 79
80. 80
81. 81
82. 82
83. 83
84. 84
85. 85
86. 86
87. 87
88. 88
89. 89
90. 90
91. 91
92. 92
93. 93
94. 94
95. 95
96. 96
97. 97
98. 98
99. 99
100. 100
101. 101
102. 102
103. 103
104. 104
105. 105
106. 106
107. 107
108. 108
109. 109
110. 110
111. 111
112. 112
113. 113
114. 114
115. 115
116. 116
117. 117

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.
//
// If the message fits within the requested size without compression,
// Truncate will set the message's Compress attribute to false. It is
// the caller's responsibility to set it back to true if they wish to
// compress the payload regardless of size.
//
// The TC bit will be set if any records were excluded from the message.
// If the TC bit is already set on the message it will be retained.
// TC indicates 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 (MinMsgSize) bytes must be treated as equal to 512 bytes.
	//
	// For ease of use, we impose that restriction here.
	if size < MinMsgSize {
		size = MinMsgSize
	}

	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 {
		_, numExtra = truncateLoop(dns.Extra, size, l, compression)
	}

	// See the function documentation for when we set this.
	dns.Truncated = 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)
}