WFE: Check NewOrder rate limits

wfe2/wfe.go

@@ -2055,6 +2055,75 @@ func (wfe *WebFrontEndImpl) orderToOrderJSON(request *http.Request, order *corep
 	return respObj
 }
 
+// checkNewOrderLimits checks whether sufficient limit quota exists for the
+// creation of a new order. If so, that quota is spent. If an error is
+// encountered during the check, it is logged but not returned.
+//
+// TODO(#5545): For now we're simply exercising the new rate limiter codepath.
+// This should eventually return a berrors.RateLimit error containing the retry
+// after duration among other information available in the ratelimits.Decision.
+func (wfe *WebFrontEndImpl) checkNewOrderLimits(ctx context.Context, regId int64, names []string) []ratelimits.Transaction {
+	if wfe.limiter == nil && wfe.txnBuilder == nil {
+		// Limiter is disabled.
+		return nil
+	}
+
+	logTxnErr := func(err error, limit ratelimits.Name) {
+		// TODO(#5545): Once key-value rate limits are authoritative this log
+		// line should be removed in favor of returning the error.
+		wfe.log.Errf("constructing rate limit transaction for %s rate limit: %s", limit, err)
+	}
+
+	var transactions []ratelimits.Transaction
+	txn, err := wfe.txnBuilder.OrdersPerAccountTransaction(regId)
+	if err != nil {
+		logTxnErr(err, ratelimits.NewOrdersPerAccount)
+		return nil
+	}
+	transactions = append(transactions, txn)
+
+	txn, err = wfe.txnBuilder.FailedAuthorizationsPerAccountCheckOnlyTransaction(regId)
+	if err != nil {
+		logTxnErr(err, ratelimits.FailedAuthorizationsPerAccount)
+		return nil
+	}
+	transactions = append(transactions, txn)
+
+	txns, err := wfe.txnBuilder.CertificatesPerDomainTransactions(regId, names)
+	if err != nil {
+		logTxnErr(err, ratelimits.CertificatesPerDomain)
+		return nil
+	}
+	transactions = append(transactions, txns...)
+
+	txn, err = wfe.txnBuilder.CertificatesPerFQDNSetTransaction(names)
+	if err != nil {
+		logTxnErr(err, ratelimits.CertificatesPerFQDNSet)
+		return nil
+	}
+	transactions = append(transactions, txn)
+
+	_, err = wfe.limiter.BatchSpend(ctx, transactions)
+	if err != nil {
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			return nil
+		}
+		wfe.log.Errf("checking limits for newOrder: %s", err)
+		return nil
+	}
+	return transactions
+}
+
+func (wfe *WebFrontEndImpl) refundNewOrderLimits(ctx context.Context, transactions []ratelimits.Transaction) {
+	if wfe.limiter == nil || wfe.txnBuilder == nil {
+		return
+	}
+	_, err := wfe.limiter.BatchRefund(ctx, transactions)
+	if err != nil {
+		wfe.log.Errf("refunding newOrder limits: %s", err)
+	}
+}
+
 // NewOrder is used by clients to create a new order object and a set of
 // authorizations to fulfill for issuance.
 func (wfe *WebFrontEndImpl) NewOrder(
@@ -2129,6 +2198,19 @@ func (wfe *WebFrontEndImpl) NewOrder(
 
 	logEvent.DNSNames = names
 
+	// TODO(#5545): Spending and Refunding can be async until these rate limits
+	// are authoritative. This saves us from adding latency to each request.
+	var ratelimitTxns []ratelimits.Transaction
+	var newOrderSuccessful bool
+	go func() {
+		ratelimitTxns = wfe.checkNewOrderLimits(ctx, acct.ID, names)
+	}()
+	defer func() {
+		if !newOrderSuccessful && ratelimitTxns != nil {
+			go wfe.refundNewOrderLimits(ctx, ratelimitTxns)
+		}
+	}()
+
 	order, err := wfe.ra.NewOrder(ctx, &rapb.NewOrderRequest{
 		RegistrationID: acct.ID,
 		Names:          names,
@@ -2150,6 +2232,7 @@ func (wfe *WebFrontEndImpl) NewOrder(
 		wfe.sendError(response, logEvent, probs.ServerInternal("Error marshaling order"), err)
 		return
 	}
+	newOrderSuccessful = true
 }
 
 // GetOrder is used to retrieve a existing order object