package noncebalancer import ( "errors" "sync" "github.com/letsencrypt/boulder/nonce" "google.golang.org/grpc/balancer" "google.golang.org/grpc/balancer/base" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" ) const ( // Name is the name used to register the nonce balancer with the gRPC // runtime. Name = "nonce" // SRVResolverScheme is the scheme used to invoke an instance of the SRV // resolver which will use the noncebalancer to pick backends. It would be // ideal to export this from the SRV resolver package but that package is // internal. SRVResolverScheme = "nonce-srv" ) // ErrNoBackendsMatchPrefix indicates that no backends were found which match // the nonce prefix provided in the RPC context. This can happen when the // provided nonce is stale, valid but the backend has since been removed from // the balancer, or valid but the backend has not yet been added to the // balancer. // // In any case, when the WFE receives this error it will return a badNonce error // to the ACME client. var ErrNoBackendsMatchPrefix = status.New(codes.Unavailable, "no backends match the nonce prefix") var errMissingPrefixCtxKey = errors.New("nonce.PrefixCtxKey value required in RPC context") var errMissingHMACKeyCtxKey = errors.New("nonce.HMACKeyCtxKey value required in RPC context") var errInvalidPrefixCtxKeyType = errors.New("nonce.PrefixCtxKey value in RPC context must be a string") var errInvalidHMACKeyCtxKeyType = errors.New("nonce.HMACKeyCtxKey value in RPC context must be a byte slice") // Balancer implements the base.PickerBuilder interface. It's used to create new // balancer.Picker instances. It should only be used by nonce-service clients. type Balancer struct{} // Compile-time assertion that *Balancer implements the base.PickerBuilder // interface. var _ base.PickerBuilder = (*Balancer)(nil) // Build implements the base.PickerBuilder interface. It is called by the gRPC // runtime when the balancer is first initialized and when the set of backend // (SubConn) addresses changes. func (b *Balancer) Build(buildInfo base.PickerBuildInfo) balancer.Picker { if len(buildInfo.ReadySCs) == 0 { // The Picker must be rebuilt if there are no backends available. return base.NewErrPicker(balancer.ErrNoSubConnAvailable) } return &Picker{ backends: buildInfo.ReadySCs, } } // Picker implements the balancer.Picker interface. It picks a backend (SubConn) // based on the nonce prefix contained in each request's Context. type Picker struct { backends map[balancer.SubConn]base.SubConnInfo prefixToBackend map[string]balancer.SubConn prefixToBackendOnce sync.Once } // Compile-time assertion that *Picker implements the balancer.Picker interface. var _ balancer.Picker = (*Picker)(nil) // Pick implements the balancer.Picker interface. It is called by the gRPC // runtime for each RPC message. It is responsible for picking a backend // (SubConn) based on the context of each RPC message. func (p *Picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) { if len(p.backends) == 0 { // This should never happen, the Picker should only be built when there // are backends available. return balancer.PickResult{}, balancer.ErrNoSubConnAvailable } // Get the HMAC key from the RPC context. hmacKeyVal := info.Ctx.Value(nonce.HMACKeyCtxKey{}) if hmacKeyVal == nil { // This should never happen. return balancer.PickResult{}, errMissingHMACKeyCtxKey } hmacKey, ok := hmacKeyVal.([]byte) if !ok { // This should never happen. return balancer.PickResult{}, errInvalidHMACKeyCtxKeyType } p.prefixToBackendOnce.Do(func() { // First call to Pick with a new Picker. prefixToBackend := make(map[string]balancer.SubConn) for sc, scInfo := range p.backends { scPrefix := nonce.DerivePrefix(scInfo.Address.Addr, hmacKey) prefixToBackend[scPrefix] = sc } p.prefixToBackend = prefixToBackend }) // Get the destination prefix from the RPC context. destPrefixVal := info.Ctx.Value(nonce.PrefixCtxKey{}) if destPrefixVal == nil { // This should never happen. return balancer.PickResult{}, errMissingPrefixCtxKey } destPrefix, ok := destPrefixVal.(string) if !ok { // This should never happen. return balancer.PickResult{}, errInvalidPrefixCtxKeyType } sc, ok := p.prefixToBackend[destPrefix] if !ok { // No backend SubConn was found for the destination prefix. return balancer.PickResult{}, ErrNoBackendsMatchPrefix.Err() } return balancer.PickResult{SubConn: sc}, nil } func init() { balancer.Register( base.NewBalancerBuilder(Name, &Balancer{}, base.Config{}), ) }