Add possibility to set read/write buffers separately

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")
 	}
 

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,
@@ -170,6 +168,8 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 		Secret:          secret,
 		AdTag:           adtag,
 		SecureMode:      secureMode,
+		ReadBufferSize:  int(readBufferSize),
+		WriteBufferSize: int(writeBufferSize),
 	}
 
 	if statsdIP != "" {

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,

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)
 	}
 }

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")
 	}