9seconds 7 лет назад
Родитель
Сommit
a99c5b2442
3 измененных файлов: 24 добавлений и 12 удалений
  1. 8
    6
      wrappers/mtproto_abridged.go
  2. 8
    2
      wrappers/mtproto_intermediate.go
  3. 8
    4
      wrappers/mtproto_proxy.go

+ 8
- 6
wrappers/mtproto_abridged.go Просмотреть файл

28
 }
28
 }
29
 
29
 
30
 func (m *MTProtoAbridged) Read() ([]byte, error) {
30
 func (m *MTProtoAbridged) Read() ([]byte, error) {
31
+	defer func() {
32
+		m.readCounter++
33
+	}()
34
+
31
 	m.logger.Debugw("Read packet",
35
 	m.logger.Debugw("Read packet",
32
 		"simple_ack", m.opts.ReadHacks.SimpleAck,
36
 		"simple_ack", m.opts.ReadHacks.SimpleAck,
33
 		"quick_ack", m.opts.ReadHacks.QuickAck,
37
 		"quick_ack", m.opts.ReadHacks.QuickAck,
75
 
79
 
76
 	buf.Reset()
80
 	buf.Reset()
77
 	buf.Grow(int(msgLength32))
81
 	buf.Grow(int(msgLength32))
78
-
79
 	if _, err := io.CopyN(buf, m.conn, int64(msgLength32)); err != nil {
82
 	if _, err := io.CopyN(buf, m.conn, int64(msgLength32)); err != nil {
80
 		return nil, errors.Annotate(err, "Cannot read message")
83
 		return nil, errors.Annotate(err, "Cannot read message")
81
 	}
84
 	}
82
-	m.readCounter++
83
 
85
 
84
 	return buf.Bytes(), nil
86
 	return buf.Bytes(), nil
85
 }
87
 }
86
 
88
 
87
 func (m *MTProtoAbridged) Write(p []byte) (int, error) {
89
 func (m *MTProtoAbridged) Write(p []byte) (int, error) {
90
+	defer func() {
91
+		m.writeCounter++
92
+	}()
93
+
88
 	m.logger.Debugw("Write packet",
94
 	m.logger.Debugw("Write packet",
89
 		"length", len(p),
95
 		"length", len(p),
90
 		"simple_ack", m.opts.WriteHacks.SimpleAck,
96
 		"simple_ack", m.opts.WriteHacks.SimpleAck,
97
 	}
103
 	}
98
 
104
 
99
 	if m.opts.WriteHacks.SimpleAck {
105
 	if m.opts.WriteHacks.SimpleAck {
100
-		m.writeCounter++
101
 		return m.conn.Write(utils.ReverseBytes(p))
106
 		return m.conn.Write(utils.ReverseBytes(p))
102
 	}
107
 	}
103
 
108
 
105
 	switch {
110
 	switch {
106
 	case packetLength < mtprotoAbridgedSmallPacketLength:
111
 	case packetLength < mtprotoAbridgedSmallPacketLength:
107
 		newData := append([]byte{byte(packetLength)}, p...)
112
 		newData := append([]byte{byte(packetLength)}, p...)
108
-
109
-		m.writeCounter++
110
 		return m.conn.Write(newData)
113
 		return m.conn.Write(newData)
111
 
114
 
112
 	case packetLength < mtprotoAbridgedLargePacketLength:
115
 	case packetLength < mtprotoAbridgedLargePacketLength:
119
 		buf.Write(length24[:])
122
 		buf.Write(length24[:])
120
 		buf.Write(p)
123
 		buf.Write(p)
121
 
124
 
122
-		m.writeCounter++
123
 		return m.conn.Write(buf.Bytes())
125
 		return m.conn.Write(buf.Bytes())
124
 	}
126
 	}
125
 
127
 

+ 8
- 2
wrappers/mtproto_intermediate.go Просмотреть файл

24
 }
24
 }
25
 
25
 
26
 func (m *MTProtoIntermediate) Read() ([]byte, error) {
26
 func (m *MTProtoIntermediate) Read() ([]byte, error) {
27
+	defer func() {
28
+		m.readCounter++
29
+	}()
30
+
27
 	m.logger.Debugw("Read packet",
31
 	m.logger.Debugw("Read packet",
28
 		"simple_ack", m.opts.ReadHacks.SimpleAck,
32
 		"simple_ack", m.opts.ReadHacks.SimpleAck,
29
 		"quick_ack", m.opts.ReadHacks.QuickAck,
33
 		"quick_ack", m.opts.ReadHacks.QuickAck,
59
 	if length%4 != 0 {
63
 	if length%4 != 0 {
60
 		length -= length % 4
64
 		length -= length % 4
61
 	}
65
 	}
62
-	m.readCounter++
63
 
66
 
64
 	return buf.Bytes()[:length], nil
67
 	return buf.Bytes()[:length], nil
65
 }
68
 }
66
 
69
 
67
 func (m *MTProtoIntermediate) Write(p []byte) (int, error) {
70
 func (m *MTProtoIntermediate) Write(p []byte) (int, error) {
71
+	defer func() {
72
+		m.writeCounter++
73
+	}()
74
+
68
 	m.logger.Debugw("Write packet",
75
 	m.logger.Debugw("Write packet",
69
 		"simple_ack", m.opts.WriteHacks.SimpleAck,
76
 		"simple_ack", m.opts.WriteHacks.SimpleAck,
70
 		"quick_ack", m.opts.WriteHacks.QuickAck,
77
 		"quick_ack", m.opts.WriteHacks.QuickAck,
71
 		"counter", m.writeCounter,
78
 		"counter", m.writeCounter,
72
 	)
79
 	)
73
-	m.writeCounter++
74
 
80
 
75
 	if m.opts.ReadHacks.SimpleAck {
81
 	if m.opts.ReadHacks.SimpleAck {
76
 		return m.conn.Write(p)
82
 		return m.conn.Write(p)

+ 8
- 4
wrappers/mtproto_proxy.go Просмотреть файл

22
 }
22
 }
23
 
23
 
24
 func (m *MTProtoProxy) Read() ([]byte, error) {
24
 func (m *MTProtoProxy) Read() ([]byte, error) {
25
+	defer func() {
26
+		m.readCounter++
27
+	}()
28
+
25
 	m.logger.Debugw("Read packet",
29
 	m.logger.Debugw("Read packet",
26
 		"counter", m.readCounter,
30
 		"counter", m.readCounter,
27
 		"simple_ack", m.req.Options.WriteHacks.SimpleAck,
31
 		"simple_ack", m.req.Options.WriteHacks.SimpleAck,
32
 	if err != nil {
36
 	if err != nil {
33
 		return nil, errors.Annotate(err, "Cannot read packet")
37
 		return nil, errors.Annotate(err, "Cannot read packet")
34
 	}
38
 	}
35
-	defer func() {
36
-		m.readCounter++
37
-	}()
38
 
39
 
39
 	m.logger.Debugw("Read packet length",
40
 	m.logger.Debugw("Read packet length",
40
 		"counter", m.readCounter,
41
 		"counter", m.readCounter,
95
 }
96
 }
96
 
97
 
97
 func (m *MTProtoProxy) Write(p []byte) (int, error) {
98
 func (m *MTProtoProxy) Write(p []byte) (int, error) {
99
+	defer func() {
100
+		m.writeCounter++
101
+	}()
102
+
98
 	m.logger.Debugw("Write packet",
103
 	m.logger.Debugw("Write packet",
99
 		"length", len(p),
104
 		"length", len(p),
100
 		"counter", m.writeCounter,
105
 		"counter", m.writeCounter,
101
 		"simple_ack", m.req.Options.ReadHacks.SimpleAck,
106
 		"simple_ack", m.req.Options.ReadHacks.SimpleAck,
102
 		"quick_ack", m.req.Options.ReadHacks.QuickAck,
107
 		"quick_ack", m.req.Options.ReadHacks.QuickAck,
103
 	)
108
 	)
104
-	m.writeCounter++
105
 
109
 
106
 	header := m.req.MakeHeader(p)
110
 	header := m.req.MakeHeader(p)
107
 	if ce := m.logger.Desugar().Check(zap.DebugLevel, "RPC_PROXY_REQ header"); ce != nil {
111
 	if ce := m.logger.Desugar().Check(zap.DebugLevel, "RPC_PROXY_REQ header"); ce != nil {

Загрузка…
Отмена
Сохранить