diff --git a/wrappers/mtproto_abridged.go b/wrappers/mtproto_abridged.go index 41f4e62..5ed40d6 100644 --- a/wrappers/mtproto_abridged.go +++ b/wrappers/mtproto_abridged.go @@ -28,6 +28,10 @@ type MTProtoAbridged struct { } func (m *MTProtoAbridged) Read() ([]byte, error) { + defer func() { + m.readCounter++ + }() + m.logger.Debugw("Read packet", "simple_ack", m.opts.ReadHacks.SimpleAck, "quick_ack", m.opts.ReadHacks.QuickAck, @@ -75,16 +79,18 @@ func (m *MTProtoAbridged) Read() ([]byte, error) { buf.Reset() buf.Grow(int(msgLength32)) - if _, err := io.CopyN(buf, m.conn, int64(msgLength32)); err != nil { return nil, errors.Annotate(err, "Cannot read message") } - m.readCounter++ return buf.Bytes(), nil } func (m *MTProtoAbridged) Write(p []byte) (int, error) { + defer func() { + m.writeCounter++ + }() + m.logger.Debugw("Write packet", "length", len(p), "simple_ack", m.opts.WriteHacks.SimpleAck, @@ -97,7 +103,6 @@ func (m *MTProtoAbridged) Write(p []byte) (int, error) { } if m.opts.WriteHacks.SimpleAck { - m.writeCounter++ return m.conn.Write(utils.ReverseBytes(p)) } @@ -105,8 +110,6 @@ func (m *MTProtoAbridged) Write(p []byte) (int, error) { switch { case packetLength < mtprotoAbridgedSmallPacketLength: newData := append([]byte{byte(packetLength)}, p...) - - m.writeCounter++ return m.conn.Write(newData) case packetLength < mtprotoAbridgedLargePacketLength: @@ -119,7 +122,6 @@ func (m *MTProtoAbridged) Write(p []byte) (int, error) { buf.Write(length24[:]) buf.Write(p) - m.writeCounter++ return m.conn.Write(buf.Bytes()) } diff --git a/wrappers/mtproto_intermediate.go b/wrappers/mtproto_intermediate.go index 5ed8e00..5c29c79 100644 --- a/wrappers/mtproto_intermediate.go +++ b/wrappers/mtproto_intermediate.go @@ -24,6 +24,10 @@ type MTProtoIntermediate struct { } func (m *MTProtoIntermediate) Read() ([]byte, error) { + defer func() { + m.readCounter++ + }() + m.logger.Debugw("Read packet", "simple_ack", m.opts.ReadHacks.SimpleAck, "quick_ack", m.opts.ReadHacks.QuickAck, @@ -59,18 +63,20 @@ func (m *MTProtoIntermediate) Read() ([]byte, error) { if length%4 != 0 { length -= length % 4 } - m.readCounter++ return buf.Bytes()[:length], nil } func (m *MTProtoIntermediate) Write(p []byte) (int, error) { + defer func() { + m.writeCounter++ + }() + m.logger.Debugw("Write packet", "simple_ack", m.opts.WriteHacks.SimpleAck, "quick_ack", m.opts.WriteHacks.QuickAck, "counter", m.writeCounter, ) - m.writeCounter++ if m.opts.ReadHacks.SimpleAck { return m.conn.Write(p) diff --git a/wrappers/mtproto_proxy.go b/wrappers/mtproto_proxy.go index d226cc8..c56e9a2 100644 --- a/wrappers/mtproto_proxy.go +++ b/wrappers/mtproto_proxy.go @@ -22,6 +22,10 @@ type MTProtoProxy struct { } func (m *MTProtoProxy) Read() ([]byte, error) { + defer func() { + m.readCounter++ + }() + m.logger.Debugw("Read packet", "counter", m.readCounter, "simple_ack", m.req.Options.WriteHacks.SimpleAck, @@ -32,9 +36,6 @@ func (m *MTProtoProxy) Read() ([]byte, error) { if err != nil { return nil, errors.Annotate(err, "Cannot read packet") } - defer func() { - m.readCounter++ - }() m.logger.Debugw("Read packet length", "counter", m.readCounter, @@ -95,13 +96,16 @@ func (m *MTProtoProxy) readCloseExt(data []byte) ([]byte, error) { } func (m *MTProtoProxy) Write(p []byte) (int, error) { + defer func() { + m.writeCounter++ + }() + m.logger.Debugw("Write packet", "length", len(p), "counter", m.writeCounter, "simple_ack", m.req.Options.ReadHacks.SimpleAck, "quick_ack", m.req.Options.ReadHacks.QuickAck, ) - m.writeCounter++ header := m.req.MakeHeader(p) if ce := m.logger.Desugar().Check(zap.DebugLevel, "RPC_PROXY_REQ header"); ce != nil {