Merge pull request #156 from tdakkota/feature/pool-and-uploader

Add text formatters, add some crypto checks

internal/crypto/cipher_decrypt.go

@@ -43,7 +43,13 @@ func (c Cipher) Decrypt(k AuthKey, encrypted *EncryptedMessage) (*EncryptedMessa
 		const maxPadding = 1024
 		n := int(msg.MessageDataLen)
 		paddingLen := len(msg.MessageDataWithPadding) - n
-		if paddingLen > maxPadding {
+
+		switch {
+		case n < 0:
+			return nil, xerrors.Errorf("message length is invalid: %d less than zero", n)
+		case n%4 != 0:
+			return nil, xerrors.Errorf("message length is invalid: %d is not divisible by 4", n)
+		case paddingLen > maxPadding:
 			return nil, xerrors.Errorf("padding %d of message is too big", paddingLen)
 		}
 	}

internal/crypto/cipher_decrypt_test.go

@@ -2,9 +2,13 @@ package crypto
 
 import (
 	"bytes"
+	"io"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/gotd/td/bin"
+	"github.com/gotd/td/internal/testutil"
 )
 
 type Zero struct{}
@@ -46,3 +50,42 @@ func TestDecrypt(t *testing.T) {
 		t.Error("mismatch")
 	}
 }
+
+func TestCipher_Decrypt(t *testing.T) {
+	var key AuthKey
+	if _, err := io.ReadFull(testutil.Rand([]byte{10}), key.Value[:]); err != nil {
+		t.Fatal(err)
+	}
+
+	c := NewClientCipher(Zero{})
+	s := NewServerCipher(Zero{})
+	tests := []struct {
+		name      string
+		data      []byte
+		dataLen   int
+		expectErr bool
+	}{
+		{"NegativeLength", []byte{1, 2, 3, 4}, -1, true},
+		{"NoPadBy4", []byte{1, 2, 3}, 3, true},
+		{"Good", bytes.Repeat([]byte{1, 2, 3, 4}, 4), 16, false},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			a := require.New(t)
+			b := bin.Buffer{}
+			data := EncryptedMessageData{
+				MessageDataLen:         int32(test.dataLen),
+				MessageDataWithPadding: test.data,
+			}
+			a.NoError(s.Encrypt(key, data, &b))
+
+			_, err := c.DecryptFromBuffer(key, &b)
+			if test.expectErr {
+				a.Error(err)
+				return
+			}
+			a.NoError(err)
+		})
+	}
+}

internal/exchange/client_flow.go

@@ -7,14 +7,12 @@ import (
 	"math/big"
 
 	"go.uber.org/zap"
-
-	"github.com/gotd/td/internal/proto"
-
 	"golang.org/x/xerrors"
 
 	"github.com/gotd/td/bin"
 	"github.com/gotd/td/internal/crypto"
 	"github.com/gotd/td/internal/mt"
+	"github.com/gotd/td/internal/proto"
 )
 
 // Run runs client-side flow.
@@ -41,6 +39,7 @@ func (c ClientExchange) Run(ctx context.Context) (ClientExchangeResult, error) {
 	if res.Nonce != nonce {
 		return ClientExchangeResult{}, xerrors.New("ResPQ nonce mismatch")
 	}
+	serverNonce := res.ServerNonce
 
 	// Selecting first public key that match fingerprint.
 	var selectedPubKey *rsa.PublicKey
@@ -89,7 +88,7 @@ Loop:
 		Pq:          res.Pq,
 		Nonce:       nonce,
 		NewNonce:    newNonce,
-		ServerNonce: res.ServerNonce,
+		ServerNonce: serverNonce,
 		P:           pBytes,
 		Q:           qBytes,
 	}
