Просмотр исходного кода

Small refactorings of rpc package

tags/0.9
9seconds 7 лет назад
Родитель
Сommit
0738631ef8

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

6
 )
6
 )
7
 
7
 
8
 var (
8
 var (
9
-	RPCTagCloseExt  = []byte{0xa2, 0x34, 0xb6, 0x5e}
10
-	RPCTagProxyAns  = []byte{0x0d, 0xda, 0x03, 0x44}
11
-	RPCTagSimpleAck = []byte{0x9b, 0x40, 0xac, 0x3b}
9
+	RPCTagCloseExt     = []byte{0xa2, 0x34, 0xb6, 0x5e}
10
+	RPCTagProxyAns     = []byte{0x0d, 0xda, 0x03, 0x44}
11
+	RPCTagSimpleAck    = []byte{0x9b, 0x40, 0xac, 0x3b}
12
+	RPCTagHandshake    = []byte{0xf5, 0xee, 0x82, 0x76}
13
+	RPCTagNonce        = []byte{0xaa, 0x87, 0xcb, 0x7a}
14
+	RPCTagProxyRequest = []byte{0xee, 0xf1, 0xce, 0x36}
15
+
16
+	RPCNonceCryptoAES = []byte{0x01, 0x00, 0x00, 0x00}
17
+
18
+	RPCHandshakeFlags = []byte{0x00, 0x00, 0x00, 0x00}
19
+
20
+	RPCProxyRequestExtraSize = []byte{0x18, 0x00, 0x00, 0x00}
21
+	RPCProxyRequestProxyTag  = []byte{0xae, 0x26, 0x1e, 0xdb}
22
+
23
+	RPCHandshakeSenderPID = []byte{}
24
+	RPCHandshakePeerPID   = []byte{}
12
 )
25
 )
26
+
27
+func init() {
28
+	RPCHandshakeSenderPID = []byte("IPIPPRPDTIME")
29
+	RPCHandshakePeerPID = []byte("IPIPPRPDTIME")
30
+}

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

1
 package rpc
1
 package rpc
2
 
2
 
3
-import (
4
-	"bytes"
5
-)
6
-
7
-const (
8
-	rpcHandshakeTagLength       = 4
9
-	rpcHandshakeFlagsLength     = 4
10
-	rpcHandshakeSenderPIDLength = 12
11
-	rpcHandshakePeerPIDLength   = rpcHandshakeSenderPIDLength
12
-
13
-	rpcHandshakeRequestLength = rpcHandshakeTagLength + rpcHandshakeFlagsLength + rpcHandshakeSenderPIDLength + rpcHandshakePeerPIDLength
14
-)
15
-
16
-var (
17
-	rpcHandshakeSenderPID [rpcHandshakeSenderPIDLength]byte
18
-	rpcHandshakePeerPID   [rpcHandshakePeerPIDLength]byte
19
-
20
-	rpcHandshakeTag   = [rpcHandshakeTagLength]byte{0xf5, 0xee, 0x82, 0x76}
21
-	rpcHandshakeFlags = [rpcHandshakeFlagsLength]byte{0x00, 0x00, 0x00, 0x00}
22
-)
3
+import "bytes"
23
 
4
 
24
 type RPCHandshakeRequest struct {
5
 type RPCHandshakeRequest struct {
25
 }
6
 }
26
 
7
 
27
 func (r *RPCHandshakeRequest) Bytes() []byte {
8
 func (r *RPCHandshakeRequest) Bytes() []byte {
28
 	buf := &bytes.Buffer{}
9
 	buf := &bytes.Buffer{}
29
-	buf.Grow(rpcHandshakeRequestLength)
30
 
10
 
31
-	buf.Write(rpcHandshakeTag[:])
32
-	buf.Write(rpcHandshakeFlags[:])
33
-	buf.Write(rpcHandshakeSenderPID[:])
34
-	buf.Write(rpcHandshakePeerPID[:])
11
+	buf.Write(RPCTagHandshake)
12
+	buf.Write(RPCHandshakeFlags)
13
+	buf.Write(RPCHandshakeSenderPID)
14
+	buf.Write(RPCHandshakePeerPID)
35
 
15
 
36
 	return buf.Bytes()
16
 	return buf.Bytes()
37
 }
17
 }
