Warp/server/internal/federationapi/queue_test.go

package federationapi

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"path/filepath"
	"slices"
	"sync"
	"testing"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/federation"
	"github.com/dariuszjaniak/warp/server/internal/matrix/event"
	"github.com/dariuszjaniak/warp/server/internal/matrix/roomversion"
	"github.com/dariuszjaniak/warp/server/internal/matrix/signing"
	"github.com/dariuszjaniak/warp/server/internal/roomserver"
	"github.com/dariuszjaniak/warp/server/internal/storage"
	_ "github.com/dariuszjaniak/warp/server/internal/storage/drivers"
)

const ourName = "warp.example"

// fakeSender records transactions and fails on demand.
type fakeSender struct {
	mu   sync.Mutex
	txns map[string][]federation.Transaction
	ids  map[string]bool
	down map[string]bool
}

func newFakeSender() *fakeSender {
	return &fakeSender{txns: map[string][]federation.Transaction{}, ids: map[string]bool{}, down: map[string]bool{}}
}

func (f *fakeSender) SendTransaction(_ context.Context, destination, txnID string, txn federation.Transaction) (federation.SendTransactionResponse, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.down[destination] {
		return federation.SendTransactionResponse{}, errors.New("connection refused")
	}
	if f.ids[txnID] {
		return federation.SendTransactionResponse{}, fmt.Errorf("transaction ID %s reused", txnID)
	}
	f.ids[txnID] = true
	f.txns[destination] = append(f.txns[destination], txn)
	return federation.SendTransactionResponse{PDUs: map[string]federation.PDUResult{}}, nil
}

func (f *fakeSender) setDown(destination string, down bool) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.down[destination] = down
}

// pdus returns how many PDUs a destination has received.
func (f *fakeSender) pdus(destination string) int {
	f.mu.Lock()
	defer f.mu.Unlock()
	n := 0
	for _, txn := range f.txns[destination] {
		n += len(txn.PDUs)
	}
	return n
}

func (f *fakeSender) largestTxn(destination string) int {
	f.mu.Lock()
	defer f.mu.Unlock()
	largest := 0
	for _, txn := range f.txns[destination] {
		largest = max(largest, len(txn.PDUs))
	}
	return largest
}

func (f *fakeSender) edus(destination string) []federation.EDU {
	f.mu.Lock()
	defer f.mu.Unlock()
	var out []federation.EDU
	for _, txn := range f.txns[destination] {
		out = append(out, txn.EDUs...)
	}
	return out
}

// fakeRooms answers JoinedServers from a table.
type fakeRooms map[string][]string

func (f fakeRooms) JoinedServers(_ context.Context, roomID string) ([]string, error) {
	return f[roomID], nil
}

func openStore(t *testing.T) storage.FederationStore {
	t.Helper()
	db, err := storage.Open(t.Context(), storage.Config{
		Driver:           storage.DriverSQLite,
		ConnectionString: "file:" + filepath.ToSlash(filepath.Join(t.TempDir(), "q.db")) + "?_pragma=busy_timeout(5000)",
	}, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })
	if err := db.Migrate(t.Context()); err != nil {
		t.Fatal(err)
	}
	return db.Federation()
}

func newTestQueue(t *testing.T, store storage.FederationStore, sender Sender, rooms RoomServers, policy *federation.Policy) *Queue {
	t.Helper()
	q := NewQueue(QueueOptions{
		ServerName: ourName,
		Store:      store,
		Sender:     sender,
		Rooms:      rooms,
		Policy:     policy,
		MinBackoff: 20 * time.Millisecond,
		MaxBackoff: 100 * time.Millisecond,
		Logger:     slog.New(slog.DiscardHandler),
	})
	t.Cleanup(q.Stop)
	return q
}

var testKey = func() *signing.KeyPair {
	k, err := signing.Generate("q")
	if err != nil {
		panic(err)
	}
	return k
}()

