From d52159d98d57ca146124e63acde6b505c56197f7 Mon Sep 17 00:00:00 2001 From: tux Date: Thu, 27 Feb 2025 15:31:41 +0530 Subject: [PATCH] fix: connection drop --- go.mod | 9 +++-- go.sum | 9 ++++- internal/client/trok.go | 4 +- internal/server/trok.go | 87 +++++++++++++++++++++++++++-------------- 4 files changed, 72 insertions(+), 37 deletions(-) diff --git a/go.mod b/go.mod index ed1d1d4..65211b4 100644 --- a/go.mod +++ b/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 ) diff --git a/go.sum b/go.sum index a21aa41..b0564bf 100644 --- a/go.sum +++ b/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= diff --git a/internal/client/trok.go b/internal/client/trok.go index 74aadc7..da273b0 100644 --- a/internal/client/trok.go +++ b/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") diff --git a/internal/server/trok.go b/internal/server/trok.go index ea00ed7..f62c0fb 100644 --- a/internal/server/trok.go +++ b/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 - tunnels map[uint16]*lib.ProtocolHandler + publicConns map[string]Conn + mutex sync.Mutex } func (t *Trok) Init(port uint16) error { - t.publicConns = make(map[string]net.Conn) - t.tunnels = make(map[uint16]*lib.ProtocolHandler) + t.publicConns = make(map[string]Conn) 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] - if !ok { - log.Error().Msgf("error finding public connection") - } + t.mutex.Lock() + pc, exists := t.publicConns[m.ARG] + delete(t.publicConns, m.ARG) + t.mutex.Unlock() - go io.Copy(pConn, conn) - io.Copy(conn, pConn) -} - -func (t *Trok) PublicConnHandler(conn net.Conn) { - id, err := gonanoid.New(12) - if err != nil { - log.Error().Msgf("error generating uid: %v", err) + if !exists || time.Since(pc.timestamp) > 10*time.Second { + conn.Close() + if exists { + pc.conn.Close() + } return } - port := uint16(conn.LocalAddr().(*net.TCPAddr).Port) - tnl, ok := t.tunnels[port] - if !ok { - log.Error().Msgf("error finding tunnel connection") - return - } - tnl.WriteMessage(&lib.Message{CMD: "CNCT", ARG: id}) - - t.publicConns[id] = conn + go io.Copy(pc.conn, conn) + io.Copy(conn, pc.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 for public conn: %v", err) + return + } + + t.mutex.Lock() + t.publicConns[id] = Conn{ + conn: conn, + timestamp: time.Now(), + } + t.mutex.Unlock() + + uidChan <- id + } }