Proxy is working but unstable

client/client.go

@@ -9,4 +9,4 @@ import (
 )
 
 // Init has to initialize client connection based on given config.
-type Init func(net.Conn, *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error)
+type Init func(net.Conn, *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error)

client/direct.go

@@ -15,7 +15,7 @@ import (
 const handshakeTimeout = 10 * time.Second
 
 // DirectInit initializes client to access Telegram bypassing middleproxies.
-func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error) {
+func DirectInit(conn net.Conn, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
 	if err := config.SetSocketOptions(conn); err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot set socket options")
 	}
@@ -37,5 +37,5 @@ func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wr
 	socket := wrappers.NewTimeoutRWC(conn, conf.PublicIPv4, conf.PublicIPv6)
 	socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
 
-	return connOpts, socket, nil
+	return socket, connOpts, nil
 }

client/middle.go

@@ -9,8 +9,8 @@ import (
 	"github.com/9seconds/mtg/wrappers"
 )
 
-func MiddleInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error) {
-	opts, newConn, err := DirectInit(conn, conf)
+func MiddleInit(conn net.Conn, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
+	newConn, opts, err := DirectInit(conn, conf)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -26,5 +26,5 @@ func MiddleInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wr
 		opts.ConnectionProto = mtproto.ConnectionProtocolIPv6
 	}
 
-	return opts, newConn, nil
+	return newConn, opts, nil
 }

mtproto/connection_options.go

@@ -13,14 +13,19 @@ type ConnectionType uint8
 
 type ConnectionProtocol uint8
 
+type Hacks struct {
+	SimpleAck bool
+	QuickAck  bool
+}
+
 // ConnectionOpts presents an options, metadata on connection requested
 // by the user on handshake.
 type ConnectionOpts struct {
 	DC              int16
 	ConnectionType  ConnectionType
 	ConnectionProto ConnectionProtocol
-	QuickAck        bool
-	SimpleAck       bool
+	ReadHacks       Hacks
+	WriteHacks      Hacks
 	ClientAddr      *net.TCPAddr
 }
 

mtproto/rpc/proxy_request.go

