fix: connection drop

go.mod

@@ -2,15 +2,16 @@ module github.com/0xtux/trok
 
 go 1.22.7
 
-require github.com/google/uuid v1.6.0
+require (
+	github.com/matoous/go-nanoid/v2 v2.1.0
+	github.com/rs/zerolog v1.33.0
+	github.com/spf13/cobra v1.8.1
+)
 
 require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
-	github.com/matoous/go-nanoid/v2 v2.1.0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
-	github.com/rs/zerolog v1.33.0 // indirect
-	github.com/spf13/cobra v1.8.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	golang.org/x/sys v0.12.0 // indirect
 )

go.sum

@@ -1,8 +1,8 @@
 github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
 github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
-github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
-github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/matoous/go-nanoid/v2 v2.1.0 h1:P64+dmq21hhWdtvZfEAofnvJULaRR1Yib0+PnU669bE=
@@ -13,6 +13,8 @@ github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/
 github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
 github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
 github.com/rs/zerolog v1.33.0 h1:1cU2KZkvPxNyfgEmhHAz/1A9Bz+llsdYzklWFzgp0r8=
 github.com/rs/zerolog v1.33.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
@@ -21,9 +23,12 @@ github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM=
 github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3kD9Y=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o=
 golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

internal/client/trok.go

@@ -51,10 +51,10 @@ func (t *Trok) ControlConnHandler(conn net.Conn) {
 		switch m.CMD {
 
 		case "EHLO":
-			t.hanldeCMDEHLO(m)
+			go t.hanldeCMDEHLO(m)
 
 		case "CNCT":
-			t.handleCMDCNCT(m)
+			go t.handleCMDCNCT(m)
 
 		default:
 			log.Info().Msgf("invalid command")

internal/server/trok.go

@@ -8,21 +8,27 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"sync"
+	"time"
 
 	"github.com/0xtux/trok/internal/lib"
 	gonanoid "github.com/matoous/go-nanoid/v2"
 	"github.com/rs/zerolog/log"
 )
 
+type Conn struct {
+	conn      net.Conn
+	timestamp time.Time
+}
+
 type Trok struct {
 	controlServer TCPServer
-	publicConns   map[string]net.Conn
+	publicConns   map[string]Conn
-	tunnels       map[uint16]*lib.ProtocolHandler
+	mutex         sync.Mutex
 }
 
 func (t *Trok) Init(port uint16) error {
-	t.publicConns = make(map[string]net.Conn)
+	t.publicConns = make(map[string]Conn)
-	t.tunnels = make(map[uint16]*lib.ProtocolHandler)
 	err := t.controlServer.Init(port, "Controller")
 	return err
 }
@@ -67,8 +73,8 @@ func (t *Trok) ControlConnHandler(conn net.Conn) {
 
 func (t *Trok) handleCMDHELO(p *lib.ProtocolHandler, m *lib.Message) {
 	log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG)
-	var s TCPServer
 
+	var s TCPServer
 	err := s.Init(0, "Handler")
 	if err != nil {
 		log.Error().Msgf("error handling HELO cmd: %v", err)
@@ -76,38 +82,61 @@ func (t *Trok) handleCMDHELO(p *lib.ProtocolHandler, m *lib.Message) {
 	}
 
 	port := s.Port()
-	go s.Start(t.PublicConnHandler)
-	t.tunnels[port] = p
-
 	p.WriteMessage(&lib.Message{CMD: "EHLO", ARG: fmt.Sprintf("%d", port)})
+
+	uidChan := make(chan string)
+	defer close(uidChan)
+
+	go t.PublicConnHandler(s.listener, uidChan)
+
+	for id := range uidChan {
+		err := p.WriteMessage(&lib.Message{CMD: "CNCT", ARG: id})
+		if err != nil {
+			break
+		}
+	}
 }
 
 func (t *Trok) handleCMDACPT(conn net.Conn, m *lib.Message) {
 	log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG)
 
-	pConn, ok := t.publicConns[m.ARG]
+	t.mutex.Lock()
-	if !ok {
+	pc, exists := t.publicConns[m.ARG]
-		log.Error().Msgf("error finding public connection")
+	delete(t.publicConns, m.ARG)
+	t.mutex.Unlock()
+
+	if !exists || time.Since(pc.timestamp) > 10*time.Second {
+		conn.Close()
+		if exists {
+			pc.conn.Close()
+		}
+		return
 	}
 
-	go io.Copy(pConn, conn)
+	go io.Copy(pc.conn, conn)
-	io.Copy(conn, pConn)
+	io.Copy(conn, pc.conn)
 }
 
-func (t *Trok) PublicConnHandler(conn net.Conn) {
+func (t *Trok) PublicConnHandler(ln net.Listener, uidChan chan<- string) {
+	for {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+
 		id, err := gonanoid.New(12)
 		if err != nil {
-		log.Error().Msgf("error generating uid: %v", err)
+			log.Error().Msgf("error generating uid for public conn: %v", err)
 			return
 		}
 
-	port := uint16(conn.LocalAddr().(*net.TCPAddr).Port)
+		t.mutex.Lock()
-	tnl, ok := t.tunnels[port]
+		t.publicConns[id] = Conn{
-	if !ok {
+			conn:      conn,
-		log.Error().Msgf("error finding tunnel connection")
+			timestamp: time.Now(),
-		return
 		}
-	tnl.WriteMessage(&lib.Message{CMD: "CNCT", ARG: id})
+		t.mutex.Unlock()
 
-	t.publicConns[id] = conn
+		uidChan <- id
+	}
 }