38
 
18
 
39
-func init() {
40
-	copy(rpcHandshakeSenderPID[:], "IPIPPRPDTIME")
41
-	copy(rpcHandshakePeerPID[:], "IPIPPRPDTIME")
42
-}
43
-
44
 func NewRPCHandshakeRequest() *RPCHandshakeRequest {
19
 func NewRPCHandshakeRequest() *RPCHandshakeRequest {
45
 	return &RPCHandshakeRequest{}
20
 	return &RPCHandshakeRequest{}
46
 }
21
 }

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

6
 	"github.com/juju/errors"
6
 	"github.com/juju/errors"
7
 )
7
 )
8
 
8
 
9
-const rpcHandshakeResponseLength = rpcHandshakeRequestLength
10
-
11
 type RPCHandshakeResponse struct {
9
 type RPCHandshakeResponse struct {
12
-	Type      [rpcHandshakeTagLength]byte
13
-	Flags     [rpcHandshakeFlagsLength]byte
14
-	SenderPID [rpcHandshakeSenderPIDLength]byte
15
-	PeerPID   [rpcHandshakePeerPIDLength]byte
10
+	Type      []byte
11
+	Flags     []byte
12
+	SenderPID []byte
13
+	PeerPID   []byte
16
 }
14
 }
17
 
15
 
18
 func (r *RPCHandshakeResponse) Bytes() []byte {
16
 func (r *RPCHandshakeResponse) Bytes() []byte {
19
 	buf := &bytes.Buffer{}
17
 	buf := &bytes.Buffer{}
20
-	buf.Grow(rpcHandshakeResponseLength)
21
 
18
 
22
 	buf.Write(r.Type[:])
19
 	buf.Write(r.Type[:])
23
 	buf.Write(r.Flags[:])
20
 	buf.Write(r.Flags[:])
28
 }
25
 }
29
 
26
 
30
 func (r *RPCHandshakeResponse) Valid(req *RPCHandshakeRequest) error {
27
 func (r *RPCHandshakeResponse) Valid(req *RPCHandshakeRequest) error {
31
-	if r.Type != rpcHandshakeTag {
28
+	if !bytes.Equal(r.Type, RPCTagHandshake) {
32
 		return errors.New("Unexpected handshake tag")
29
 		return errors.New("Unexpected handshake tag")
33
 	}
30
 	}
34
-	if r.PeerPID != rpcHandshakeSenderPID {
31
+	if !bytes.Equal(r.PeerPID, RPCHandshakeSenderPID) {
35
 		return errors.New("Incorrect sender PID")
32
 		return errors.New("Incorrect sender PID")
36
 	}
33
 	}
37
 
34
 
39
 }
36
 }
40
 
37
 
41
 func NewRPCHandshakeResponse(data []byte) (*RPCHandshakeResponse, error) {
38
 func NewRPCHandshakeResponse(data []byte) (*RPCHandshakeResponse, error) {
42
-	if len(data) != rpcHandshakeResponseLength {
39
+	if len(data) != 32 {
43
 		return nil, errors.New("Incorrect handshake response length")
40
 		return nil, errors.New("Incorrect handshake response length")
44
 	}
41
 	}
45
 
42
 
46
-	resp := RPCHandshakeResponse{}
47
-	copy(resp.Type[:], data[:4])
48
-	copy(resp.Flags[:], data[4:8])
49
-	copy(resp.SenderPID[:], data[8:20])
50
-	copy(resp.PeerPID[:], data[20:])
51
-
52
-	return &resp, nil
43
+	return &RPCHandshakeResponse{
44
+		Type:      data[:4],
45
+		Flags:     data[4:8],
46
+		SenderPID: data[8:20],
47
+		PeerPID:   data[20:],
48
+	}, nil
53
 }
49
 }

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

9
 	"github.com/juju/errors"
9
 	"github.com/juju/errors"
10
 )
10
 )
11
 
11
 