// notification builds a persisted-event notification for a message, or a
// membership event about target when target is set.
func notification(t *testing.T, roomID, sender, target string, n int) roomserver.Notification {
	t.Helper()
	b := event.Builder{
		RoomID: roomID, Type: "m.room.message", Sender: sender,
		Content:    json.RawMessage(fmt.Sprintf(`{"body":"%d"}`, n)),
		PrevEvents: []string{"$prev"}, AuthEvents: []string{"$auth"}, Depth: int64(n + 2),
	}
	if target != "" {
		b.Type = event.TypeMember
		b.StateKey = &target
		b.Content = json.RawMessage(`{"membership":"leave"}`)
	}
	e, err := b.Build(time.Unix(1_700_000_000+int64(n), 0), testKey, roomversion.MustGet("11"))
	if err != nil {
		t.Fatal(err)
	}
	return roomserver.Notification{RoomID: roomID, Event: e, StreamOrdering: int64(n)}
}

func waitFor(t *testing.T, what string, check func() bool) {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		if check() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

func TestQueueSendsToTheRoomsServers(t *testing.T) {
	sender := newFakeSender()
	rooms := fakeRooms{"!r:warp.example": {ourName, "a.example", "b.example"}}
	q := newTestQueue(t, openStore(t), sender, rooms, nil)
	if err := q.Start(t.Context()); err != nil {
		t.Fatal(err)
	}

	q.EventPersisted(t.Context(), notification(t, "!r:warp.example", "@me:"+ourName, "", 1))
	// Received from b and not to be relayed: b sends its own events.
	received := notification(t, "!r:warp.example", "@x:b.example", "", 2)
	received.Origin = "b.example"
	q.EventPersisted(t.Context(), received)
	// A join handed over by send_join from c: relayed to everybody but c.
	relayed := notification(t, "!r:warp.example", "@y:c.example", "", 3)
	relayed.Origin, relayed.Relay = "c.example", true
	q.EventPersisted(t.Context(), relayed)
	// A kick of a user on d, whose server is no longer joined but must
	// hear of it.
	q.EventPersisted(t.Context(), notification(t, "!r:warp.example", "@me:"+ourName, "@z:d.example", 4))

	waitFor(t, "deliveries", func() bool {
		return sender.pdus("a.example") == 3 && sender.pdus("b.example") == 3 && sender.pdus("d.example") == 1
	})
	if n := sender.pdus("c.example"); n != 0 {
		t.Errorf("c received %d PDUs; it is not in the room and its own join was not to be echoed", n)
	}
	if n := sender.pdus(ourName); n != 0 {
		t.Errorf("the queue sent %d PDUs to this server", n)
	}
}

func TestQueueBacksOffAndRetries(t *testing.T) {
	sender := newFakeSender()
	sender.setDown("a.example", true)
	store := openStore(t)
	q := newTestQueue(t, store, sender, fakeRooms{"!r:warp.example": {"a.example"}}, nil)
	if err := q.Start(t.Context()); err != nil {
		t.Fatal(err)
	}

	for i := range 120 {
		q.EventPersisted(t.Context(), notification(t, "!r:warp.example", "@me:"+ourName, "", i))
	}
	waitFor(t, "a recorded failure", func() bool {
		d, err := store.Destination(t.Context(), "a.example")
		return err == nil && d.Failures >= 2
	})

	sender.setDown("a.example", false)
	waitFor(t, "every PDU delivered", func() bool { return sender.pdus("a.example") == 120 })
	if largest := sender.largestTxn("a.example"); largest > maxPDUsPerTxn {
		t.Errorf("a transaction carried %d PDUs, the limit is %d", largest, maxPDUsPerTxn)
	}
	waitFor(t, "the failure count to clear", func() bool {
		d, err := store.Destination(t.Context(), "a.example")
		return err == nil && d.Failures == 0 && d.LastSuccessMS > 0
	})
	left, err := store.Queued(t.Context(), "a.example", storage.QueueKindPDU, 10)
	if err != nil || len(left) != 0 {
		t.Errorf("queue after delivery = %d items, %v", len(left), err)
	}
}

func TestQueueSurvivesARestart(t *testing.T) {
	sender := newFakeSender()
	sender.setDown("a.example", true)
	store := openStore(t)
	rooms := fakeRooms{"!r:warp.example": {"a.example"}}

	first := newTestQueue(t, store, sender, rooms, nil)
	if err := first.Start(t.Context()); err != nil {
		t.Fatal(err)
	}
	for i := range 5 {
		first.EventPersisted(t.Context(), notification(t, "!r:warp.example", "@me:"+ourName, "", i))
	}
	waitFor(t, "the events to be queued", func() bool {
		items, _ := store.Queued(t.Context(), "a.example", storage.QueueKindPDU, 10)
		return len(items) == 5
	})
	first.Stop()

	// While stopped, events are dropped rather than held.
	first.EventPersisted(t.Context(), notification(t, "!r:warp.example", "@me:"+ourName, "", 99))

	sender.setDown("a.example", false)
	if err := store.SetDestination(t.Context(), storage.Destination{ServerName: "a.example"}); err != nil {
		t.Fatal(err)
	}
	second := newTestQueue(t, store, sender, rooms, nil)
	if err := second.Start(t.Context()); err != nil {
		t.Fatal(err)
	}
	waitFor(t, "the backlog after a restart", func() bool { return sender.pdus("a.example") == 5 })
	time.Sleep(50 * time.Millisecond)
	if n := sender.pdus("a.example"); n != 5 {
		t.Errorf("a received %d PDUs; the one persisted while stopped should not have been queued", n)
	}
}

func TestQueueHonoursThePolicy(t *testing.T) {
	sender := newFakeSender()
	policy := federation.NewPolicy(nil, []string{"denied.example"})
	q := newTestQueue(t, openStore(t), sender, fakeRooms{"!r:warp.example": {"ok.example", "denied.example"}}, policy)
	if err := q.Start(t.Context()); err != nil {
		t.Fatal(err)
	}
	q.EventPersisted(t.Context(), notification(t, "!r:warp.example", "@me:"+ourName, "", 1))
	waitFor(t, "delivery to the allowed server", func() bool { return sender.pdus("ok.example") == 1 })
	if n := sender.pdus("denied.example"); n != 0 {
		t.Errorf("a denied server received %d PDUs", n)
	}
}

func TestEDUsGoToTheRoomsServers(t *testing.T) {
	sender := newFakeSender()
	rooms := fakeRooms{"!r:warp.example": {ourName, "a.example"}}
	q := newTestQueue(t, openStore(t), sender, rooms, nil)
	if err := q.Start(t.Context()); err != nil {
		t.Fatal(err)
	}
	edus := NewEDUs(q, rooms)
	edus.Typing(t.Context(), "!r:warp.example", "@me:"+ourName, true, 0)
	edus.Receipt(t.Context(), "!r:warp.example", "@me:"+ourName, "m.read", "$e", "", 5)
	edus.Receipt(t.Context(), "!r:warp.example", "@me:"+ourName, "m.read.private", "$e", "", 5)
	edus.ToDevice(t.Context(), "@me:"+ourName, "m.room_key", map[string]map[string]json.RawMessage{
		"@u:a.example": {"DEV": json.RawMessage(`{"k":1}`)},
		"@v:b.example": {"*": json.RawMessage(`{"k":2}`)},
	})

	waitFor(t, "EDUs", func() bool { return len(sender.edus("a.example")) == 3 && len(sender.edus("b.example")) == 1 })
	var types []string
	for _, e := range sender.edus("a.example") {
		types = append(types, e.Type)
	}
	slices.Sort(types)
	if !slices.Equal(types, []string{federation.EDUDirectToDev, federation.EDUReceipt, federation.EDUTyping}) {
		t.Errorf("EDUs to a = %v; a private receipt must not leave the server", types)
	}
	var toDevice federation.DirectToDeviceContent
	if err := json.Unmarshal(sender.edus("b.example")[0].Content, &toDevice); err != nil {
		t.Fatal(err)
	}
	if _, ok := toDevice.Messages["@u:a.example"]; ok {
		t.Error("b was sent a's recipients")
	}
}