Try out zerolog

cmd/backfill/main.go

@@ -0,0 +1,264 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"crypto/ed25519"
+	"encoding/json"
+	"encoding/pem"
+	"flag"
+	"fmt"
+	"math/rand"
+	"os"
+	"regexp"
+	"sort"
+	"time"
+
+	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+	"github.com/tidwall/gjson"
+)
+
+var requestFrom = flag.String("from", "", "the server name that the request should originate from")
+var requestKey = flag.String("key", "matrix_key.pem", "the private key to use when signing the request")
+var requestTo = flag.String("to", "", "the server name to start backfilling from")
+var startEventID = flag.String("eventid", "", "the event ID to start backfilling from")
+var roomID = flag.String("room", "", "the room ID to backfill")
+
+// nolint: gocyclo
+func main() {
+	zerolog.TimeFieldFormat = time.RFC3339Nano
+	log.Logger = log.Output(zerolog.ConsoleWriter{
+		Out:        os.Stderr,
+		TimeFormat: "15:04:05.000",
+	})
+	flag.Parse()
+
+	if requestFrom == nil || *requestFrom == "" {
+		fmt.Println("expecting: furl -from origin.com [-key matrix_key.pem] https://path/to/url")
+		fmt.Println("supported flags:")
+		flag.PrintDefaults()
+		os.Exit(1)
+	}
+
+	if requestTo == nil || *requestTo == "" {
+		fmt.Println("expecting a non empty -to value")
+		flag.PrintDefaults()
+		os.Exit(1)
+	}
+	if roomID == nil || *roomID == "" {
+		fmt.Println("expecting a non empty -room value")
+		flag.PrintDefaults()
+		os.Exit(1)
+	}
+	if startEventID == nil || *startEventID == "" {
+		fmt.Println("expecting a non empty -eventid value")
+		flag.PrintDefaults()
+		os.Exit(1)
+	}
+
+	data, err := os.ReadFile(*requestKey)
+	if err != nil {
+		log.Fatal().
+			Err(err).
+			Msg("failed to read file")
+	}
+
+	var privateKey ed25519.PrivateKey
+	keyBlock, _ := pem.Decode(data)
+	if keyBlock == nil {
+		panic("keyBlock is nil")
+	}
+	if keyBlock.Type == "MATRIX PRIVATE KEY" {
+		_, privateKey, err = ed25519.GenerateKey(bytes.NewReader(keyBlock.Bytes))
+		if err != nil {
+			panic(err)
+		}
+	} else {
+		panic("unexpected key block")
+	}
+
+	serverName := gomatrixserverlib.ServerName(*requestFrom)
+	client := gomatrixserverlib.NewFederationClient(
+		[]*gomatrixserverlib.SigningIdentity{
+			{
+				ServerName: serverName,
+				KeyID:      gomatrixserverlib.KeyID(keyBlock.Headers["Key-ID"]),
+				PrivateKey: privateKey,
+			},
+		},
+		gomatrixserverlib.WithKeepAlives(true),
+	)
+
+	b := &backfiller{
+		FedClient: client,
+		servers: map[gomatrixserverlib.ServerName]struct{}{
+			gomatrixserverlib.ServerName(*requestTo): {},
+		},
+		preferServer: gomatrixserverlib.ServerName(*requestTo),
+	}
+
+	ctx := context.Background()
+	eventID := *startEventID
+	start := time.Now()
+	seenEvents := make(map[string]struct{})
+
+	defer func() {
+		log.Debug().
+			TimeDiff("duration", time.Now(), start).
+			Int("events", len(seenEvents)).
+			Msg("Finished backfilling")
+	}()
+	f, err := os.Create(tokenise(*roomID) + "_backfill.json")
+	if err != nil {
+		log.Fatal().
+			Err(err).
+			Msg("failed to create JSON file")
+	}
+	defer f.Close() // nolint: errcheck
+
+	encoder := json.NewEncoder(f)
+
+	for {
+		log.Debug().
+			Int("events", len(seenEvents)).
+			Str("event_id", eventID).
+			Msg("requesting event")
+		evs, err := gomatrixserverlib.RequestBackfill(ctx, serverName, b, &nopJSONVerifier{}, *roomID, "9", []string{eventID}, 100)
+		if err != nil && len(evs) == 0 {
+			log.Err(err).Msg("failed to backfill, retrying")
+			continue
+		}
+		var createSeen bool
+		sort.Sort(headeredEvents(evs))
+		for _, x := range evs {
+			if _, ok := seenEvents[x.EventID()]; ok {
+				continue
+			}
+
+			sender := gomatrixserverlib.ServerName(gjson.GetBytes(x.JSON(), "origin").Str)
+			if sender != "" && sender != serverName {
+				b.servers[sender] = struct{}{}
+			}
+
+			if x.Type() == "m.room.message" {
+				x.Redact()
+			}
+
+			// The following ensures we preserve the "_event_id" field
+			err = encoder.Encode(x)
+			if err != nil {
+				log.Fatal().
+					Err(err).
+					Msg("failed to write to file")
+			}
+			if x.Type() == gomatrixserverlib.MRoomCreate {
+				createSeen = true
+			}
+		}
+		// We've reached the beginning of the room
+		if createSeen {
+			log.Debug().
+				Int("events", len(seenEvents)).
+				Msg("Reached beginning of the room, existing")
+			return
+		}
+
+		// Remember the event ID before trying to find a new one
+		beforeEvID := eventID
+		for _, x := range evs {
+			if _, ok := seenEvents[x.EventID()]; ok {
+				continue
+			}
+			if x.EventID() == beforeEvID {
+				continue
+			}
+			eventID = x.EventID()
+			break
+		}
+		if beforeEvID == eventID {
+			log.Debug().
+				Msg("no new eventID found in backfill response")
+			return
+		}
+		// Finally store which events we've already seen
+		for _, x := range evs {
+			seenEvents[x.EventID()] = struct{}{}
+		}
+		time.Sleep(time.Second) // don't hit remotes to hard
+	}
+}
+
+type headeredEvents []*gomatrixserverlib.HeaderedEvent
+
+func (h headeredEvents) Len() int {
+	return len(h)
+}
+
+func (h headeredEvents) Less(i, j int) bool {
+	return h[i].Depth() < h[j].Depth()
+}
+
+func (h headeredEvents) Swap(i, j int) {
+	h[i], h[j] = h[j], h[i]
+}
+
+type backfiller struct {
+	FedClient    *gomatrixserverlib.FederationClient
+	servers      map[gomatrixserverlib.ServerName]struct{}
+	preferServer gomatrixserverlib.ServerName
+}
+
+func (b backfiller) StateIDsBeforeEvent(ctx context.Context, event *gomatrixserverlib.HeaderedEvent) ([]string, error) {
+	return []string{}, nil
+}
+
+func (b backfiller) StateBeforeEvent(ctx context.Context, roomVer gomatrixserverlib.RoomVersion, event *gomatrixserverlib.HeaderedEvent, eventIDs []string) (map[string]*gomatrixserverlib.Event, error) {
+	return nil, nil
+}
+
+func (b backfiller) Backfill(ctx context.Context, origin, server gomatrixserverlib.ServerName, roomID string, limit int, fromEventIDs []string) (gomatrixserverlib.Transaction, error) {
+	return b.FedClient.Backfill(ctx, origin, server, roomID, limit, fromEventIDs)
+}
+
+func (b backfiller) ServersAtEvent(ctx context.Context, roomID, eventID string) []gomatrixserverlib.ServerName {
+	servers := make([]gomatrixserverlib.ServerName, 0, len(b.servers)+1)
+	for v := range b.servers {
+		if v == b.preferServer { // will be added to the front anyway
+			continue
+		}
+		servers = append(servers, v)
+	}
+	rand.Shuffle(len(servers), func(i, j int) {
+		servers[i], servers[j] = servers[j], servers[i]
+	})
+
+	// always prefer specified server
+	servers = append([]gomatrixserverlib.ServerName{b.preferServer}, servers...)
+
+	if len(servers) > 5 {
+		servers = servers[:5]
+	}
+	return servers
+}
+
+func (b backfiller) ProvideEvents(roomVer gomatrixserverlib.RoomVersion, eventIDs []string) ([]*gomatrixserverlib.Event, error) {
+	return []*gomatrixserverlib.Event{}, nil
+}
+
+var safeCharacters = regexp.MustCompile("[^A-Za-z0-9$]+")
+
+func tokenise(str string) string {
+	return safeCharacters.ReplaceAllString(str, "_")
+}
+
+// NopJSONVerifier is a JSONVerifier that verifies nothing and returns no errors.
+type nopJSONVerifier struct {
+	// this verifier verifies nothing
+}
+
+func (t *nopJSONVerifier) VerifyJSONs(ctx context.Context, requests []gomatrixserverlib.VerifyJSONRequest) ([]gomatrixserverlib.VerifyJSONResult, error) {
+	result := make([]gomatrixserverlib.VerifyJSONResult, len(requests))
+	return result, nil
+}

