Skip to content

Commit

Permalink
Keep unfit RTs b/c they might be useful for others
Browse files Browse the repository at this point in the history
Mark last_modification_time for RT
  • Loading branch information
rayluo committed Mar 3, 2021
1 parent 78e9ccf commit 26f0689
Show file tree
Hide file tree
Showing 4 changed files with 20 additions and 7 deletions.
13 changes: 11 additions & 2 deletions msal/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -918,11 +918,20 @@ def _acquire_token_silent_by_finding_specific_refresh_token(
client = self._build_client(self.client_credential, authority)

response = None # A distinguishable value to mean cache is empty
for entry in matches:
for entry in sorted( # Since unfit RTs would not be aggressively removed,
# we start from newer RTs which are more likely fit.
matches,
key=lambda e: int(e.get("last_modification_time", "0")),
reverse=True):
logger.debug("Cache attempts an RT")
response = client.obtain_token_by_refresh_token(
entry, rt_getter=lambda token_item: token_item["secret"],
on_removing_rt=rt_remover or self.token_cache.remove_rt,
on_removing_rt=(rt_remover or self.token_cache.remove_rt)
if # we can remove a RT when a single scope is an exact match
len(scopes) == 1
and set(entry.get("target", "").split()) <= set(scopes)
else # otherwise keep the RT as it might work for a subset of scopes
lambda rt_item: None, # No OP
on_obtaining_tokens=lambda event: self.token_cache.add(dict(
event,
environment=authority.instance,
Expand Down
9 changes: 6 additions & 3 deletions msal/token_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,9 +148,9 @@ def __add(self, event, now=None):
target = ' '.join(event.get("scope", [])) # Per schema, we don't sort it

with self._lock:
now = int(time.time() if now is None else now)

if access_token:
now = int(time.time() if now is None else now)
expires_in = int( # AADv1-like endpoint returns a string
response.get("expires_in", 3599))
ext_expires_in = int( # AADv1-like endpoint returns a string
Expand Down Expand Up @@ -212,6 +212,7 @@ def __add(self, event, now=None):
"environment": environment,
"client_id": event.get("client_id"),
"target": target, # Optional per schema though
"last_modification_time": str(now), # Optional. Schema defines it as a string.
}
if "foci" in response:
rt["family_id"] = response["foci"]
Expand Down Expand Up @@ -249,8 +250,10 @@ def remove_rt(self, rt_item):

def update_rt(self, rt_item, new_rt):
assert rt_item.get("credential_type") == self.CredentialType.REFRESH_TOKEN
return self.modify(
self.CredentialType.REFRESH_TOKEN, rt_item, {"secret": new_rt})
return self.modify(self.CredentialType.REFRESH_TOKEN, rt_item, {
"secret": new_rt,
"last_modification_time": str(int(time.time())), # Optional. Schema defines it as a string.
})

def remove_at(self, at_item):
assert at_item.get("credential_type") == self.CredentialType.ACCESS_TOKEN
Expand Down
3 changes: 1 addition & 2 deletions tests/test_application.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,8 +73,7 @@ def setUp(self):
self.client_id, authority=self.authority_url, token_cache=self.cache)

def test_cache_empty_will_be_returned_as_None(self):
self.assertEqual(
None, self.app.acquire_token_silent(['cache_miss'], self.account))
self.app.token_cache = msal.SerializableTokenCache() # Reset it to empty
self.assertEqual(
None, self.app.acquire_token_silent_with_error(['cache_miss'], self.account))

Expand Down
2 changes: 2 additions & 0 deletions tests/test_token_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ def testAddByAad(self):
'credential_type': 'RefreshToken',
'environment': 'login.example.com',
'home_account_id': "uid.utid",
'last_modification_time': '1000',
'secret': 'a refresh token',
'target': 's2 s1 s3',
},
Expand Down Expand Up @@ -157,6 +158,7 @@ def testAddByAdfs(self):
'credential_type': 'RefreshToken',
'environment': 'fs.msidlab8.com',
'home_account_id': "subject",
'last_modification_time': "1000",
'secret': 'a refresh token',
'target': 's2 s1 s3',
},
Expand Down

0 comments on commit 26f0689

Please sign in to comment.