From adba59f55285a2ac23a662342565a1d877891319 Mon Sep 17 00:00:00 2001
From: 0xTux <0xtux@pm.me>
Date: Thu, 24 Oct 2024 00:54:34 +0530
Subject: [PATCH] feat: tunnel connection

---
 internal/client/trok.go | 31 +++++++++++++++++++-----
 internal/server/trok.go | 52 +++++++++++++++++++++++++++++++++++++----
 2 files changed, 73 insertions(+), 10 deletions(-)

diff --git a/internal/client/trok.go b/internal/client/trok.go
index 926768c..74aadc7 100644
--- a/internal/client/trok.go
+++ b/internal/client/trok.go
@@ -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)
 }
diff --git a/internal/server/trok.go b/internal/server/trok.go
index 293d47b..ea00ed7 100644
--- a/internal/server/trok.go
+++ b/internal/server/trok.go
@@ -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
 }