Small refactorings of rpc package

mtproto/rpc/rpc.go

@@ -6,7 +6,25 @@ const (
 )
 
 var (
-	RPCTagCloseExt  = []byte{0xa2, 0x34, 0xb6, 0x5e}
-	RPCTagProxyAns  = []byte{0x0d, 0xda, 0x03, 0x44}
-	RPCTagSimpleAck = []byte{0x9b, 0x40, 0xac, 0x3b}
+	RPCTagCloseExt     = []byte{0xa2, 0x34, 0xb6, 0x5e}
+	RPCTagProxyAns     = []byte{0x0d, 0xda, 0x03, 0x44}
+	RPCTagSimpleAck    = []byte{0x9b, 0x40, 0xac, 0x3b}
+	RPCTagHandshake    = []byte{0xf5, 0xee, 0x82, 0x76}
+	RPCTagNonce        = []byte{0xaa, 0x87, 0xcb, 0x7a}
+	RPCTagProxyRequest = []byte{0xee, 0xf1, 0xce, 0x36}
+
+	RPCNonceCryptoAES = []byte{0x01, 0x00, 0x00, 0x00}
+
+	RPCHandshakeFlags = []byte{0x00, 0x00, 0x00, 0x00}
+
+	RPCProxyRequestExtraSize = []byte{0x18, 0x00, 0x00, 0x00}
+	RPCProxyRequestProxyTag  = []byte{0xae, 0x26, 0x1e, 0xdb}
+
+	RPCHandshakeSenderPID = []byte{}
+	RPCHandshakePeerPID   = []byte{}
 )
+
+func init() {
+	RPCHandshakeSenderPID = []byte("IPIPPRPDTIME")
+	RPCHandshakePeerPID = []byte("IPIPPRPDTIME")
+}

mtproto/rpc/rpc_handshake_request.go

@@ -1,46 +1,21 @@
 package rpc
 
-import (
-	"bytes"
-)
-
-const (
-	rpcHandshakeTagLength       = 4
-	rpcHandshakeFlagsLength     = 4
-	rpcHandshakeSenderPIDLength = 12
-	rpcHandshakePeerPIDLength   = rpcHandshakeSenderPIDLength
-
-	rpcHandshakeRequestLength = rpcHandshakeTagLength + rpcHandshakeFlagsLength + rpcHandshakeSenderPIDLength + rpcHandshakePeerPIDLength
-)
-
-var (
-	rpcHandshakeSenderPID [rpcHandshakeSenderPIDLength]byte
-	rpcHandshakePeerPID   [rpcHandshakePeerPIDLength]byte
-
-	rpcHandshakeTag   = [rpcHandshakeTagLength]byte{0xf5, 0xee, 0x82, 0x76}
-	rpcHandshakeFlags = [rpcHandshakeFlagsLength]byte{0x00, 0x00, 0x00, 0x00}
-)
+import "bytes"
 
 type RPCHandshakeRequest struct {
 }
 
 func (r *RPCHandshakeRequest) Bytes() []byte {
 	buf := &bytes.Buffer{}
-	buf.Grow(rpcHandshakeRequestLength)
 
-	buf.Write(rpcHandshakeTag[:])
-	buf.Write(rpcHandshakeFlags[:])
-	buf.Write(rpcHandshakeSenderPID[:])
-	buf.Write(rpcHandshakePeerPID[:])
+	buf.Write(RPCTagHandshake)
+	buf.Write(RPCHandshakeFlags)
+	buf.Write(RPCHandshakeSenderPID)
+	buf.Write(RPCHandshakePeerPID)
 
 	return buf.Bytes()
 }
 
-func init() {
-	copy(rpcHandshakeSenderPID[:], "IPIPPRPDTIME")
-	copy(rpcHandshakePeerPID[:], "IPIPPRPDTIME")
-}
-
 func NewRPCHandshakeRequest() *RPCHandshakeRequest {
 	return &RPCHandshakeRequest{}
 }

mtproto/rpc/rpc_handshake_response.go

