From 0ce325f908b0c9b7d93933f5c0ad52c528105145 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Fri, 27 Jul 2018 14:26:14 +0300 Subject: [PATCH] Add possibility to set read/write buffers separately --- client/direct.go | 10 +++------- config/config.go | 38 +++++++++++++++++++------------------- main.go | 16 +++++++++++++++- proxy/proxy.go | 10 ++++++---- telegram/dialer.go | 10 +++------- 5 files changed, 46 insertions(+), 38 deletions(-) diff --git a/client/direct.go b/client/direct.go index de8b66f..01730bd 100644 --- a/client/direct.go +++ b/client/direct.go @@ -12,11 +12,7 @@ import ( "github.com/9seconds/mtg/wrappers" ) -const ( - handshakeTimeout = 10 * time.Second - readBufferSize = 64 * 1024 - writeBufferSize = 64 * 1024 -) +const handshakeTimeout = 10 * time.Second // DirectInit initializes client connection for proxy which connects to // Telegram directly. @@ -25,10 +21,10 @@ func DirectInit(socket net.Conn, connID string, conf *config.Config) (wrappers.W if err := tcpSocket.SetNoDelay(false); err != nil { return nil, nil, errors.Annotate(err, "Cannot disable NO_DELAY to client socket") } - if err := tcpSocket.SetReadBuffer(readBufferSize); err != nil { + if err := tcpSocket.SetReadBuffer(conf.ReadBufferSize); err != nil { return nil, nil, errors.Annotate(err, "Cannot set read buffer size of client socket") } - if err := tcpSocket.SetWriteBuffer(writeBufferSize); err != nil { + if err := tcpSocket.SetWriteBuffer(conf.WriteBufferSize); err != nil { return nil, nil, errors.Annotate(err, "Cannot set write buffer size of client socket") } diff --git a/config/config.go b/config/config.go index f96b56f..6d0e641 100644 --- a/config/config.go +++ b/config/config.go @@ -11,18 +11,15 @@ import ( statsd "gopkg.in/alexcesaro/statsd.v2" ) -// Buffer sizes define internal socket buffer sizes. -const ( - BufferWriteSize = 32 * 1024 - BufferReadSize = 32 * 1024 -) - // Config represents common configuration of mtg. type Config struct { Debug bool Verbose bool SecureMode bool + ReadBufferSize int + WriteBufferSize int + BindPort uint16 PublicIPv4Port uint16 PublicIPv6Port uint16 @@ -114,6 +111,7 @@ func getAddr(host fmt.Stringer, port uint16) string { // fetches data from external sources. Parameters passed to this // function, should come from command line arguments. func NewConfig(debug, verbose bool, // nolint: gocyclo + writeBufferSize, readBufferSize uint32, bindIP, publicIPv4, publicIPv6, statsIP net.IP, bindPort, publicIPv4Port, publicIPv6Port, statsPort, statsdPort uint16, statsdIP, statsdNetwork, statsdPrefix, statsdTagsFormat string, @@ -157,19 +155,21 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo } conf := &Config{ - Debug: debug, - Verbose: verbose, - BindIP: bindIP, - BindPort: bindPort, - PublicIPv4: publicIPv4, - PublicIPv4Port: publicIPv4Port, - PublicIPv6: publicIPv6, - PublicIPv6Port: publicIPv6Port, - StatsIP: statsIP, - StatsPort: statsPort, - Secret: secret, - AdTag: adtag, - SecureMode: secureMode, + Debug: debug, + Verbose: verbose, + BindIP: bindIP, + BindPort: bindPort, + PublicIPv4: publicIPv4, + PublicIPv4Port: publicIPv4Port, + PublicIPv6: publicIPv6, + PublicIPv6Port: publicIPv6Port, + StatsIP: statsIP, + StatsPort: statsPort, + Secret: secret, + AdTag: adtag, + SecureMode: secureMode, + ReadBufferSize: int(readBufferSize), + WriteBufferSize: int(writeBufferSize), } if statsdIP != "" { diff --git a/main.go b/main.go index 72eee60..5b294cd 100644 --- a/main.go +++ b/main.go @@ -49,6 +49,19 @@ var ( Default("3128"). Uint16() + writeBufferSize = app.Flag("write-buffer", + "Write buffer size in bytes. You can think about it as a buffer from client to Telegram."). + Short('w'). + Envar("MTG_BUFFER_WRITE"). + Default("65536"). + Uint32() + readBufferSize = app.Flag("read-buffer", + "Read buffer size in bytes. You can think about it as a buffer from Telegram to client."). + Short('r'). + Envar("MTG_BUFFER_READ"). + Default("131072"). + Uint32() + publicIPv4 = app.Flag("public-ipv4", "Which IPv4 address is public."). Short('4'). @@ -117,7 +130,7 @@ var ( func init() { rand.Seed(time.Now().UTC().UnixNano()) app.Version(version) - + app.HelpFlag.Short('h') } func main() { // nolint: gocyclo @@ -129,6 +142,7 @@ func main() { // nolint: gocyclo } conf, err := config.NewConfig(*debug, *verbose, + *writeBufferSize, *readBufferSize, *bindIP, *publicIPv4, *publicIPv6, *statsIP, *bindPort, *publicIPv4Port, *publicIPv6Port, *statsPort, *statsdPort, *statsdIP, *statsdNetwork, *statsdPrefix, *statsdTagsFormat, diff --git a/proxy/proxy.go b/proxy/proxy.go index b8866b9..7a26bb7 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -83,8 +83,8 @@ func (p *Proxy) accept(conn net.Conn) { } else { clientStream := clientConn.(wrappers.StreamReadWriteCloser) serverStream := serverConn.(wrappers.StreamReadWriteCloser) - go p.directPipe(clientStream, serverStream, wait) - go p.directPipe(serverStream, clientStream, wait) + go p.directPipe(clientStream, serverStream, wait, p.conf.ReadBufferSize) + go p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize) } wait.Wait() @@ -130,14 +130,16 @@ func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser, } } -func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.WriteCloser, wait *sync.WaitGroup) { +func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.WriteCloser, + wait *sync.WaitGroup, bufferSize int) { defer func() { src.Close() // nolint: errcheck dst.Close() // nolint: errcheck wait.Done() }() - if _, err := io.Copy(dst, src); err != nil { + buffer := make([]byte, bufferSize) + if _, err := io.CopyBuffer(dst, src, buffer); err != nil { src.Logger().Warnw("Cannot pump sockets", "error", err) } } diff --git a/telegram/dialer.go b/telegram/dialer.go index 58b4ff5..9958403 100644 --- a/telegram/dialer.go +++ b/telegram/dialer.go @@ -10,11 +10,7 @@ import ( "github.com/9seconds/mtg/wrappers" ) -const ( - telegramDialTimeout = 10 * time.Second - readBufferSize = 64 * 1024 - writeBufferSize = 64 * 1024 -) +const telegramDialTimeout = 10 * time.Second type tgDialer struct { net.Dialer @@ -32,10 +28,10 @@ func (t *tgDialer) dial(addr string) (net.Conn, error) { if err = tcpSocket.SetNoDelay(true); err != nil { return nil, errors.Annotate(err, "Cannot set NO_DELAY to Telegram") } - if err = tcpSocket.SetReadBuffer(readBufferSize); err != nil { + if err = tcpSocket.SetReadBuffer(t.conf.WriteBufferSize); err != nil { return nil, errors.Annotate(err, "Cannot set read buffer size on telegram socket") } - if err = tcpSocket.SetWriteBuffer(writeBufferSize); err != nil { + if err = tcpSocket.SetWriteBuffer(t.conf.ReadBufferSize); err != nil { return nil, errors.Annotate(err, "Cannot set write buffer size on telegram socket") }