9seconds 7 лет назад
Родитель
Сommit
6fa5f31ca8

+ 0
- 32
mtproto/bufferpool.go Просмотреть файл

1
-package mtproto
2
-
3
-import (
4
-	"bytes"
5
-	"sync"
6
-)
7
-
8
-const bufferPoolSize = 4 * 1024
9
-
10
-var bufferPool sync.Pool
11
-
12
-func GetBuffer() *bytes.Buffer {
13
-	buf := bufferPool.Get().(*bytes.Buffer)
14
-	buf.Reset()
15
-
16
-	return buf
17
-}
18
-
19
-func ReturnBuffer(buf *bytes.Buffer) {
20
-	bufferPool.Put(buf)
21
-}
22
-
23
-func init() {
24
-	bufferPool = sync.Pool{
25
-		New: func() interface{} {
26
-			buf := &bytes.Buffer{}
27
-			buf.Grow(bufferPoolSize)
28
-
29
-			return buf
30
-		},
31
-	}
32
-}

+ 0
- 6
mtproto/rpc/rpc.go Просмотреть файл

1
 package rpc
1
 package rpc
2
 
2
 
3
-import "bytes"
4
-
5
 const (
3
 const (
6
 	RPCNonceSeqNo     = -2
4
 	RPCNonceSeqNo     = -2
7
 	RPCHandshakeSeqNo = -1
5
 	RPCHandshakeSeqNo = -1
8
 )
6
 )
9
-
10
-type RPC interface {
11
-	Bytes() *bytes.Buffer
12
-}

+ 2
- 2
mtproto/rpc/rpc_handshake_request.go Просмотреть файл

25
 type RPCHandshakeRequest struct {
25
 type RPCHandshakeRequest struct {
26
 }
26
 }
27
 
27
 
28
-func (r *RPCHandshakeRequest) Bytes() *bytes.Buffer {
28
+func (r *RPCHandshakeRequest) Bytes() []byte {
29
 	buf := &bytes.Buffer{}
29
 	buf := &bytes.Buffer{}
30
 	buf.Grow(rpcHandshakeRequestLength)
30
 	buf.Grow(rpcHandshakeRequestLength)
31
 
31
 
34
 	buf.Write(rpcHandshakeSenderPID[:])
34
 	buf.Write(rpcHandshakeSenderPID[:])
35
 	buf.Write(rpcHandshakePeerPID[:])
35
 	buf.Write(rpcHandshakePeerPID[:])
36
 
36
 
37
-	return buf
37
+	return buf.Bytes()
38
 }
38
 }
39
 
39
 
40
 func init() {
40
 func init() {

+ 2
- 2
mtproto/rpc/rpc_handshake_response.go Просмотреть файл

15
 	PeerPID   [rpcHandshakePeerPIDLength]byte
15
 	PeerPID   [rpcHandshakePeerPIDLength]byte
16
 }
16
 }
17
 
17
 
18
-func (r *RPCHandshakeResponse) Bytes() *bytes.Buffer {
18
+func (r *RPCHandshakeResponse) Bytes() []byte {
19
 	buf := &bytes.Buffer{}
19
 	buf := &bytes.Buffer{}
20
 	buf.Grow(rpcHandshakeResponseLength)
20
 	buf.Grow(rpcHandshakeResponseLength)
21
 
21
 
24
 	buf.Write(r.SenderPID[:])
24
 	buf.Write(r.SenderPID[:])
25
 	buf.Write(r.PeerPID[:])
25
 	buf.Write(r.PeerPID[:])
26
 
26
 
27
-	return buf
27
+	return buf.Bytes()
28
 }
28
 }
29
 
29
 
30
 func (r *RPCHandshakeResponse) Valid(req *RPCHandshakeRequest) error {
30
 func (r *RPCHandshakeResponse) Valid(req *RPCHandshakeRequest) error {

+ 2
- 2
mtproto/rpc/rpc_nonce_request.go Просмотреть файл

31
 	Nonce       [rpcNonceLength]byte
31
 	Nonce       [rpcNonceLength]byte
32
 }
32
 }
33
 
33
 
34
-func (r *RPCNonceRequest) Bytes() *bytes.Buffer {
34
+func (r *RPCNonceRequest) Bytes() []byte {
35
 	buf := &bytes.Buffer{}
35
 	buf := &bytes.Buffer{}
36
 	buf.Grow(rpcNonceRequestLength)
36
 	buf.Grow(rpcNonceRequestLength)
37
 
37
 
41
 	buf.Write(r.CryptoTS[:])
41
 	buf.Write(r.CryptoTS[:])
42
 	buf.Write(r.Nonce[:])
42
 	buf.Write(r.Nonce[:])
43
 
43
 
44
-	return buf
44
+	return buf.Bytes()
45
 }
45
 }
46
 
46
 
47
 func NewRPCNonceRequest(proxySecret []byte) (*RPCNonceRequest, error) {
47
 func NewRPCNonceRequest(proxySecret []byte) (*RPCNonceRequest, error) {

+ 2
- 2
mtproto/rpc/rpc_nonce_response.go Просмотреть файл

15
 	Crypto  [rpcNonceCryptoAESLength]byte
15
 	Crypto  [rpcNonceCryptoAESLength]byte
16
 }
16
 }