@@ -6,18 +6,15 @@ import (
 	"github.com/juju/errors"
 )
 
-const rpcHandshakeResponseLength = rpcHandshakeRequestLength
-
 type RPCHandshakeResponse struct {
-	Type      [rpcHandshakeTagLength]byte
-	Flags     [rpcHandshakeFlagsLength]byte
-	SenderPID [rpcHandshakeSenderPIDLength]byte
-	PeerPID   [rpcHandshakePeerPIDLength]byte
+	Type      []byte
+	Flags     []byte
+	SenderPID []byte
+	PeerPID   []byte
 }
 
 func (r *RPCHandshakeResponse) Bytes() []byte {
 	buf := &bytes.Buffer{}
-	buf.Grow(rpcHandshakeResponseLength)
 
 	buf.Write(r.Type[:])
 	buf.Write(r.Flags[:])
@@ -28,10 +25,10 @@ func (r *RPCHandshakeResponse) Bytes() []byte {
 }
 
 func (r *RPCHandshakeResponse) Valid(req *RPCHandshakeRequest) error {
-	if r.Type != rpcHandshakeTag {
+	if !bytes.Equal(r.Type, RPCTagHandshake) {
 		return errors.New("Unexpected handshake tag")
 	}
-	if r.PeerPID != rpcHandshakeSenderPID {
+	if !bytes.Equal(r.PeerPID, RPCHandshakeSenderPID) {
 		return errors.New("Incorrect sender PID")
 	}
 
@@ -39,15 +36,14 @@ func (r *RPCHandshakeResponse) Valid(req *RPCHandshakeRequest) error {
 }
 
 func NewRPCHandshakeResponse(data []byte) (*RPCHandshakeResponse, error) {
-	if len(data) != rpcHandshakeResponseLength {
+	if len(data) != 32 {
 		return nil, errors.New("Incorrect handshake response length")
 	}
 
-	resp := RPCHandshakeResponse{}
-	copy(resp.Type[:], data[:4])
-	copy(resp.Flags[:], data[4:8])
-	copy(resp.SenderPID[:], data[8:20])
-	copy(resp.PeerPID[:], data[20:])
-
-	return &resp, nil
+	return &RPCHandshakeResponse{
+		Type:      data[:4],
+		Flags:     data[4:8],
+		SenderPID: data[8:20],
+		PeerPID:   data[20:],
+	}, nil
 }

mtproto/rpc/rpc_nonce_request.go

@@ -9,52 +9,36 @@ import (
 	"github.com/juju/errors"
 )
 
-const (
-	rpcNonceLength            = 16
-	rpcNonceKeySelectorLength = 4
-	rpcNonceCryptoTSLength    = 4
-	rpcNonceTagLength         = 4
-	rpcNonceCryptoAESLength   = 4
-
-	rpcNonceRequestLength = rpcNonceTagLength + rpcNonceKeySelectorLength + rpcNonceCryptoAESLength + rpcNonceCryptoTSLength + rpcNonceLength
-)
-
-var (
-	rpcNonceTag          = [rpcNonceTagLength]byte{0xaa, 0x87, 0xcb, 0x7a}
-	rpcNonceCryptoAESTag = [rpcNonceCryptoAESLength]byte{0x01, 0x00, 0x00, 0x00}
-)
-
 type RPCNonceRequest struct {
-	KeySelector [rpcNonceKeySelectorLength]byte
-	CryptoTS    [rpcNonceCryptoTSLength]byte
-	Nonce       [rpcNonceLength]byte
+	KeySelector []byte
+	CryptoTS    []byte
+	Nonce       []byte
 }
 
 func (r *RPCNonceRequest) Bytes() []byte {
 	buf := &bytes.Buffer{}
-	buf.Grow(rpcNonceRequestLength)
 
-	buf.Write(rpcNonceTag[:])
-	buf.Write(r.KeySelector[:])
-	buf.Write(rpcNonceCryptoAESTag[:])
-	buf.Write(r.CryptoTS[:])
-	buf.Write(r.Nonce[:])
+	buf.Write(RPCTagNonce)
+	buf.Write(r.KeySelector)
+	buf.Write(RPCNonceCryptoAES)
+	buf.Write(r.CryptoTS)
+	buf.Write(r.Nonce)
 
 	return buf.Bytes()
 }
 
 func NewRPCNonceRequest(proxySecret []byte) (*RPCNonceRequest, error) {
-	var nonce [rpcNonceLength]byte
-	var keySelector [rpcNonceKeySelectorLength]byte
-	var cryptoTS [rpcNonceCryptoTSLength]byte
+	nonce := make([]byte, 16)
+	keySelector := make([]byte, 4)
+	cryptoTS := make([]byte, 4)
 
-	if _, err := rand.Read(nonce[:]); err != nil {
+	if _, err := rand.Read(nonce); err != nil {
 		return nil, errors.Annotate(err, "Cannot generate nonce")
 	}
-	copy(keySelector[:], proxySecret)
+	copy(keySelector, proxySecret)
 
 	timestamp := time.Now().Truncate(time.Second).Unix() % 4294967296 // 256 ^ 4 - do not know how to name
-	binary.LittleEndian.PutUint32(cryptoTS[:], uint32(timestamp))
+	binary.LittleEndian.PutUint32(cryptoTS, uint32(timestamp))
 
 	return &RPCNonceRequest{
 		KeySelector: keySelector,

mtproto/rpc/rpc_nonce_response.go

@@ -6,36 +6,33 @@ import (
 	"github.com/juju/errors"
 )
 
-const rpcNonceResponseLength = rpcNonceRequestLength
-
 type RPCNonceResponse struct {
 	RPCNonceRequest
 
-	RPCType [rpcNonceTagLength]byte
-	Crypto  [rpcNonceCryptoAESLength]byte
+	RPCType []byte
+	Crypto  []byte
 }
 
 func (r *RPCNonceResponse) Bytes() []byte {
 	buf := &bytes.Buffer{}
-	buf.Grow(rpcNonceResponseLength)
 
-	buf.Write(r.RPCType[:])
-	buf.Write(r.KeySelector[:])
-	buf.Write(r.Crypto[:])
-	buf.Write(r.CryptoTS[:])
-	buf.Write(r.Nonce[:])
+	buf.Write(r.RPCType)
+	buf.Write(r.KeySelector)
+	buf.Write(r.Crypto)
+	buf.Write(r.CryptoTS)
+	buf.Write(r.Nonce)
 
 	return buf.Bytes()
 }
 
 func (r *RPCNonceResponse) Valid(req *RPCNonceRequest) error {
-	if r.RPCType != rpcNonceTag {
+	if !bytes.Equal(r.RPCType, RPCTagNonce) {
 		return errors.New("Unexpected RPC type")
 	}
-	if r.Crypto != rpcNonceCryptoAESTag {
+	if !bytes.Equal(r.Crypto, RPCNonceCryptoAES) {
 		return errors.New("Unexpected crypto type")
 	}
-	if r.KeySelector != req.KeySelector {
+	if !bytes.Equal(r.KeySelector, req.KeySelector) {
 		return errors.New("Unexpected key selector")
 	}
 
@@ -43,16 +40,17 @@ func (r *RPCNonceResponse) Valid(req *RPCNonceRequest) error {
 }
 
 func NewRPCNonceResponse(data []byte) (*RPCNonceResponse, error) {
-	if len(data) != rpcNonceResponseLength {
+	if len(data) != 32 {
 		return nil, errors.New("Unexpected message length")
 	}
 
-	resp := RPCNonceResponse{}
-	copy(resp.RPCType[:], data[:4])
-	copy(resp.KeySelector[:], data[4:8])
-	copy(resp.Crypto[:], data[8:12])
-	copy(resp.CryptoTS[:], data[12:16])
-	copy(resp.Nonce[:], data[16:])
-
-	return &resp, nil
+	return &RPCNonceResponse{
+		RPCNonceRequest: RPCNonceRequest{
+			KeySelector: data[4:8],
+			CryptoTS:    data[12:16],
+			Nonce:       data[16:],
+		},
+		RPCType: data[:4],
+		Crypto:  data[8:12],
+	}, nil
 }

mtproto/rpc/rpc_proxy_request.go

@@ -11,22 +11,11 @@ import (
 	"github.com/9seconds/mtg/mtproto"
 )
 
-const (
-	rpcProxyRequestConnectionIDLength = 8
-	rpcProxyRequestIPPortLength       = 16 + 4
-)
-
-var (
-	rpcProxyRequestTag       = []byte{0xee, 0xf1, 0xce, 0x36}
-	rpcProxyRequestExtraSize = []byte{0x18, 0x00, 0x00, 0x00}
-	rpcProxyRequestProxyTag  = []byte{0xae, 0x26, 0x1e, 0xdb}
-)
-
 type RPCProxyRequest struct {
 	Flags        RPCProxyRequestFlags
-	ConnectionID [rpcProxyRequestConnectionIDLength]byte
-	OurIPPort    [rpcProxyRequestIPPortLength]byte
-	ClientIPPort [rpcProxyRequestIPPortLength]byte
+	ConnectionID []byte
+	OurIPPort    []byte
+	ClientIPPort []byte
 	ADTag        []byte
 	Options      *mtproto.ConnectionOpts
 }
@@ -43,13 +32,13 @@ func (r *RPCProxyRequest) Bytes(message []byte) []byte {
 		flags |= RPCProxyRequestFlagsEncrypted
 	}
 
-	buf.Write(rpcProxyRequestTag)
+	buf.Write(RPCTagProxyRequest)
 	buf.Write(flags.Bytes())
 	buf.Write(r.ConnectionID[:])
 	buf.Write(r.ClientIPPort[:])
 	buf.Write(r.OurIPPort[:])
-	buf.Write(rpcProxyRequestExtraSize)
-	buf.Write(rpcProxyRequestProxyTag)
+	buf.Write(RPCProxyRequestExtraSize)
+	buf.Write(RPCProxyRequestProxyTag)
 	buf.WriteByte(byte(len(r.ADTag)))
 	buf.Write(r.ADTag)
 	buf.Write(bytes.Repeat([]byte{0x00}, buf.Len()%4))
@@ -69,23 +58,26 @@ func NewRPCProxyRequest(clientAddr, ownAddr *net.TCPAddr, opts *mtproto.Connecti
 	}
 
 	request := RPCProxyRequest{
-		Flags:   flags,
-		ADTag:   adTag,
-		Options: opts,
+		Flags:        flags,
+		ADTag:        adTag,
+		Options:      opts,
+		ConnectionID: make([]byte, 8),
+		ClientIPPort: make([]byte, 16+4),
+		OurIPPort:    make([]byte, 16+4),
 	}
 
-	if _, err := rand.Read(request.ConnectionID[:]); err != nil {
+	if _, err := rand.Read(request.ConnectionID); err != nil {
 		return nil, errors.Annotate(err, "Cannot generate connection ID")
 	}
 
-	port := make([]byte, 4)
+	port := [4]byte{}
 	copy(request.ClientIPPort[:16], clientAddr.IP.To16())
-	binary.LittleEndian.PutUint32(port, uint32(clientAddr.Port))
-	copy(request.ClientIPPort[16:], port)
+	binary.LittleEndian.PutUint32(port[:], uint32(clientAddr.Port))
+	copy(request.ClientIPPort[16:], port[:])
 
 	copy(request.OurIPPort[:16], ownAddr.IP.To16())
-	binary.LittleEndian.PutUint32(port, uint32(ownAddr.Port))
-	copy(request.OurIPPort[16:], port)
+	binary.LittleEndian.PutUint32(port[:], uint32(ownAddr.Port))
+	copy(request.OurIPPort[16:], port[:])
 
 	return &request, nil
 }