add auth to liteclient

liteclient/client.go

@@ -53,7 +53,7 @@ func OptionWorkersPerConnection(n int) Options {
 		}
 		connFirst := c.connections[0]
 		for i := 0; i < n-1; i++ {
-			conn, err := NewConnection(context.Background(), connFirst.peerPublicKey, connFirst.host)
+			conn, err := NewConnection(context.Background(), connFirst.peerPublicKey, connFirst.host, connFirst.authKey)
 			if err != nil {
 				slog.Warn("liteclient clone connection error", err)
 				continue

liteclient/connection.go

@@ -2,15 +2,22 @@ package liteclient
 
 import (
 	"context"
+	"crypto/ed25519"
+	"crypto/rand"
 	"encoding/binary"
-	mrand "math/rand"
+	"fmt"
+	"github.com/tonkeeper/tongo/tl"
 	"sync"
 	"time"
 )
 
 const (
-	magicTCPPing = 0x4d082b9a //crc32(tcp.ping random_id:long = tcp.Pong)
-	magicTCPPong = 0xdc69fb03 //crc32(tcp.pong random_id:long = tcp.Pong)
+	magicTCPPing                     = 0x4d082b9a //crc32(tcp.ping random_id:long = tcp.Ping)
+	magicTCPPong                     = 0xdc69fb03 //crc32(tcp.pong random_id:long = tcp.Pong)
+	magicTcpAuthentificate           = 0x445bab12 //crc32(tcp.authentificate nonce:bytes = tcp.Message)
+	magicTcpAuthentificationNonce    = 0xe35d4ab6 //crc32(tcp.authentificationNonce nonce:bytes = tcp.Message)
+	magicTcpAuthentificationComplete = 0xf7ad9ea6 //crc32(tcp.authentificationComplete key:PublicKey signature:bytes = tcp.Message)
+	magicPubKey                      = 0xc6b41348 //crc32(pub.ed25519#c6b41348 key:int256 = PublicKey)
 )
 
 type ConnectionStatus int
@@ -23,25 +30,40 @@ const (
 	reconnectTimeout = 10 * time.Second
 )
 
+const (
+	ClientNonceSize    = 32
+	MaxServerNonceSize = 512
+)
+
 type Connection struct {
-	peerPublicKey []byte
-	host          string
-	resp          chan Packet
+	peerPublicKey    []byte
+	host             string
+	resp             chan Packet
+	authKey          ed25519.PrivateKey
+	authCompleteChan chan error // Closes when auth is complete or error is sent
 
 	// mu protects all fields below.
 	mu           sync.Mutex
 	status       ConnectionStatus
 	econn        *encryptedConn
 	pings        map[uint64]time.Time
 	avgRoundTrip time.Duration
+	nonce        []byte
 }
 
-func NewConnection(ctx context.Context, peerPublicKey []byte, host string) (*Connection, error) {
+func NewConnection(ctx context.Context, peerPublicKey []byte, host string, authKeys ...ed25519.PrivateKey) (*Connection, error) {
+	if len(authKeys) > 1 {
+		return nil, fmt.Errorf("expected 0 or 1 authentication key, got: %d", len(authKeys))
+	}
 	c := Connection{
-		host:          host,
-		peerPublicKey: peerPublicKey,
-		resp:          make(chan Packet),
-		status:        Connecting,
+		host:             host,
+		peerPublicKey:    peerPublicKey,
+		resp:             make(chan Packet),
+		status:           Connecting,
+		authCompleteChan: make(chan error),
+	}
+	if len(authKeys) == 1 {
+		c.authKey = authKeys[0]
 	}
 	if err := c.setupEncryptedConnection(ctx); err != nil {
 		return nil, err
@@ -56,12 +78,41 @@ func (c *Connection) setupEncryptedConnection(ctx context.Context) error {
 		return err
 	}
 	c.mu.Lock()
-	defer c.mu.Unlock()
-
 	c.econn = econn
-	c.status = Connected
 	c.pings = make(map[uint64]time.Time, 5)
+
 	go c.reader(econn.handleIncomingPackets())
+
+	if c.authKey == nil {
+		c.status = Connected
+		c.mu.Unlock()
+		return nil
+	}
+
+	c.mu.Unlock()
+
+	if err := c.sendAuthRequest(); err != nil {
+		c.econn.close()
+		return fmt.Errorf("failed to send auth request: %w", err)
+	}
+
+	timer := time.NewTimer(10 * time.Second)
+	defer timer.Stop()
+
+	select {
+	case authErr := <-c.authCompleteChan:
+		if authErr != nil {
+			err = authErr
+		}
+	case <-timer.C:
+		err = fmt.Errorf("authentication timeout")
+	case <-ctx.Done():
+		err = ctx.Err()
+	}
+	if err != nil {
+		c.econn.close()
+		return err
+	}
 	return nil
 }
 
@@ -77,6 +128,7 @@ func (c *Connection) reconnect() {
 
 	for {
 		if err := c.setupEncryptedConnection(context.Background()); err != nil {
+			fmt.Printf("error reconnecting to %s: %s\n", c.host, err)
 			time.Sleep(1 * time.Second)
 			continue
 		}
@@ -116,6 +168,14 @@ func (c *Connection) reader(packetCh chan Packet) {
 				// the next iteration of this for loop will restart reconnect timeout.
 				continue
 			}
+
+			if p.MagicType() == magicTcpAuthentificationNonce {
+				if err := c.handleAuthResponse(p); err != nil {
+					fmt.Printf("failed to handle auth nonce: %s\n", err)
+				}
+				continue
+			}
+
 			c.resp <- p
 
 		case <-time.After(reconnectTimeout):
@@ -150,7 +210,9 @@ func (c *Connection) ping() {
 	binary.LittleEndian.PutUint32(ping[:4], magicTCPPing)
 	for {
 		time.Sleep(time.Second * 3)
-		mrand.Read(ping[4:])
+		if _, err := rand.Read(ping[4:]); err != nil {
+			panic(err) // impossible if source of randomness is correct
+		}
 		p, err := NewPacket(ping)
 		if err != nil {
 			panic(err) // impossible if NewPacket function is correct
@@ -203,3 +265,85 @@ func (c *Connection) Status() ConnectionStatus {
 	defer c.mu.Unlock()
 	return c.status
 }
+
+func (c *Connection) sendAuthRequest() error {
+	nonce := make([]byte, ClientNonceSize)
+	if _, err := rand.Read(nonce); err != nil {
+		return fmt.Errorf("error generating nonce: %w", err)
+	}
+
+	c.mu.Lock()
+	c.nonce = nonce
+	c.mu.Unlock()
+
+	payload := make([]byte, 4)
+	binary.LittleEndian.PutUint32(payload, magicTcpAuthentificate)
+	payload = append(payload, tl.EncodeLength(len(nonce))...)
+	payload = append(payload, nonce...)
+	payload = alignBytes(payload)
+
+	p, err := NewPacket(payload)
+	if err != nil {
+		return fmt.Errorf("failed to create auth packet: %w", err)
+	}
+	b := p.marshal()
+	return c.econn.send(b)
+}
+
+func (c *Connection) handleAuthResponse(p Packet) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.status != Connecting {
+		return fmt.Errorf("received unexpected auth packet")
+	}
+
+	err := c.sendAuthComplete(p)
+	if err == nil {
+		c.status = Connected
+	}
+	c.authCompleteChan <- err
+	return nil
+}
+
+func (c *Connection) sendAuthComplete(received Packet) error {
+	if c.authKey == nil {
+		return fmt.Errorf("no auth key available")
+	}
+
+	if len(received.Payload) < 37 {
+		return fmt.Errorf("too short payload")
+	}
+	length, data, err := decodeLength(received.Payload[4:])
+	if err != nil {
+		return err
+	}
+	if len(data) < length {
+		return fmt.Errorf("payload is smaller than should be according to length")
+	}
+	nonce := data[:length]
+	if len(nonce) > MaxServerNonceSize {
+		return fmt.Errorf("too long nonce")
+	}
+
+	signature := ed25519.Sign(c.authKey, append(append([]byte{}, c.nonce...), nonce...))
+
+	payload := make([]byte, 4)
+	binary.LittleEndian.PutUint32(payload, magicTcpAuthentificationComplete)
+	binary.LittleEndian.PutUint32(payload, magicPubKey)
+	pubKey := c.authKey.Public().(ed25519.PublicKey)
+	payload = append(payload, pubKey[:]...)
+	payload = append(payload, tl.EncodeLength(len(signature))...)
+	payload = append(payload, signature...)
+	payload = alignBytes(payload)
+
+	p, err := NewPacket(payload)
+	if err != nil {
+		return fmt.Errorf("failed to create packet: %w", err)
+	}
+	b := p.marshal()
+	if err := c.econn.send(b); err != nil {
+		return fmt.Errorf("failed to send auth sign request, err: %w", err)
+	}
+	return nil
+}