From cb31ac762c404f8edb6edd2e985666009f9952fe Mon Sep 17 00:00:00 2001
From: tux <t@tux.rs>
Date: Fri, 12 Dec 2025 02:18:26 +0530
Subject: [PATCH] fix(server): tcp tunnel missing initial bytes

---
 internal/lib/protocol.go |  4 ++++
 internal/server/trok.go  | 10 +++++-----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/internal/lib/protocol.go b/internal/lib/protocol.go
index 87acd76..43bde60 100644
--- a/internal/lib/protocol.go
+++ b/internal/lib/protocol.go
@@ -29,6 +29,10 @@ func InitProtocolHandler(conn net.Conn) *ProtocolHandler {
 	}
 }
 
+func (p *ProtocolHandler) Reader() *bufio.Reader {
+	return p.reader
+}
+
 func (p *ProtocolHandler) ReadMessage() (*Message, error) {
 	data, err := p.reader.ReadString('\n')
 	if err != nil {
diff --git a/internal/server/trok.go b/internal/server/trok.go
index d4fdfa1..b32bccd 100644
--- a/internal/server/trok.go
+++ b/internal/server/trok.go
@@ -67,7 +67,7 @@ func (t *Trok) ControlConnHandler(conn net.Conn) {
 			t.handleCMDHELO(p, m)
 
 		case "ACPT":
-			t.handleCMDACPT(conn, m)
+			t.handleCMDACPT(p, conn, m)
 
 		default:
 			log.Info().Msgf("invalid command")
@@ -101,7 +101,7 @@ func (t *Trok) handleCMDHELO(p *lib.ProtocolHandler, m *lib.Message) {
 	}
 }
 
-func (t *Trok) handleCMDACPT(conn net.Conn, m *lib.Message) {
+func (t *Trok) handleCMDACPT(p *lib.ProtocolHandler, conn net.Conn, m *lib.Message) {
 	log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG)
 
 	t.mutex.Lock()
@@ -117,7 +117,7 @@ func (t *Trok) handleCMDACPT(conn net.Conn, m *lib.Message) {
 		return
 	}
 
-	t.Bind(pc.conn, conn)
+	t.Bind(pc.conn, conn, p.Reader())
 }
 
 func (t *Trok) PublicConnHandler(ln net.Listener, uidChan chan<- string) {
@@ -144,9 +144,9 @@ func (t *Trok) PublicConnHandler(ln net.Listener, uidChan chan<- string) {
 	}
 }
 
-func (t *Trok) Bind(src, dst net.Conn) {
+func (t *Trok) Bind(src net.Conn, dst net.Conn, dstReader io.Reader) {
 	defer src.Close()
 	defer dst.Close()
-	go io.Copy(src, dst)
+	go io.Copy(src, dstReader)
 	io.Copy(dst, src)
 }