From 7a005516107d54eb021b653a2c310495dff83331 Mon Sep 17 00:00:00 2001
From: Brian Olson <brian.olson@gmail.com>
Date: Wed, 8 Jan 2025 13:49:27 -0500
Subject: [PATCH] RELAY_ALLOW_PDS_PROXY -allow-pds-proxy

---
 bgs/bgs.go         | 6 +++---
 cmd/bigsky/main.go | 7 +++++++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/bgs/bgs.go b/bgs/bgs.go
index 1de4fcdf6..30697c640 100644
--- a/bgs/bgs.go
+++ b/bgs/bgs.go
@@ -125,8 +125,8 @@ type BGSConfig struct {
 	ConcurrencyPerPDS    int64
 	MaxQueuePerPDS       int64
 	NumCompactionWorkers int
-
-	VerboseAPILog bool
+	VerboseAPILog        bool
+	AllowPDSProxies      bool
 
 	// NextCrawlers gets forwarded POST /xrpc/com.atproto.sync.requestCrawl
 	NextCrawlers []*url.URL
@@ -946,7 +946,7 @@ func (bgs *BGS) handleFedEvent(ctx context.Context, host *models.PDS, env *event
 			return fmt.Errorf("rebase was true in event seq:%d,host:%s", evt.Seq, host.Host)
 		}
 
-		if host.ID != u.PDS && u.PDS != 0 {
+		if host.ID != u.PDS && u.PDS != 0 && !bgs.config.AllowPDSProxies {
 			bgs.log.Warn("received event for repo from different pds than expected", "repo", evt.Repo, "expPds", u.PDS, "gotPds", host.Host)
 			// Flush any cached DID documents for this user
 			bgs.didr.FlushCacheFor(env.RepoCommit.Repo)
diff --git a/cmd/bigsky/main.go b/cmd/bigsky/main.go
index 4b56678d1..a566869c0 100644
--- a/cmd/bigsky/main.go
+++ b/cmd/bigsky/main.go
@@ -101,6 +101,12 @@ func run(args []string) error {
 			Value:   "prod",
 			EnvVars: []string{"RELAY_CRAWL_NONSSL"},
 		},
+		&cli.BoolFlag{
+			Name:    "allow-pds-proxy",
+			Usage:   "allow getting PDS data from a non-canonical source (e.g. another relay)",
+			Value:   true,
+			EnvVars: []string{"RELAY_ALLOW_PDS_PROXY"},
+		},
 		&cli.BoolFlag{
 			Name:    "spidering",
 			Value:   false,
@@ -493,6 +499,7 @@ func runBigsky(cctx *cli.Context) error {
 	default:
 		return fmt.Errorf("crawl-insecure-ws/RELAY_CRAWL_NONSSL exepected true|false|prod, got %s", cctx.String("crawl-insecure-ws"))
 	}
+	bgsConfig.AllowPDSProxies = cctx.Bool("allow-pds-proxy")
 	bgsConfig.CompactInterval = cctx.Duration("compact-interval")
 	bgsConfig.ConcurrencyPerPDS = cctx.Int64("concurrency-per-pds")
 	bgsConfig.MaxQueuePerPDS = cctx.Int64("max-queue-per-pds")