From 52889273a9cdd503dd65954176bbaf114d34952e Mon Sep 17 00:00:00 2001
From: Jaz Volpert <ericvolp12@gmail.com>
Date: Sat, 23 Sep 2023 03:08:16 +0000
Subject: [PATCH] Support PLC rate limiters and DID Web resolution limiter
 functions

---
 atproto/identity/base_directory.go |  5 +++++
 atproto/identity/did.go            | 14 ++++++++++++++
 cmd/palomar/main.go                |  8 ++++++++
 3 files changed, 27 insertions(+)

diff --git a/atproto/identity/base_directory.go b/atproto/identity/base_directory.go
index 40577eac3..15cec3e24 100644
--- a/atproto/identity/base_directory.go
+++ b/atproto/identity/base_directory.go
@@ -7,12 +7,17 @@ import (
 	"net/http"
 
 	"github.com/bluesky-social/indigo/atproto/syntax"
+	"golang.org/x/time/rate"
 )
 
 // The zero value ('BaseDirectory{}') is a usable Directory.
 type BaseDirectory struct {
 	// if non-empty, this string should have URL method, hostname, and optional port; it should not have a path or trailing slash
 	PLCURL string
+	// If not nil, this limiter will be used to rate-limit requests to the PLCURL
+	PLCLimiter *rate.Limiter
+	// If not nil, this function will be called inline with DID Web lookups, and can be used to limit the number of requests to a given hostname
+	DIDWebLimitFunc func(ctx context.Context, hostname string) error
 	// HTTP client used for did:web, did:plc, and HTTP (well-known) handle resolution
 	HTTPClient http.Client
 	// DNS resolver used for DNS handle resolution. Calling code can use a custom Dialer to query against a specific DNS server, or re-implement the interface for even more control over the resolution process
diff --git a/atproto/identity/did.go b/atproto/identity/did.go
index bca976f5c..b1af03ff5 100644
--- a/atproto/identity/did.go
+++ b/atproto/identity/did.go
@@ -67,6 +67,13 @@ func (d *BaseDirectory) ResolveDIDWeb(ctx context.Context, did syntax.DID) (*DID
 
 	// TODO: use a more robust client
 	// TODO: allow ctx to specify unsafe http:// resolution, for testing?
+
+	if d.DIDWebLimitFunc != nil {
+		if err := d.DIDWebLimitFunc(ctx, hostname); err != nil {
+			return nil, fmt.Errorf("did:web limit func returned an error for (%s): %w", hostname, err)
+		}
+	}
+
 	resp, err := http.Get("https://" + hostname + "/.well-known/did.json")
 	// look for NXDOMAIN
 	var dnsErr *net.DNSError
@@ -101,6 +108,13 @@ func (d *BaseDirectory) ResolveDIDPLC(ctx context.Context, did syntax.DID) (*DID
 	if plcURL == "" {
 		plcURL = DefaultPLCURL
 	}
+
+	if d.PLCLimiter != nil {
+		if err := d.PLCLimiter.Wait(ctx); err != nil {
+			return nil, fmt.Errorf("failed to wait for PLC limiter: %w", err)
+		}
+	}
+
 	resp, err := http.Get(plcURL + "/" + did.String())
 	if err != nil {
 		return nil, fmt.Errorf("failed did:plc directory resolution: %w", err)
diff --git a/cmd/palomar/main.go b/cmd/palomar/main.go
index 1b8641c79..f26d241ad 100644
--- a/cmd/palomar/main.go
+++ b/cmd/palomar/main.go
@@ -11,6 +11,7 @@ import (
 	"time"
 
 	_ "github.com/joho/godotenv/autoload"
+	"golang.org/x/time/rate"
 
 	"github.com/bluesky-social/indigo/atproto/identity"
 	"github.com/bluesky-social/indigo/search"
@@ -132,6 +133,12 @@ var runCmd = &cli.Command{
 			Value:   20,
 			EnvVars: []string{"PALOMAR_INDEX_MAX_CONCURRENCY"},
 		},
+		&cli.IntFlag{
+			Name:    "plc-rate-limit",
+			Usage:   "max number of requests per second to PLC registry",
+			Value:   100,
+			EnvVars: []string{"PALOMAR_PLC_RATE_LIMIT"},
+		},
 	},
 	Action: func(cctx *cli.Context) error {
 		logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
@@ -155,6 +162,7 @@ var runCmd = &cli.Command{
 			HTTPClient: http.Client{
 				Timeout: time.Second * 15,
 			},
+			PLCLimiter:            rate.NewLimiter(rate.Limit(cctx.Int("plc-rate-limit")), 1),
 			TryAuthoritativeDNS:   true,
 			SkipDNSDomainSuffixes: []string{".bsky.social"},
 		}