From a9e3e93ad2c152f79a96214e6f604ab57f2e3344 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 15 May 2025 14:31:11 +0100 Subject: [PATCH 1/4] feat(go/adbc/driver/flightsql): add SSL root certs to oauth --- .../flightsql/flightsql_adbc_server_test.go | 22 ++++--- .../driver/flightsql/flightsql_database.go | 4 +- go/adbc/driver/flightsql/flightsql_oauth.go | 62 +++++++++++++++++-- 3 files changed, 74 insertions(+), 14 deletions(-) diff --git a/go/adbc/driver/flightsql/flightsql_adbc_server_test.go b/go/adbc/driver/flightsql/flightsql_adbc_server_test.go index 1261d6becc..e28583f4e6 100644 --- a/go/adbc/driver/flightsql/flightsql_adbc_server_test.go +++ b/go/adbc/driver/flightsql/flightsql_adbc_server_test.go @@ -124,7 +124,7 @@ func (suite *ServerBasedTests) TearDownSuite() { suite.s.Shutdown() } -func (suite *ServerBasedTests) generateCertOption() grpc.ServerOption { +func (suite *ServerBasedTests) generateCertOption() (*tls.Config, string) { // Generate a self-signed certificate in-process for testing privKey, err := rsa.GenerateKey(rand.Reader, 2048) suite.Require().NoError(err) @@ -156,9 +156,10 @@ func (suite *ServerBasedTests) generateCertOption() grpc.ServerOption { suite.Require().NoError(err) tlsConfig := &tls.Config{Certificates: []tls.Certificate{cert}} - tlsCreds := credentials.NewTLS(tlsConfig) + return tlsConfig, string(certBytes) + // tlsCreds := credentials.NewTLS(tlsConfig) - return grpc.Creds(tlsCreds) + // return grpc.Creds(tlsCreds), string(certBytes) } func (suite *ServerBasedTests) openAndExecuteQuery(query string) { @@ -343,6 +344,7 @@ type OAuthTests struct { oauthServer *httptest.Server mockOAuthServer *MockOAuthServer + pemCert string } // MockOAuthServer simulates an OAuth 2.0 server for testing @@ -421,12 +423,18 @@ func oauthTestUnary(ctx context.Context, req interface{}, info *grpc.UnaryServer } func (suite *OAuthTests) SetupSuite() { + + tlsConfig, pemCertString := suite.generateCertOption() + suite.pemCert = pemCertString + suite.mockOAuthServer = &MockOAuthServer{} - suite.oauthServer = httptest.NewServer(http.HandlerFunc(suite.mockOAuthServer.handleTokenRequest)) + suite.oauthServer = httptest.NewUnstartedServer(http.HandlerFunc(suite.mockOAuthServer.handleTokenRequest)) + suite.oauthServer.TLS = tlsConfig + suite.oauthServer.StartTLS() suite.setupFlightServer(&AuthnTestServer{}, []flight.ServerMiddleware{ {Unary: oauthTestUnary}, - }, suite.generateCertOption()) + }, grpc.Creds(credentials.NewTLS(tlsConfig))) } func (suite *OAuthTests) TearDownSuite() { @@ -451,7 +459,7 @@ func (suite *OAuthTests) TestTokenExchangeFlow() { driver.OptionKeySubjectToken: "test-subject-token", driver.OptionKeySubjectTokenType: "urn:ietf:params:oauth:token-type:jwt", driver.OptionKeyTokenURI: suite.oauthServer.URL, - driver.OptionSSLSkipVerify: adbc.OptionValueEnabled, + driver.OptionSSLRootCerts: suite.pemCert, }) suite.Require().NoError(err) @@ -465,7 +473,7 @@ func (suite *OAuthTests) TestClientCredentialsFlow() { driver.OptionKeyClientId: "test-client", driver.OptionKeyClientSecret: "test-secret", driver.OptionKeyTokenURI: suite.oauthServer.URL, - driver.OptionSSLSkipVerify: adbc.OptionValueEnabled, + driver.OptionSSLRootCerts: suite.pemCert, }) suite.Require().NoError(err) diff --git a/go/adbc/driver/flightsql/flightsql_database.go b/go/adbc/driver/flightsql/flightsql_database.go index 9ff89d5f57..d8070c2513 100644 --- a/go/adbc/driver/flightsql/flightsql_database.go +++ b/go/adbc/driver/flightsql/flightsql_database.go @@ -183,9 +183,9 @@ func (d *databaseImpl) SetOptions(cnOptions map[string]string) error { var err error switch flow { case ClientCredentials: - d.oauthToken, err = newClientCredentials(cnOptions) + d.oauthToken, err = newClientCredentials(cnOptions, &tlsConfig) case TokenExchange: - d.oauthToken, err = newTokenExchangeFlow(cnOptions) + d.oauthToken, err = newTokenExchangeFlow(cnOptions, &tlsConfig) default: return adbc.Error{ Msg: fmt.Sprintf("oauth flow not implemented: %s", flow), diff --git a/go/adbc/driver/flightsql/flightsql_oauth.go b/go/adbc/driver/flightsql/flightsql_oauth.go index 707590a0df..52dc268046 100644 --- a/go/adbc/driver/flightsql/flightsql_oauth.go +++ b/go/adbc/driver/flightsql/flightsql_oauth.go @@ -19,13 +19,42 @@ package flightsql import ( "context" + "crypto/tls" "fmt" + "net/http" "golang.org/x/oauth2" "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/oauth" ) +// TokenSource supplies PerRPCCredentials from an oauth2.TokenSource. +type FlightTokenSource struct { + oauth2.TokenSource +} + +// GetRequestMetadata gets the request metadata as a map from a TokenSource. +func (ts FlightTokenSource) GetRequestMetadata(ctx context.Context, _ ...string) (map[string]string, error) { + token, err := ts.Token() + if err != nil { + return nil, err + } + // ri, _ := credentials.RequestInfoFromContext(ctx) + // if err = credentials.CheckSecurityLevel(ri.AuthInfo, credentials.PrivacyAndIntegrity); err != nil { + // return nil, fmt.Errorf("unable to transfer TokenSource PerRPCCredentials: %v", err) + // } + return map[string]string{ + "authorization": token.Type() + " " + token.AccessToken, + }, nil +} + +// RequireTransportSecurity indicates whether the credentials requires transport security. +func (ts FlightTokenSource) RequireTransportSecurity() bool { + return false +} + +// Bit flags for different OAuth authentication methods. Enables multiple authentication methods to be +// specified simultaneaously if needed const ( ClientCredentials = "client_credentials" TokenExchange = "token_exchange" @@ -54,6 +83,16 @@ var ( } ) +const ( + ttPrefix = "urn:ietf:params:oauth:token-type:" + TokenTypeAccessToken = ttPrefix + "access_token" + TokenTypeRefreshToken = ttPrefix + "refresh_token" + TokenTypeIdToken = ttPrefix + "id_token" + TokenTypeSaml1 = ttPrefix + "saml1" + TokenTypeSaml2 = ttPrefix + "saml2" + TokenTypeJWT = ttPrefix + "jwt" +) + func parseOAuthOptions(options map[string]string, paramMap map[string]oAuthOption, flowName string) (map[string]string, error) { params := map[string]string{} @@ -69,16 +108,29 @@ func parseOAuthOptions(options map[string]string, paramMap map[string]oAuthOptio return params, nil } -func exchangeToken(conf *oauth2.Config, codeOptions []oauth2.AuthCodeOption) (credentials.PerRPCCredentials, error) { +func exchangeToken(conf *oauth2.Config, codeOptions []oauth2.AuthCodeOption, tlsConfig *tls.Config) (credentials.PerRPCCredentials, error) { ctx := context.Background() + + if tlsConfig != nil { + // Set the HTTP client with custom TLS config in the context + httpClient := &http.Client{ + Transport: &http.Transport{ + TLSClientConfig: tlsConfig, + }, + } + + ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient) + } + tok, err := conf.Exchange(ctx, "", codeOptions...) if err != nil { return nil, err } + // return &FlightTokenSource{TokenSource: conf.TokenSource(ctx, tok)}, nil return &oauth.TokenSource{TokenSource: conf.TokenSource(ctx, tok)}, nil } -func newClientCredentials(options map[string]string) (credentials.PerRPCCredentials, error) { +func newClientCredentials(options map[string]string, tlsConfig *tls.Config) (credentials.PerRPCCredentials, error) { codeOptions := []oauth2.AuthCodeOption{ // Required value for client credentials requests as specified in https://datatracker.ietf.org/doc/html/rfc6749#section-4.4.2 oauth2.SetAuthURLParam("grant_type", "client_credentials"), @@ -101,10 +153,10 @@ func newClientCredentials(options map[string]string) (credentials.PerRPCCredenti conf.Scopes = []string{scopes} } - return exchangeToken(conf, codeOptions) + return exchangeToken(conf, codeOptions, tlsConfig) } -func newTokenExchangeFlow(options map[string]string) (credentials.PerRPCCredentials, error) { +func newTokenExchangeFlow(options map[string]string, tlsConfig *tls.Config) (credentials.PerRPCCredentials, error) { tokenURI, ok := options[OptionKeyTokenURI] if !ok { return nil, fmt.Errorf("token exchange grant requires %s", OptionKeyTokenURI) @@ -147,5 +199,5 @@ func newTokenExchangeFlow(options map[string]string) (credentials.PerRPCCredenti } } - return exchangeToken(conf, codeOptions) + return exchangeToken(conf, codeOptions, tlsConfig) } From c8a4a9c0d770f69bef93d1deb016444458898dac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 16 May 2025 13:39:28 +0100 Subject: [PATCH 2/4] chore(go/adbc/driver/flightsql): remove code used for testing --- go/adbc/driver/flightsql/flightsql_oauth.go | 26 --------------------- 1 file changed, 26 deletions(-) diff --git a/go/adbc/driver/flightsql/flightsql_oauth.go b/go/adbc/driver/flightsql/flightsql_oauth.go index 52dc268046..1288d0696e 100644 --- a/go/adbc/driver/flightsql/flightsql_oauth.go +++ b/go/adbc/driver/flightsql/flightsql_oauth.go @@ -28,31 +28,6 @@ import ( "google.golang.org/grpc/credentials/oauth" ) -// TokenSource supplies PerRPCCredentials from an oauth2.TokenSource. -type FlightTokenSource struct { - oauth2.TokenSource -} - -// GetRequestMetadata gets the request metadata as a map from a TokenSource. -func (ts FlightTokenSource) GetRequestMetadata(ctx context.Context, _ ...string) (map[string]string, error) { - token, err := ts.Token() - if err != nil { - return nil, err - } - // ri, _ := credentials.RequestInfoFromContext(ctx) - // if err = credentials.CheckSecurityLevel(ri.AuthInfo, credentials.PrivacyAndIntegrity); err != nil { - // return nil, fmt.Errorf("unable to transfer TokenSource PerRPCCredentials: %v", err) - // } - return map[string]string{ - "authorization": token.Type() + " " + token.AccessToken, - }, nil -} - -// RequireTransportSecurity indicates whether the credentials requires transport security. -func (ts FlightTokenSource) RequireTransportSecurity() bool { - return false -} - // Bit flags for different OAuth authentication methods. Enables multiple authentication methods to be // specified simultaneaously if needed const ( @@ -126,7 +101,6 @@ func exchangeToken(conf *oauth2.Config, codeOptions []oauth2.AuthCodeOption, tls if err != nil { return nil, err } - // return &FlightTokenSource{TokenSource: conf.TokenSource(ctx, tok)}, nil return &oauth.TokenSource{TokenSource: conf.TokenSource(ctx, tok)}, nil } From 9e7bcd824c92ddcbd96c099a336fffaef8a2df3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 16 May 2025 13:43:09 +0100 Subject: [PATCH 3/4] feat(go/adbc/driver/flightsql): set tls config for oauth http client calls --- go/adbc/driver/flightsql/flightsql_oauth.go | 30 +++++++++++++-------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/go/adbc/driver/flightsql/flightsql_oauth.go b/go/adbc/driver/flightsql/flightsql_oauth.go index 1288d0696e..24ad3f1f6a 100644 --- a/go/adbc/driver/flightsql/flightsql_oauth.go +++ b/go/adbc/driver/flightsql/flightsql_oauth.go @@ -83,20 +83,24 @@ func parseOAuthOptions(options map[string]string, paramMap map[string]oAuthOptio return params, nil } -func exchangeToken(conf *oauth2.Config, codeOptions []oauth2.AuthCodeOption, tlsConfig *tls.Config) (credentials.PerRPCCredentials, error) { +func createOAuthContext(tlsConfig *tls.Config) context.Context { ctx := context.Background() - if tlsConfig != nil { - // Set the HTTP client with custom TLS config in the context - httpClient := &http.Client{ - Transport: &http.Transport{ - TLSClientConfig: tlsConfig, - }, - } + if tlsConfig == nil { + return ctx + } - ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient) + // Create a custom HTTP client with TLS config to use for oauth calls + httpClient := &http.Client{ + Transport: &http.Transport{ + TLSClientConfig: tlsConfig, + }, } + return context.WithValue(ctx, oauth2.HTTPClient, httpClient) +} + +func exchangeToken(ctx context.Context, conf *oauth2.Config, codeOptions []oauth2.AuthCodeOption) (credentials.PerRPCCredentials, error) { tok, err := conf.Exchange(ctx, "", codeOptions...) if err != nil { return nil, err @@ -105,6 +109,8 @@ func exchangeToken(conf *oauth2.Config, codeOptions []oauth2.AuthCodeOption, tls } func newClientCredentials(options map[string]string, tlsConfig *tls.Config) (credentials.PerRPCCredentials, error) { + ctx := createOAuthContext(tlsConfig) + codeOptions := []oauth2.AuthCodeOption{ // Required value for client credentials requests as specified in https://datatracker.ietf.org/doc/html/rfc6749#section-4.4.2 oauth2.SetAuthURLParam("grant_type", "client_credentials"), @@ -127,10 +133,12 @@ func newClientCredentials(options map[string]string, tlsConfig *tls.Config) (cre conf.Scopes = []string{scopes} } - return exchangeToken(conf, codeOptions, tlsConfig) + return exchangeToken(ctx, conf, codeOptions) } func newTokenExchangeFlow(options map[string]string, tlsConfig *tls.Config) (credentials.PerRPCCredentials, error) { + ctx := createOAuthContext(tlsConfig) + tokenURI, ok := options[OptionKeyTokenURI] if !ok { return nil, fmt.Errorf("token exchange grant requires %s", OptionKeyTokenURI) @@ -173,5 +181,5 @@ func newTokenExchangeFlow(options map[string]string, tlsConfig *tls.Config) (cre } } - return exchangeToken(conf, codeOptions, tlsConfig) + return exchangeToken(ctx, conf, codeOptions) } From 3f347bf27ee503c5bfb520a2382645d2356a10a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 16 May 2025 13:52:13 +0100 Subject: [PATCH 4/4] fix(go/adbc/driver/flightsql): remove extra unnecessary code --- .../driver/flightsql/flightsql_adbc_server_test.go | 3 --- go/adbc/driver/flightsql/flightsql_oauth.go | 12 ------------ 2 files changed, 15 deletions(-) diff --git a/go/adbc/driver/flightsql/flightsql_adbc_server_test.go b/go/adbc/driver/flightsql/flightsql_adbc_server_test.go index e28583f4e6..d7d4c73b90 100644 --- a/go/adbc/driver/flightsql/flightsql_adbc_server_test.go +++ b/go/adbc/driver/flightsql/flightsql_adbc_server_test.go @@ -157,9 +157,6 @@ func (suite *ServerBasedTests) generateCertOption() (*tls.Config, string) { suite.Require().NoError(err) tlsConfig := &tls.Config{Certificates: []tls.Certificate{cert}} return tlsConfig, string(certBytes) - // tlsCreds := credentials.NewTLS(tlsConfig) - - // return grpc.Creds(tlsCreds), string(certBytes) } func (suite *ServerBasedTests) openAndExecuteQuery(query string) { diff --git a/go/adbc/driver/flightsql/flightsql_oauth.go b/go/adbc/driver/flightsql/flightsql_oauth.go index 24ad3f1f6a..32217202ff 100644 --- a/go/adbc/driver/flightsql/flightsql_oauth.go +++ b/go/adbc/driver/flightsql/flightsql_oauth.go @@ -28,8 +28,6 @@ import ( "google.golang.org/grpc/credentials/oauth" ) -// Bit flags for different OAuth authentication methods. Enables multiple authentication methods to be -// specified simultaneaously if needed const ( ClientCredentials = "client_credentials" TokenExchange = "token_exchange" @@ -58,16 +56,6 @@ var ( } ) -const ( - ttPrefix = "urn:ietf:params:oauth:token-type:" - TokenTypeAccessToken = ttPrefix + "access_token" - TokenTypeRefreshToken = ttPrefix + "refresh_token" - TokenTypeIdToken = ttPrefix + "id_token" - TokenTypeSaml1 = ttPrefix + "saml1" - TokenTypeSaml2 = ttPrefix + "saml2" - TokenTypeJWT = ttPrefix + "jwt" -) - func parseOAuthOptions(options map[string]string, paramMap map[string]oAuthOption, flowName string) (map[string]string, error) { params := map[string]string{}