clientv3: move auth token credential to "credentials" package

Signed-off-by: Gyuho Lee <leegyuho@amazon.com>
This commit is contained in:
Gyuho Lee 2019-07-22 00:11:17 -07:00
parent db61ee106c
commit 3dc00ab615
2 changed files with 27 additions and 39 deletions

View File

@ -86,9 +86,8 @@ type Client struct {
// Username is a user name for authentication. // Username is a user name for authentication.
Username string Username string
// Password is a password for authentication. // Password is a password for authentication.
Password string Password string
// tokenCred is an instance of WithPerRPCCredentials()'s argument authTokenBundle credentials.Bundle
tokenCred *authTokenCredential
callOpts []grpc.CallOption callOpts []grpc.CallOption
@ -193,23 +192,6 @@ func (c *Client) autoSync() {
} }
} }
type authTokenCredential struct {
token string
tokenMu *sync.RWMutex
}
func (cred authTokenCredential) RequireTransportSecurity() bool {
return false
}
func (cred authTokenCredential) GetRequestMetadata(ctx context.Context, s ...string) (map[string]string, error) {
cred.tokenMu.RLock()
defer cred.tokenMu.RUnlock()
return map[string]string{
rpctypes.TokenFieldNameGRPC: cred.token,
}, nil
}
func (c *Client) processCreds(scheme string) (creds grpccredentials.TransportCredentials) { func (c *Client) processCreds(scheme string) (creds grpccredentials.TransportCredentials) {
creds = c.creds creds = c.creds
switch scheme { switch scheme {
@ -316,10 +298,7 @@ func (c *Client) getToken(ctx context.Context) error {
continue continue
} }
c.tokenCred.tokenMu.Lock() c.authTokenBundle.UpdateAuthToken(resp.Token)
c.tokenCred.token = resp.Token
c.tokenCred.tokenMu.Unlock()
return nil return nil
} }
@ -343,9 +322,7 @@ func (c *Client) dial(target string, creds grpccredentials.TransportCredentials,
} }
if c.Username != "" && c.Password != "" { if c.Username != "" && c.Password != "" {
c.tokenCred = &authTokenCredential{ c.authTokenBundle = credentials.NewBundle(credentials.Config{})
tokenMu: &sync.RWMutex{},
}
ctx, cancel := c.ctx, func() {} ctx, cancel := c.ctx, func() {}
if c.cfg.DialTimeout > 0 { if c.cfg.DialTimeout > 0 {
@ -362,7 +339,7 @@ func (c *Client) dial(target string, creds grpccredentials.TransportCredentials,
return nil, err return nil, err
} }
} else { } else {
opts = append(opts, grpc.WithPerRPCCredentials(c.tokenCred)) opts = append(opts, grpc.WithPerRPCCredentials(c.authTokenBundle.PerRPCCredentials()))
} }
cancel() cancel()
} }

View File

@ -29,14 +29,19 @@ import (
// Config defines gRPC credential configuration. // Config defines gRPC credential configuration.
type Config struct { type Config struct {
TLSConfig *tls.Config TLSConfig *tls.Config
AuthToken string }
// Bundle defines gRPC credential interface.
type Bundle interface {
grpccredentials.Bundle
UpdateAuthToken(token string)
} }
// NewBundle constructs a new gRPC credential bundle. // NewBundle constructs a new gRPC credential bundle.
func NewBundle(cfg Config) grpccredentials.Bundle { func NewBundle(cfg Config) Bundle {
return &bundle{ return &bundle{
tc: newTransportCredential(cfg.TLSConfig), tc: newTransportCredential(cfg.TLSConfig),
rc: newPerRPCCredential(cfg.AuthToken), rc: newPerRPCCredential(),
} }
} }
@ -125,14 +130,7 @@ type perRPCCredential struct {
authTokenMu sync.RWMutex authTokenMu sync.RWMutex
} }
func newPerRPCCredential(authToken string) *perRPCCredential { func newPerRPCCredential() *perRPCCredential { return &perRPCCredential{} }
if authToken == "" {
return nil
}
return &perRPCCredential{
authToken: authToken,
}
}
func (rc *perRPCCredential) RequireTransportSecurity() bool { return false } func (rc *perRPCCredential) RequireTransportSecurity() bool { return false }
@ -142,3 +140,16 @@ func (rc *perRPCCredential) GetRequestMetadata(ctx context.Context, s ...string)
rc.authTokenMu.RUnlock() rc.authTokenMu.RUnlock()
return map[string]string{rpctypes.TokenFieldNameGRPC: authToken}, nil return map[string]string{rpctypes.TokenFieldNameGRPC: authToken}, nil
} }
func (b *bundle) UpdateAuthToken(token string) {
if b.rc == nil {
return
}
b.rc.UpdateAuthToken(token)
}
func (rc *perRPCCredential) UpdateAuthToken(token string) {
rc.authTokenMu.Lock()
rc.authToken = token
rc.authTokenMu.Unlock()
}