Add support for AD federated authentication login sequence

.gitignore

@@ -1,2 +1,8 @@
 /.idea
 /.connstr
+.vscode
+.terraform
+*.tfstate*
+*.log
+*.swp
+*~

accesstokenconnector.go

@@ -6,19 +6,8 @@ import (
 	"context"
 	"database/sql/driver"
 	"errors"
-	"fmt"
 )
 
-var _ driver.Connector = &accessTokenConnector{}
-
-// accessTokenConnector wraps Connector and injects a
-// fresh access token when connecting to the database
-type accessTokenConnector struct {
-	Connector
-
-	accessTokenProvider func() (string, error)
-}
-
 // NewAccessTokenConnector creates a new connector from a DSN and a token provider.
 // The token provider func will be called when a new connection is requested and should return a valid access token.
 // The returned connector may be used with sql.OpenDB.
@@ -32,20 +21,26 @@ func NewAccessTokenConnector(dsn string, tokenProvider func() (string, error)) (
 		return nil, err
 	}
 
-	c := &accessTokenConnector{
-		Connector:           *conn,
-		accessTokenProvider: tokenProvider,
+	conn.FederatedAuthenticationProvider = &accessTokenProvider{
+		tokenProvider: tokenProvider,
 	}
-	return c, nil
+
+	return conn, nil
 }
 
