tux/trok
1
0
Fork 0
mirror of https://github.com/tuxdotrs/trok.git synced 2025-07-06 11:56:34 +05:30
Code Issues Packages Projects Releases Wiki Activity

feat: tunnel connection

Browse Source
This commit is contained in:
0xTux
2024-10-24 00:54:34 +05:30
parent 9dd5ec0ada
commit adba59f552
2 changed files with 73 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
internal/client/trok.go
View File

@ -5,6 +5,8 @@ Copyright © 2024 tux <0xtux@pm.me>
package client
import (
"fmt"
"io"
"net"
"github.com/0xtux/trok/internal/lib"
@ -49,10 +51,10 @@ func (t *Trok) ControlConnHandler(conn net.Conn) {
switch m.CMD {
case "EHLO":
t.hanldeCMDEHLO(p, m)
t.hanldeCMDEHLO(m)
case "CNCT":
t.handleCMDCNCT(p, m)
t.handleCMDCNCT(m)
default:
log.Info().Msgf("invalid command")
@ -60,12 +62,29 @@ func (t *Trok) ControlConnHandler(conn net.Conn) {
}
}
func (t *Trok) hanldeCMDEHLO(p *lib.ProtocolHandler, m *lib.Message) {
func (t *Trok) hanldeCMDEHLO(m *lib.Message) {
log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG)
p.WriteMessage(m)
}
func (t *Trok) handleCMDCNCT(p *lib.ProtocolHandler, m *lib.Message) {
func (t *Trok) handleCMDCNCT(m *lib.Message) {
log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG)
p.WriteMessage(m)
var upstream TCPClient
var downstream TCPClient
err := upstream.Init(3000, "UpStream")
if err != nil {
log.Error().Msgf("can't connect to upstream socket: %v", err)
return
}
err = downstream.Init(1421, "DownStream")
if err != nil {
log.Error().Msgf("can't connect to downstream socket: %v", err)
return
}
downstream.conn.Write([]byte(fmt.Sprintf("ACPT %s\n", m.ARG)))
go io.Copy(upstream.conn, downstream.conn)
io.Copy(downstream.conn, upstream.conn)
}

52
internal/server/trok.go
View File

@ -5,17 +5,24 @@ Copyright © 2024 tux <0xtux@pm.me>
package server
import (
"fmt"
"io"
"net"
"github.com/0xtux/trok/internal/lib"
gonanoid "github.com/matoous/go-nanoid/v2"
"github.com/rs/zerolog/log"
)
type Trok struct {
controlServer TCPServer
publicConns map[string]net.Conn
tunnels map[uint16]*lib.ProtocolHandler
}
func (t *Trok) Init(port uint16) error {
t.publicConns = make(map[string]net.Conn)
t.tunnels = make(map[uint16]*lib.ProtocolHandler)
err := t.controlServer.Init(port, "Controller")
return err
}
@ -50,7 +57,7 @@ func (t *Trok) ControlConnHandler(conn net.Conn) {
t.handleCMDHELO(p, m)
case "ACPT":
t.handleCMDACPT(p, m)
t.handleCMDACPT(conn, m)
default:
log.Info().Msgf("invalid command")
@ -60,10 +67,47 @@ 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)
p.WriteMessage(m)
var s TCPServer
err := s.Init(0, "Handler")
if err != nil {
log.Error().Msgf("error handling HELO cmd: %v", err)
return
}
port := s.Port()
go s.Start(t.PublicConnHandler)
t.tunnels[port] = p
p.WriteMessage(&lib.Message{CMD: "EHLO", ARG: fmt.Sprintf("%d", port)})
}
func (t *Trok) handleCMDACPT(p *lib.ProtocolHandler, m *lib.Message) {
func (t *Trok) handleCMDACPT(conn net.Conn, m *lib.Message) {
log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG)
p.WriteMessage(m)
pConn, ok := t.publicConns[m.ARG]
if !ok {
log.Error().Msgf("error finding public connection")
}
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)
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
}