Proxy is working but unstable

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
 }