Class: Legate::Auth::TokenManager

Inherits:
Object
  • Object
show all
Defined in:
lib/legate/auth/token_manager.rb

Overview

TokenManager is responsible for managing the lifecycle of authentication tokens. It provides a centralized system for token acquisition, refresh, and invalidation. This class works with the TokenStore for persistence and the various authentication schemes for token operations.

Constant Summary collapse

DEFAULT_CONFIG =

Default configuration values

{
  refresh_buffer: 60,       # Seconds before expiration to trigger refresh
  retry_max_attempts: 3,    # Maximum number of refresh retry attempts
  retry_delay: 2,           # Initial delay between retries (seconds)
  retry_backoff: 1.5,       # Backoff multiplier for subsequent retries
  auto_refresh: true,       # Whether to automatically refresh tokens
  background_refresh: false # Whether to refresh tokens in background
}.freeze

Instance Method Summary collapse

Constructor Details

#initialize(token_store, config = {}) ⇒ TokenManager

Initialize a new TokenManager

Parameters:

  • token_store (Legate::Auth::TokenStore) —

    The token store for persistence

  • config (Hash) (defaults to: {}) —

    Configuration options



27
28
29
30
31
32
33
34
35
36
37
 
# File 'lib/legate/auth/token_manager.rb', line 27

def initialize(token_store, config = {})
  @token_store = token_store
  @config = DEFAULT_CONFIG.merge(config)
  @callbacks = {
    before_expiry: [],
    refresh_success: [],
    refresh_failure: [],
    invalidated: []
  }
  @lock = Mutex.new
end

Instance Method Details

#get_token(scheme, credential, force_refresh: false) ⇒ Legate::Auth::ExchangedCredential?

Get a token for the given scheme and credential

Parameters:

Returns:

Raises:

  • (ArgumentError) 


44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
 
# File 'lib/legate/auth/token_manager.rb', line 44

def get_token(scheme, credential, force_refresh: false)
  raise ArgumentError, 'Scheme must be an Legate::Auth::Scheme' unless scheme.is_a?(Legate::Auth::Scheme)

  cache_key = generate_cache_key(scheme, credential)

  # Use a mutex to prevent race conditions during token retrieval/refresh
  @lock.synchronize do
    # Try to get the token from the store
    token = @token_store.get(cache_key)

    # If no token, refresh it
    return refresh_token(scheme, credential, nil, cache_key) if token.nil?

    # Check if a scheme supports refresh
    supports_refresh = scheme.respond_to?(:supports_refresh?) && scheme.supports_refresh?
    is_refreshable = token.respond_to?(:refreshable?) && token.refreshable?

    # If force refresh is requested, refresh the token regardless of expiration
    if force_refresh
      # Check if the token is refreshable before attempting to refresh
      return refresh_token(scheme, credential, token, cache_key) if supports_refresh && is_refreshable

      # For tokens that aren't refreshable, create a new one
      invalidate_token(cache_key)
      new_token = exchange_token(scheme, credential)
      if new_token
        @token_store.store(cache_key, new_token)
        trigger_callback(:refresh_success, new_token, scheme, credential)
        return new_token
      end

    end

    # Check if token needs refresh based on expiration
    if needs_refresh?(token)
      # Only try to refresh if the scheme supports it and the token is refreshable
      return refresh_token(scheme, credential, token, cache_key) if supports_refresh && is_refreshable

      # If not refreshable, invalidate and return nil
      invalidate_token(cache_key)
      return nil

    end

    # Check if token is approaching expiration and trigger callback
    if approaching_expiration?(token)
      trigger_callback(:before_expiry, token, scheme, credential)

      # If auto_refresh is enabled and we can refresh, do it
      return refresh_token(scheme, credential, token, cache_key) if @config[:auto_refresh] && supports_refresh && is_refreshable
    end

    token
  end
end

#invalidate_token(cache_key) ⇒ Boolean

Invalidate a token, removing it from the store

Parameters:

  • cache_key (String) —

    The cache key for the token

Returns:

  • (Boolean) —

    True if the token was invalidated



171
172
173
174
175
 
# File 'lib/legate/auth/token_manager.rb', line 171

def invalidate_token(cache_key)
  result = @token_store.clear(cache_key)
  trigger_callback(:invalidated, nil, nil, nil, cache_key: cache_key) if result
  result
end

#on(event, &callback) ⇒ self

Register a callback for token lifecycle events

