Skip to content
This repository was archived by the owner on Oct 23, 2023. It is now read-only.

Commit c1c892a

Browse files
pmahindrakar-osspmahindrakar-oss
authored
Adding configurable audience property for flyte clients (#329)
* Adding configurable audience property for flyte clients Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> * changed the const audience to audienceKey Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> * fixed unit tests Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> * fixed unit test Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> * nit Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> * feedback Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> * refactored unit tests Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> * Added UseAudienceFromAdmin property to force pull audience from admin config. Default is false and expects clients to pass it Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> * Added test for expected number of calls to the public admin endpoint Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> * fixed the tests Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com> Signed-off-by: pmahindrakar-oss <prafulla.mahindrakar@gmail.com>
1 parent 9fbac98 commit c1c892a

6 files changed

Lines changed: 155 additions & 36 deletions

File tree

‎clients/go/admin/auth_interceptor_test.go‎

Lines changed: 24 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -13,22 +13,18 @@ import (
1313
"sync"
1414
"testing"
1515

16-
"github.com/flyteorg/flytestdlib/logger"
17-
18-
"k8s.io/apimachinery/pkg/util/rand"
19-
20-
mocks2 "github.com/flyteorg/flyteidl/clients/go/admin/mocks"
21-
"github.com/stretchr/testify/mock"
22-
23-
service2 "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
24-
"github.com/flyteorg/flytestdlib/config"
25-
2616
"github.com/stretchr/testify/assert"
27-
28-
"github.com/flyteorg/flyteidl/clients/go/admin/cache/mocks"
17+
"github.com/stretchr/testify/mock"
2918
"google.golang.org/grpc"
3019
"google.golang.org/grpc/codes"
3120
"google.golang.org/grpc/status"
21+
"k8s.io/apimachinery/pkg/util/rand"
22+
23+
"github.com/flyteorg/flyteidl/clients/go/admin/cache/mocks"
24+
adminMocks "github.com/flyteorg/flyteidl/clients/go/admin/mocks"
25+
"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
26+
"github.com/flyteorg/flytestdlib/config"
27+
"github.com/flyteorg/flytestdlib/logger"
3228
)
3329

3430
// authMetadataServer is a fake AuthMetadataServer that takes in an AuthMetadataServer implementation (usually one
@@ -39,15 +35,15 @@ type authMetadataServer struct {
3935
port int
4036
grpcServer *grpc.Server
4137
netListener net.Listener
42-
impl service2.AuthMetadataServiceServer
38+
impl service.AuthMetadataServiceServer
4339
lck *sync.RWMutex
4440
}
4541

46-
func (s authMetadataServer) GetOAuth2Metadata(ctx context.Context, in *service2.OAuth2MetadataRequest) (*service2.OAuth2MetadataResponse, error) {
42+
func (s authMetadataServer) GetOAuth2Metadata(ctx context.Context, in *service.OAuth2MetadataRequest) (*service.OAuth2MetadataResponse, error) {
4743
return s.impl.GetOAuth2Metadata(ctx, in)
4844
}
4945

50-
func (s authMetadataServer) GetPublicClientConfig(ctx context.Context, in *service2.PublicClientAuthConfigRequest) (*service2.PublicClientAuthConfigResponse, error) {
46+
func (s authMetadataServer) GetPublicClientConfig(ctx context.Context, in *service.PublicClientAuthConfigRequest) (*service.PublicClientAuthConfigResponse, error) {
5147
return s.impl.GetPublicClientConfig(ctx, in)
5248
}
5349

@@ -84,7 +80,7 @@ func (s *authMetadataServer) Start(_ context.Context) error {
8480
}
8581

8682
grpcS := grpc.NewServer()
87-
service2.RegisterAuthMetadataServiceServer(grpcS, s)
83+
service.RegisterAuthMetadataServiceServer(grpcS, s)
8884
go func() {
8985
_ = grpcS.Serve(lis)
9086
//assert.NoError(s.t, err)
@@ -106,7 +102,7 @@ func (s *authMetadataServer) Close() {
106102
s.s.Close()
107103
}
108104

109-
func newAuthMetadataServer(t testing.TB, port int, impl service2.AuthMetadataServiceServer) *authMetadataServer {
105+
func newAuthMetadataServer(t testing.TB, port int, impl service.AuthMetadataServiceServer) *authMetadataServer {
110106
return &authMetadataServer{
111107
port: port,
112108
t: t,
@@ -132,13 +128,13 @@ func Test_newAuthInterceptor(t *testing.T) {
132128
}))
133129

134130
port := rand.IntnRange(10000, 60000)
135-
m := &mocks2.AuthMetadataServiceServer{}
136-
m.OnGetOAuth2MetadataMatch(mock.Anything, mock.Anything).Return(&service2.OAuth2MetadataResponse{
131+
m := &adminMocks.AuthMetadataServiceServer{}
132+
m.OnGetOAuth2MetadataMatch(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{
137133
AuthorizationEndpoint: fmt.Sprintf("http://localhost:%d/oauth2/authorize", port),
138134
TokenEndpoint: fmt.Sprintf("http://localhost:%d/oauth2/token", port),
139135
JwksUri: fmt.Sprintf("http://localhost:%d/oauth2/jwks", port),
140136
}, nil)
141-
m.OnGetPublicClientConfigMatch(mock.Anything, mock.Anything).Return(&service2.PublicClientAuthConfigResponse{
137+
m.OnGetPublicClientConfigMatch(mock.Anything, mock.Anything).Return(&service.PublicClientAuthConfigResponse{
142138
Scopes: []string{"all"},
143139
}, nil)
144140
s := newAuthMetadataServer(t, port, m)
@@ -171,7 +167,7 @@ func Test_newAuthInterceptor(t *testing.T) {
171167
}))
172168

173169
port := rand.IntnRange(10000, 60000)
174-
m := &mocks2.AuthMetadataServiceServer{}
170+
m := &adminMocks.AuthMetadataServiceServer{}
175171
s := newAuthMetadataServer(t, port, m)
176172
ctx := context.Background()
177173
assert.NoError(t, s.Start(ctx))
@@ -201,13 +197,13 @@ func Test_newAuthInterceptor(t *testing.T) {
201197
}))
202198

203199
port := rand.IntnRange(10000, 60000)
204-
m := &mocks2.AuthMetadataServiceServer{}
205-
m.OnGetOAuth2MetadataMatch(mock.Anything, mock.Anything).Return(&service2.OAuth2MetadataResponse{
200+
m := &adminMocks.AuthMetadataServiceServer{}
201+
m.OnGetOAuth2MetadataMatch(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{
206202
AuthorizationEndpoint: fmt.Sprintf("http://localhost:%d/oauth2/authorize", port),
207203
TokenEndpoint: fmt.Sprintf("http://localhost:%d/oauth2/token", port),
208204
JwksUri: fmt.Sprintf("http://localhost:%d/oauth2/jwks", port),
209205
}, nil)
210-
m.OnGetPublicClientConfigMatch(mock.Anything, mock.Anything).Return(&service2.PublicClientAuthConfigResponse{
206+
m.OnGetPublicClientConfigMatch(mock.Anything, mock.Anything).Return(&service.PublicClientAuthConfigResponse{
211207
Scopes: []string{"all"},
212208
}, nil)
213209

@@ -237,8 +233,8 @@ func Test_newAuthInterceptor(t *testing.T) {
237233

238234
func TestMaterializeCredentials(t *testing.T) {
239235
port := rand.IntnRange(10000, 60000)
240-
t.Run("No public client config or oauth2 metadata endpoint lookup", func(t *testing.T) {
241-
m := &mocks2.AuthMetadataServiceServer{}
236+
t.Run("No oauth2 metadata endpoint or Public client config lookup", func(t *testing.T) {
237+
m := &adminMocks.AuthMetadataServiceServer{}
242238
m.OnGetOAuth2MetadataMatch(mock.Anything, mock.Anything).Return(nil, errors.New("unexpected call to get oauth2 metadata"))
243239
m.OnGetPublicClientConfigMatch(mock.Anything, mock.Anything).Return(nil, errors.New("unexpected call to get public client config"))
244240
s := newAuthMetadataServer(t, port, m)
@@ -256,12 +252,13 @@ func TestMaterializeCredentials(t *testing.T) {
256252
AuthType: AuthTypeClientSecret,
257253
TokenURL: fmt.Sprintf("http://localhost:%d/api/v1/token", port),
258254
Scopes: []string{"all"},
255+
Audience: "http://localhost:30081",
259256
AuthorizationHeader: "authorization",
260257
}, &mocks.TokenCache{}, f)
261258
assert.NoError(t, err)
262259
})
263260
t.Run("Failed to fetch client metadata", func(t *testing.T) {
264-
m := &mocks2.AuthMetadataServiceServer{}
261+
m := &adminMocks.AuthMetadataServiceServer{}
265262
m.OnGetOAuth2MetadataMatch(mock.Anything, mock.Anything).Return(nil, errors.New("unexpected call to get oauth2 metadata"))
266263
failedPublicClientConfigLookup := errors.New("expected err")
267264
m.OnGetPublicClientConfigMatch(mock.Anything, mock.Anything).Return(nil, failedPublicClientConfigLookup)

‎clients/go/admin/config.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,8 @@ type Config struct {
5353
ClientSecretLocation string `json:"clientSecretLocation" pflag:",File containing the client secret"`
5454
ClientSecretEnvVar string `json:"clientSecretEnvVar" pflag:",Environment variable containing the client secret"`
5555
Scopes []string `json:"scopes" pflag:",List of scopes to request"`
56+
UseAudienceFromAdmin bool `json:"useAudienceFromAdmin" pflag:",Use Audience configured from admins public endpoint config."`
57+
Audience string `json:"audience" pflag:",Audience to use when initiating OAuth2 authorization requests."`
5658

5759
// There are two ways to get the token URL. If the authorization server url is provided, the client will try to use RFC 8414 to
5860
// try to get the token URL. Or it can be specified directly through TokenURL config.

‎clients/go/admin/config_flags.go‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎clients/go/admin/config_flags_test.go‎

Lines changed: 28 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎clients/go/admin/token_source_provider.go‎

Lines changed: 29 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"fmt"
66
"io/ioutil"
7+
"net/url"
78
"os"
89
"strings"
910
"sync"
@@ -22,6 +23,10 @@ import (
2223
"github.com/flyteorg/flytestdlib/logger"
2324
)
2425

26+
const (
27+
audienceKey = "audience"
28+
)
29+
2530
// TokenSourceProvider defines the interface needed to provide a TokenSource that is used to
2631
// create a client with authentication enabled.
2732
type TokenSourceProvider interface {
@@ -46,15 +51,24 @@ func NewTokenSourceProvider(ctx context.Context, cfg *Config, tokenCache cache.T
4651
}
4752

4853
scopes := cfg.Scopes
49-
if len(scopes) == 0 {
50-
clientMetadata, err := authClient.GetPublicClientConfig(ctx, &service.PublicClientAuthConfigRequest{})
54+
audienceValue := cfg.Audience
55+
56+
if len(scopes) == 0 || cfg.UseAudienceFromAdmin {
57+
publicClientConfig, err := authClient.GetPublicClientConfig(ctx, &service.PublicClientAuthConfigRequest{})
5158
if err != nil {
5259
return nil, fmt.Errorf("failed to fetch client metadata. Error: %v", err)
5360
}
54-
scopes = clientMetadata.Scopes
61+
// Update scopes from publicClientConfig
62+
if len(scopes) == 0 {
63+
scopes = publicClientConfig.Scopes
64+
}
65+
// Update audience from publicClientConfig
66+
if cfg.UseAudienceFromAdmin {
67+
audienceValue = publicClientConfig.Audience
68+
}
5569
}
5670

57-
tokenProvider, err = NewClientCredentialsTokenSourceProvider(ctx, cfg, scopes, tokenURL)
71+
tokenProvider, err = NewClientCredentialsTokenSourceProvider(ctx, cfg, scopes, tokenURL, audienceValue)
5872
if err != nil {
5973
return nil, err
6074
}
@@ -152,7 +166,7 @@ type ClientCredentialsTokenSourceProvider struct {
152166
TokenRefreshWindow time.Duration
153167
}
154168

155-
func NewClientCredentialsTokenSourceProvider(ctx context.Context, cfg *Config, scopes []string, tokenURL string) (TokenSourceProvider, error) {
169+
func NewClientCredentialsTokenSourceProvider(ctx context.Context, cfg *Config, scopes []string, tokenURL string, audience string) (TokenSourceProvider, error) {
156170
var secret string
157171
if len(cfg.ClientSecretEnvVar) > 0 {
158172
secret = os.Getenv(cfg.ClientSecretEnvVar)
@@ -164,13 +178,19 @@ func NewClientCredentialsTokenSourceProvider(ctx context.Context, cfg *Config, s
164178
}
165179
secret = string(secretBytes)
166180
}
181+
endpointParams := url.Values{}
182+
if len(audience) > 0 {
183+
endpointParams = url.Values{audienceKey: {audience}}
184+
}
167185
secret = strings.TrimSpace(secret)
168186
return ClientCredentialsTokenSourceProvider{
169187
ccConfig: clientcredentials.Config{
170-
ClientID: cfg.ClientID,
171-
ClientSecret: secret,
172-
TokenURL: tokenURL,
173-
Scopes: scopes},
188+
ClientID: cfg.ClientID,
189+
ClientSecret: secret,
190+
TokenURL: tokenURL,
191+
Scopes: scopes,
192+
EndpointParams: endpointParams,
193+
},
174194
TokenRefreshWindow: cfg.TokenRefreshWindow.Duration}, nil
175195
}
176196

‎clients/go/admin/token_source_test.go‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,16 @@ package admin
22

33
import (
44
"context"
5+
"net/url"
56
"testing"
67

78
"github.com/stretchr/testify/assert"
9+
"github.com/stretchr/testify/mock"
810
"golang.org/x/oauth2"
11+
12+
tokenCacheMocks "github.com/flyteorg/flyteidl/clients/go/admin/cache/mocks"
13+
adminMocks "github.com/flyteorg/flyteidl/clients/go/admin/mocks"
14+
"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
915
)
1016

1117
type DummyTestTokenSource struct {
@@ -25,3 +31,67 @@ func TestNewTokenSource(t *testing.T) {
2531
assert.NoError(t, err)
2632
assert.Equal(t, "Bearer abc", metadata["test"])
2733
}
34+
35+
func TestNewTokenSourceProvider(t *testing.T) {
36+
ctx := context.Background()
37+
tests := []struct {
38+
name string
39+
audienceCfg string
40+
scopesCfg []string
41+
useAudienceFromAdmin bool
42+
clientConfigResponse service.PublicClientAuthConfigResponse
43+
expectedAudience string
44+
expectedScopes []string
45+
expectedCallsPubEndpoint int
46+
}{
47+
{
48+
name: "audience from client config",
49+
audienceCfg: "clientConfiguredAud",
50+
scopesCfg: []string{"all"},
51+
clientConfigResponse: service.PublicClientAuthConfigResponse{},
52+
expectedAudience: "clientConfiguredAud",
53+
expectedScopes: []string{"all"},
54+
expectedCallsPubEndpoint: 0,
55+
},
56+
{
57+
name: "audience from public client response",
58+
audienceCfg: "clientConfiguredAud",
59+
useAudienceFromAdmin: true,
60+
scopesCfg: []string{"all"},
61+
clientConfigResponse: service.PublicClientAuthConfigResponse{Audience: "AdminConfiguredAud", Scopes: []string{}},
62+
expectedAudience: "AdminConfiguredAud",
63+
expectedScopes: []string{"all"},
64+
expectedCallsPubEndpoint: 1,
65+
},
66+
67+
{
68+
name: "audience from client with useAudience from admin false",
69+
audienceCfg: "clientConfiguredAud",
70+
useAudienceFromAdmin: false,
71+
scopesCfg: []string{"all"},
72+
clientConfigResponse: service.PublicClientAuthConfigResponse{Audience: "AdminConfiguredAud", Scopes: []string{}},
73+
expectedAudience: "clientConfiguredAud",
74+
expectedScopes: []string{"all"},
75+
expectedCallsPubEndpoint: 0,
76+
},
77+
}
78+
for _, test := range tests {
79+
cfg := GetConfig(ctx)
80+
tokenCache := &tokenCacheMocks.TokenCache{}
81+
metadataClient := &adminMocks.AuthMetadataServiceClient{}
82+
metadataClient.OnGetOAuth2MetadataMatch(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{}, nil)
83+
metadataClient.OnGetPublicClientConfigMatch(mock.Anything, mock.Anything).Return(&test.clientConfigResponse, nil)
84+
cfg.AuthType = AuthTypeClientSecret
85+
cfg.Audience = test.audienceCfg
86+
cfg.Scopes = test.scopesCfg
87+
cfg.UseAudienceFromAdmin = test.useAudienceFromAdmin
88+
flyteTokenSource, err := NewTokenSourceProvider(ctx, cfg, tokenCache, metadataClient)
89+
assert.True(t, metadataClient.AssertNumberOfCalls(t, "GetPublicClientConfig", test.expectedCallsPubEndpoint))
90+
assert.NoError(t, err)
91+
assert.NotNil(t, flyteTokenSource)
92+
clientCredSourceProvider, ok := flyteTokenSource.(ClientCredentialsTokenSourceProvider)
93+
assert.True(t, ok)
94+
assert.Equal(t, test.expectedScopes, clientCredSourceProvider.ccConfig.Scopes)
95+
assert.Equal(t, url.Values{audienceKey: {test.expectedAudience}}, clientCredSourceProvider.ccConfig.EndpointParams)
96+
}
97+
}

0 commit comments

Comments
 (0)