github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/sink/kafka/oauth2_token_provider.go (about)

     1  // Copyright 2023 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package kafka
    15  
    16  import (
    17  	"context"
    18  	"net/url"
    19  
    20  	"github.com/IBM/sarama"
    21  	"github.com/pingcap/errors"
    22  	"golang.org/x/oauth2"
    23  	"golang.org/x/oauth2/clientcredentials"
    24  )
    25  
    26  // tsokenProvider is a user-defined callback for generating
    27  // access tokens for SASL/OAUTHBEARER auth.
    28  type tokenProvider struct {
    29  	tokenSource oauth2.TokenSource
    30  }
    31  
    32  var _ sarama.AccessTokenProvider = (*tokenProvider)(nil)
    33  
    34  // Token implements the sarama.AccessTokenProvider interface.
    35  // Token returns an access token. The implementation should ensure token
    36  // reuse so that multiple calls at connect time do not create multiple
    37  // tokens. The implementation should also periodically refresh the token in
    38  // order to guarantee that each call returns an unexpired token.  This
    39  // method should not block indefinitely--a timeout error should be returned
    40  // after a short period of inactivity so that the broker connection logic
    41  // can log debugging information and retry.
    42  func (t *tokenProvider) Token() (*sarama.AccessToken, error) {
    43  	token, err := t.tokenSource.Token()
    44  	if err != nil {
    45  		// Errors will result in Sarama retrying the broker connection and logging
    46  		// the transient error, with a Broker connection error surfacing after retry
    47  		// attempts have been exhausted.
    48  		return nil, err
    49  	}
    50  
    51  	return &sarama.AccessToken{Token: token.AccessToken}, nil
    52  }
    53  
    54  func newTokenProvider(ctx context.Context, o *Options) (sarama.AccessTokenProvider, error) {
    55  	// grant_type is by default going to be set to 'client_credentials' by the
    56  	// clientcredentials library as defined by the spec, however non-compliant
    57  	// auth server implementations may want a custom type
    58  	var endpointParams url.Values
    59  	if o.SASL.OAuth2.GrantType != "" {
    60  		if endpointParams == nil {
    61  			endpointParams = url.Values{}
    62  		}
    63  		endpointParams.Set("grant_type", o.SASL.OAuth2.GrantType)
    64  	}
    65  
    66  	// audience is an optional parameter that can be used to specify the
    67  	// intended audience of the token.
    68  	if o.SASL.OAuth2.Audience != "" {
    69  		if endpointParams == nil {
    70  			endpointParams = url.Values{}
    71  		}
    72  		endpointParams.Set("audience", o.SASL.OAuth2.Audience)
    73  	}
    74  
    75  	tokenURL, err := url.Parse(o.SASL.OAuth2.TokenURL)
    76  	if err != nil {
    77  		return nil, errors.Trace(err)
    78  	}
    79  
    80  	cfg := clientcredentials.Config{
    81  		ClientID:       o.SASL.OAuth2.ClientID,
    82  		ClientSecret:   o.SASL.OAuth2.ClientSecret,
    83  		TokenURL:       tokenURL.String(),
    84  		EndpointParams: endpointParams,
    85  		Scopes:         o.SASL.OAuth2.Scopes,
    86  	}
    87  	return &tokenProvider{
    88  		tokenSource: cfg.TokenSource(ctx),
    89  	}, nil
    90  }