From 9317445e3716213a85839a75136de8ba1d0d23c6 Mon Sep 17 00:00:00 2001
From: fregie <xiaohao950830@live.com>
Date: Wed, 16 Jun 2021 16:58:48 +0800
Subject: [PATCH] feature: support specify buffer size in config

---
 proxy/config.go | 10 ++++++----
 proxy/proxy.go  | 31 +++++++++++++++----------------
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/proxy/config.go b/proxy/config.go
index 692f5c697..44bd2c8ef 100644
--- a/proxy/config.go
+++ b/proxy/config.go
@@ -3,15 +3,17 @@ package proxy
 import "github.com/p4gefau1t/trojan-go/config"
 
 type Config struct {
-	RunType  string `json:"run_type" yaml:"run-type"`
-	LogLevel int    `json:"log_level" yaml:"log-level"`
-	LogFile  string `json:"log_file" yaml:"log-file"`
+	RunType         string `json:"run_type" yaml:"run-type"`
+	LogLevel        int    `json:"log_level" yaml:"log-level"`
+	LogFile         string `json:"log_file" yaml:"log-file"`
+	RelayBufferSize int    `json:"relay_buffer_size" yaml:"relay_buffer_size"`
 }
 
 func init() {
 	config.RegisterConfigCreator(Name, func() interface{} {
 		return &Config{
-			LogLevel: 1,
+			LogLevel:        1,
+			RelayBufferSize: 4 * 1024,
 		}
 	})
 }
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 7c333c38e..9c93ed07f 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"io"
 	"math/rand"
-	"net"
 	"os"
 	"strings"
 
@@ -16,16 +15,13 @@ import (
 
 const Name = "PROXY"
 
-const (
-	MaxPacketSize = 1024 * 8
-)
-
 // Proxy relay connections and packets
 type Proxy struct {
-	sources []tunnel.Server
-	sink    tunnel.Client
-	ctx     context.Context
-	cancel  context.CancelFunc
+	sources         []tunnel.Server
+	sink            tunnel.Client
+	ctx             context.Context
+	cancel          context.CancelFunc
+	relayBufferSize int
 }
 
 func (p *Proxy) Run() error {
@@ -68,8 +64,9 @@ func (p *Proxy) relayConnLoop() {
 					}
 					defer outbound.Close()
 					errChan := make(chan error, 2)
-					copyConn := func(a, b net.Conn) {
-						_, err := io.Copy(a, b)
+					copyConn := func(dst io.Writer, src io.Reader) {
+						buffer := make([]byte, p.relayBufferSize)
+						_, err := io.CopyBuffer(dst, src, buffer)
 						errChan <- err
 					}
 					go copyConn(inbound, outbound)
@@ -116,7 +113,7 @@ func (p *Proxy) relayPacketLoop() {
 					errChan := make(chan error, 2)
 					copyPacket := func(a, b tunnel.PacketConn) {
 						for {
-							buf := make([]byte, MaxPacketSize)
+							buf := make([]byte, p.relayBufferSize)
 							n, metadata, err := a.ReadWithMetadata(buf)
 							if err != nil {
 								errChan <- err
@@ -151,11 +148,13 @@ func (p *Proxy) relayPacketLoop() {
 }
 
 func NewProxy(ctx context.Context, cancel context.CancelFunc, sources []tunnel.Server, sink tunnel.Client) *Proxy {
+	cfg := config.FromContext(ctx, Name).(*Config)
 	return &Proxy{
-		sources: sources,
-		sink:    sink,
-		ctx:     ctx,
-		cancel:  cancel,
+		sources:         sources,
+		sink:            sink,
+		ctx:             ctx,
+		cancel:          cancel,
+		relayBufferSize: cfg.RelayBufferSize,
 	}
 }