Parameters:

  • event (Symbol) —

    The event to register for (:before_expiry, :refresh_success, :refresh_failure, :invalidated)

  • callback (Proc) —

    The callback to execute

Returns:

  • (self)

Raises:

  • (ArgumentError) 


216
217
218
219
220
221
 
# File 'lib/legate/auth/token_manager.rb', line 216

def on(event, &callback)
  raise ArgumentError, "Unknown event: #{event}. Valid events: #{@callbacks.keys.join(', ')}" unless @callbacks.key?(event)

  @callbacks[event] << callback
  self
end

#refresh_token(scheme, credential, token = nil, cache_key = nil) ⇒ Legate::Auth::ExchangedCredential?

Explicitly refresh a token

Parameters:

Returns:

Raises:

  • (ArgumentError) 


105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
 
# File 'lib/legate/auth/token_manager.rb', line 105

def refresh_token(scheme, credential, token = nil, cache_key = nil)
  raise ArgumentError, 'Scheme must be an Legate::Auth::Scheme' unless scheme.is_a?(Legate::Auth::Scheme)

  cache_key ||= generate_cache_key(scheme, credential)

  # If we don't have a token and it's an oauth or service account scheme,
  # we need to authenticate from scratch
  if token.nil?
    if %i[oauth2 oidc service_account].include?(scheme.scheme_type)
      # For these schemes, we need a complete authentication flow
      # which can't be handled here - return nil to indicate need for full auth
      return nil
    end

    # For other schemes, we can simply apply the credential
    begin
      # Basic auth, API key, etc. - create a new token directly
      token = exchange_token(scheme, credential)
      if token
        @token_store.store(cache_key, token)
        trigger_callback(:refresh_success, token, scheme, credential)
      end
      return token
    rescue Legate::Auth::Error => e
      Legate.logger.error("Failed to create token: #{e.message}")
      trigger_callback(:refresh_failure, nil, scheme, credential, error: e)
      return nil
    end
  end

  # Token exists - attempt to refresh it if scheme supports refresh
  supports_refresh = scheme.respond_to?(:supports_refresh?) && scheme.supports_refresh?
  is_refreshable = token.respond_to?(:refreshable?) && token.refreshable?

  if supports_refresh && is_refreshable
    begin
      refreshed = scheme.refresh_token(token, credential)
      if refreshed
        @token_store.store(cache_key, refreshed)
        trigger_callback(:refresh_success, refreshed, scheme, credential)
        return refreshed
      else
        # Handle the case where refresh_token returns nil but doesn't raise an error
        Legate.logger.error('Failed to refresh token: refresh_token returned nil')
        trigger_callback(:refresh_failure, token, scheme, credential, error: nil)
        return nil
      end
    rescue Legate::Auth::TokenRefreshError => e
      Legate.logger.error("Failed to refresh token: #{e.message}")
      trigger_callback(:refresh_failure, token, scheme, credential, error: e)
      return nil
    end
  end

  # Scheme doesn't support refresh or token isn't refreshable
  # Return existing token if it's not expired
  return token unless token.expired?

  # Otherwise, invalidate it
  invalidate_token(cache_key)
  nil
end

#revoke_token(scheme, credential, token) ⇒ Boolean

Revoke a token with the authentication provider

Parameters:

Returns:

  • (Boolean) —

    True if the token was revoked

Raises:

  • (ArgumentError) 


182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
 
# File 'lib/legate/auth/token_manager.rb', line 182

def revoke_token(scheme, credential, token)
  raise ArgumentError, 'Scheme must be an Legate::Auth::Scheme' unless scheme.is_a?(Legate::Auth::Scheme)
  raise ArgumentError, 'Token must be an ExchangedCredential' unless token.is_a?(Legate::Auth::ExchangedCredential)

  # Check if scheme supports revocation
  unless scheme.respond_to?(:revoke_token)
    Legate.logger.warn("Scheme #{scheme.scheme_type} does not support token revocation")
    return false
  end

  begin
    # Attempt to revoke the token
    result = scheme.revoke_token(token, credential)

    # Invalidate the token in our store if revocation succeeded
    if result
      cache_key = generate_cache_key(scheme, credential)
      invalidate_token(cache_key)
    end

    result
  rescue Legate::Auth::Error => e
    Legate.logger.error("Failed to revoke token: #{e.message}")
    false
  rescue NotImplementedError => e
    Legate.logger.error("#{e.message}")
    false
  end
end