feat: support --dns command (#154)

tssh/args.go

@@ -81,6 +81,7 @@ type sshArgs struct {
 	Client         bool        `arg:"--client" help:"force trzsz run as a client on the jump server"`
 	Debug          bool        `arg:"--debug" help:"verbose mode for debugging, same as ssh's -vvv"`
 	Zmodem         bool        `arg:"--zmodem" help:"enable zmodem lrzsz ( rz / sz ) feature"`
+	Dns            string      `arg:"--dns" placeholder:"[udp://|tcp://]host[:port]" help:"custom DNS server"`
 	Udp            bool        `arg:"--udp" help:"ssh over UDP like mosh (default mode: QUIC)"`
 	TsshdPath      string      `arg:"--tsshd-path" placeholder:"path" help:"[udp] tsshd absolute path on the server"`
 	NewHost        bool        `arg:"--new-host" help:"[tools] add new host to configuration"`

tssh/dns.go

@@ -0,0 +1,88 @@
+/*
+MIT License
+
+Copyright (c) 2023-2024 The Trzsz SSH Authors.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+package tssh
+
+import (
+	"context"
+	"net"
+	"net/url"
+	"strings"
+)
+
+// setDNS sets the net.DefaultResolver to use the given DNS server.
+func setDNS(dns string) {
+
+	network, dns, err := resolveDnsAddress(dns)
+	if err != nil {
+		return
+
+	}
+
+	net.DefaultResolver = &net.Resolver{
+		PreferGo: true,
+		Dial: func(ctx context.Context, _, addr string) (net.Conn, error) {
+			debug("use custom DNS: %s://%s", network, dns)
+			var d net.Dialer
+			return d.DialContext(ctx, network, dns)
+		},
+	}
+
+}
+
+func resolveDnsAddress(dns string) (string, string, error) {
+
+	var preParseDns string
+	if !strings.Contains(dns, "://") {
+		preParseDns = "udp://" + dns
+	} else {
+		preParseDns = dns
+	}
+
+	svrParse, err := url.Parse(preParseDns)
+	if err != nil {
+		warning("parse dns [%s] failed: %v", dns, err)
+		return "", "", err
+
+	}
+
+	var network string
+	switch strings.ToLower(svrParse.Scheme) {
+	case "tcp":
+		network = "tcp"
+	default:
+		network = "udp"
+	}
+
+	host, port, err := net.SplitHostPort(svrParse.Host)
+	if err != nil {
+		// If no port is specified, use default port 53
+		host = svrParse.Host
+		port = "53"
+	}
+
+	dns = net.JoinHostPort(host, port)
+	return network, dns, nil
+
+}

tssh/dns_test.go

@@ -0,0 +1,72 @@
+/*
+MIT License
+
+Copyright (c) 2023-2024 The Trzsz SSH Authors.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+package tssh
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDNS(t *testing.T) {
+
+	assert := assert.New(t)
+	assertDestEqual := func(waitParseDns, expectedDns string) {
+
+		t.Helper()
+		network, dns, err := resolveDnsAddress(waitParseDns)
+		assert.Nil(err)
+		assert.Equal(expectedDns, fmt.Sprintf("%s://%s", network, dns))
+
+	}
+
+	assertDestNotNil := func(preParseDns string) {
+		t.Helper()
+		_, _, err := resolveDnsAddress(preParseDns)
+		assert.NotNil(err)
+
+	}
+
+	assertDestNotNil("ab cd")
+	assertDestNotNil("udp://ab:cd")
+
+	assertDestEqual("8.8.8.8", "udp://8.8.8.8:53")
+	assertDestEqual("8.8.8.8:53", "udp://8.8.8.8:53")
+	assertDestEqual("udp://8.8.8.8", "udp://8.8.8.8:53")
+	assertDestEqual("udp://8.8.8.8:53", "udp://8.8.8.8:53")
+	assertDestEqual("tcp://8.8.8.8", "tcp://8.8.8.8:53")
+	assertDestEqual("tcp://8.8.8.8:53", "tcp://8.8.8.8:53")
+	assertDestEqual("udp://8.8.8.8:5300", "udp://8.8.8.8:5300")
+
+	assertDestEqual("2001:4860:4860::8888", "udp://[2001:4860:4860::8888]:53")
+	assertDestEqual("[2001:4860:4860::8888]:53", "udp://[2001:4860:4860::8888]:53")
+	assertDestEqual("udp://2001:4860:4860::8888", "udp://[2001:4860:4860::8888]:53")
+	assertDestEqual("udp://[2001:4860:4860::8888]:53", "udp://[2001:4860:4860::8888]:53")
+	assertDestEqual("tcp://2001:4860:4860::8888", "tcp://[2001:4860:4860::8888]:53")
+	assertDestEqual("tcp://[2001:4860:4860::8888]:53", "tcp://[2001:4860:4860::8888]:53")
+	assertDestEqual("udp://[2001:4860:4860::8888]:5300", "udp://[2001:4860:4860::8888]:5300")
+
+}

tssh/main.go

@@ -193,6 +193,10 @@ func TsshMain(argv []string) int {
 	args.Destination = dest
 	args.originalDest = dest
 
+	if args.Dns != "" {
+		setDNS(args.Dns)
+	}
+
 	// start ssh program
 	var code int
 	code, err = sshStart(&args)