go.mod

@@ -22,7 +22,7 @@ require (
 	github.com/matrix-org/dugong v0.0.0-20210921133753-66e6b1c67e2e
 	github.com/matrix-org/go-sqlite3-js v0.0.0-20220419092513-28aa791a1c91
 	github.com/matrix-org/gomatrix v0.0.0-20220926102614-ceba4d9f7530
-	github.com/matrix-org/gomatrixserverlib v0.0.0-20230131183213-122f1e0e3fa1
+	github.com/matrix-org/gomatrixserverlib v0.0.0-20230307091615-c09404be6e15
 	github.com/matrix-org/pinecone v0.11.1-0.20230210171230-8c3b24f2649a
 	github.com/matrix-org/util v0.0.0-20221111132719-399730281e66
 	github.com/mattn/go-sqlite3 v1.14.15
@@ -94,6 +94,7 @@ require (
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
 	github.com/klauspost/compress v1.15.11 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.16 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 // indirect
 	github.com/minio/highwayhash v1.0.2 // indirect
@@ -118,6 +119,7 @@ require (
 	github.com/quic-go/qtls-go1-20 v0.1.0 // indirect
 	github.com/quic-go/quic-go v0.32.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20220927061507-ef77025ab5aa // indirect
+	github.com/rs/zerolog v1.29.0 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.1 // indirect
 	go.etcd.io/bbolt v1.3.6 // indirect

go.sum

@@ -123,6 +123,7 @@ github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDk
 github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc=
 github.com/codeclysm/extract v2.2.0+incompatible h1:q3wyckoA30bhUSiwdQezMqVhwd8+WGE64/GL//LtUhI=
 github.com/codeclysm/extract v2.2.0+incompatible/go.mod h1:2nhFMPHiU9At61hz+12bfrlpXSUrOnK+wR+KlGO4Uks=
+github.com/coreos/go-systemd/v22 v22.3.3-0.20220203105225-a9a7ef127534/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/creack/pty v1.1.11/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -188,6 +189,7 @@ github.com/gobwas/pool v0.2.0/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6Wezm
 github.com/gobwas/ws v1.0.2 h1:CoAavW/wd/kulfZmSIBt6p24n4j7tHgNVCjsfHVNUbo=
 github.com/gobwas/ws v1.0.2/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM=
 github.com/goccy/go-json v0.9.11 h1:/pAaQDLHEoCq/5FFmSKBswWmK6H0e8g4159Kc/X/nqk=
+github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
@@ -323,11 +325,19 @@ github.com/matrix-org/gomatrix v0.0.0-20220926102614-ceba4d9f7530 h1:kHKxCOLcHH8
 github.com/matrix-org/gomatrix v0.0.0-20220926102614-ceba4d9f7530/go.mod h1:/gBX06Kw0exX1HrwmoBibFA98yBk/jxKpGVeyQbff+s=
 github.com/matrix-org/gomatrixserverlib v0.0.0-20230131183213-122f1e0e3fa1 h1:JSw0nmjMrgBmoM2aQsa78LTpI5BnuD9+vOiEQ4Qo0qw=
 github.com/matrix-org/gomatrixserverlib v0.0.0-20230131183213-122f1e0e3fa1/go.mod h1:Mtifyr8q8htcBeugvlDnkBcNUy5LO8OzUoplAf1+mb4=
+github.com/matrix-org/gomatrixserverlib v0.0.0-20230301150509-5f7577968b8f h1:gJQy+K/OYxAycj5rs0/7U2R4/Cx+XfkdERRLjIOQmZ4=
+github.com/matrix-org/gomatrixserverlib v0.0.0-20230301150509-5f7577968b8f/go.mod h1:Mtifyr8q8htcBeugvlDnkBcNUy5LO8OzUoplAf1+mb4=
+github.com/matrix-org/gomatrixserverlib v0.0.0-20230307091615-c09404be6e15 h1:7ywZ3mH1EUWo9SvkikCz99EmYixjBNAAKRDnwXpJS3Q=
+github.com/matrix-org/gomatrixserverlib v0.0.0-20230307091615-c09404be6e15/go.mod h1:Mtifyr8q8htcBeugvlDnkBcNUy5LO8OzUoplAf1+mb4=
 github.com/matrix-org/pinecone v0.11.1-0.20230210171230-8c3b24f2649a h1:awrPDf9LEFySxTLKYBMCiObelNx/cBuv/wzllvCCH3A=
 github.com/matrix-org/pinecone v0.11.1-0.20230210171230-8c3b24f2649a/go.mod h1:HchJX9oKMXaT2xYFs0Ha/6Zs06mxLU8k6F1ODnrGkeQ=
 github.com/matrix-org/util v0.0.0-20221111132719-399730281e66 h1:6z4KxomXSIGWqhHcfzExgkH3Z3UkIXry4ibJS4Aqz2Y=
 github.com/matrix-org/util v0.0.0-20221111132719-399730281e66/go.mod h1:iBI1foelCqA09JJgPV0FYz4qA5dUXYOxMi57FxKBdd4=
+github.com/mattn/go-colorable v0.1.12/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4=
+github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
+github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
 github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
+github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
 github.com/mattn/go-isatty v0.0.16 h1:bq3VjFmv/sOjHtdEhmkEV4x1AJtvUvOJ2PFAZ5+peKQ=
 github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
 github.com/mattn/go-sqlite3 v1.14.15 h1:vfoHhTN1af61xCRSWzFIWzx2YskyMTwHLrExkBOjvxI=
@@ -432,6 +442,9 @@ github.com/remyoudompheng/bigfft v0.0.0-20220927061507-ef77025ab5aa/go.mod h1:qq
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
 github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8=
 github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
+github.com/rs/xid v1.4.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
+github.com/rs/zerolog v1.29.0 h1:Zes4hju04hjbvkVkOhdl2HpZa+0PmVwigmo8XoORE5w=
+github.com/rs/zerolog v1.29.0/go.mod h1:NILgTygv/Uej1ra5XxGf82ZFSLk58MFGAUS2o6usyD0=
 github.com/ryszard/goskiplist v0.0.0-20150312221310-2dfbae5fcf46/go.mod h1:uAQ5PCi+MFsC7HjREoAz1BU+Mq60+05gifQSsHSDG/8=
 github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
@@ -651,6 +664,8 @@ golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210603081109-ebe580a85c40/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220114195835-da31bd327af9/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=