From 2b58bd9876d5faa8600fe58c3c23ab6288cac38f Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Mon, 15 Aug 2022 00:32:32 +0300
Subject: [PATCH] lib/promscrape/discovery/dns: add support for resolving MX
 records

See https://github.com/prometheus/prometheus/pull/10099
---
 docs/CHANGELOG.md                   |  1 +
 docs/sd_configs.md                  |  3 +-
 lib/promscrape/discovery/dns/dns.go | 54 ++++++++++++++++++++++++++++-
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index 4ab54883e..f6e197561 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -18,6 +18,7 @@ The following tip changes can be tested by building VictoriaMetrics components f
 * SECURITY: [vmalert](https://docs.victoriametrics.com/vmalert.html): do not expose `-remoteWrite.url`, `-remoteRead.url` and `-datasource.url` command-line flag values in logs and at `http://vmalert:8880/flags` page by default, since they may contain sensitive data such as auth keys. This aligns `vmalert` behaviour with [vmagent](https://docs.victoriametrics.com/vmagent.html), which doesn't expose `-remoteWrite.url` command-line flag value in logs and at `http://vmagent:8429/flags` page by default. Specify `-remoteWrite.showURL`, `-remoteRead.showURL` and `-datasource.showURL` command-line flags for showing values for the corresponding `-*.url` flags in logs. Thanks to @mble for [the pull request](https://github.com/VictoriaMetrics/VictoriaMetrics/pull/2965).
 
 * FEATURE: [VictoriaMetrics cluster](https://docs.victoriametrics.com/Cluster-VictoriaMetrics.html): improve performance for heavy queries on systems with many CPU cores.
+* FEATURE: [vmagent](https://docs.victoriametrics.com/vmagent.html): add support for MX record types in [dns_sd_configs](https://docs.victoriametrics.com/sd_configs.html#dns_sd_configs) in the same way as Prometheus 2.38 [does](https://github.com/prometheus/prometheus/pull/10099).
 
 * BUGFIX: prevent from excess CPU usage when the storage enters [read-only mode](https://docs.victoriametrics.com/Cluster-VictoriaMetrics.html#readonly-mode).
 
diff --git a/docs/sd_configs.md b/docs/sd_configs.md
index 0f80e2620..f3281e5ca 100644
--- a/docs/sd_configs.md
+++ b/docs/sd_configs.md
@@ -215,7 +215,7 @@ scrape_configs:
   - names: ["...", "..."]
 
     # type is an optional type of DNS query to perform.
-    # Supported values are: SRV, A, or AAAA.
+    # Supported values are: SRV, A, AAAA or MX.
     # By default SRV is used.
     # type: ...
 
@@ -228,6 +228,7 @@ The following meta labels are available on discovered targets during [relabeling
 * `__meta_dns_name`: the record name that produced the discovered target.
 * `__meta_dns_srv_record_target`: the target field of the SRV record
 * `__meta_dns_srv_record_port`: the port field of the SRV record
+* `__meta_dns_mx_record_target`: the target field of the MX record.
 
 
 ## docker_sd_configs
diff --git a/lib/promscrape/discovery/dns/dns.go b/lib/promscrape/discovery/dns/dns.go
index d6a791b1a..4dde4b9cc 100644
--- a/lib/promscrape/discovery/dns/dns.go
+++ b/lib/promscrape/discovery/dns/dns.go
@@ -45,6 +45,9 @@ func (sdc *SDConfig) GetLabels(baseDir string) ([]map[string]string, error) {
 	case "SRV":
 		ms := getSRVAddrLabels(ctx, sdc)
 		return ms, nil
+	case "MX":
+		ms := getMXAddrLabels(ctx, sdc)
+		return ms, nil
 	case "A", "AAAA":
 		return getAAddrLabels(ctx, sdc, typ)
 	default:
@@ -57,6 +60,45 @@ func (sdc *SDConfig) MustStop() {
 	// nothing to do
 }
 
+func getMXAddrLabels(ctx context.Context, sdc *SDConfig) []map[string]string {
+	port := 25
+	if sdc.Port != nil {
+		port = *sdc.Port
+	}
+	type result struct {
+		name string
+		mx   []*net.MX
+		err  error
+	}
+	ch := make(chan result, len(sdc.Names))
+	for _, name := range sdc.Names {
+		go func(name string) {
+			mx, err := resolver.LookupMX(ctx, name)
+			ch <- result{
+				name: name,
+				mx:   mx,
+				err:  err,
+			}
+		}(name)
+	}
+	var ms []map[string]string
+	for range sdc.Names {
+		r := <-ch
+		if r.err != nil {
+			logger.Errorf("error in MX lookup for %q; skipping it; error: %s", r.name, r.err)
+			continue
+		}
+		for _, mx := range r.mx {
+			target := mx.Host
+			for strings.HasSuffix(target, ".") {
+				target = target[:len(target)-1]
+			}
+			ms = appendMXLabels(ms, r.name, target, port)
+		}
+	}
+	return ms
+}
+
 func getSRVAddrLabels(ctx context.Context, sdc *SDConfig) []map[string]string {
 	type result struct {
 		name string
@@ -94,7 +136,7 @@ func getSRVAddrLabels(ctx context.Context, sdc *SDConfig) []map[string]string {
 
 func getAAddrLabels(ctx context.Context, sdc *SDConfig, lookupType string) ([]map[string]string, error) {
 	if sdc.Port == nil {
-		return nil, fmt.Errorf("missing `port` in `dns_sd_config`")
+		return nil, fmt.Errorf("missing `port` in `dns_sd_config` for `type: %s`", lookupType)
 	}
 	port := *sdc.Port
 	type result struct {
@@ -131,6 +173,16 @@ func getAAddrLabels(ctx context.Context, sdc *SDConfig, lookupType string) ([]ma
 	return ms, nil
 }
 
+func appendMXLabels(ms []map[string]string, name, target string, port int) []map[string]string {
+	addr := discoveryutils.JoinHostPort(target, port)
+	m := map[string]string{
+		"__address__":                 addr,
+		"__meta_dns_name":             name,
+		"__meta_dns_mx_record_target": target,
+	}
+	return append(ms, m)
+}
+
 func appendAddrLabels(ms []map[string]string, name, target string, port int) []map[string]string {
 	addr := discoveryutils.JoinHostPort(target, port)
 	m := map[string]string{