From 88faeb7195354cf7c64c5bfd452b08b1e7ce7a8d Mon Sep 17 00:00:00 2001 From: 9seconds Date: Fri, 6 Jul 2018 10:21:04 +0300 Subject: [PATCH] Proxy is working but unstable --- client/client.go | 2 +- client/direct.go | 4 +-- client/middle.go | 6 ++-- mtproto/connection_options.go | 9 ++++-- mtproto/rpc/proxy_request.go | 2 +- mtproto/wrappers/abridged.go | 13 ++++---- mtproto/wrappers/intermediate.go | 7 ++--- mtproto/wrappers/proxy_request.go | 38 ++++++++++++------------ proxy/copy_pool.go | 18 ------------ proxy/server.go | 49 ++++++++++++++++++++++++++----- 10 files changed, 83 insertions(+), 65 deletions(-) delete mode 100644 proxy/copy_pool.go diff --git a/client/client.go b/client/client.go index cb6b291..32e23a9 100644 --- a/client/client.go +++ b/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) diff --git a/client/direct.go b/client/direct.go index 98b518d..3ce1eed 100644 --- a/client/direct.go +++ b/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 } diff --git a/client/middle.go b/client/middle.go index 2197cec..f8877fc 100644 --- a/client/middle.go +++ b/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 } diff --git a/mtproto/connection_options.go b/mtproto/connection_options.go index 585e0f7..0285488 100644 --- a/mtproto/connection_options.go +++ b/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 } diff --git a/mtproto/rpc/proxy_request.go b/mtproto/rpc/proxy_request.go index 994af28..e36a2ec 100644 --- a/mtproto/rpc/proxy_request.go +++ b/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 } diff --git a/mtproto/wrappers/abridged.go b/mtproto/wrappers/abridged.go index 9c22785..1850220 100644 --- a/mtproto/wrappers/abridged.go +++ b/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)) } diff --git a/mtproto/wrappers/intermediate.go b/mtproto/wrappers/intermediate.go index d80eabe..a22e56d 100644 --- a/mtproto/wrappers/intermediate.go +++ b/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) } diff --git a/mtproto/wrappers/proxy_request.go b/mtproto/wrappers/proxy_request.go index 7d0adbe..2641715 100644 --- a/mtproto/wrappers/proxy_request.go +++ b/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 } diff --git a/proxy/copy_pool.go b/proxy/copy_pool.go deleted file mode 100644 index 23477aa..0000000 --- a/proxy/copy_pool.go +++ /dev/null @@ -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 - }, - } -} diff --git a/proxy/server.go b/proxy/server.go index 314b3d9..9b43613 100644 --- a/proxy/server.go +++ b/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.