17
 
17
 
18
-func (r *RPCNonceResponse) Bytes() *bytes.Buffer {
18
+func (r *RPCNonceResponse) Bytes() []byte {
19
 	buf := &bytes.Buffer{}
19
 	buf := &bytes.Buffer{}
20
 	buf.Grow(rpcNonceResponseLength)
20
 	buf.Grow(rpcNonceResponseLength)
21
 
21
 
25
 	buf.Write(r.CryptoTS[:])
25
 	buf.Write(r.CryptoTS[:])
26
 	buf.Write(r.Nonce[:])
26
 	buf.Write(r.Nonce[:])
27
 
27
 
28
-	return buf
28
+	return buf.Bytes()
29
 }
29
 }
30
 
30
 
31
 func (r *RPCNonceResponse) Valid(req *RPCNonceRequest) error {
31
 func (r *RPCNonceResponse) Valid(req *RPCNonceRequest) error {

+ 3
- 3
mtproto/rpc/rpc_proxy_request.go Просмотреть файл

32
 	Extras       *mtproto.Extras
32
 	Extras       *mtproto.Extras
33
 }
33
 }
34
 
34
 
35
-func (r *RPCProxyRequest) Bytes() *bytes.Buffer {
36
-	buf := mtproto.GetBuffer()
35
+func (r *RPCProxyRequest) Bytes() []byte {
36
+	buf := &bytes.Buffer{}
37
 
37
 
38
 	flags := r.Flags
38
 	flags := r.Flags
39
 	if r.Extras.QuickAck {
39
 	if r.Extras.QuickAck {
62
 		buf.Write(messageBytes)
62
 		buf.Write(messageBytes)
63
 	}
63
 	}
64
 
64
 
65
-	return buf
65
+	return buf.Bytes()
66
 }
66
 }
67
 
67
 
