prometheus

1
// Copyright 2019 The Prometheus Authors
2
// Licensed under the Apache License, Version 2.0 (the "License");
3
// you may not use this file except in compliance with the License.
4
// You may obtain a copy of the License at
5
//
6
// http://www.apache.org/licenses/LICENSE-2.0
7
//
8
// Unless required by applicable law or agreed to in writing, software
9
// distributed under the License is distributed on an "AS IS" BASIS,
10
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11
// See the License for the specific language governing permissions and
12
// limitations under the License.
13

14
package dns
15

16
import (
17
	"context"
18
	"fmt"
19
	"net"
20
	"testing"
21
	"time"
22

23
	"github.com/go-kit/log"
24
	"github.com/miekg/dns"
25
	"github.com/prometheus/client_golang/prometheus"
26
	"github.com/prometheus/common/model"
27
	"github.com/stretchr/testify/require"
28
	"go.uber.org/goleak"
29
	"gopkg.in/yaml.v2"
30

31
	"github.com/prometheus/prometheus/discovery"
32
	"github.com/prometheus/prometheus/discovery/targetgroup"
33
)
34

35
func TestMain(m *testing.M) {
36
	goleak.VerifyTestMain(m)
37
}
38

39
func TestDNS(t *testing.T) {
40
	testCases := []struct {
41
		name   string
42
		config SDConfig
43
		lookup func(name string, qtype uint16, logger log.Logger) (*dns.Msg, error)
44

45
		expected []*targetgroup.Group
46
	}{
47
		{
48
			name: "A record query with error",
49
			config: SDConfig{
50
				Names:           []string{"web.example.com."},
51
				RefreshInterval: model.Duration(time.Minute),
52
				Port:            80,
53
				Type:            "A",
54
			},
55
			lookup: func(name string, qtype uint16, logger log.Logger) (*dns.Msg, error) {
56
				return nil, fmt.Errorf("some error")
57
			},
58
			expected: []*targetgroup.Group{},
59
		},
60
		{
61
			name: "A record query",
62
			config: SDConfig{
63
				Names:           []string{"web.example.com."},
64
				RefreshInterval: model.Duration(time.Minute),
65
				Port:            80,
66
				Type:            "A",
67
			},
68
			lookup: func(name string, qtype uint16, logger log.Logger) (*dns.Msg, error) {
69
				return &dns.Msg{
70
						Answer: []dns.RR{
71
							&dns.A{A: net.IPv4(192, 0, 2, 2)},
72
						},
73
					},
74
					nil
75
			},
76
			expected: []*targetgroup.Group{
77
				{
78
					Source: "web.example.com.",
79
					Targets: []model.LabelSet{
80
						{
81
							"__address__":                  "192.0.2.2:80",
82
							"__meta_dns_name":              "web.example.com.",
83
							"__meta_dns_srv_record_target": "",
84
							"__meta_dns_srv_record_port":   "",
85
							"__meta_dns_mx_record_target":  "",
86
							"__meta_dns_ns_record_target":  "",
87
						},
88
					},
89
				},
90
			},
91
		},
92
		{
93
			name: "AAAA record query",
94
			config: SDConfig{
95
				Names:           []string{"web.example.com."},
96
				RefreshInterval: model.Duration(time.Minute),
97
				Port:            80,
98
				Type:            "AAAA",
99
			},
100
			lookup: func(name string, qtype uint16, logger log.Logger) (*dns.Msg, error) {
101
				return &dns.Msg{
102
						Answer: []dns.RR{
103
							&dns.AAAA{AAAA: net.IPv6loopback},
104
						},
105
					},
106
					nil
107
			},
108
			expected: []*targetgroup.Group{
109
				{
110
					Source: "web.example.com.",
111
					Targets: []model.LabelSet{
112
						{
113
							"__address__":                  "[::1]:80",
114
							"__meta_dns_name":              "web.example.com.",
115
							"__meta_dns_srv_record_target": "",
116
							"__meta_dns_srv_record_port":   "",
117
							"__meta_dns_mx_record_target":  "",
118
							"__meta_dns_ns_record_target":  "",
119
						},
120
					},
121
				},
122
			},
123
		},
124
		{
125
			name: "SRV record query",
126
			config: SDConfig{
127
				Names:           []string{"_mysql._tcp.db.example.com."},
128
				Type:            "SRV",
129
				RefreshInterval: model.Duration(time.Minute),
130
			},
131
			lookup: func(name string, qtype uint16, logger log.Logger) (*dns.Msg, error) {
132
				return &dns.Msg{
133
						Answer: []dns.RR{
134
							&dns.SRV{Port: 3306, Target: "db1.example.com."},
135
							&dns.SRV{Port: 3306, Target: "db2.example.com."},
136
						},
137
					},
138
					nil
139
			},
140
			expected: []*targetgroup.Group{
141
				{
142
					Source: "_mysql._tcp.db.example.com.",
143
					Targets: []model.LabelSet{
144
						{
145
							"__address__":                  "db1.example.com:3306",
146
							"__meta_dns_name":              "_mysql._tcp.db.example.com.",
147
							"__meta_dns_srv_record_target": "db1.example.com.",
148
							"__meta_dns_srv_record_port":   "3306",
149
							"__meta_dns_mx_record_target":  "",
150
							"__meta_dns_ns_record_target":  "",
151
						},
152
						{
153
							"__address__":                  "db2.example.com:3306",
154
							"__meta_dns_name":              "_mysql._tcp.db.example.com.",
155
							"__meta_dns_srv_record_target": "db2.example.com.",
156
							"__meta_dns_srv_record_port":   "3306",
157
							"__meta_dns_mx_record_target":  "",
158
							"__meta_dns_ns_record_target":  "",
159
						},
160
					},
161
				},
162
			},
163
		},
164
		{
165
			name: "SRV record query with unsupported resource records",
166
			config: SDConfig{
167
				Names:           []string{"_mysql._tcp.db.example.com."},
168
				RefreshInterval: model.Duration(time.Minute),
169
			},
170
			lookup: func(name string, qtype uint16, logger log.Logger) (*dns.Msg, error) {
171
				return &dns.Msg{
172
						Answer: []dns.RR{
173
							&dns.SRV{Port: 3306, Target: "db1.example.com."},
174
							&dns.TXT{Txt: []string{"this should be discarded"}},
175
						},
176
					},
177
					nil
178
			},
179
			expected: []*targetgroup.Group{
180
				{
181
					Source: "_mysql._tcp.db.example.com.",
182
					Targets: []model.LabelSet{
183
						{
184
							"__address__":                  "db1.example.com:3306",
185
							"__meta_dns_name":              "_mysql._tcp.db.example.com.",
186
							"__meta_dns_srv_record_target": "db1.example.com.",
187
							"__meta_dns_srv_record_port":   "3306",
188
							"__meta_dns_mx_record_target":  "",
189
							"__meta_dns_ns_record_target":  "",
190
						},
191
					},
192
				},
193
			},
194
		},
195
		{
196
			name: "SRV record query with empty answer (NXDOMAIN)",
197
			config: SDConfig{
198
				Names:           []string{"_mysql._tcp.db.example.com."},
199
				RefreshInterval: model.Duration(time.Minute),
200
			},
201
			lookup: func(name string, qtype uint16, logger log.Logger) (*dns.Msg, error) {
202
				return &dns.Msg{}, nil
203
			},
204
			expected: []*targetgroup.Group{
205
				{
206
					Source: "_mysql._tcp.db.example.com.",
207
				},
208
			},
209
		},
210
		{
211
			name: "MX record query",
212
			config: SDConfig{
213
				Names:           []string{"example.com."},
214
				Type:            "MX",
215
				Port:            25,
216
				RefreshInterval: model.Duration(time.Minute),
217
			},
218
			lookup: func(name string, qtype uint16, logger log.Logger) (*dns.Msg, error) {
219
				return &dns.Msg{
220
						Answer: []dns.RR{
221
							&dns.MX{Preference: 0, Mx: "smtp1.example.com."},
222
							&dns.MX{Preference: 10, Mx: "smtp2.example.com."},
223
						},
224
					},
225
					nil
226
			},
227
			expected: []*targetgroup.Group{
228
				{
229
					Source: "example.com.",
230
					Targets: []model.LabelSet{
231
						{
232
							"__address__":                  "smtp1.example.com:25",
233
							"__meta_dns_name":              "example.com.",
234
							"__meta_dns_srv_record_target": "",
235
							"__meta_dns_srv_record_port":   "",
236
							"__meta_dns_mx_record_target":  "smtp1.example.com.",
237
							"__meta_dns_ns_record_target":  "",
238
						},
239
						{
240
							"__address__":                  "smtp2.example.com:25",
241
							"__meta_dns_name":              "example.com.",
242
							"__meta_dns_srv_record_target": "",
243
							"__meta_dns_srv_record_port":   "",
244
							"__meta_dns_mx_record_target":  "smtp2.example.com.",
245
							"__meta_dns_ns_record_target":  "",
246
						},
247
					},
248
				},
249
			},
250
		},
251
	}
252

253
	for _, tc := range testCases {
254
		tc := tc
255
		t.Run(tc.name, func(t *testing.T) {
256
			t.Parallel()
257

258
			reg := prometheus.NewRegistry()
259
			refreshMetrics := discovery.NewRefreshMetrics(reg)
260
			metrics := tc.config.NewDiscovererMetrics(reg, refreshMetrics)
261
			require.NoError(t, metrics.Register())
262

263
			sd, err := NewDiscovery(tc.config, nil, metrics)
264
			require.NoError(t, err)
265
			sd.lookupFn = tc.lookup
266

267
			tgs, err := sd.refresh(context.Background())
268
			require.NoError(t, err)
269
			require.Equal(t, tc.expected, tgs)
270

271
			metrics.Unregister()
272
		})
273
	}
274
}
275

