ScuroNeko/mtg
1
0
Fork 0
mirror of https://github.com/ScuroNeko/mtg.git synced 2026-08-31 12:54:01 +03:00
Code Issues Packages Projects Releases Wiki Activity

Add possibility to set read/write buffers separately

Browse Source
This commit is contained in:
9seconds
2018-07-27 14:26:14 +03:00
parent b86ad07ba9
commit 0ce325f908
5 changed files with 46 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/direct.go
+3 -7
View File
@@ -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
+19 -19
View File
@@ -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 != "" {
main.go
+15 -1
View File
@@ -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
+6 -4
View File
@@ -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
+3 -7
View File
@@ -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")
}