From e13c7dc7b71f5d6d07587fff13274316e2fa15f1 Mon Sep 17 00:00:00 2001
From: Lonny Wong <lonnywong@qq.com>
Date: Sat, 27 Jul 2024 15:52:56 +0800
Subject: [PATCH] support custom udp port range

---
 README.cn.md |  3 ++-
 README.en.md |  3 ++-
 tssh/udp.go  | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/README.cn.md b/README.cn.md
index 661dc35..c64269b 100644
--- a/README.cn.md
+++ b/README.cn.md
@@ -715,6 +715,7 @@ trzsz-ssh ( tssh ) 设计为 ssh 客户端的直接替代品，提供与 openssh
   ```
   Host xxx
       #!! UdpMode yes
+      #!! UdpPort 61000-62000
       #!! TsshdPath ~/go/bin/tsshd
   ```
 
@@ -722,7 +723,7 @@ trzsz-ssh ( tssh ) 设计为 ssh 客户端的直接替代品，提供与 openssh
 
 - `tssh` 会先作为一个 ssh 客户端正常登录到服务器上，然后在服务器上启动一个新的 `tsshd` 进程。
 
-- `tsshd` 进程会随机侦听一个 61000 到 62000 之间的 UDP 端口，并将其端口和密钥通过 ssh 通道发回给 `tssh` 进程。登录的 ssh 连接会被关闭，然后 `tssh` 进程通过 UDP 与 `tsshd` 进程通讯。
+- `tsshd` 进程会随机侦听一个 61000 到 62000 之间的 UDP 端口（可通过 `UdpPort` 配置自定义），并将其端口和密钥通过 ssh 通道发回给`tssh`进程。登录的 ssh 连接会被关闭，然后`tssh`进程通过 UDP 与`tsshd` 进程通讯。
 
 - `tsshd` 支持 `QUIC` 协议和 `KCP` 协议（默认是 `QUIC` 协议），可以命令行指定（如 `-oUdpMode=KCP`），或如下配置：
 
diff --git a/README.en.md b/README.en.md
index 038829f..66577df 100644
--- a/README.en.md
+++ b/README.en.md
@@ -717,6 +717,7 @@ trzsz-ssh ( tssh ) is an ssh client designed as a drop-in replacement for the op
   ```
   Host xxx
       #!! UdpMode yes
+      #!! UdpPort 61000-62000
       #!! TsshdPath ~/go/bin/tsshd
   ```
 
@@ -724,7 +725,7 @@ trzsz-ssh ( tssh ) is an ssh client designed as a drop-in replacement for the op
 
 - The `tssh` will first login to the server normally as an ssh client, and then run a new `tsshd` process on the server.
 
-- The `tsshd` process listens on a random udp port between 61000 and 62000, and sends its port number and a secret key back to the `tssh` process over the ssh channel. The ssh connection is then shut down, and the `tssh` process communicates with the `tsshd` process over udp.
+- The `tsshd` process listens on a random udp port between 61000 and 62000 (can be customized by `UdpPort`), and sends its port number and a secret key back to the `tssh`process over the ssh channel. The ssh connection is then shut down, and the`tssh`process communicates with the`tsshd` process over udp.
 
 - The `tsshd` supports `QUIC` protocol and `KCP` protocol (the default is `QUIC`), which can be specified on the command line (such as `-oUdpMode=KCP`), or configured as follows:
 
diff --git a/tssh/udp.go b/tssh/udp.go
index 090b496..29e32c0 100644
--- a/tssh/udp.go
+++ b/tssh/udp.go
@@ -35,6 +35,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+	"unicode"
 
 	"github.com/google/shlex"
 	"github.com/trzsz/tsshd/tsshd"
@@ -851,6 +852,37 @@ func getTsshdCommand(args *sshArgs, udpMode int) string {
 		buf.WriteString(" --kcp")
 	}
 
+	if udpPort := getExOptionConfig(args, "UdpPort"); udpPort != "" {
+		ports := strings.FieldsFunc(udpPort, func(c rune) bool {
+			return unicode.IsSpace(c) || c == ',' || c == '-'
+		})
+		if len(ports) == 1 {
+			port, err := strconv.Atoi(ports[0])
+			if err != nil {
+				warning("UdpPort %s is invalid: %v", udpPort, err)
+			} else {
+				buf.WriteString(fmt.Sprintf(" --port %d", port))
+			}
+		} else if len(ports) == 2 {
+			for {
+				lowPort, err := strconv.Atoi(ports[0])
+				if err != nil {
+					warning("UdpPort %s is invalid: %v", udpPort, err)
+					break
+				}
+				highPort, err := strconv.Atoi(ports[1])
+				if err != nil {
+					warning("UdpPort %s is invalid: %v", udpPort, err)
+					break
+				}
+				buf.WriteString(fmt.Sprintf(" --port %d-%d", lowPort, highPort))
+				break // nolint:all
+			}
+		} else {
+			warning("UdpPort %s is invalid", udpPort)
+		}
+	}
+
 	return buf.String()
 }