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

     1  package noncebalancer
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"google.golang.org/grpc/balancer"
     8  	"google.golang.org/grpc/balancer/base"
     9  	"google.golang.org/grpc/resolver"
    10  
    11  	"github.com/letsencrypt/boulder/nonce"
    12  	"github.com/letsencrypt/boulder/test"
    13  )
    14  
    15  func TestPickerPicksCorrectBackend(t *testing.T) {
    16  	_, p, subConns := setupTest(false)
    17  	prefix := nonce.DerivePrefix(subConns[0].addrs[0].Addr, []byte("Kala namak"))
    18  
    19  	testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, "HNmOnt8w")
    20  	testCtx = context.WithValue(testCtx, nonce.HMACKeyCtxKey{}, []byte(prefix))
    21  	info := balancer.PickInfo{Ctx: testCtx}
    22  
    23  	gotPick, err := p.Pick(info)
    24  	test.AssertNotError(t, err, "Pick failed")
    25  	test.AssertDeepEquals(t, subConns[0], gotPick.SubConn)
    26  }
    27  
    28  func TestPickerMissingPrefixInCtx(t *testing.T) {
    29  	_, p, subConns := setupTest(false)
    30  	prefix := nonce.DerivePrefix(subConns[0].addrs[0].Addr, []byte("Kala namak"))
    31  
    32  	testCtx := context.WithValue(context.Background(), nonce.HMACKeyCtxKey{}, []byte(prefix))
    33  	info := balancer.PickInfo{Ctx: testCtx}
    34  
    35  	gotPick, err := p.Pick(info)
    36  	test.AssertErrorIs(t, err, errMissingPrefixCtxKey)
    37  	test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
    38  }
    39  
    40  func TestPickerInvalidPrefixInCtx(t *testing.T) {
    41  	_, p, _ := setupTest(false)
    42  
    43  	testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, 9)
    44  	testCtx = context.WithValue(testCtx, nonce.HMACKeyCtxKey{}, []byte("foobar"))
    45  	info := balancer.PickInfo{Ctx: testCtx}
    46  
    47  	gotPick, err := p.Pick(info)
    48  	test.AssertErrorIs(t, err, errInvalidPrefixCtxKeyType)
    49  	test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
    50  }
    51  
    52  func TestPickerMissingHMACKeyInCtx(t *testing.T) {
    53  	_, p, _ := setupTest(false)
    54  
    55  	testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, "HNmOnt8w")
    56  	info := balancer.PickInfo{Ctx: testCtx}
    57  
    58  	gotPick, err := p.Pick(info)
    59  	test.AssertErrorIs(t, err, errMissingHMACKeyCtxKey)
    60  	test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
    61  }
    62  
    63  func TestPickerInvalidHMACKeyInCtx(t *testing.T) {
    64  	_, p, _ := setupTest(false)
    65  
    66  	testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, "HNmOnt8w")
    67  	testCtx = context.WithValue(testCtx, nonce.HMACKeyCtxKey{}, 9)
    68  	info := balancer.PickInfo{Ctx: testCtx}
    69  
    70  	gotPick, err := p.Pick(info)
    71  	test.AssertErrorIs(t, err, errInvalidHMACKeyCtxKeyType)
    72  	test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
    73  }
    74  
    75  func TestPickerNoMatchingSubConnAvailable(t *testing.T) {
    76  	_, p, subConns := setupTest(false)
    77  	prefix := nonce.DerivePrefix(subConns[0].addrs[0].Addr, []byte("Kala namak"))
    78  
    79  	testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, "rUsTrUin")
    80  	testCtx = context.WithValue(testCtx, nonce.HMACKeyCtxKey{}, []byte(prefix))
    81  	info := balancer.PickInfo{Ctx: testCtx}
    82  
    83  	gotPick, err := p.Pick(info)
    84  	test.AssertErrorIs(t, err, ErrNoBackendsMatchPrefix.Err())
    85  	test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
    86  }
    87  
    88  func TestPickerNoSubConnsAvailable(t *testing.T) {
    89  	b, p, _ := setupTest(true)
    90  	b.Build(base.PickerBuildInfo{})
    91  	info := balancer.PickInfo{Ctx: context.Background()}
    92  
    93  	gotPick, err := p.Pick(info)
    94  	test.AssertErrorIs(t, err, balancer.ErrNoSubConnAvailable)
    95  	test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
    96  }
    97  
    98  func setupTest(noSubConns bool) (*pickerBuilder, balancer.Picker, []*subConn) {
    99  	var subConns []*subConn
   100  	bi := base.PickerBuildInfo{
   101  		ReadySCs: make(map[balancer.SubConn]base.SubConnInfo),
   102  	}
   103  
   104  	sc := &subConn{}
   105  	addr := resolver.Address{Addr: "10.77.77.77:8080"}
   106  	sc.UpdateAddresses([]resolver.Address{addr})
   107  
   108  	if !noSubConns {
   109  		bi.ReadySCs[sc] = base.SubConnInfo{Address: addr}
   110  		subConns = append(subConns, sc)
   111  	}
   112  
   113  	b := &pickerBuilder{}
   114  	p := b.Build(bi)
   115  	return b, p, subConns
   116  }
   117  
   118  // subConn is a test mock which implements the balancer.SubConn interface.
   119  type subConn struct {
   120  	balancer.SubConn
   121  	addrs []resolver.Address
   122  }
   123  
   124  func (s *subConn) UpdateAddresses(addrs []resolver.Address) {
   125  	s.addrs = addrs
   126  }