@@ -105,7 +104,7 @@ Loop:
 	}
 	reqDHParams := &mt.ReqDHParamsRequest{
 		Nonce:                nonce,
-		ServerNonce:          res.ServerNonce,
+		ServerNonce:          serverNonce,
 		P:                    pBytes,
 		Q:                    qBytes,
 		PublicKeyFingerprint: crypto.RSAFingerprint(selectedPubKey),
@@ -138,8 +137,11 @@ Loop:
 		if p.Nonce != nonce {
 			return ClientExchangeResult{}, xerrors.New("ServerDHParamsOk nonce mismatch")
 		}
+		if p.ServerNonce != serverNonce {
+			return ClientExchangeResult{}, xerrors.New("ServerDHParamsOk server nonce mismatch")
+		}
 
-		key, iv := crypto.TempAESKeys(newNonce.BigInt(), res.ServerNonce.BigInt())
+		key, iv := crypto.TempAESKeys(newNonce.BigInt(), serverNonce.BigInt())
 		// Decrypting inner data.
 		data, err := crypto.DecryptExchangeAnswer(p.EncryptedAnswer, key, iv)
 		if err != nil {
@@ -151,6 +153,12 @@ Loop:
 		if err := innerData.Decode(b); err != nil {
 			return ClientExchangeResult{}, err
 		}
+		if innerData.Nonce != nonce {
+			return ClientExchangeResult{}, xerrors.New("ServerDHInnerData nonce mismatch")
+		}
+		if innerData.ServerNonce != serverNonce {
+			return ClientExchangeResult{}, xerrors.New("ServerDHInnerData server nonce mismatch")
+		}
 
 		dhPrime := big.NewInt(0).SetBytes(innerData.DhPrime)
 		g := big.NewInt(int64(innerData.G))
@@ -215,6 +223,13 @@ Loop:
 		}
 		switch v := dhSetRes.(type) {
 		case *mt.DhGenOk: // dh_gen_ok#3bcbf734
+			if v.Nonce != nonce {
+				return ClientExchangeResult{}, xerrors.New("DhGenOk nonce mismatch")
+			}
+			if v.ServerNonce != serverNonce {
+				return ClientExchangeResult{}, xerrors.New("DhGenOk server nonce mismatch")
+			}
+
 			var key crypto.Key
 			authKey.FillBytes(key[:])
 			authKeyID := key.ID()

internal/exchange/client_flow_test.go

@@ -0,0 +1,47 @@
+package exchange
+
+import (
+	"context"
+	"crypto/rsa"
+	"math/rand"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap/zaptest"
+
+	"github.com/gotd/td/internal/tdsync"
+	"github.com/gotd/td/transport"
+)
+
+func TestExchangeTimeout(t *testing.T) {
+	a := require.New(t)
+
+	reader := rand.New(rand.NewSource(1))
+	key, err := rsa.GenerateKey(reader, 2048)
+	a.NoError(err)
+	log := zaptest.NewLogger(t)
+
+	i := transport.Intermediate(nil)
+	client, _ := i.Pipe()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	grp := tdsync.NewCancellableGroup(ctx)
+	grp.Go(func(groupCtx context.Context) error {
+		_, err := NewExchanger(client).
+			WithLogger(log.Named("client")).
+			WithRand(reader).
+			WithTimeout(1 * time.Second).
+			Client([]*rsa.PublicKey{&key.PublicKey}).
+			Run(groupCtx)
+		return err
+	})
+
+	err = grp.Wait()
+	if err, ok := err.(net.Error); !ok || !err.Timeout() {
+		require.NoError(t, err)
+	}
+}

internal/exchange/flow_test.go

@@ -5,7 +5,6 @@ import (
 	"crypto/rsa"
 	"fmt"
 	"math/rand"
-	"net"
 	"testing"
 	"time"
 
@@ -53,37 +52,6 @@ func TestExchange(t *testing.T) {
 	require.NoError(t, grp.Wait())
 }
 
-func TestExchangeTimeout(t *testing.T) {
-	a := require.New(t)
-
-	reader := rand.New(rand.NewSource(1))
-	key, err := rsa.GenerateKey(reader, 2048)
-	a.NoError(err)
-	log := zaptest.NewLogger(t)
-
-	i := transport.Intermediate(nil)
-	client, _ := i.Pipe()
-
-	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
-	defer cancel()
-
-	grp := tdsync.NewCancellableGroup(ctx)
-	grp.Go(func(groupCtx context.Context) error {
-		_, err := NewExchanger(client).
-			WithLogger(log.Named("client")).
-			WithRand(reader).
-			WithTimeout(1 * time.Second).
-			Client([]*rsa.PublicKey{&key.PublicKey}).
-			Run(groupCtx)
-		return err
-	})
-
-	err = grp.Wait()
-	if err, ok := err.(net.Error); !ok || !err.Timeout() {
-		require.NoError(t, err)
-	}
-}
-
 func TestExchangeCorpus(t *testing.T) {
 	k := testutil.RSAPrivateKey()
 

telegram/internal/tgtest/loop.go

@@ -3,16 +3,12 @@ package tgtest
 import (
 	"context"
 	"encoding/binary"
-	"io"
-	"net"
 
-	"go.uber.org/zap"
 	"golang.org/x/xerrors"
 
-	"github.com/gotd/td/internal/mt"
-
 	"github.com/gotd/td/bin"
 	"github.com/gotd/td/internal/crypto"
+	"github.com/gotd/td/internal/mt"
 	"github.com/gotd/td/transport"
 )
 
@@ -23,19 +19,9 @@ type Session struct {
 
 func (s *Server) rpcHandle(ctx context.Context, conn *connection) error {
 	var b bin.Buffer
-	var key crypto.AuthKey
 	for {
 		b.Reset()
 		if err := conn.Recv(ctx, &b); err != nil {
-			var syscallErr *net.OpError
-			// TODO(tdakkota): Find a better way to detect forcibly closed connection.
-			if xerrors.Is(err, io.EOF) || (xerrors.As(err, &syscallErr) && syscallErr.Op == "read") {
-				// Client disconnected.
-				s.users.deleteConnection(key)
-				s.log.Info("Read failed, closing loop.", zap.Error(err))
-				return nil
-			}
-
 			return err
 		}
 
@@ -101,11 +87,6 @@ func (s *Server) handle(session Session, msgID int64, in *bin.Buffer) error {
 	}
 
 	if err := s.handler.OnMessage(session, msgID, in); err != nil {
-		// Client disconnected during write.
-		var syscallErr *net.OpError
-		if xerrors.As(err, &syscallErr) && syscallErr.Op == "write" {
-			return nil
-		}
 		return xerrors.Errorf("failed to call handler: %w", err)
 	}
 

telegram/internal/tgtest/server.go

@@ -5,9 +5,11 @@ import (
 	"context"
 	"crypto/rand"
 	"crypto/rsa"
+	"io"
 	"net"
 
 	"go.uber.org/zap"
+	"golang.org/x/xerrors"
 
 	"github.com/gotd/td/bin"
 	"github.com/gotd/td/clock"
@@ -105,6 +107,13 @@ func (s *Server) serve() error {
 	return s.server.Serve(s.ctx, func(ctx context.Context, conn transport.Conn) error {
 		err := s.serveConn(ctx, conn)
 		if err != nil {
+			// Client disconnected.
+			var syscallErr *net.OpError
+			if xerrors.Is(err, io.EOF) ||
+				xerrors.As(err, &syscallErr) && syscallErr.Op == "write" ||
+				syscallErr.Op == "read" {
+				return nil
+			}
 			s.log.Info("Serving handler error", zap.Error(err))
 		}
 		return err

telegram/message/doc.go

@@ -0,0 +1,2 @@
+// Package message contains some useful utilities for creating Telegram messages.
+package message