From 1757750e580a0ff03fec877ba792a69825a3b720 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Mon, 15 Apr 2024 18:53:33 -0400 Subject: [PATCH] Added unmarshal attribute for expires_in for device flow auth token (#4319) Signed-off-by: Eduardo Apolinario --- flyteidl/clients/go/admin/deviceflow/payload.go | 3 ++- flyteidl/clients/go/admin/deviceflow/token_orchestrator.go | 6 ++++-- .../clients/go/admin/deviceflow/token_orchestrator_test.go | 2 ++ 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/flyteidl/clients/go/admin/deviceflow/payload.go b/flyteidl/clients/go/admin/deviceflow/payload.go index 656e1a88d8..38e8b5502f 100644 --- a/flyteidl/clients/go/admin/deviceflow/payload.go +++ b/flyteidl/clients/go/admin/deviceflow/payload.go @@ -40,5 +40,6 @@ type DeviceAccessTokenRequest struct { type DeviceAccessTokenResponse struct { oauth2.Token - Error string `json:"error"` + Error string `json:"error"` + ExpiresIn int64 `json:"expires_in"` // relative seconds from now } diff --git a/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go b/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go index fc4f1d3786..e934ec6949 100644 --- a/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go +++ b/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go @@ -125,12 +125,14 @@ func (t TokenOrchestrator) PollTokenEndpoint(ctx context.Context, tokReq DeviceA // Unmarshalled response if it contains an error then check if we need to increase the polling interval if len(tokResp.Error) > 0 { if tokResp.Error == errSlowDown || tokResp.Error == errAuthPending { - pollInterval = pollInterval * 2 - + logger.Debugf(ctx, "going to poll again due to error %v", tokResp.Error) } else { return nil, fmt.Errorf("oauth error : %v", tokResp.Error) } } else { + if secs := tokResp.ExpiresIn; secs > 0 { + tokResp.Token.Expiry = time.Now().Add(time.Duration(secs) * time.Second) + } // Got the auth token in the response and save it in the cache err = t.TokenCache.SaveToken(&tokResp.Token) // Saving into the cache is only considered to be a warning in this case. diff --git a/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go b/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go index d5a22c523a..5c1dc5f2bd 100644 --- a/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go +++ b/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go @@ -85,6 +85,7 @@ func TestFetchFromAuthFlow(t *testing.T) { Token: oauth2.Token{ AccessToken: "access_token", }, + ExpiresIn: 300, } darBytes, err := json.Marshal(dar) assert.Nil(t, err) @@ -119,5 +120,6 @@ func TestFetchFromAuthFlow(t *testing.T) { assert.Nil(t, err) assert.NotNil(t, authToken) assert.Equal(t, "access_token", authToken.AccessToken) + assert.True(t, authToken.Expiry.After(time.Now().Add(time.Second*200))) }) }