@@ -24,7 +24,7 @@ func (r *ProxyRequest) Bytes(message []byte) []byte {
 	buf := &bytes.Buffer{}
 
 	flags := r.Flags
-	if r.Options.QuickAck {
+	if r.Options.ReadHacks.QuickAck {
 		flags |= proxyRequestFlagsQuickAck
 	}
 

mtproto/wrappers/abridged.go

@@ -28,9 +28,6 @@ type AbridgedReadWriteCloserWithAddr struct {
 
 func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 	return a.BufferedRead(p, func() error {
-		a.opts.QuickAck = false
-		a.opts.SimpleAck = false
-
 		buf := &bytes.Buffer{}
 		buf.Grow(3)
 
@@ -41,7 +38,7 @@ func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 		buf.Reset()
 
 		if msgLength >= abridgedQuickAckLength {
-			a.opts.QuickAck = true
+			a.opts.ReadHacks.QuickAck = true
 			msgLength -= 0x80
 		}
 
@@ -56,9 +53,13 @@ func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 		}
 		msgLength32 *= 4
 
-		if _, err := io.CopyN(a.Buffer, a.conn, int64(msgLength32)); err != nil {
+		buf.Reset()
+		buf.Grow(int(msgLength32))
+
+		if _, err := io.CopyN(buf, a.conn, int64(msgLength32)); err != nil {
 			return errors.Annotate(err, "Cannot read message")
 		}
+		a.Buffer.Write(buf.Bytes())
 
 		return nil
 	})
@@ -68,7 +69,7 @@ func (a *AbridgedReadWriteCloserWithAddr) Write(p []byte) (int, error) {
 	if len(p)%4 != 0 {
 		return 0, errors.Errorf("Incorrect packet length %d", len(p))
 	}
-	if a.opts.SimpleAck {
+	if a.opts.WriteHacks.SimpleAck {
 		return a.conn.Write(reverseBytes(p))
 	}
 

mtproto/wrappers/intermediate.go

@@ -23,9 +23,6 @@ type IntermediateReadWriteCloserWithAddr struct {
 
 func (i *IntermediateReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 	return i.BufferedRead(p, func() error {
-		i.opts.QuickAck = false
-		i.opts.SimpleAck = false
-
 		buf := &bytes.Buffer{}
 		buf.Grow(4)
 
@@ -37,7 +34,7 @@ func (i *IntermediateReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 		buf.Grow(int(length))
 
 		if length > intermediateQuickAckLength {
-			i.opts.QuickAck = true
+			i.opts.ReadHacks.QuickAck = true
 			length -= intermediateQuickAckLength
 		}
 
@@ -58,7 +55,7 @@ func (i *IntermediateReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 }
 
 func (i *IntermediateReadWriteCloserWithAddr) Write(p []byte) (int, error) {
-	if i.opts.SimpleAck {
+	if i.opts.WriteHacks.SimpleAck {
 		return i.conn.Write(p)
 	}
 

mtproto/wrappers/proxy_request.go

@@ -22,23 +22,21 @@ type ProxyRequestReadWriteCloserWithAddr struct {
 
 func (p *ProxyRequestReadWriteCloserWithAddr) Read(buf []byte) (int, error) {
 	return p.BufferedRead(buf, func() error {
-		ansBuf := &bytes.Buffer{}
-		ansBuf.Grow(4)
-
-		if _, err := io.CopyN(ansBuf, p.conn, 4); err != nil {
+		ans := make([]byte, 4)
+		if _, err := io.ReadFull(p.conn, ans); err != nil {
 			return errors.Annotate(err, "Cannot read RPC tag")
 		}
 
 		switch {
-		case bytes.Equal(ansBuf.Bytes(), rpc.TagCloseExt):
-			return p.readCloseExt()
-		case bytes.Equal(ansBuf.Bytes(), rpc.TagProxyAns):
+		case bytes.Equal(ans, rpc.TagProxyAns):
 			return p.readProxyAns(buf)
-		case bytes.Equal(ansBuf.Bytes(), rpc.TagSimpleAck):
+		case bytes.Equal(ans, rpc.TagSimpleAck):
 			return p.readSimpleAck()
+		case bytes.Equal(ans, rpc.TagCloseExt):
+			return p.readCloseExt()
 		}
 
-		return errors.Errorf("Unknown RPC answer %s", ansBuf.Bytes())
+		return errors.Errorf("Unknown RPC answer %v", ans)
 	})
 }
 
@@ -46,21 +44,21 @@ func (p *ProxyRequestReadWriteCloserWithAddr) readCloseExt() error {
 	return errors.New("Connection has been closed remotely")
 }
 
-func (p *ProxyRequestReadWriteCloserWithAddr) readProxyAns(buf []byte) error {
-	if _, err := io.CopyN(ioutil.Discard, p.conn, 8+4); err != nil {
+func (p *ProxyRequestReadWriteCloserWithAddr) readProxyAns(buf []byte) (err error) {
+	if _, err = io.CopyN(ioutil.Discard, p.conn, 8+4); err != nil {
 		return errors.Annotate(err, "Cannot skip flags and connid")
 	}
 
-	for {
-		n, err := p.conn.Read(buf)
+	n := len(buf)
+	preBuffer := &bytes.Buffer{}
+	for n == len(buf) {
+		n, err = p.conn.Read(buf)
 		if err != nil {
 			return errors.Annotate(err, "Cannot read proxy answer")
 		}
-		if n == 0 {
-			break
-		}
-		p.Buffer.Write(buf[:n])
+		preBuffer.Write(buf[:n])
 	}
+	p.Buffer.Write(preBuffer.Bytes())
 
 	return nil
 }
@@ -69,10 +67,12 @@ func (p *ProxyRequestReadWriteCloserWithAddr) readSimpleAck() error {
 	if _, err := io.CopyN(ioutil.Discard, p.conn, 8); err != nil {
 		return errors.Annotate(err, "Cannot skip connid")
 	}
-	if _, err := io.CopyN(p.Buffer, p.conn, 4); err != nil {
+
+	ackData := make([]byte, 4)
+	if _, err := io.ReadFull(p.conn, ackData); err != nil {
 		return errors.Annotate(err, "Cannot read simple ack")
 	}
-	p.req.Options.SimpleAck = true
+	p.Buffer.Write(ackData)
 
 	return nil
 }

proxy/copy_pool.go

@@ -1,18 +0,0 @@
-package proxy
-
-import (
-	"sync"
-
-	"github.com/9seconds/mtg/config"
-)
-
-var copyPool sync.Pool
-
-func init() {
-	copyPool = sync.Pool{
-		New: func() interface{} {
-			data := make([]byte, config.BufferSizeCopy)
-			return &data
-		},
-	}
-}

proxy/server.go

@@ -85,8 +85,32 @@ func (s *Server) accept(conn net.Conn) {
 	wait := &sync.WaitGroup{}
 	wait.Add(2)
 
-	go s.pipe(clientConn, tgConn, wait)
-	go s.pipe(tgConn, clientConn, wait)
+	go func() {
+		defer wait.Done()
+
+		for {
+			connOpts.ReadHacks.QuickAck = false
+			connOpts.ReadHacks.SimpleAck = false
+			if err := s.pump(clientConn, tgConn, socketID, "client"); err != nil {
+				s.logger.Infow("Client stream is aborted",
+					"socketid", socketID, "error", err)
+				return
+			}
+		}
+	}()
+	go func() {
+		defer wait.Done()
+
+		for {
+			connOpts.WriteHacks.QuickAck = false
+			connOpts.WriteHacks.SimpleAck = false
+			if err := s.pump(tgConn, clientConn, socketID, "telegram"); err != nil {
+				s.logger.Infow("Telegram stream is aborted",
+					"socketid", socketID, "error", err)
+				return
+			}
+		}
+	}()
 
 	<-ctx.Done()
 	wait.Wait()
@@ -98,7 +122,7 @@ func (s *Server) accept(conn net.Conn) {
 }
 
 func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc, conn net.Conn, socketID string) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
-	connOpts, socket, err := s.clientInit(conn, s.conf)
+	socket, connOpts, err := s.clientInit(conn, s.conf)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot init client connection")
 	}
@@ -128,13 +152,22 @@ func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFun
 	return conn, nil
 }
 
-func (s *Server) pipe(dst io.Writer, src io.Reader, wait *sync.WaitGroup) {
-	defer wait.Done()
+func (s *Server) pump(src io.Reader, dst io.Writer, socketID, name string) (err error) {
+	copyBuf := make([]byte, 1024*1024*2)
 
-	buf := copyPool.Get().(*[]byte)
-	defer copyPool.Put(buf)
+	n := config.BufferSizeCopy
+	for n == config.BufferSizeCopy {
+		n, err = src.Read(copyBuf)
+		if err != nil {
+			break
+		}
+		_, err = dst.Write(copyBuf[:n])
+		if err != nil {
+			break
+		}
+	}
 
-	io.CopyBuffer(dst, src, *buf) // nolint: errcheck
+	return
 }
 
 // NewServer creates new instance of MTPROTO proxy.