12
-const (
13
-	rpcNonceLength            = 16
14
-	rpcNonceKeySelectorLength = 4
15
-	rpcNonceCryptoTSLength    = 4
16
-	rpcNonceTagLength         = 4
17
-	rpcNonceCryptoAESLength   = 4
18
-
19
-	rpcNonceRequestLength = rpcNonceTagLength + rpcNonceKeySelectorLength + rpcNonceCryptoAESLength + rpcNonceCryptoTSLength + rpcNonceLength
20
-)
21
-
22
-var (
23
-	rpcNonceTag          = [rpcNonceTagLength]byte{0xaa, 0x87, 0xcb, 0x7a}
24
-	rpcNonceCryptoAESTag = [rpcNonceCryptoAESLength]byte{0x01, 0x00, 0x00, 0x00}
25
-)
26
-
27
 type RPCNonceRequest struct {
12
 type RPCNonceRequest struct {
28
-	KeySelector [rpcNonceKeySelectorLength]byte
29
-	CryptoTS    [rpcNonceCryptoTSLength]byte
30
-	Nonce       [rpcNonceLength]byte
13
+	KeySelector []byte
14
+	CryptoTS    []byte
15
+	Nonce       []byte
31
 }
16
 }
32
 
17
 
33
 func (r *RPCNonceRequest) Bytes() []byte {
18
 func (r *RPCNonceRequest) Bytes() []byte {
34
 	buf := &bytes.Buffer{}
19
 	buf := &bytes.Buffer{}
35
-	buf.Grow(rpcNonceRequestLength)
36
 
20
 
37
-	buf.Write(rpcNonceTag[:])
38
-	buf.Write(r.KeySelector[:])
39
-	buf.Write(rpcNonceCryptoAESTag[:])
40
-	buf.Write(r.CryptoTS[:])
41
-	buf.Write(r.Nonce[:])
21
+	buf.Write(RPCTagNonce)
22
+	buf.Write(r.KeySelector)
23
+	buf.Write(RPCNonceCryptoAES)
24
+	buf.Write(r.CryptoTS)
25
+	buf.Write(r.Nonce)
42
 
26
 
43
 	return buf.Bytes()
27
 	return buf.Bytes()
44
 }
28
 }
45
 
29
 
