mirror of
https://github.com/tuxdotrs/trok.git
synced 2025-07-06 11:56:34 +05:30
fix: connection drop
This commit is contained in:
9
go.mod
9
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
|
||||
)
|
||||
|
||||
9
go.sum
9
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=
|
||||
|
||||
@ -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")
|
||||
|
||||
@ -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()
|
||||
|
||||
if !exists || time.Since(pc.timestamp) > 10*time.Second {
|
||||
conn.Close()
|
||||
if exists {
|
||||
pc.conn.Close()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
go io.Copy(pConn, conn)
|
||||
io.Copy(conn, pConn)
|
||||
go io.Copy(pc.conn, conn)
|
||||
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)
|
||||
tnl, ok := t.tunnels[port]
|
||||
if !ok {
|
||||
log.Error().Msgf("error finding tunnel connection")
|
||||
return
|
||||
t.mutex.Lock()
|
||||
t.publicConns[id] = Conn{
|
||||
conn: conn,
|
||||
timestamp: time.Now(),
|
||||
}
|
||||
tnl.WriteMessage(&lib.Message{CMD: "CNCT", ARG: id})
|
||||
t.mutex.Unlock()
|
||||
|
||||
t.publicConns[id] = conn
|
||||
uidChan <- id
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user