credentials/alts: Change ALTS hsConn to a map (#3741)

This commit is contained in:
Ryan Kim
2020-07-15 14:08:37 -06:00
committed by GitHub
parent 46d42f139d
commit 9fcde86ebe
2 changed files with 39 additions and 20 deletions

View File

@ -27,9 +27,12 @@ import (
)
var (
// hsConn represents a connection to hypervisor handshaker service.
hsConn *grpc.ClientConn
mu sync.Mutex
// mu guards hsConnMap and hsDialer.
mu sync.Mutex
// hsConn represents a mapping from a hypervisor handshaker service address
// to a corresponding connection to a hypervisor handshaker service
// instance.
hsConnMap = make(map[string]*grpc.ClientConn)
// hsDialer will be reassigned in tests.
hsDialer = grpc.Dial
)
@ -41,7 +44,8 @@ func Dial(hsAddress string) (*grpc.ClientConn, error) {
mu.Lock()
defer mu.Unlock()
if hsConn == nil {
hsConn, ok := hsConnMap[hsAddress]
if !ok {
// Create a new connection to the handshaker service. Note that
// this connection stays open until the application is closed.
var err error
@ -49,6 +53,7 @@ func Dial(hsAddress string) (*grpc.ClientConn, error) {
if err != nil {
return nil, err
}
hsConnMap[hsAddress] = hsConn
}
return hsConn, nil
}

View File

@ -25,8 +25,8 @@ import (
)
const (
// The address is irrelevant in this test.
testAddress = "some_address"
testAddress1 = "some_address_1"
testAddress2 = "some_address_2"
)
func TestDial(t *testing.T) {
@ -40,30 +40,44 @@ func TestDial(t *testing.T) {
}
}()
// Ensure that hsConn is nil at first.
hsConn = nil
// First call to Dial, it should create set hsConn.
conn1, err := Dial(testAddress)
// First call to Dial, it should create a connection to the server running
// at the given address.
conn1, err := Dial(testAddress1)
if err != nil {
t.Fatalf("first call to Dial failed: %v", err)
t.Fatalf("first call to Dial(%v) failed: %v", testAddress1, err)
}
if conn1 == nil {
t.Fatal("first call to Dial(_)=(nil, _), want not nil")
t.Fatalf("first call to Dial(%v)=(nil, _), want not nil", testAddress1)
}
if got, want := hsConn, conn1; got != want {
t.Fatalf("hsConn=%v, want %v", got, want)
if got, want := hsConnMap[testAddress1], conn1; got != want {
t.Fatalf("hsConnMap[%v]=%v, want %v", testAddress1, got, want)
}
// Second call to Dial should return conn1 above.
conn2, err := Dial(testAddress)
conn2, err := Dial(testAddress1)
if err != nil {
t.Fatalf("second call to Dial(_) failed: %v", err)
t.Fatalf("second call to Dial(%v) failed: %v", testAddress1, err)
}
if got, want := conn2, conn1; got != want {
t.Fatalf("second call to Dial(_)=(%v, _), want (%v,. _)", got, want)
t.Fatalf("second call to Dial(%v)=(%v, _), want (%v,. _)", testAddress1, got, want)
}
if got, want := hsConn, conn1; got != want {
t.Fatalf("hsConn=%v, want %v", got, want)
if got, want := hsConnMap[testAddress1], conn1; got != want {
t.Fatalf("hsConnMap[%v]=%v, want %v", testAddress1, got, want)
}
// Third call to Dial using a different address should create a new
// connection.
conn3, err := Dial(testAddress2)
if err != nil {
t.Fatalf("third call to Dial(%v) failed: %v", testAddress2, err)
}
if conn3 == nil {
t.Fatalf("third call to Dial(%v)=(nil, _), want not nil", testAddress2)
}
if got, want := hsConnMap[testAddress2], conn3; got != want {
t.Fatalf("hsConnMap[%v]=%v, want %v", testAddress2, got, want)
}
if got, want := conn2 == conn3, false; got != want {
t.Fatalf("(conn2==conn3)=%v, want %v", got, want)
}
}