github.com/letsencrypt/boulder@v0.20251208.0/cmd/admin/overrides_import.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"flag"
     7  	"fmt"
     8  	"sync"
     9  
    10  	"github.com/letsencrypt/boulder/ratelimits"
    11  	sapb "github.com/letsencrypt/boulder/sa/proto"
    12  	"google.golang.org/protobuf/types/known/durationpb"
    13  )
    14  
    15  type subcommandImportOverrides struct {
    16  	file        string
    17  	parallelism int
    18  }
    19  
    20  func (*subcommandImportOverrides) Desc() string { return "Push overrides to SA" }
    21  
    22  func (c *subcommandImportOverrides) Flags(f *flag.FlagSet) {
    23  	f.StringVar(&c.file, "file", "", "path to YAML file containing rate limit overrides")
    24  	f.IntVar(&c.parallelism, "parallelism", 10, "the number of concurrent RPCs to send to the SA (default: 10)")
    25  }
    26  
    27  func (c *subcommandImportOverrides) Run(ctx context.Context, a *admin) error {
    28  	if c.file == "" {
    29  		return errors.New("--file is required")
    30  	}
    31  	if c.parallelism <= 0 {
    32  		return errors.New("--parallelism must be greater than 0")
    33  	}
    34  	overrides, err := ratelimits.LoadOverridesByBucketKey(c.file)
    35  	if err != nil {
    36  		return err
    37  	}
    38  	var overrideCount = len(overrides)
    39  
    40  	work := make(chan *sapb.RateLimitOverride, overrideCount)
    41  	for k, ov := range overrides {
    42  		work <- &sapb.RateLimitOverride{
    43  			LimitEnum: int64(ov.Name),
    44  			BucketKey: k,
    45  			Comment:   ov.Comment,
    46  			Period:    durationpb.New(ov.Period.Duration),
    47  			Count:     ov.Count,
    48  			Burst:     ov.Burst,
    49  		}
    50  	}
    51  	close(work)
    52  
    53  	type result struct {
    54  		ov  *sapb.RateLimitOverride
    55  		err error
    56  	}
    57  	results := make(chan result, c.parallelism)
    58  
    59  	var wg sync.WaitGroup
    60  	for i := 0; i < c.parallelism; i++ {
    61  		wg.Go(func() {
    62  			for ov := range work {
    63  				_, err := a.sac.AddRateLimitOverride(ctx, &sapb.AddRateLimitOverrideRequest{Override: ov})
    64  				results <- result{ov: ov, err: err}
    65  			}
    66  		})
    67  	}
    68  
    69  	var errorCount int
    70  	for range overrideCount {
    71  		result := <-results
    72  		if result.err != nil {
    73  			a.log.AuditErrf("failed to add override: key=%q limit=%d: %s", result.ov.BucketKey, result.ov.LimitEnum, result.err)
    74  			errorCount++
    75  		}
    76  	}
    77  
    78  	wg.Wait()
    79  	close(results)
    80  
    81  	if errorCount > 0 {
    82  		return fmt.Errorf("%d out of %d overrides failed to be added, see log message(s) for more details", errorCount, overrideCount)
    83  	}
    84  	a.log.Infof("Successfully added %d overrides", overrideCount)
    85  	return nil
    86  }