276
func TestSDConfigUnmarshalYAML(t *testing.T) {
277
	marshal := func(c SDConfig) []byte {
278
		d, err := yaml.Marshal(c)
279
		if err != nil {
280
			panic(err)
281
		}
282
		return d
283
	}
284

285
	unmarshal := func(d []byte) func(interface{}) error {
286
		return func(o interface{}) error {
287
			return yaml.Unmarshal(d, o)
288
		}
289
	}
290

291
	cases := []struct {
292
		name      string
293
		input     SDConfig
294
		expectErr bool
295
	}{
296
		{
297
			name: "valid srv",
298
			input: SDConfig{
299
				Names: []string{"a.example.com", "b.example.com"},
300
				Type:  "SRV",
301
			},
302
			expectErr: false,
303
		},
304
		{
305
			name: "valid a",
306
			input: SDConfig{
307
				Names: []string{"a.example.com", "b.example.com"},
308
				Type:  "A",
309
				Port:  5300,
310
			},
311
			expectErr: false,
312
		},
313
		{
314
			name: "valid aaaa",
315
			input: SDConfig{
316
				Names: []string{"a.example.com", "b.example.com"},
317
				Type:  "AAAA",
318
				Port:  5300,
319
			},
320
			expectErr: false,
321
		},
322
		{
323
			name: "invalid a without port",
324
			input: SDConfig{
325
				Names: []string{"a.example.com", "b.example.com"},
326
				Type:  "A",
327
			},
328
			expectErr: true,
329
		},
330
		{
331
			name: "invalid aaaa without port",
332
			input: SDConfig{
333
				Names: []string{"a.example.com", "b.example.com"},
334
				Type:  "AAAA",
335
			},
336
			expectErr: true,
337
		},
338
		{
339
			name: "invalid empty names",
340
			input: SDConfig{
341
				Names: []string{},
342
				Type:  "AAAA",
343
			},
344
			expectErr: true,
345
		},
346
		{
347
			name: "invalid unknown dns type",
348
			input: SDConfig{
349
				Names: []string{"a.example.com", "b.example.com"},
350
				Type:  "PTR",
351
			},
352
			expectErr: true,
353
		},
354
	}
355

356
	for _, c := range cases {
357
		t.Run(c.name, func(t *testing.T) {
358
			var config SDConfig
359
			d := marshal(c.input)
360
			err := config.UnmarshalYAML(unmarshal(d))
361
			require.Equal(t, c.expectErr, err != nil)
362
		})
363
	}
364
}
365