boulder/grpc/noncebalancer/noncebalancer.go

130 lines
4.5 KiB
Go

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{}),
)
}