46
 func NewRPCNonceRequest(proxySecret []byte) (*RPCNonceRequest, error) {
30
 func NewRPCNonceRequest(proxySecret []byte) (*RPCNonceRequest, error) {
47
-	var nonce [rpcNonceLength]byte
48
-	var keySelector [rpcNonceKeySelectorLength]byte
49
-	var cryptoTS [rpcNonceCryptoTSLength]byte
31
+	nonce := make([]byte, 16)
32
+	keySelector := make([]byte, 4)
33
+	cryptoTS := make([]byte, 4)
50
 
34
 
51
-	if _, err := rand.Read(nonce[:]); err != nil {
35
+	if _, err := rand.Read(nonce); err != nil {
52
 		return nil, errors.Annotate(err, "Cannot generate nonce")
36
 		return nil, errors.Annotate(err, "Cannot generate nonce")
53
 	}
37
 	}
54
-	copy(keySelector[:], proxySecret)
38
+	copy(keySelector, proxySecret)
55
 
39
 
56
 	timestamp := time.Now().Truncate(time.Second).Unix() % 4294967296 // 256 ^ 4 - do not know how to name
40
 	timestamp := time.Now().Truncate(time.Second).Unix() % 4294967296 // 256 ^ 4 - do not know how to name
57
-	binary.LittleEndian.PutUint32(cryptoTS[:], uint32(timestamp))
41
+	binary.LittleEndian.PutUint32(cryptoTS, uint32(timestamp))
58
 
42
 
59
 	return &RPCNonceRequest{
43
 	return &RPCNonceRequest{
60
 		KeySelector: keySelector,
44
 		KeySelector: keySelector,

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

6
 	"github.com/juju/errors"
6
 	"github.com/juju/errors"
7
 )
7
 )
8
 
8
 
9
-const rpcNonceResponseLength = rpcNonceRequestLength
10
-
11
 type RPCNonceResponse struct {
9
 type RPCNonceResponse struct {
12
 	RPCNonceRequest
10
 	RPCNonceRequest
13
 
11
 
14
-	RPCType [rpcNonceTagLength]byte
15
-	Crypto  [rpcNonceCryptoAESLength]byte
12
+	RPCType []byte
13
+	Crypto  []byte
16
 }
14
 }
17
 
15
 
18
 func (r *RPCNonceResponse) Bytes() []byte {
16
 func (r *RPCNonceResponse) Bytes() []byte {
19
 	buf := &bytes.Buffer{}
17
 	buf := &bytes.Buffer{}
20
-	buf.Grow(rpcNonceResponseLength)
21
 
18
 
22
-	buf.Write(r.RPCType[:])
23
-	buf.Write(r.KeySelector[:])
24
-	buf.Write(r.Crypto[:])
25
-	buf.Write(r.CryptoTS[:])
26
-	buf.Write(r.Nonce[:])
19
+	buf.Write(r.RPCType)
20
+	buf.Write(r.KeySelector)
21
+	buf.Write(r.Crypto)
22
+	buf.Write(r.CryptoTS)
23
+	buf.Write(r.Nonce)
27
 
24
 
28
 	return buf.Bytes()
25
 	return buf.Bytes()
29
 }
26
 }
30
 
27
 
31
 func (r *RPCNonceResponse) Valid(req *RPCNonceRequest) error {
28
 func (r *RPCNonceResponse) Valid(req *RPCNonceRequest) error {
32
-	if r.RPCType != rpcNonceTag {
29
+	if !bytes.Equal(r.RPCType, RPCTagNonce) {
33
 		return errors.New("Unexpected RPC type")
30
 		return errors.New("Unexpected RPC type")
34
 	}
31
 	}
35
-	if r.Crypto != rpcNonceCryptoAESTag {
32
+	if !bytes.Equal(r.Crypto, RPCNonceCryptoAES) {
36
 		return errors.New("Unexpected crypto type")
33
 		return errors.New("Unexpected crypto type")
37
 	}
34
 	}
38
-	if r.KeySelector != req.KeySelector {
35
+	if !bytes.Equal(r.KeySelector, req.KeySelector) {
39
 		return errors.New("Unexpected key selector")
36
 		return errors.New("Unexpected key selector")
40
 	}
37
 	}
41
 
38
 
43
 }
40
 }
44
 
41
 
45
 func NewRPCNonceResponse(data []byte) (*RPCNonceResponse, error) {
42
 func NewRPCNonceResponse(data []byte) (*RPCNonceResponse, error) {
46
-	if len(data) != rpcNonceResponseLength {
43
+	if len(data) != 32 {
47
 		return nil, errors.New("Unexpected message length")
44
 		return nil, errors.New("Unexpected message length")
48
 	}
45
 	}
49
 
46
 
50
-	resp := RPCNonceResponse{}
51
-	copy(resp.RPCType[:], data[:4])
52
-	copy(resp.KeySelector[:], data[4:8])
53
-	copy(resp.Crypto[:], data[8:12])
54
-	copy(resp.CryptoTS[:], data[12:16])
55
-	copy(resp.Nonce[:], data[16:])
56
-
57
-	return &resp, nil
47
+	return &RPCNonceResponse{
48
+		RPCNonceRequest: RPCNonceRequest{
49
+			KeySelector: data[4:8],
50
+			CryptoTS:    data[12:16],
51
+			Nonce:       data[16:],
52
+		},
53
+		RPCType: data[:4],
54
+		Crypto:  data[8:12],
55
+	}, nil
58
 }
56
 }

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

11
 	"github.com/9seconds/mtg/mtproto"
11
 	"github.com/9seconds/mtg/mtproto"
12
 )
12
 )
13
 
13
 
