From f586506180cbe16dafeb286ff766715f68851a5b Mon Sep 17 00:00:00 2001
From: tux <t@tux.rs>
Date: Tue, 27 Jan 2026 17:31:19 +0530
Subject: [PATCH] feat: add cli option to pass custom remote port

---
 cmd/tcp.go              | 10 ++++++++--
 internal/client/init.go |  4 ++--
 internal/client/trok.go | 16 ++++++++++++++--
 internal/server/trok.go |  3 ++-
 4 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/cmd/tcp.go b/cmd/tcp.go
index 3beb78c..b0a6740 100644
--- a/cmd/tcp.go
+++ b/cmd/tcp.go
@@ -16,18 +16,24 @@ var tcpCmd = &cobra.Command{
 	Short: "Start TCP proxy",
 	Args:  cobra.ExactArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
-		port := args[0]
+		localPort := args[0]
 
 		serverAddr, err := cmd.Flags().GetString("server")
 		if err != nil {
 			panic(err)
 		}
 
-		client.Start(serverAddr, fmt.Sprintf(":%s", port))
+		remotePort, err := cmd.Flags().GetString("port")
+		if err != nil {
+			panic(err)
+		}
+
+		client.Start(serverAddr, fmt.Sprintf(":%s", localPort), remotePort)
 	},
 }
 
 func init() {
 	rootCmd.AddCommand(tcpCmd)
 	tcpCmd.Flags().StringP("server", "s", "trok.cloud:1337", "Remote server address")
+	tcpCmd.Flags().StringP("port", "p", "0", "port on the remote server")
 }
diff --git a/internal/client/init.go b/internal/client/init.go
index 38c3d3a..506f592 100644
--- a/internal/client/init.go
+++ b/internal/client/init.go
@@ -11,8 +11,8 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-func Start(serverAddr, localAddr string) {
-	trok, err := NewTrokClient(serverAddr, localAddr)
+func Start(serverAddr, localAddr, remotePort string) {
+	trok, err := NewTrokClient(serverAddr, localAddr, remotePort)
 	if err != nil {
 		log.Fatal().Msgf("failed init trok %v", err)
 	}
diff --git a/internal/client/trok.go b/internal/client/trok.go
index 4288da6..8fabd5d 100644
--- a/internal/client/trok.go
+++ b/internal/client/trok.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"os"
 	"strings"
 
 	"github.com/rs/zerolog/log"
@@ -18,15 +19,17 @@ type Trok struct {
 	controlClient *TCPClient
 	serverAddr    string
 	localAddr     string
+	remotePort    string
 }
 
-func NewTrokClient(serverAddr, localAddr string) (*Trok, error) {
+func NewTrokClient(serverAddr, localAddr, remotePort string) (*Trok, error) {
 	controlClient, err := NewTCPClient(serverAddr, "Controller")
 
 	return &Trok{
 		controlClient: controlClient,
 		serverAddr:    serverAddr,
 		localAddr:     localAddr,
+		remotePort:    remotePort,
 	}, err
 }
 
@@ -45,7 +48,7 @@ func (t *Trok) Stop() {
 func (t *Trok) ControlConnHandler(conn net.Conn) {
 	p := lib.InitProtocolHandler(conn)
 
-	p.WriteMessage(&lib.Message{CMD: "HELO", ARG: "Trok"})
+	p.WriteMessage(&lib.Message{CMD: "HELO", ARG: t.remotePort})
 
 	for {
 		m, err := p.ReadMessage()
@@ -66,6 +69,9 @@ func (t *Trok) ControlConnHandler(conn net.Conn) {
 		case "CNCT":
 			go t.handleCMDCNCT(m)
 
+		case "ERR":
+			go t.handleCMDERR(m, conn)
+
 		default:
 			log.Info().Msgf("invalid command")
 		}
@@ -96,3 +102,9 @@ func (t *Trok) handleCMDCNCT(m *lib.Message) {
 	go io.Copy(upStream.conn, downStream.conn)
 	io.Copy(downStream.conn, upStream.conn)
 }
+
+func (t *Trok) handleCMDERR(m *lib.Message, conn net.Conn) {
+	log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG)
+	conn.Close()
+	os.Exit(1)
+}
diff --git a/internal/server/trok.go b/internal/server/trok.go
index b32bccd..cb197d3 100644
--- a/internal/server/trok.go
+++ b/internal/server/trok.go
@@ -79,9 +79,10 @@ func (t *Trok) handleCMDHELO(p *lib.ProtocolHandler, m *lib.Message) {
 	log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG)
 
 	var s TCPServer
-	err := s.Init(":", "Handler")
+	err := s.Init(fmt.Sprintf(":%s", m.ARG), "Handler")
 	if err != nil {
 		log.Error().Msgf("error handling HELO cmd: %v", err)
+		p.WriteMessage(&lib.Message{CMD: "ERR", ARG: "ADDRESS_UNAVAILABLE"})
 		return
 	}