feat(go/adbc/driver/flightsql): add SSL root certs to oauth

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,7 @@ func (suite *ServerBasedTests) generateCertOption() grpc.ServerOption {
 
 	suite.Require().NoError(err)
 	tlsConfig := &tls.Config{Certificates: []tls.Certificate{cert}}
-	tlsCreds := credentials.NewTLS(tlsConfig)
-
-	return grpc.Creds(tlsCreds)
+	return tlsConfig, string(certBytes)
 }
 
 func (suite *ServerBasedTests) openAndExecuteQuery(query string) {
@@ -343,6 +341,7 @@ type OAuthTests struct {
 
 	oauthServer     *httptest.Server
 	mockOAuthServer *MockOAuthServer
+	pemCert         string
 }
 
 // MockOAuthServer simulates an OAuth 2.0 server for testing
@@ -421,12 +420,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 +456,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 +470,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)
 

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),

go/adbc/driver/flightsql/flightsql_oauth.go

@@ -19,7 +19,9 @@ package flightsql
 
 import (
 	"context"
+	"crypto/tls"
 	"fmt"
+	"net/http"
 
 	"golang.org/x/oauth2"
 	"google.golang.org/grpc/credentials"
@@ -69,16 +71,34 @@ 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 createOAuthContext(tlsConfig *tls.Config) context.Context {
 	ctx := context.Background()
+
+	if tlsConfig == nil {
+		return ctx
+	}
+
+	// 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
 	}
 	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) {
+	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"),
@@ -101,10 +121,12 @@ func newClientCredentials(options map[string]string) (credentials.PerRPCCredenti
 		conf.Scopes = []string{scopes}
 	}
 
-	return exchangeToken(conf, codeOptions)
+	return exchangeToken(ctx, conf, codeOptions)
 }
 
-func newTokenExchangeFlow(options map[string]string) (credentials.PerRPCCredentials, error) {
+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)
@@ -147,5 +169,5 @@ func newTokenExchangeFlow(options map[string]string) (credentials.PerRPCCredenti
 		}
 	}
 
-	return exchangeToken(conf, codeOptions)
+	return exchangeToken(ctx, conf, codeOptions)
 }