perf: sync protocol and core hardening updates

internal/mtprotoedge/exchange_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"crypto/rsa"
+	"encoding/binary"
 	"errors"
 	"net"
 	"testing"
@@ -106,6 +107,214 @@ func TestKeyExchange(t *testing.T) {
 	}
 }
 
+type authKeySaveContextObservation struct {
+	hasDeadline bool
+	deadline    time.Time
+}
+
+type observingAuthKeyStore struct {
+	store.AuthKeyStore
+	saveContext chan authKeySaveContextObservation
+}
+
+func (s *observingAuthKeyStore) Save(ctx context.Context, key store.AuthKeyData) error {
+	deadline, hasDeadline := ctx.Deadline()
+	select {
+	case s.saveContext <- authKeySaveContextObservation{hasDeadline: hasDeadline, deadline: deadline}:
+	default:
+	}
+	return s.AuthKeyStore.Save(ctx, key)
+}
+
+type gatedAuthKeyStore struct {
+	store.AuthKeyStore
+	entered chan store.AuthKeyData
+	release chan struct{}
+	saveErr error
+}
+
+type ownershipFrameConn struct {
+	transport.Conn
+	frame []byte
+}
+
+func (c *ownershipFrameConn) Recv(_ context.Context, b *bin.Buffer) error {
+	b.ResetTo(c.frame)
+	return nil
+}
+
+func TestExchangeEncryptedReplayTransfersFrameOwnership(t *testing.T) {
+	backing := make([]byte, 64)
+	copy(backing[:8], []byte{1, 2, 3, 4, 5, 6, 7, 8})
+	conn := &ownershipFrameConn{frame: backing}
+	ex := serverExchangeCompat{conn: conn, timeout: time.Second}
+	var b bin.Buffer
+	err := ex.readUnencrypted(context.Background(), &b, &compatReqPQ{})
+	var encrypted *exchange.UnexpectedEncryptedError
+	if !errors.As(err, &encrypted) {
+		t.Fatalf("read encrypted frame err = %v, want UnexpectedEncryptedError", err)
+	}
+	if len(encrypted.Frame) != len(backing) || &encrypted.Frame[0] != &backing[0] {
+		t.Fatal("encrypted replay copied the received frame instead of transferring ownership")
+	}
+	if b.Buf != nil {
+		t.Fatal("exchange buffer retained transferred encrypted frame backing")
+	}
+}
+
+func (s *gatedAuthKeyStore) Save(ctx context.Context, key store.AuthKeyData) error {
+	select {
+	case s.entered <- key:
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+	select {
+	case <-s.release:
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+	if s.saveErr != nil {
+		return s.saveErr
+	}
+	return s.AuthKeyStore.Save(ctx, key)
+}
+
+// TestKeyExchangeDoesNotAcknowledgeBeforeAuthKeyCommit pins the protocol commit
+// boundary: while durable Save is blocked, the client must not receive DhGenOk
+// and therefore must not report a successful exchange.
+func TestKeyExchangeDoesNotAcknowledgeBeforeAuthKeyCommit(t *testing.T) {
+	base := memory.NewAuthKeyStore()
+	keys := &gatedAuthKeyStore{
+		AuthKeyStore: base,
+		entered:      make(chan store.AuthKeyData, 1),
+		release:      make(chan struct{}, 1),
+	}
+	addr, pub, _ := startTestServer(t, Options{DC: 2, AuthKeys: keys})
+	conn := dialTransportOnly(t, addr)
+
+	type exchangeOutcome struct {
+		result exchange.ClientExchangeResult
+		err    error
+	}
+	outcome := make(chan exchangeOutcome, 1)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	go func() {
+		result, err := exchange.NewExchanger(conn, 2).
+			WithRand(rand.Reader).
+			Client([]exchange.PublicKey{pub}).
+			Run(ctx)
+		outcome <- exchangeOutcome{result: result, err: err}
+	}()
+
+	var pending store.AuthKeyData
+	select {
+	case pending = <-keys.entered:
+	case <-time.After(5 * time.Second):
+		t.Fatal("AuthKeyStore.Save was not reached")
+	}
+	defer func() {
+		select {
+		case keys.release <- struct{}{}:
+		default:
+		}
+	}()
+
+	select {
+	case got := <-outcome:
+		t.Fatalf("client exchange completed before auth key commit: err=%v", got.err)
+	case <-time.After(150 * time.Millisecond):
+	}
+	if _, found, err := base.Get(context.Background(), pending.ID); err != nil {
+		t.Fatalf("Get before commit: %v", err)
+	} else if found {
+		t.Fatal("auth key became visible while durable Save was blocked")
+	}
+
+	keys.release <- struct{}{}
+	select {
+	case got := <-outcome:
+		if got.err != nil {
+			t.Fatalf("client exchange after commit: %v", got.err)
+		}
+		if got.result.AuthKey.ID != pending.ID {
+			t.Fatalf("committed auth key id = %x, client got %x", pending.ID, got.result.AuthKey.ID)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("client exchange did not finish after auth key commit")
+	}
+	if _, found, err := base.Get(context.Background(), pending.ID); err != nil {
+		t.Fatalf("Get after commit: %v", err)
+	} else if !found {
+		t.Fatal("auth key is not durable after successful client exchange")
+	}
+}
+
+// TestKeyExchangeAuthKeyCommitFailureWithholdsDhGenOk proves the failure side
+// of the same invariant. The client must not observe success if storage rejects
+// the key; the server closes this exchange and lets the client retry cleanly.
+func TestKeyExchangeAuthKeyCommitFailureWithholdsDhGenOk(t *testing.T) {
+	base := memory.NewAuthKeyStore()
+	keys := &gatedAuthKeyStore{
+		AuthKeyStore: base,
+		entered:      make(chan store.AuthKeyData, 1),
+		release:      make(chan struct{}, 1),
+		saveErr:      errors.New("injected auth key persistence failure"),
+	}
+	keys.release <- struct{}{}
+	addr, pub, _ := startTestServer(t, Options{DC: 2, AuthKeys: keys})
+	conn := dialTransportOnly(t, addr)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	_, err := exchange.NewExchanger(conn, 2).
+		WithRand(rand.Reader).
+		Client([]exchange.PublicKey{pub}).
+		Run(ctx)
+	if err == nil {
+		t.Fatal("client exchange succeeded even though auth key commit failed")
+	}
+
+	select {
+	case attempted := <-keys.entered:
+		if _, found, getErr := base.Get(context.Background(), attempted.ID); getErr != nil {
+			t.Fatalf("Get failed key: %v", getErr)
+		} else if found {
+			t.Fatal("failed auth key commit became visible")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("AuthKeyStore.Save was not attempted")
+	}
+}
+
+func TestKeyExchangeAuthKeySaveUsesHandshakeDeadline(t *testing.T) {
+	const handshakeMax = 10 * time.Second
+	observed := make(chan authKeySaveContextObservation, 1)
+	keys := &observingAuthKeyStore{
+		AuthKeyStore: memory.NewAuthKeyStore(),
+		saveContext:  observed,
+	}
+	addr, pub, _ := startTestServer(t, Options{
+		DC:                   2,
+		AuthKeys:             keys,
+		HandshakeMaxDuration: handshakeMax,
+	})
+
+	_, _, _ = dialHandshake(t, addr, 2, pub)
+	select {
+	case got := <-observed:
+		if !got.hasDeadline {
+			t.Fatal("AuthKeyStore.Save context has no handshake deadline")
+		}
+		remaining := time.Until(got.deadline)
+		if remaining <= 0 || remaining > handshakeMax {
+			t.Fatalf("AuthKeyStore.Save deadline remaining = %v, want (0, %v]", remaining, handshakeMax)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("AuthKeyStore.Save was not called")
+	}
+}
+
 func TestKeyExchangeAcceptsAndroidMediaTempNegativeDC(t *testing.T) {
 	const dc = 2
 	addr, pub, srv := startTestServer(t, Options{DC: dc})
@@ -227,6 +436,82 @@ func TestKeyExchangeIgnoresUnencryptedMsgsAck(t *testing.T) {
 	}
 }
 
+func TestBufferedExchangePushTransfersFrameOwnershipWithoutCopy(t *testing.T) {
+	backing := make([]byte, 64)
+	for i := range backing {
+		backing[i] = byte(i)
+	}
+	source := &bin.Buffer{Buf: backing}
+	buffered := newBufferedConn(nil)
+	buffered.push(source)
+	if source.Buf != nil {
+		t.Fatal("push retained ownership in the source buffer")
+	}
+
+	var got bin.Buffer
+	if err := buffered.Recv(context.Background(), &got); err != nil {
+		t.Fatalf("Recv pending frame: %v", err)
+	}
+	if len(got.Buf) != len(backing) || &got.Buf[0] != &backing[0] {
+		t.Fatal("pending frame was copied instead of transferring its backing")
+	}
+	if len(buffered.pending) != 0 || cap(buffered.pending) != 0 {
+		t.Fatalf("consumed pending ownership retained: len=%d cap=%d", len(buffered.pending), cap(buffered.pending))
+	}
+}
+
+func TestBufferedExchangeLargeTrailingMsgsAckReleasesFrameBeforeNextRecv(t *testing.T) {
+	encodeUnencrypted := func(msg bin.Encoder, msgID int64) []byte {
+		var payload bin.Buffer
+		if err := msg.Encode(&payload); err != nil {
+			t.Fatalf("encode payload: %v", err)
+		}
+		var frame bin.Buffer
+		if err := (tgproto.UnencryptedMessage{MessageID: msgID, MessageData: payload.Raw()}).Encode(&frame); err != nil {
+			t.Fatalf("encode unencrypted frame: %v", err)
+		}
+		return frame.Copy()
+	}
+	intermediate := func(frame []byte) []byte {
+		packet := make([]byte, bin.Word+len(frame))
+		binary.LittleEndian.PutUint32(packet, uint32(len(frame)))
+		copy(packet[bin.Word:], frame)
+		return packet
+	}
+
+	// Make the ignored ack larger than the per-codec retained-buffer threshold. The following
+	// small req_pq frame forces bufferedConn to cross the next-Recv ownership boundary while the
+	// same destination bin.Buffer is reused.
+	ids := make([]int64, 300_000)
+	for i := range ids {
+		ids[i] = int64(i + 1)
+	}
+	ackFrame := encodeUnencrypted(&mt.MsgsAck{MsgIDs: ids}, 4)
+	reqFrame := encodeUnencrypted(&mt.ReqPqMultiRequest{}, 8)
+	packet := append(intermediate(ackFrame), intermediate(reqFrame)...)
+	budget := newInboundFrameBudget(2 * int64(len(ackFrame)))
+	conn, _ := newFrameBudgetTestTransport(packet, &quickAckIntermediateCodec{}, budget)
+	buffered := newBufferedConn(conn)
+
+	var got bin.Buffer
+	if err := buffered.Recv(context.Background(), &got); err != nil {
+		t.Fatalf("Recv after large msgs_ack: %v", err)
+	}
+	if id, ok := unencryptedPayloadID(&got); !ok || id != mt.ReqPqMultiRequestTypeID {
+		t.Fatalf("returned frame type = 0x%x ok=%v, want req_pq_multi", id, ok)
+	}
+	if used, want := budget.usedBytes(), 2*int64(len(reqFrame)); used != want {
+		t.Fatalf("inbound budget after skipped ack = %d, want only next frame %d", used, want)
+	}
+	if cap(got.Buf) >= len(ackFrame)/2 {
+		t.Fatalf("large ignored ack backing retained by next frame: cap=%d ack=%d", cap(got.Buf), len(ackFrame))
+	}
+	conn.releaseInboundFrame()
+	if used := budget.usedBytes(); used != 0 {
+		t.Fatalf("inbound budget after final ownership release = %d, want 0", used)
+	}
+}
+
 type ackingExchangeConn struct {
 	transport.Conn
 	t *testing.T