Merge remote-tracking branch 'origin/master' into stable

9seconds · Mar 31, 2020 · e5b9c84 · e5b9c84
2 parents 8ccf61c + ad7acee
commit e5b9c84
Show file tree

Hide file tree

Showing 7 changed files with 57 additions and 10 deletions.
diff --git 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
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
 	"net"
 
 	"github.com/alecthomas/units"
@@ -104,6 +105,52 @@ 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 * 2 * math.Log(float64(c.MultiplexPerConnection))
+	newValue = math.Ceil(newValue)
+	newValue = math.Max(fvalue, newValue)
+
+	return int(newValue)
+}
+
 type Opt struct {
 	Option OptionType
 	Value  interface{}

diff --git a/go.mod b/go.mod
@@ -14,7 +14,7 @@ require (
 	golang.org/x/crypto v0.0.0-20200323165209-0ec3e9974c59
 	golang.org/x/lint v0.0.0-20200302205851-738671d3881b // indirect
 	golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e // indirect
-	golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd
+	golang.org/x/sys v0.0.0-20200327173247-9dae0f8f5775
 	golang.org/x/tools v0.0.0-20200319210407-521f4a0cd458 // indirect
 	gopkg.in/alecthomas/kingpin.v2 v2.2.6
 	honnef.co/go/tools v0.0.1-2020.1.3 // indirect

diff --git a/go.sum b/go.sum
@@ -156,6 +156,8 @@ golang.org/x/sys v0.0.0-20200106162015-b016eb3dc98e/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20200122134326-e047566fdf82/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd h1:xhmwyvizuTgC2qz7ZlMluP20uW+C3Rm0FD/WLDX8884=
 golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20200327173247-9dae0f8f5775 h1:TC0v2RSO1u2kn1ZugjrFXkRZAEaqMN/RW+OTZkBzmLE=
+golang.org/x/sys v0.0.0-20200327173247-9dae0f8f5775/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/tools v0.0.0-20190621195816-6e04913cbbac/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc=

diff --git 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
@@ -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
@@ -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)
 	}