From b5a8d8b804aaa6e28b4b3b3483d23b3253b3907c Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Mon, 30 Mar 2020 16:18:30 +0300
Subject: [PATCH] Add possibility to scale middleproxy buffers

---
 cli/proxy.go      |  2 +-
 config/config.go  | 45 +++++++++++++++++++++++++++++++++++++++++++++
 proxy/proxy.go    |  4 ++--
 telegram/base.go  |  2 +-
 utils/init_tcp.go |  8 +++-----
 5 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/cli/proxy.go b/cli/proxy.go
index fe15e39..a42a780 100644
--- a/cli/proxy.go
+++ b/cli/proxy.go
@@ -50,7 +50,7 @@ func Proxy() error { // nolint: funlen
 
 	zap.S().Debugw("Configuration", "config", config.Printable())
 
-	if len(config.C.AdTag) > 0 {
+	if config.C.MiddleProxyMode() {
 		zap.S().Infow("Use middle proxy connection to Telegram")
 
 		diff, err := ntp.Fetch()
diff --git a/config/config.go b/config/config.go
index 874a50b..aa13937 100644
--- a/config/config.go
+++ b/config/config.go
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
 	"net"
 
 	"github.com/alecthomas/units"
@@ -104,6 +105,50 @@ type Config struct {
 	AdTag  []byte `json:"adtag"`
 }
 
+func (c *Config) ClientReadBuffer() int {
+	return c.ReadBuffer
+}
+
+func (c *Config) ClientWriteBuffer() int {
+	return c.WriteBuffer
+}
+
+func (c *Config) MiddleProxyMode() bool {
+	return len(c.AdTag) > 0
+}
+
+func (c *Config) ProxyReadBuffer() int {
+	value := c.ReadBuffer
+
+	if c.MiddleProxyMode() {
+		value = c.adjustProxyValue(value)
+	}
+
+	return value
+}
+
+func (c *Config) ProxyWriteBuffer() int {
+	value := c.WriteBuffer
+
+	if c.MiddleProxyMode() {
+		value = c.adjustProxyValue(value)
+	}
+
+	return value
+}
+
+func (c *Config) adjustProxyValue(value int) int {
+	if c.MultiplexPerConnection == 0 {
+		return value
+	}
+
+	fvalue := float64(value)
+	newValue := fvalue * 1.5 * math.Log2(float64(c.MultiplexPerConnection))
+	newValue = math.Max(fvalue, math.Ceil(newValue))
+
+	return int(newValue)
+}
+
 type Opt struct {
 	Option OptionType
 	Value  interface{}
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 8a8202e..7bc462d 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -51,7 +51,7 @@ func (p *Proxy) accept(conn net.Conn) {
 	connID := conntypes.NewConnID()
 	logger := p.Logger.With("connection_id", connID)
 
-	if err := utils.InitTCP(conn); err != nil {
+	if err := utils.InitTCP(conn, config.C.ClientReadBuffer(), config.C.ClientWriteBuffer()); err != nil {
 		logger.Errorw("Cannot initialize client TCP connection", "error", err)
 		return
 	}
@@ -90,7 +90,7 @@ func (p *Proxy) accept(conn net.Conn) {
 
 	err = nil
 
-	if len(config.C.AdTag) > 0 {
+	if config.C.MiddleProxyMode() {
 		middleConnection(req)
 	} else {
 		err = directConnection(req)
diff --git a/telegram/base.go b/telegram/base.go
index bb27ff0..d454983 100644
--- a/telegram/base.go
+++ b/telegram/base.go
@@ -37,7 +37,7 @@ func (b *baseTelegram) dial(dc conntypes.DC,
 			continue
 		}
 
-		if err := utils.InitTCP(conn); err != nil {
+		if err := utils.InitTCP(conn, config.C.ProxyReadBuffer(), config.C.ProxyWriteBuffer()); err != nil {
 			b.logger.Infow("Cannot initialize TCP socket", "address", addr, "error", err)
 			continue
 		}
diff --git a/utils/init_tcp.go b/utils/init_tcp.go
index d9166c1..c8d8b65 100644
--- a/utils/init_tcp.go
+++ b/utils/init_tcp.go
@@ -4,24 +4,22 @@ import (
 	"fmt"
 	"net"
 	"time"
-
-	"github.com/9seconds/mtg/config"
 )
 
 const tcpKeepAlivePingPeriod = 2 * time.Second
 
-func InitTCP(conn net.Conn) error {
+func InitTCP(conn net.Conn, readBufferSize int, writeBufferSize int) error {
 	tcpConn := conn.(*net.TCPConn)
 
 	if err := tcpConn.SetNoDelay(true); err != nil {
 		return fmt.Errorf("cannot set TCP_NO_DELAY: %w", err)
 	}
 
-	if err := tcpConn.SetReadBuffer(config.C.ReadBuffer); err != nil {
+	if err := tcpConn.SetReadBuffer(readBufferSize); err != nil {
 		return fmt.Errorf("cannot set read buffer size: %w", err)
 	}
 
-	if err := tcpConn.SetWriteBuffer(config.C.WriteBuffer); err != nil {
+	if err := tcpConn.SetWriteBuffer(writeBufferSize); err != nil {
 		return fmt.Errorf("cannot set write buffer size: %w", err)
 	}