14
-const (
15
-	rpcProxyRequestConnectionIDLength = 8
16
-	rpcProxyRequestIPPortLength       = 16 + 4
17
-)
18
-
19
-var (
20
-	rpcProxyRequestTag       = []byte{0xee, 0xf1, 0xce, 0x36}
21
-	rpcProxyRequestExtraSize = []byte{0x18, 0x00, 0x00, 0x00}
22
-	rpcProxyRequestProxyTag  = []byte{0xae, 0x26, 0x1e, 0xdb}
23
-)
24
-
25
 type RPCProxyRequest struct {
14
 type RPCProxyRequest struct {
26
 	Flags        RPCProxyRequestFlags
15
 	Flags        RPCProxyRequestFlags
27
-	ConnectionID [rpcProxyRequestConnectionIDLength]byte
28
-	OurIPPort    [rpcProxyRequestIPPortLength]byte
29
-	ClientIPPort [rpcProxyRequestIPPortLength]byte
16
+	ConnectionID []byte
17
+	OurIPPort    []byte
18
+	ClientIPPort []byte
30
 	ADTag        []byte
19
 	ADTag        []byte
31
 	Options      *mtproto.ConnectionOpts
20
 	Options      *mtproto.ConnectionOpts
32
 }
21
 }
43
 		flags |= RPCProxyRequestFlagsEncrypted
32
 		flags |= RPCProxyRequestFlagsEncrypted
44
 	}
33
 	}
45
 
34
 
46
-	buf.Write(rpcProxyRequestTag)
35
+	buf.Write(RPCTagProxyRequest)
47
 	buf.Write(flags.Bytes())
36
 	buf.Write(flags.Bytes())
48
 	buf.Write(r.ConnectionID[:])
37
 	buf.Write(r.ConnectionID[:])
49
 	buf.Write(r.ClientIPPort[:])
38
 	buf.Write(r.ClientIPPort[:])
50
 	buf.Write(r.OurIPPort[:])
39
 	buf.Write(r.OurIPPort[:])
51
-	buf.Write(rpcProxyRequestExtraSize)
52
-	buf.Write(rpcProxyRequestProxyTag)
40
+	buf.Write(RPCProxyRequestExtraSize)
41
+	buf.Write(RPCProxyRequestProxyTag)
53
 	buf.WriteByte(byte(len(r.ADTag)))
42
 	buf.WriteByte(byte(len(r.ADTag)))
54
 	buf.Write(r.ADTag)
43
 	buf.Write(r.ADTag)
55
 	buf.Write(bytes.Repeat([]byte{0x00}, buf.Len()%4))
44
 	buf.Write(bytes.Repeat([]byte{0x00}, buf.Len()%4))
69
 	}
58
 	}
70
 
59
 
71
 	request := RPCProxyRequest{
60
 	request := RPCProxyRequest{
72
-		Flags:   flags,
73
-		ADTag:   adTag,
74
-		Options: opts,
61
+		Flags:        flags,
62
+		ADTag:        adTag,
63
+		Options:      opts,
64
+		ConnectionID: make([]byte, 8),
65
+		ClientIPPort: make([]byte, 16+4),
66
+		OurIPPort:    make([]byte, 16+4),
75
 	}
67
 	}
76
 
68
 
77
-	if _, err := rand.Read(request.ConnectionID[:]); err != nil {
69
+	if _, err := rand.Read(request.ConnectionID); err != nil {
78
 		return nil, errors.Annotate(err, "Cannot generate connection ID")
70
 		return nil, errors.Annotate(err, "Cannot generate connection ID")
79
 	}
71
 	}
80
 
72
 
81
-	port := make([]byte, 4)
73
+	port := [4]byte{}
82
 	copy(request.ClientIPPort[:16], clientAddr.IP.To16())
74
 	copy(request.ClientIPPort[:16], clientAddr.IP.To16())
83
-	binary.LittleEndian.PutUint32(port, uint32(clientAddr.Port))
84
-	copy(request.ClientIPPort[16:], port)
75
+	binary.LittleEndian.PutUint32(port[:], uint32(clientAddr.Port))
76
+	copy(request.ClientIPPort[16:], port[:])
85
 
77
 
86
 	copy(request.OurIPPort[:16], ownAddr.IP.To16())
78
 	copy(request.OurIPPort[:16], ownAddr.IP.To16())
87
-	binary.LittleEndian.PutUint32(port, uint32(ownAddr.Port))
88
-	copy(request.OurIPPort[16:], port)
79
+	binary.LittleEndian.PutUint32(port[:], uint32(ownAddr.Port))
80
+	copy(request.OurIPPort[16:], port[:])
89
 
81
 
90
 	return &request, nil
82
 	return &request, nil
91
 }
83
 }

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