From 96df7d7d20ee4b84698d5758c55314211d69038d Mon Sep 17 00:00:00 2001
From: Stanislav Chzhen <s.chzhen@adguard.com>
Date: Wed, 20 Dec 2023 20:46:24 +0300
Subject: [PATCH] all: imp code

---
 proxy/upstreams.go        | 10 +++++++---
 upstream/upstream.go      | 24 ++++++++++++++----------
 upstream/upstream_test.go | 14 ++++++++------
 3 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/proxy/upstreams.go b/proxy/upstreams.go
index 7f469cd37..39c0d83e3 100644
--- a/proxy/upstreams.go
+++ b/proxy/upstreams.go
@@ -13,6 +13,8 @@ import (
 	"golang.org/x/exp/slices"
 )
 
+// errWrongUpstreamSpecs is returned by [ParseUpstreamsConfig] when upstreams
+// configuration is invalid.
 const errWrongUpstreamSpecs errors.Error = "wrong upstream specification"
 
 // UpstreamConfig is a wrapper for a list of default upstreams, a map of
@@ -37,7 +39,8 @@ type UpstreamConfig struct {
 var _ io.Closer = (*UpstreamConfig)(nil)
 
 // ParseUpstreamsConfig returns UpstreamConfig and error if upstreams
-// configuration is invalid.
+// configuration is invalid.  It also skips empty lines and comments (lines
+// begining with `#`).
 //
 // # Simple upstreams
 //
@@ -158,9 +161,9 @@ type configParser struct {
 }
 
 // parse returns UpstreamConfig and error if upstreams configuration is invalid.
-func (p *configParser) parse(conf []string) (c *UpstreamConfig, err error) {
+func (p *configParser) parse(lines []string) (c *UpstreamConfig, err error) {
 	var errs []error
-	for i, l := range conf {
+	for i, l := range lines {
 		if err = p.parseLine(i, l); err != nil {
 			errs = append(errs, &ParseError{Idx: i, err: err})
 		}
@@ -261,6 +264,7 @@ func (p *configParser) specifyUpstream(domains []string, u string, idx int) (err
 	if len(domains) == 0 {
 		p.upstreams = append(p.upstreams, dnsUpstream)
 
+		// TODO(s.chzhen):  Logs without index.
 		log.Debug("dnsproxy: upstream at index %d: %s", idx, addr)
 	} else {
 		p.includeToReserved(dnsUpstream, domains)
diff --git a/upstream/upstream.go b/upstream/upstream.go
index cb2e3f79a..c3c2c1adf 100644
--- a/upstream/upstream.go
+++ b/upstream/upstream.go
@@ -11,7 +11,6 @@ import (
 	"net"
 	"net/netip"
 	"net/url"
-	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -150,6 +149,7 @@ const (
 // AddressToUpstream converts addr to an Upstream using the specified options.
 // addr can be either a URL, or a plain address, either a domain name or an IP.
 //
+//   - 1.2.3.4 or 1.2.3.4:4321 for plain DNS using IP address;
 //   - udp://5.3.5.3:53 or 5.3.5.3:53 for plain DNS using IP address;
 //   - udp://name.server:53 or name.server:53 for plain DNS using domain name;
 //   - tcp://5.3.5.3:53 for plain DNS-over-TCP using IP address;
@@ -175,22 +175,26 @@ func AddressToUpstream(addr string, opts *Options) (u Upstream, err error) {
 		opts = &Options{}
 	}
 
+	_, err = netip.ParseAddr(addr)
+	if err == nil {
+		return urlToUpstream(&url.URL{Scheme: "udp", Host: addr}, opts)
+	}
+
+	_, err = netip.ParseAddrPort(addr)
+	if err == nil {
+		return urlToUpstream(&url.URL{Scheme: "udp", Host: addr}, opts)
+	}
+
 	var uu *url.URL
 	if strings.Contains(addr, "://") {
-		// Parse as URL.
 		uu, err = url.Parse(addr)
 		if err != nil {
 			return nil, fmt.Errorf("failed to parse %s: %w", addr, err)
 		}
 	} else {
-		// Probably, plain UDP upstream defined by address or address:port.
-		_, port, splitErr := net.SplitHostPort(addr)
-		if splitErr == nil {
-			// Validate port.
-			_, err = strconv.ParseUint(port, 10, 16)
-			if err != nil {
-				return nil, fmt.Errorf("invalid address %s: %w", addr, err)
-			}
+		err = netutil.ValidateHostname(addr)
+		if err != nil {
+			return nil, fmt.Errorf("invalid address %s: %w", addr, err)
 		}
 
 		uu = &url.URL{
diff --git a/upstream/upstream_test.go b/upstream/upstream_test.go
index f56569253..7907f1f54 100644
--- a/upstream/upstream_test.go
+++ b/upstream/upstream_test.go
@@ -271,15 +271,17 @@ func TestAddressToUpstream_bads(t *testing.T) {
 		wantErrMsg: "unsupported url scheme: asdf",
 	}, {
 		addr: "12345.1.1.1:1234567",
-		wantErrMsg: `invalid address 12345.1.1.1:1234567: ` +
-			`strconv.ParseUint: parsing "1234567": value out of range`,
+		wantErrMsg: `invalid address 12345.1.1.1:1234567: bad hostname ` +
+			`"12345.1.1.1:1234567": bad top-level domain name label "1:1234567": ` +
+			`bad top-level domain name label rune ':'`,
 	}, {
 		addr: ":1234567",
-		wantErrMsg: `invalid address :1234567: ` +
-			`strconv.ParseUint: parsing "1234567": value out of range`,
+		wantErrMsg: `invalid address :1234567: bad hostname ":1234567": bad top-level ` +
+			`domain name label ":1234567": bad top-level domain name label rune ':'`,
 	}, {
-		addr:       "host:",
-		wantErrMsg: `invalid address host:: strconv.ParseUint: parsing "": invalid syntax`,
+		addr: "host:",
+		wantErrMsg: `invalid address host:: bad hostname "host:": bad top-level ` +
+			`domain name label "host:": bad top-level domain name label rune ':'`,
 	}}
 
 	for _, tc := range testCases {