feat: add cli option to pass custom remote port

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)
 	}

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)
+}

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
 	}