From 28ae482495c02966fe968dc9b73cbce78d6a8ef7 Mon Sep 17 00:00:00 2001
From: Samantha <hello@entropy.cat>
Date: Tue, 16 Jan 2024 16:29:11 -0500
Subject: [PATCH] Address comments.

---
 wfe2/wfe.go | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/wfe2/wfe.go b/wfe2/wfe.go
index 2bb061edc97b..1902cc3df25f 100644
--- a/wfe2/wfe.go
+++ b/wfe2/wfe.go
@@ -672,24 +672,27 @@ func (wfe *WebFrontEndImpl) checkNewAccountLimits(ctx context.Context, ip net.IP
 }
 
 // refundNewAccountLimits is typically called when a new account creation fails.
-// It refunds the limit quota consumed by the request, allowing the caller to
-// retry immediately. If an error is encountered during the refund, it is logged
-// but not returned.
-func (wfe *WebFrontEndImpl) refundNewAccountLimits(ctx context.Context, ip net.IP) {
+// It waits for limitCheck channel to close and then refunds the limit quota
+// consumed by the request, allowing the caller to retry immediately. If an
+// error is encountered during the refund, it is logged but not returned.
+func (wfe *WebFrontEndImpl) refundNewAccountLimits(ctx context.Context, limitCheck <-chan struct{}, ip net.IP) {
 	if wfe.limiter == nil && wfe.txnBuilder == nil {
 		// Limiter is disabled.
 		return
 	}
 
 	warn := func(err error, limit ratelimits.Name) {
-		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
-			return
-		}
-		// TODO(#5545): Once key-value rate limits are authoritative this log
-		// line should be removed in favor of returning the error.
+		// Log the error for debugging purposes.
 		wfe.log.Warningf("refunding %s rate limit: %s", limit, err)
 	}
 
+	select {
+	case <-ctx.Done():
+		return
+	case <-limitCheck:
+		// Continue with refunding the limits.
+	}
+
 	txn, err := wfe.txnBuilder.RegistrationsPerIPAddressTransaction(ip)
 	if err != nil {
 		warn(err, ratelimits.NewRegistrationsPerIPAddress)
@@ -842,22 +845,19 @@ func (wfe *WebFrontEndImpl) NewAccount(
 
 	// TODO(#5545): Spending and Refunding can be async until these rate limits
 	// are authoritative. This saves us from adding latency to each request.
-	doneCheckingLimits := make(chan struct{})
+	limitCheck := make(chan struct{})
 	var newRegistrationSuccessful bool
 	var errIsRateLimit bool
 
 	go func() {
 		// Close the channel on goroutine completion.
-		defer close(doneCheckingLimits)
+		defer close(limitCheck)
 		wfe.checkNewAccountLimits(ctx, ip)
 	}()
 
 	defer func() {
 		if !newRegistrationSuccessful && !errIsRateLimit {
-			// Wait for the rate limit check to complete before attempting to refund
-			// the limits. If the check failed, we don't want to refund anything.
-			<-doneCheckingLimits
-			go wfe.refundNewAccountLimits(ctx, ip)
+			go wfe.refundNewAccountLimits(ctx, limitCheck, ip)
 		}
 	}()
 
@@ -865,6 +865,9 @@ func (wfe *WebFrontEndImpl) NewAccount(
 	acctPB, err := wfe.ra.NewRegistration(ctx, &reg)
 	if err != nil {
 		if errors.Is(err, berrors.RateLimit) {
+			// Request was denied by a legacy rate limit. In this error case we
+			// do not want to refund the quota consumed by the request because
+			// repeated requests would result in unearned refunds.
 			errIsRateLimit = true
 		}
 		if errors.Is(err, berrors.Duplicate) {