Change all callsites so requestToken just takes type+value

deephaven · Aug 16, 2024 · e1e98a0 · e1e98a0
1 parent 2db1e60
commit e1e98a0
Showing 1 changed file with 16 additions and 15 deletions.
diff --git a/go/pkg/client/tokens.go b/go/pkg/client/tokens.go
@@ -38,8 +38,19 @@ func withAuthToken(ctx context.Context, token []byte) context.Context {
 }
 
 // requestToken requests a new token from flight.
-func requestToken(handshakeClient flight.FlightService_HandshakeClient, handshakeReq *flight.HandshakeRequest) ([]byte, error) {
-	err := handshakeClient.Send(handshakeReq)
+func requestToken(handshakeClient flight.FlightService_HandshakeClient, authType string, authToken []byte) ([]byte, error) {
+
+	war := sessionpb2.WrappedAuthenticationRequest{
+		Type:    authType,
+		Payload: authToken,
+	}
+	payload, err := proto.Marshal(&war)
+	if err != nil {
+		return nil, err
+	}
+	handshakeReq := flight.HandshakeRequest{Payload: []byte(payload)}
+
+	err = handshakeClient.Send(&handshakeReq)
 
 	if err != nil {
 		return nil, err
@@ -124,23 +135,13 @@ func (tr *tokenManager) Close() error {
 // "user:password"; when auth_type is DefaultAuth, it will be ignored; when auth_type is a custom-built
 // authenticator, it must conform to the specific requirement of the authenticator.
 func newTokenManager(ctx context.Context, fs *flightStub, cfg configpb2.ConfigServiceClient, authType string, authToken string) (*tokenManager, error) {
-	authString := makeAuthString(authType, authToken)
-
 	handshakeClient, err := fs.handshake(ctx)
 
 	if err != nil {
 		return nil, err
 	}
 
-	war := sessionpb2.WrappedAuthenticationRequest{
-		Type:    authType,
-		Payload: []byte(authToken),
-	}
-	payload, err := proto.Marshal(&war)
-	if err != nil {
-		return nil, err
-	}
-	tkn, err := requestToken(handshakeClient, &flight.HandshakeRequest{Payload: []byte(payload)})
+	tkn, err := requestToken(handshakeClient, authType, []byte(authToken))
 
 	if err != nil {
 		return nil, err
@@ -184,10 +185,10 @@ func newTokenManager(ctx context.Context, fs *flightStub, cfg configpb2.ConfigSe
 				var tkn []byte
 
 				if err == nil {
-					tkn, err = requestToken(handshakeClient, &flight.HandshakeRequest{Payload: oldToken})
+					tkn, err = requestToken(handshakeClient, "Bearer", oldToken)
 				} else {
 					log.Println("Old token has an error during token update.  Attempting to acquire a fresh token.  err=", err)
-					tkn, err = requestToken(handshakeClient, &flight.HandshakeRequest{Payload: []byte(authString)})
+					tkn, err = requestToken(handshakeClient, authType, []byte(authToken))
 				}
 
 				if err != nil {