github.com/letsencrypt/boulder@v0.20251208.0/grpc/noncebalancer/noncebalancer.go (about)

     1  package noncebalancer
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  
     7  	"github.com/letsencrypt/boulder/nonce"
     8  
     9  	"google.golang.org/grpc/balancer"
    10  	"google.golang.org/grpc/balancer/base"
    11  	"google.golang.org/grpc/codes"
    12  	"google.golang.org/grpc/status"
    13  )
    14  
    15  const (
    16  	// Name is the name used to register the nonce balancer with the gRPC
    17  	// runtime.
    18  	Name = "nonce"
    19  
    20  	// SRVResolverScheme is the scheme used to invoke an instance of the SRV
    21  	// resolver which will use the noncebalancer to pick backends. It would be
    22  	// ideal to export this from the SRV resolver package but that package is
    23  	// internal.
    24  	SRVResolverScheme = "nonce-srv"
    25  )
    26  
    27  // ErrNoBackendsMatchPrefix indicates that no backends were found which match
    28  // the nonce prefix provided in the RPC context. This can happen when the
    29  // provided nonce is stale, valid but the backend has since been removed from
    30  // the balancer, or valid but the backend has not yet been added to the
    31  // balancer.
    32  //
    33  // In any case, when the WFE receives this error it will return a badNonce error
    34  // to the ACME client. Note that the WFE uses exact pointer comparison to
    35  // detect that the status it receives is this exact status object, so don't
    36  // wrap this with fmt.Errorf when returning it.
    37  var ErrNoBackendsMatchPrefix = status.New(codes.Unavailable, "no backends match the nonce prefix")
    38  var errMissingPrefixCtxKey = errors.New("nonce.PrefixCtxKey value required in RPC context")
    39  var errMissingHMACKeyCtxKey = errors.New("nonce.HMACKeyCtxKey value required in RPC context")
    40  var errInvalidPrefixCtxKeyType = errors.New("nonce.PrefixCtxKey value in RPC context must be a string")
    41  var errInvalidHMACKeyCtxKeyType = errors.New("nonce.HMACKeyCtxKey value in RPC context must be a byte slice")
    42  
    43  // pickerBuilder implements the base.PickerBuilder interface. It's used to
    44  // create new Picker instances. It should only be used by nonce-service clients.
    45  type pickerBuilder struct{}
    46  
    47  // Build implements the base.PickerBuilder interface. It is called by the gRPC
    48  // runtime when the balancer is first initialized and when the set of backend
    49  // (SubConn) addresses changes.
    50  func (b *pickerBuilder) Build(buildInfo base.PickerBuildInfo) balancer.Picker {
    51  	if len(buildInfo.ReadySCs) == 0 {
    52  		// The Picker must be rebuilt if there are no backends available.
    53  		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
    54  	}
    55  	return &picker{
    56  		backends: buildInfo.ReadySCs,
    57  	}
    58  }
    59  
    60  // picker implements the balancer.Picker interface. It picks a backend (SubConn)
    61  // based on the nonce prefix contained in each request's Context.
    62  type picker struct {
    63  	backends            map[balancer.SubConn]base.SubConnInfo
    64  	prefixToBackend     map[string]balancer.SubConn
    65  	prefixToBackendOnce sync.Once
    66  }
    67  
    68  // Pick implements the balancer.Picker interface. It is called by the gRPC
    69  // runtime for each RPC message. It is responsible for picking a backend
    70  // (SubConn) based on the context of each RPC message.
    71  func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
    72  	if len(p.backends) == 0 {
    73  		// This should never happen, the Picker should only be built when there
    74  		// are backends available.
    75  		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
    76  	}
    77  
    78  	// Get the HMAC key from the RPC context.
    79  	hmacKeyVal := info.Ctx.Value(nonce.HMACKeyCtxKey{})
    80  	if hmacKeyVal == nil {
    81  		// This should never happen.
    82  		return balancer.PickResult{}, errMissingHMACKeyCtxKey
    83  	}
    84  	hmacKey, ok := hmacKeyVal.([]byte)
    85  	if !ok {
    86  		// This should never happen.
    87  		return balancer.PickResult{}, errInvalidHMACKeyCtxKeyType
    88  	}
    89  
    90  	p.prefixToBackendOnce.Do(func() {
    91  		// First call to Pick with a new Picker.
    92  		prefixToBackend := make(map[string]balancer.SubConn)
    93  		for sc, scInfo := range p.backends {
    94  			scPrefix := nonce.DerivePrefix(scInfo.Address.Addr, hmacKey)
    95  			prefixToBackend[scPrefix] = sc
    96  		}
    97  		p.prefixToBackend = prefixToBackend
    98  	})
    99  
   100  	// Get the destination prefix from the RPC context.
   101  	destPrefixVal := info.Ctx.Value(nonce.PrefixCtxKey{})
   102  	if destPrefixVal == nil {
   103  		// This should never happen.
   104  		return balancer.PickResult{}, errMissingPrefixCtxKey
   105  	}
   106  	destPrefix, ok := destPrefixVal.(string)
   107  	if !ok {
   108  		// This should never happen.
   109  		return balancer.PickResult{}, errInvalidPrefixCtxKeyType
   110  	}
   111  
   112  	sc, ok := p.prefixToBackend[destPrefix]
   113  	if !ok {
   114  		// No backend SubConn was found for the destination prefix.
   115  		return balancer.PickResult{}, ErrNoBackendsMatchPrefix.Err()
   116  	}
   117  	return balancer.PickResult{SubConn: sc}, nil
   118  }
   119  
   120  func init() {
   121  	balancer.Register(
   122  		base.NewBalancerBuilder(Name, &pickerBuilder{}, base.Config{}),
   123  	)
   124  }