-// Connect returns a new database connection
-func (c *accessTokenConnector) Connect(ctx context.Context) (driver.Conn, error) {
-	var err error
-	c.Connector.params.fedAuthAccessToken, err = c.accessTokenProvider()
-	if err != nil {
-		return nil, fmt.Errorf("mssql: error retrieving access token: %+v", err)
-	}
+type accessTokenProvider struct {
+	tokenProvider func() (string, error)
+}
+
+func (p *accessTokenProvider) ConfigureProvider(fa *FederatedAuthenticationState) error {
+	fa.FedAuthLibrary = FedAuthLibrarySecurityToken
+	return nil
+}
+
+func (p *accessTokenProvider) ProvideSecurityToken(ctx context.Context) (string, error) {
+	return p.tokenProvider()
+}
 
-	return c.Connector.Connect(ctx)
+func (p *accessTokenProvider) ProvideActiveDirectoryToken(ctx context.Context, serverSPN, stsURL string) (string, error) {
+	return p.tokenProvider()
 }

accesstokenconnector_test.go

@@ -30,21 +30,21 @@ func TestNewAccessTokenConnector(t *testing.T) {
 				dsn:           dsn,
 				tokenProvider: tp},
 			want: func(c driver.Connector) error {
-				tc, ok := c.(*accessTokenConnector)
+				tc, ok := c.(*Connector)
 				if !ok {
-					return fmt.Errorf("Expected driver to be of type *accessTokenConnector, but got %T", c)
+					return fmt.Errorf("Expected driver to be of type *Connector, but got %T", c)
 				}
-				p := tc.Connector.params
+				p := tc.params
 				if p.database != "db" {
 					return fmt.Errorf("expected params.database=db, but got %v", p.database)
 				}
 				if p.host != "server.database.windows.net" {
 					return fmt.Errorf("expected params.host=server.database.windows.net, but got %v", p.host)
 				}
-				if tc.accessTokenProvider == nil {
-					return fmt.Errorf("Expected tokenProvider to not be nil")
+				if tc.FederatedAuthenticationProvider == nil {
+					return fmt.Errorf("Expected federated authentication provider to not be nil")
 				}
-				t, err := tc.accessTokenProvider()
+				t, err := tc.FederatedAuthenticationProvider.ProvideSecurityToken(context.TODO())
 				if t != "token" || err != nil {
 					return fmt.Errorf("Unexpected results from tokenProvider: %v, %v", t, err)
 				}

conn_str.go

@@ -37,7 +37,8 @@ type connectParams struct {
 	failOverPartner           string
 	failOverPort              uint64
 	packetSize                uint16
-	fedAuthAccessToken        string
+	fedAuthWorkflow           string
+	aadClientCertPath         string
 }
 
 // default packet size for TDS buffer
@@ -232,6 +233,13 @@ func parseConnectParams(dsn string) (connectParams, error) {
 		}
 	}
 
+	p.fedAuthWorkflow, ok = params["fedauth"]
+	if ok && p.disableEncryption {
+		f := "Encryption must not be disabled for federated authentication: encrypt='%s'"
+		return p, fmt.Errorf(f, encrypt)
+	}
+	p.aadClientCertPath, _ = params["clientcertpath"]
+
 	return p, nil
 }
 
@@ -247,8 +255,8 @@ func (p connectParams) toUrl() *url.URL {
 	}
 	res := url.URL{
 		Scheme: "sqlserver",
-		Host: p.host,
-		User: url.UserPassword(p.user, p.password),
+		Host:   p.host,
+		User:   url.UserPassword(p.user, p.password),
 	}
 	if p.instance != "" {
 		res.Path = p.instance

conn_str_test.go

@@ -34,6 +34,9 @@ func TestInvalidConnectionString(t *testing.T) {
 		// URL mode
 		"sqlserver://\x00",
 		"sqlserver://host?key=value1&key=value2", // duplicate keys
+
+		// cannot use federated authentication when encryption is disabled
+		"encrypt=disable;fedauth=ActiveDirectoryMSI",
 	}
 	for _, connStr := range connStrings {
 		_, err := parseConnectParams(connStr)
@@ -78,6 +81,10 @@ func TestValidConnectionString(t *testing.T) {
 		{"log=64;packet size=8192", func(p connectParams) bool { return p.logFlags == 64 && p.packetSize == 8192 }},
 		{"log=64;packet size=48000", func(p connectParams) bool { return p.logFlags == 64 && p.packetSize == 32767 }},
 
+		// federated authentication workflow
+		{"fedauth=ActiveDirectoryPassword", func(p connectParams) bool { return p.fedAuthWorkflow == "ActiveDirectoryPassword" }},
+		{"clientcertpath=client.pem", func(p connectParams) bool { return p.aadClientCertPath == "client.pem" }},
+
 		// those are supported currently, but maybe should not be
 		{"someparam", func(p connectParams) bool { return true }},
 		{";;=;", func(p connectParams) bool { return true }},

fedauth.go

@@ -0,0 +1,94 @@
+package mssql
+
+import (
+	"context"
+)
+
+// Federated authentication library affects the login data structure and message sequence.
+const (
+	// FedAuthLibraryLiveIDCompactToken specifies the Microsoft Live ID Compact Token authentication scheme
+	FedAuthLibraryLiveIDCompactToken = 0x00
+
+	// FedAuthLibrarySecurityToken specifies a token-based authentication where the token is available
+	// without additional information provided during the login sequence.
+	FedAuthLibrarySecurityToken = 0x01
+
+	// FedAuthLibraryADAL specifies a token-based authentication where a token is obtained during the
+	// login sequence using the server SPN and STS URL provided by the server during login.
+	FedAuthLibraryADAL = 0x02
+
+	// FedAuthLibraryReserved is used to indicate that no federated authentication scheme applies.
+	FedAuthLibraryReserved = 0x7F
+)
+
+// Federated authentication ADAL workflow affects the mechanism used to authenticate.
+const (
+	// FedAuthADALWorkflowPassword uses a username/password to obtain a token from Active Directory
+	FedAuthADALWorkflowPassword = 0x01
+
+	// FedAuthADALWorkflowPassword uses the Windows identity to obtain a token from Active Directory
+	FedAuthADALWorkflowIntegrated = 0x02
+
+	// FedAuthADALWorkflowMSI uses the managed identity service to obtain a token
+	FedAuthADALWorkflowMSI = 0x03
+)
+
+type FederatedAuthenticationState struct {
+	// FedAuthWorkflow captures the "fedauth" connection parameter
+	FedAuthWorkflow string
+
+	// UserName is initially set to the user id connection parameter.
+	// The federated authentication configurer can modify this value to
+	// change what is sent in the login packet.
+	UserName string
+
+	// Password is initially set to the user id connection parameter.
+	// The federated authentication configurer can modify this value to
+	// change what is sent in the login packet.
+	Password string
+
+	// Password is initially set to the client cert path connection parameter.
+	ClientCertPath string
+
+	// FedAuthLibrary is populated by the federated authentication provider.
+	FedAuthLibrary int
+
+	// ADALWorkflow is populated by the federated authentication provider.
+	ADALWorkflow byte
+
+	// FedAuthEcho is populated from the prelogin response
+	FedAuthEcho bool
+
+	// FedAuthToken is populated during login with the value from the provider.
+	FedAuthToken string
+
+	// Nonce is populated during login with the value from the provider.
+	Nonce []byte
+
+	// Signature is populated during login with the value from the server.
+	Signature []byte
+}
+
+// FederatedAuthenticationProvider implementations use the connection string
+// parameters to determine the library and workflow, if any, and obtain tokens
+// during the login sequence.
+type FederatedAuthenticationProvider interface {
+	// Configure accepts the incoming connection parameters and determines
+	// the values for the authentication library and ADAL workflow.
+	ConfigureProvider(fedAuth *FederatedAuthenticationState) error
+
+	// ProvideActiveDirectoryToken implementations are called during federated
+	// authentication login sequences where the server provides a service
+	// principal name and security token service endpoint that should be used
+	// to obtain the token. Implementations should contact the security token
+	// service specified and obtain the appropriate token, or return an error
+	// to indicate why a token is not available.
+	ProvideActiveDirectoryToken(ctx context.Context, serverSPN, stsURL string) (string, error)
+
+	// ProvideSecurityToken implementations are called during federated
+	// authentication security token login sequences at the point when the
+	// security token is required.  The string returned should be the access
+	// token to supply to the server, otherwise an error can be returned to
+	// indicate why a token is not available.
+	ProvideSecurityToken(ctx context.Context) (string, error)
+}

mssql.go

@@ -58,6 +58,7 @@ func (d *Driver) OpenConnector(dsn string) (*Connector, error) {
 	if err != nil {
 		return nil, err
 	}
+
 	return &Connector{
 		params: params,
 		driver: d,
@@ -126,6 +127,10 @@ type Connector struct {
 	// Dialer sets a custom dialer for all network operations.
 	// If Dialer is not set, normal net dialers are used.
 	Dialer Dialer
+
+	// FederatedAuthenticationProvider handles choosing the parameters
+	// and obtaining tokens for federated authentication login scenarios.
+	FederatedAuthenticationProvider FederatedAuthenticationProvider
 }
 
 type Dialer interface {
@@ -148,7 +153,7 @@ type Conn struct {
 	processQueryText bool
 	connectionGood   bool
 
-	outs         map[string]interface{}
+	outs map[string]interface{}
 }
 
 func (c *Conn) checkBadConn(err error) error {
@@ -653,9 +658,9 @@ func (s *Stmt) processExec(ctx context.Context) (res driver.Result, err error) {
 }
 
 type Rows struct {
-	stmt    *Stmt
-	cols    []columnStruct
-	reader  *tokenProcessor
+	stmt     *Stmt
+	cols     []columnStruct
+	reader   *tokenProcessor
 	nextCols []columnStruct
 
 	cancel func()
@@ -669,7 +674,7 @@ func (rc *Rows) Close() error {
 	for {
 		tok, err := rc.reader.nextToken()
 		if err == nil {
-			if  tok == nil {
+			if tok == nil {
 				return nil
 			} else {
 				// continue consuming tokens