68
 func NewRPCProxyRequest(connectionType mtproto.ConnectionType, local, remote *net.TCPAddr, adTag []byte, extras *mtproto.Extras) (*RPCProxyRequest, error) {
68
 func NewRPCProxyRequest(connectionType mtproto.ConnectionType, local, remote *net.TCPAddr, adTag []byte, extras *mtproto.Extras) (*RPCProxyRequest, error) {

+ 0
- 48
mtproto/rwc.go Просмотреть файл

1
-package mtproto
2
-
3
-import (
4
-	"bytes"
5
-	"io"
6
-)
7
-
8
-type BytesRWC interface {
9
-	Write(*bytes.Buffer) (int, error)
10
-	Read([]byte) (int, error)
11
-	Close() error
12
-}
13
-
14
-type StartBytesRWC struct {
15
-	conn BytesRWC
16
-}
17
-
18
-func (s *StartBytesRWC) Write(p []byte) (int, error) {
19
-	buf := GetBuffer()
20
-	buf.Write(p)
21
-	defer ReturnBuffer(buf)
22
-
23
-	return s.conn.Write(buf)
24
-}
25
-
26
-func (s *StartBytesRWC) Read(p []byte) (int, error) {
27
-	return s.conn.Read(p)
28
-}
29
-
30
-func (s *StartBytesRWC) Close() error {
31
-	return s.conn.Close()
32
-}
33
-
34
-type FinishBytesRWC struct {
35
-	conn io.ReadWriteCloser
36
-}
37
-
38
-func (f *FinishBytesRWC) Write(buf *bytes.Buffer) (int, error) {
39
-	return f.conn.Write(buf.Bytes())
40
-}
41
-
42
-func (f *FinishBytesRWC) Read(p []byte) (int, error) {
43
-	return f.conn.Read(p)
44
-}
45
-
46
-func (f *FinishBytesRWC) Close() error {
47
-	return f.conn.Close()
48
-}

mtproto/crypt/crypt.go → mtproto/wrappers/crypt.go Просмотреть файл

1
-package crypt
1
+package wrappers
2
 
2
 
3
 import (
3
 import (
4
 	"bytes"
4
 	"bytes"

+ 11
- 20
mtproto/wrappers/frame.go Просмотреть файл

8
 	"io"
8
 	"io"
9
 
9
 
10
 	"github.com/juju/errors"
10
 	"github.com/juju/errors"
11
-
12
-	"github.com/9seconds/mtg/mtproto"
13
 )
11
 )
14
 
12
 
15
 // Frame: { MessageLength(4) | SequenceNumber(4) | Message(???) | CRC32(4) [| padding(4), ...] }
13
 // Frame: { MessageLength(4) | SequenceNumber(4) | Message(???) | CRC32(4) [| padding(4), ...] }
21
 var frameRWCPadding = [4]byte{0x04, 0x00, 0x00, 0x00}
19
 var frameRWCPadding = [4]byte{0x04, 0x00, 0x00, 0x00}
22
 
20
 
23
 type FrameRWC struct {
21
 type FrameRWC struct {
24
-	conn mtproto.BytesRWC
22
+	conn io.ReadWriteCloser
25
 
23
 
26
 	readSeqNo  int32
24
 	readSeqNo  int32
27
 	writeSeqNo int32
25
 	writeSeqNo int32
28
 	readBuf    *bytes.Buffer
26
 	readBuf    *bytes.Buffer
29
 }
27
 }
30
 
28
 
31
-func (f *FrameRWC) Write(buf *bytes.Buffer) (int, error) {
32
-	writeBuf := mtproto.GetBuffer()
33
-	defer mtproto.ReturnBuffer(writeBuf)
29
+func (f *FrameRWC) Write(buf []byte) (int, error) {
30
+	writeBuf := &bytes.Buffer{}
34
 
31
 
35
 	// 4 - len bytes
32
 	// 4 - len bytes
36
 	// 4 - seq bytes
33
 	// 4 - seq bytes
37
 	// . - message
34
 	// . - message
38
 	// 4 - crc32
35
 	// 4 - crc32
39
-	messageLength := 4 + 4 + buf.Len() + 4
36
+	messageLength := 4 + 4 + len(buf) + 4
40
 	paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize
37
 	paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize
41
 	writeBuf.Grow(messageLength + paddingLength)
38
 	writeBuf.Grow(messageLength + paddingLength)
42
 
39
 
43
 	binary.Write(writeBuf, binary.LittleEndian, uint32(messageLength))
40
 	binary.Write(writeBuf, binary.LittleEndian, uint32(messageLength))
44
 	binary.Write(writeBuf, binary.LittleEndian, f.writeSeqNo)
41
 	binary.Write(writeBuf, binary.LittleEndian, f.writeSeqNo)
45
-	writeBuf.Write(buf.Bytes())
42
+	writeBuf.Write(buf)
46
 	f.writeSeqNo++
43
 	f.writeSeqNo++
47
 
44
 
48
 	checksum := crc32.ChecksumIEEE(writeBuf.Bytes())
45
 	checksum := crc32.ChecksumIEEE(writeBuf.Bytes())
49
 	binary.Write(writeBuf, binary.LittleEndian, checksum)
46
 	binary.Write(writeBuf, binary.LittleEndian, checksum)
50
 	writeBuf.Write(bytes.Repeat(frameRWCPadding[:], paddingLength/4))
47
 	writeBuf.Write(bytes.Repeat(frameRWCPadding[:], paddingLength/4))
51
 
48
 
52
-	return f.conn.Write(writeBuf)
49
+	_, err := f.conn.Write(writeBuf.Bytes())
50
+	return len(buf), err
53
 }
51
 }
54
 
52
 
55
 func (f *FrameRWC) Read(p []byte) (int, error) {
53
 func (f *FrameRWC) Read(p []byte) (int, error) {
57
 		return f.flush(p)
55
 		return f.flush(p)
58
 	}
56
 	}
59
 
57
 
60
-	buf := mtproto.GetBuffer()
61
-	defer mtproto.ReturnBuffer(buf)
62
-
58
+	buf := &bytes.Buffer{}
63
 	for {
59
 	for {
64
 		buf.Reset()
60
 		buf.Reset()
65
 		if _, err := io.CopyN(buf, f.conn, 4); err != nil {
61
 		if _, err := io.CopyN(buf, f.conn, 4); err != nil {
103
 }
99
 }
104
 
100
 
105
 func (f *FrameRWC) Close() error {
101
 func (f *FrameRWC) Close() error {
106
-	defer mtproto.ReturnBuffer(f.readBuf)
107
 	return f.conn.Close()
102
 	return f.conn.Close()
108
 }
103
 }
109
 
104
 
118
 	if sizeToRead == f.readBuf.Len() {
113
 	if sizeToRead == f.readBuf.Len() {
119
 		f.readBuf.Reset()
114
 		f.readBuf.Reset()
120
 	} else {
115
 	} else {
121
-		newBuf := mtproto.GetBuffer()
122
-		newBuf.Write(data[sizeToRead:])
123
-
124
-		mtproto.ReturnBuffer(f.readBuf)
125
-		f.readBuf = newBuf
116
+		f.readBuf = bytes.NewBuffer(data[sizeToRead:])
126
 	}
117
 	}
127
 
118
 
128
 	return sizeToRead, nil
119
 	return sizeToRead, nil
129
 }
120
 }
130
 
121
 
131
-func NewFrameRWC(conn mtproto.BytesRWC, seqNo int32) mtproto.BytesRWC {
122
+func NewFrameRWC(conn io.ReadWriteCloser, seqNo int32) io.ReadWriteCloser {
132
 	return &FrameRWC{
123
 	return &FrameRWC{
133
 		conn:       conn,
124
 		conn:       conn,
134
 		readSeqNo:  seqNo,
125
 		readSeqNo:  seqNo,
135
 		writeSeqNo: seqNo,
126
 		writeSeqNo: seqNo,
136
-		readBuf:    mtproto.GetBuffer(),
127
+		readBuf:    &bytes.Buffer{},
137
 	}
128
 	}
138
 }
129
 }

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