From f0efa4697ee92466a871eaa985fb4b6c98fa8e8a Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Fri, 26 Mar 2021 12:15:26 +0300
Subject: [PATCH] Faketls works now

---
 mtglib/internal/faketls/conn.go        | 70 ++++++++++++++++++++++++++
 mtglib/internal/faketls/record/init.go |  2 +
 mtglib/internal/relay/conn.go          |  6 +--
 mtglib/proxy.go                        | 13 +++--
 4 files changed, 82 insertions(+), 9 deletions(-)
 create mode 100644 mtglib/internal/faketls/conn.go

diff --git a/mtglib/internal/faketls/conn.go b/mtglib/internal/faketls/conn.go
new file mode 100644
index 0000000..9962cbb
--- /dev/null
+++ b/mtglib/internal/faketls/conn.go
@@ -0,0 +1,70 @@
+package faketls
+
+import (
+	"bytes"
+	"fmt"
+	"net"
+
+	"github.com/9seconds/mtg/v2/mtglib/internal/faketls/record"
+)
+
+type Conn struct {
+	net.Conn
+
+	readBuffer bytes.Buffer
+}
+
+func (c *Conn) Read(p []byte) (int, error) {
+	if n, _ := c.readBuffer.Read(p); n > 0 {
+		return n, nil
+	}
+
+	rec := record.AcquireRecord()
+	defer record.ReleaseRecord(rec)
+
+	for {
+		if err := rec.Read(c.Conn); err != nil {
+			return 0, err // nolint: wrapcheck
+		}
+
+		switch rec.Type { // nolint: exhaustive
+		case record.TypeChangeCipherSpec:
+		case record.TypeApplicationData:
+			rec.Payload.WriteTo(&c.readBuffer)
+
+			n, err := c.readBuffer.Read(p)
+
+			return n, err
+		default:
+			return 0, fmt.Errorf("unsupported record type %v", rec.Type)
+		}
+	}
+}
+
+func (c *Conn) Write(p []byte) (int, error) {
+	rec := record.AcquireRecord()
+	defer record.ReleaseRecord(rec)
+
+	rec.Type = record.TypeApplicationData
+	rec.Version = record.Version12
+	written := 0
+
+	for len(p) > 0 {
+		chunkSize := record.TLSMaxRecordSize
+		if chunkSize > len(p) {
+			chunkSize = len(p)
+		}
+
+		rec.Payload.Reset()
+		rec.Payload.Write(p[:chunkSize])
+
+		if err := rec.Dump(c.Conn); err != nil {
+			return 0, err
+		}
+
+		written += chunkSize
+		p = p[chunkSize:]
+	}
+
+	return written, nil
+}
diff --git a/mtglib/internal/faketls/record/init.go b/mtglib/internal/faketls/record/init.go
index 1e543aa..764c29e 100644
--- a/mtglib/internal/faketls/record/init.go
+++ b/mtglib/internal/faketls/record/init.go
@@ -2,6 +2,8 @@ package record
 
 import "fmt"
 
+const TLSMaxRecordSize = 65535 // max uint16
+
 type Type uint8
 
 const (
diff --git a/mtglib/internal/relay/conn.go b/mtglib/internal/relay/conn.go
index 211a32a..8f4e89d 100644
--- a/mtglib/internal/relay/conn.go
+++ b/mtglib/internal/relay/conn.go
@@ -9,11 +9,10 @@ type conn struct {
 }
 
 func (c conn) Read(p []byte) (int, error) {
-	ctx := c.relay.ctx
 	n, err := c.ReadWriteCloser.Read(p)
 
 	select {
-	case <-ctx.Done():
+	case <-c.relay.ctx.Done():
 	case c.relay.tickChannel <- struct{}{}:
 	}
 
@@ -21,11 +20,10 @@ func (c conn) Read(p []byte) (int, error) {
 }
 
 func (c conn) Write(p []byte) (int, error) {
-	ctx := c.relay.ctx
 	n, err := c.ReadWriteCloser.Write(p)
 
 	select {
-	case <-ctx.Done():
+	case <-c.relay.ctx.Done():
 	case c.relay.tickChannel <- struct{}{}:
 	}
 
diff --git a/mtglib/proxy.go b/mtglib/proxy.go
index 7bf9f7c..6766c00 100644
--- a/mtglib/proxy.go
+++ b/mtglib/proxy.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"net"
 	"sync"
 	"time"
@@ -59,7 +58,7 @@ func (p *Proxy) ServeConn(conn net.Conn) {
 		ctx.logger.Info("Stream has been finished")
 	}()
 
-	if err := p.doFakeTLSHandshake(ctx, ctx.clientConn); err != nil {
+	if err := p.doFakeTLSHandshake(ctx); err != nil {
 		p.logger.InfoError("faketls handshake is failed", err)
 
 		return
@@ -122,11 +121,11 @@ func (p *Proxy) Shutdown() {
 	p.workerPool.Release()
 }
 
-func (p *Proxy) doFakeTLSHandshake(ctx *streamContext, conn io.ReadWriter) error {
+func (p *Proxy) doFakeTLSHandshake(ctx *streamContext) error {
 	rec := record.AcquireRecord()
 	defer record.ReleaseRecord(rec)
 
-	if err := rec.Read(conn); err != nil {
+	if err := rec.Read(ctx.clientConn); err != nil {
 		return fmt.Errorf("cannot read client hello: %w", err)
 	}
 
@@ -145,10 +144,14 @@ func (p *Proxy) doFakeTLSHandshake(ctx *streamContext, conn io.ReadWriter) error
 		return fmt.Errorf("anti replay attack from %s", ctx.ClientIP().String())
 	}
 
-	if err := faketls.SendWelcomePacket(conn, p.secret.Key[:], hello); err != nil {
+	if err := faketls.SendWelcomePacket(ctx.clientConn, p.secret.Key[:], hello); err != nil {
 		return fmt.Errorf("cannot send a welcome packet: %w", err)
 	}
 
+	ctx.clientConn = &faketls.Conn{
+		Conn: ctx.clientConn,
+	}
+
 	return nil
 }