Repository: denisenkom/go-mssqldb
Pull request #551 (2020-01-30 10:57:35)

Add support for NTLMv2
Fix for https://github.com/denisenkom/go-mssqldb/issues/146

codecov[bot]: # [Codecov](https://codecov.io/gh/denisenkom/go-mssqldb/pull/551?src=pr&el=h1) Report
> Merging [#551](https://codecov.io/gh/denisenkom/go-mssqldb/pull/551?src=pr&el=desc) into [master](https://codecov.io/gh/denisenkom/go-mssqldb/commit/1d7a30a10f73a69b1adb731b0ab12be3da43a833?src=pr&el=desc) will **not change** coverage.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/denisenkom/go-mssqldb/pull/551/graphs/tree.svg?width=650&token=MoMB2u5avp&height=150&src=pr)](https://codecov.io/gh/denisenkom/go-mssqldb/pull/551?src=pr&el=tree)

```diff
@@          Coverage Diff           @@
##           master    #551   +/-   ##
======================================
  Coverage    68.9%   68.9%           
======================================
  Files          22      22           
  Lines        5068    5068           
======================================
  Hits         3492    3492           
  Misses       1370    1370           
  Partials      206     206
```



------

[Continue to review full report at Codecov](https://codecov.io/gh/denisenkom/go-mssqldb/pull/551?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/denisenkom/go-mssqldb/pull/551?src=pr&el=footer). Last update [1d7a30a...8c414af](https://codecov.io/gh/denisenkom/go-mssqldb/pull/551?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/README.md b/README.md
index b655176..46e4eea 100644
--- a/README.md
+++ b/README.md
@@ -106,6 +106,26 @@ Other supported formats are listed below.
   * `odbc:server=localhost;user id=sa;password={foo{bar}` // Literal `{`, password is "foo{bar"
   * `odbc:server=localhost;user id=sa;password={foo}}bar}` // Escaped `} with `}}`, password is "foo}bar"
 
+### Azure Active Directory authentication - preview
+
+The configuration of functionality might change in the future.
+
+Azure Active Directory (AAD) access tokens are relatively short lived and need to be 
+valid when a new connection is made. Authentication is supported using a callback func that
+provides a fresh and valid token using a connector:
+``` golang
+conn, err := mssql.NewAccessTokenConnector(
+  "Server=test.database.windows.net;Database=testdb",
+  tokenProvider)
+if err != nil {
+	// handle errors in DSN
+}
+db := sql.OpenDB(conn)
+```
+Where `tokenProvider` is a function that returns a fresh access token or an error. None of these statements
+actually trigger the retrieval of a token, this happens when the first statment is issued and a connection
+is created.
+
 ## Executing Stored Procedures
 
 To run a stored procedure, set the query text to the procedure name:
diff --git a/accesstokenconnector.go b/accesstokenconnector.go
new file mode 100644
index 0000000..8dbe509
--- /dev/null
+++ b/accesstokenconnector.go
@@ -0,0 +1,51 @@
+// +build go1.10
+
+package mssql
+
+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.
+func NewAccessTokenConnector(dsn string, tokenProvider func() (string, error)) (driver.Connector, error) {
+	if tokenProvider == nil {
+		return nil, errors.New("mssql: tokenProvider cannot be nil")
+	}
+
+	conn, err := NewConnector(dsn)
+	if err != nil {
+		return nil, err
+	}
+
+	c := &accessTokenConnector{
+		Connector:           *conn,
+		accessTokenProvider: tokenProvider,
+	}
+	return c, 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)
+	}
+
+	return c.Connector.Connect(ctx)
+}
diff --git a/conn_str.go b/conn_str.go
index 4ff54b8..26ac50f 100644
--- a/conn_str.go
+++ b/conn_str.go
@@ -37,6 +37,7 @@ type connectParams struct {
 	failOverPartner           string
 	failOverPort              uint64
 	packetSize                uint16
+	fedAuthAccessToken        string
 }
 
 func parseConnectParams(dsn string) (connectParams, error) {
diff --git a/examples/azuread-accesstoken/README.md b/examples/azuread-accesstoken/README.md
new file mode 100644
index 0000000..cb20a76
--- /dev/null
+++ b/examples/azuread-accesstoken/README.md
@@ -0,0 +1,9 @@
+## Azure Managed Identity example
+
+This example shows how Azure Managed Identity can be used to access SQL Azure. Take note of the
+trust boundary before using MSI to prevent exposure of the tokens outside of the trust boundary.
+
+This example can only be run from a Azure Virtual Machine with Managed Identity configured.
+You can follow the steps from this tutorial to turn on managed identity for your VM and grant the
+VM access to a SQL Azure database:
+https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/tutorial-windows-vm-access-sql
diff --git a/examples/azuread-accesstoken/go.mod b/examples/azuread-accesstoken/go.mod
new file mode 100644
index 0000000..eb113f3
--- /dev/null
+++ b/examples/azuread-accesstoken/go.mod
@@ -0,0 +1,8 @@
+module github.com/denisenkom/go-mssqldb/examples/azure-ad-accesstoken
+
+go 1.13
+
+require (
+	github.com/Azure/go-autorest/autorest/adal v0.8.1
+	github.com/denisenkom/go-mssqldb v0.0.0-20191128021309-1d7a30a10f73
+)
diff --git a/examples/azuread-accesstoken/go.sum b/examples/azuread-accesstoken/go.sum
new file mode 100644
index 0000000..17148ff
--- /dev/null
+++ b/examples/azuread-accesstoken/go.sum
@@ -0,0 +1,30 @@
+github.com/Azure/go-autorest v13.3.2+incompatible h1:VxzPyuhtnlBOzc4IWCZHqpyH2d+QMLQEuy3wREyY4oc=
+github.com/Azure/go-autorest/autorest v0.9.0 h1:MRvx8gncNaXJqOoLmhNjUAKh33JJF8LyxPhomEtOsjs=
+github.com/Azure/go-autorest/autorest v0.9.0/go.mod h1:xyHB1BMZT0cuDHU7I0+g046+BFDTQ8rEZB0s4Yfa6bI=
+github.com/Azure/go-autorest/autorest v0.9.4 h1:1cM+NmKw91+8h5vfjgzK4ZGLuN72k87XVZBWyGwNjUM=
+github.com/Azure/go-autorest/autorest/adal v0.5.0/go.mod h1:8Z9fGy2MpX0PvDjB1pEgQTmVqjGhiHBW7RJJEciWzS0=
+github.com/Azure/go-autorest/autorest/adal v0.8.1 h1:pZdL8o72rK+avFWl+p9nE8RWi1JInZrWJYlnpfXJwHk=
+github.com/Azure/go-autorest/autorest/adal v0.8.1/go.mod h1:ZjhuQClTqx435SRJ2iMlOxPYt3d2C/T/7TiQCVZSn3Q=
+github.com/Azure/go-autorest/autorest/date v0.1.0/go.mod h1:plvfp3oPSKwf2DNjlBjWF/7vwR+cUD/ELuzDCXwHUVA=
+github.com/Azure/go-autorest/autorest/date v0.2.0 h1:yW+Zlqf26583pE43KhfnhFcdmSWlm5Ew6bxipnr/tbM=
+github.com/Azure/go-autorest/autorest/date v0.2.0/go.mod h1:vcORJHLJEh643/Ioh9+vPmf1Ij9AEBM5FuBIXLmIy0g=
+github.com/Azure/go-autorest/autorest/mocks v0.1.0/go.mod h1:OTyCOPRA2IgIlWxVYxBee2F5Gr4kF2zd2J5cFRaIDN0=
+github.com/Azure/go-autorest/autorest/mocks v0.2.0/go.mod h1:OTyCOPRA2IgIlWxVYxBee2F5Gr4kF2zd2J5cFRaIDN0=
+github.com/Azure/go-autorest/autorest/mocks v0.3.0/go.mod h1:a8FDP3DYzQ4RYfVAxAN3SVSiiO77gL2j2ronKKP0syM=
+github.com/Azure/go-autorest/logger v0.1.0/go.mod h1:oExouG+K6PryycPJfVSxi/koC6LSNgds39diKLz7Vrc=
+github.com/Azure/go-autorest/tracing v0.5.0 h1:TRn4WjSnkcSy5AEG3pnbtFSwNtwzjr4VYyQflFE619k=
+github.com/Azure/go-autorest/tracing v0.5.0/go.mod h1:r/s2XiOKccPW3HrqB+W0TQzfbtp2fGCgRFtBroKn4Dk=
+github.com/denisenkom/go-mssqldb v0.0.0-20191128021309-1d7a30a10f73 h1:OGNva6WhsKst5OZf7eZOklDztV3hwtTHovdrLHV+MsA=
+github.com/denisenkom/go-mssqldb v0.0.0-20191128021309-1d7a30a10f73/go.mod h1:xbL0rPBG9cCiLr28tMa8zpbdarY27NDyej4t/EjAShU=
+github.com/dgrijalva/jwt-go v3.2.0+incompatible h1:7qlOGliEKZXTDg6OTjfoBKDXWrumCAMpl/TFQ4/5kLM=
+github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
+github.com/golang-sql/civil v0.0.0-20190719163853-cb61b32ac6fe h1:lXe2qZdvpiX5WZkZR4hgp4KJVfY3nMkvmwbVkpv1rVY=
+github.com/golang-sql/civil v0.0.0-20190719163853-cb61b32ac6fe/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20190325154230-a5d413f7728c/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20191206172530-e9b2fee46413 h1:ULYEB3JvPRE/IfO+9uO7vKV/xzVTO7XPAwm8xbf4w2g=
+golang.org/x/crypto v0.0.0-20191206172530-e9b2fee46413/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
+golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
diff --git a/examples/azuread-accesstoken/managed_identity.go b/examples/azuread-accesstoken/managed_identity.go
new file mode 100644
index 0000000..0adf99c
--- /dev/null
+++ b/examples/azuread-accesstoken/managed_identity.go
@@ -0,0 +1,82 @@
+package main
+
+import (
+	"database/sql"
+	"flag"
+	"fmt"
+	"log"
+
+	"github.com/Azure/go-autorest/autorest/adal"
+	mssql "github.com/denisenkom/go-mssqldb"
+)
+
+var (
+	debug    = flag.Bool("debug", false, "enable debugging")
+	server   = flag.String("server", "", "the database server")
+	database = flag.String("database", "", "the database")
+)
+
+func main() {
+	flag.Parse()
+
+	if *debug {
+		fmt.Printf(" server:%s\n", *server)
+		fmt.Printf(" database:%s\n", *database)
+	}
+
+	if *server == "" {
+		log.Fatal("Server name cannot be left empty")
+	}
+
+	if *database == "" {
+		log.Fatal("Database name cannot be left empty")
+	}
+
+	connString := fmt.Sprintf("Server=%s;Database=%s", *server, *database)
+	if *debug {
+		fmt.Printf(" connString:%s\n", connString)
+	}
+
+	tokenProvider, err := getMSITokenProvider()
+	if err != nil {
+		log.Fatal("Error creating token provider for system assigned Azure Managed Identity:", err.Error())
+	}
+
+	connector, err := mssql.NewAccessTokenConnector(
+		connString, tokenProvider)
+	if err != nil {
+		log.Fatal("Connector creation failed:", err.Error())
+	}
+	conn := sql.OpenDB(connector)
+	defer conn.Close()
+
+	row := conn.QueryRow("select 1, 'abc'")
+	var somenumber int64
+	var somechars string
+	err = row.Scan(&somenumber, &somechars)
+	if err != nil {
+		log.Fatal("Scan failed:", err.Error())
+	}
+	fmt.Printf("somenumber:%d\n", somenumber)
+	fmt.Printf("somechars:%s\n", somechars)
+
+	fmt.Printf("bye\n")
+}
+
+func getMSITokenProvider() (func() (string, error), error) {
+	msiEndpoint, err := adal.GetMSIEndpoint()
+	if err != nil {
+		return nil, err
+	}
+	msi, err := adal.NewServicePrincipalTokenFromMSI(
+		msiEndpoint, "https://database.windows.net/")
+	if err != nil {
+		return nil, err
+	}
+
+	return func() (string, error) {
+		msi.EnsureFresh()
+		token := msi.OAuthToken()
+		return token, nil
+	}, nil
+}
diff --git a/ntlm.go b/ntlm.go
index 7c0cc4f..3ba48ff 100644
--- a/ntlm.go
+++ b/ntlm.go
@@ -4,11 +4,14 @@ package mssql
 
 import (
 	"crypto/des"
+	"crypto/hmac"
 	"crypto/md5"
 	"crypto/rand"
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"strings"
+	"time"
 	"unicode/utf16"
 
 	"golang.org/x/crypto/md4"
@@ -198,86 +201,209 @@ func ntlmSessionResponse(clientNonce [8]byte, serverChallenge [8]byte, password
 	return response(hash, passwordHash)
 }
 
-func (auth *ntlmAuth) NextBytes(bytes []byte) ([]byte, error) {
-	if string(bytes[0:8]) != "NTLMSSP\x00" {
-		return nil, errorNTLM
+func ntlmHashNoPadding(val string) []byte {
+	hash := make([]byte, 16)
+	h := md4.New()
+	h.Write(utf16le(val))
+	h.Sum(hash[:0])
+
+	return hash
+}
+
+func hmacMD5(passwordHash, data []byte) []byte {
+	hmacEntity := hmac.New(md5.New, passwordHash)
+	hmacEntity.Write(data)
+
+	return hmacEntity.Sum(nil)
+}
+
+func getNTLMv2AndLMv2ResponsePayloads(target, username, password string, challenge, nonce [8]byte, targetInfoFields []byte, timestamp time.Time) (ntlmV2Payload, lmV2Payload []byte) {
+	// NTLMv2 response payload: http://davenport.sourceforge.net/ntlm.html#theNtlmv2Response
+
+	ntlmHash := ntlmHashNoPadding(password)
+	usernameAndTargetBytes := utf16le(strings.ToUpper(username) + target)
+	ntlmV2Hash := hmacMD5(ntlmHash, usernameAndTargetBytes)
+	targetInfoLength := len(targetInfoFields)
+	blob := make([]byte, 32+targetInfoLength)
+	binary.BigEndian.PutUint32(blob[:4], 0x01010000)
+	binary.BigEndian.PutUint32(blob[4:8], 0x00000000)
+	binary.BigEndian.PutUint64(blob[8:16], uint64(timestamp.UnixNano()))
+	copy(blob[16:24], nonce[:])
+	binary.BigEndian.PutUint32(blob[24:28], 0x00000000)
+	copy(blob[28:], targetInfoFields)
+	binary.BigEndian.PutUint32(blob[28+targetInfoLength:], 0x00000000)
+	challengeLength := len(challenge)
+	blobLength := len(blob)
+	challengeAndBlob := make([]byte, challengeLength + blobLength)
+	copy(challengeAndBlob[:challengeLength], challenge[:])
+	copy(challengeAndBlob[challengeLength:], blob)
+	hashedChallenge := hmacMD5(ntlmV2Hash, challengeAndBlob)
+	ntlmV2Payload = append(hashedChallenge, blob...)
+
+	// LMv2 response payload: http://davenport.sourceforge.net/ntlm.html#theLmv2Response
+	ntlmV2hash := hmacMD5(ntlmHash, usernameAndTargetBytes)
+	challengeAndNonce := make([]byte, 16)
+	copy(challengeAndNonce[:8], challenge[:])
+	copy(challengeAndNonce[8:], nonce[:])
+	hashedChallenge = hmacMD5(ntlmV2hash, challengeAndNonce)
+	lmV2Payload = append(hashedChallenge, nonce[:]...)
+
+	return
+}
+
+func negotiateExtendedSessionSecurity(flags uint32, message []byte, challenge [8]byte, username, password string) (lm, nt []byte, err error) {
+	nonce := clientChallenge()
+
+	// Official specification: https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-nlmp/b38c36ed-2804-4868-a9ff-8dd3182128e4
+	// Unofficial walk through referenced by https://www.freetds.org/userguide/domains.htm: http://davenport.sourceforge.net/ntlm.html
+	if (flags & _NEGOTIATE_TARGET_INFO) != 0 {
+		targetInfoFields, target, err := getNTLMv2TargetInfoFields(message)
+		if err != nil {
+			return lm, nt, err
+		}
+
+		nt, lm = getNTLMv2AndLMv2ResponsePayloads(target, username, password, challenge, nonce, targetInfoFields, time.Now())
+
+		return lm, nt, nil
 	}
-	if binary.LittleEndian.Uint32(bytes[8:12]) != _CHALLENGE_MESSAGE {
-		return nil, errorNTLM
+
+	var lm_bytes [24]byte
+	copy(lm_bytes[:8], nonce[:])
+	lm = lm_bytes[:]
+	nt_bytes := ntlmSessionResponse(nonce, challenge, password)
+	nt = nt_bytes[:]
+
+	return lm, nt, nil
+}
+
+func getNTLMv2TargetInfoFields(type2Message []byte) (info []byte, target string, err error) {
+	type2MessageError := "mssql: while parsing NTLMv2 type 2 message, length %d too small for offset %d"
+	type2MessageLength := len(type2Message)
+	if type2MessageLength < 20 {
+		return nil, target, fmt.Errorf(type2MessageError, type2MessageLength, 20)
 	}
-	flags := binary.LittleEndian.Uint32(bytes[20:24])
-	var challenge [8]byte
-	copy(challenge[:], bytes[24:32])
 
-	var lm, nt []byte
-	if (flags & _NEGOTIATE_EXTENDED_SESSIONSECURITY) != 0 {
-		nonce := clientChallenge()
-		var lm_bytes [24]byte
-		copy(lm_bytes[:8], nonce[:])
-		lm = lm_bytes[:]
-		nt_bytes := ntlmSessionResponse(nonce, challenge, auth.Password)
-		nt = nt_bytes[:]
-	} else {
-		lm_bytes := lmResponse(challenge, auth.Password)
-		lm = lm_bytes[:]
-		nt_bytes := ntResponse(challenge, auth.Password)
-		nt = nt_bytes[:]
+	targetNameAllocated := binary.LittleEndian.Uint16(type2Message[14:16])
+	targetNameOffset := binary.LittleEndian.Uint32(type2Message[16:20])
+	endOfOffset := int(targetNameOffset + uint32(targetNameAllocated))
+	if type2MessageLength < endOfOffset {
+		return nil, target, fmt.Errorf(type2MessageError, type2MessageLength, endOfOffset)
+	}
+
+	target, err = ucs22str(type2Message[targetNameOffset:targetNameOffset+uint32(targetNameAllocated)])
+	if err != nil {
+		return nil, target, err
+	}
+
+	targetInformationAllocated := binary.LittleEndian.Uint16(type2Message[42:44])
+	targetInformationDataOffset := binary.LittleEndian.Uint32(type2Message[44:48])
+	endOfOffset = int(targetInformationDataOffset + uint32(targetInformationAllocated))
+	if type2MessageLength < endOfOffset {
+		return nil, target, fmt.Errorf(type2MessageError, type2MessageLength, endOfOffset)
 	}
+
+	targetInformationBytes := make([]byte, targetInformationAllocated)
+	copy(targetInformationBytes, type2Message[targetInformationDataOffset:targetInformationDataOffset+uint32(targetInformationAllocated)])
+
+	return targetInformationBytes, target, nil
+}
+
+func buildNTLMResponsePayload(lm, nt []byte, flags uint32, domain, workstation, username string) ([]byte, error) {
 	lm_len := len(lm)
 	nt_len := len(nt)
-
-	domain16 := utf16le(auth.Domain)
+	domain16 := utf16le(domain)
 	domain_len := len(domain16)
-	user16 := utf16le(auth.UserName)
+	user16 := utf16le(username)
 	user_len := len(user16)
-	workstation16 := utf16le(auth.Workstation)
+	workstation16 := utf16le(workstation)
 	workstation_len := len(workstation16)
-
 	msg := make([]byte, 88+lm_len+nt_len+domain_len+user_len+workstation_len)
 	copy(msg, []byte("NTLMSSP\x00"))
 	binary.LittleEndian.PutUint32(msg[8:], _AUTHENTICATE_MESSAGE)
+
 	// Lm Challenge Response Fields
 	binary.LittleEndian.PutUint16(msg[12:], uint16(lm_len))
 	binary.LittleEndian.PutUint16(msg[14:], uint16(lm_len))
 	binary.LittleEndian.PutUint32(msg[16:], 88)
+
 	// Nt Challenge Response Fields
 	binary.LittleEndian.PutUint16(msg[20:], uint16(nt_len))
 	binary.LittleEndian.PutUint16(msg[22:], uint16(nt_len))
 	binary.LittleEndian.PutUint32(msg[24:], uint32(88+lm_len))
+
 	// Domain Name Fields
 	binary.LittleEndian.PutUint16(msg[28:], uint16(domain_len))
 	binary.LittleEndian.PutUint16(msg[30:], uint16(domain_len))
 	binary.LittleEndian.PutUint32(msg[32:], uint32(88+lm_len+nt_len))
+
 	// User Name Fields
 	binary.LittleEndian.PutUint16(msg[36:], uint16(user_len))
 	binary.LittleEndian.PutUint16(msg[38:], uint16(user_len))
 	binary.LittleEndian.PutUint32(msg[40:], uint32(88+lm_len+nt_len+domain_len))
+
 	// Workstation Fields
 	binary.LittleEndian.PutUint16(msg[44:], uint16(workstation_len))
 	binary.LittleEndian.PutUint16(msg[46:], uint16(workstation_len))
 	binary.LittleEndian.PutUint32(msg[48:], uint32(88+lm_len+nt_len+domain_len+user_len))
+
 	// Encrypted Random Session Key Fields
 	binary.LittleEndian.PutUint16(msg[52:], 0)
 	binary.LittleEndian.PutUint16(msg[54:], 0)
 	binary.LittleEndian.PutUint32(msg[56:], uint32(88+lm_len+nt_len+domain_len+user_len+workstation_len))
+
 	// Negotiate Flags
 	binary.LittleEndian.PutUint32(msg[60:], flags)
+
 	// Version
 	binary.LittleEndian.PutUint32(msg[64:], 0)
 	binary.LittleEndian.PutUint32(msg[68:], 0)
+
 	// MIC
 	binary.LittleEndian.PutUint32(msg[72:], 0)
 	binary.LittleEndian.PutUint32(msg[76:], 0)
 	binary.LittleEndian.PutUint32(msg[88:], 0)
 	binary.LittleEndian.PutUint32(msg[84:], 0)
+
 	// Payload
 	copy(msg[88:], lm)
 	copy(msg[88+lm_len:], nt)
 	copy(msg[88+lm_len+nt_len:], domain16)
 	copy(msg[88+lm_len+nt_len+domain_len:], user16)
 	copy(msg[88+lm_len+nt_len+domain_len+user_len:], workstation16)
+
 	return msg, nil
 }
 
+func (auth *ntlmAuth) NextBytes(bytes []byte) ([]byte, error) {
+	signature := string(bytes[0:8])
+	if signature != "NTLMSSP\x00" {
+		return nil, errorNTLM
+	}
+
+	messageTypeIndicator := binary.LittleEndian.Uint32(bytes[8:12])
+	if messageTypeIndicator != _CHALLENGE_MESSAGE {
+		return nil, errorNTLM
+	}
+
+	var challenge [8]byte
+	copy(challenge[:], bytes[24:32])
+	flags := binary.LittleEndian.Uint32(bytes[20:24])
+	if (flags & _NEGOTIATE_EXTENDED_SESSIONSECURITY) != 0 {
+		lm, nt, err := negotiateExtendedSessionSecurity(flags, bytes, challenge, auth.UserName, auth.Password)
+		if err != nil {
+			return nil, err
+		}
+
+		return buildNTLMResponsePayload(lm, nt, flags, auth.Domain, auth.Workstation, auth.UserName)
+	}
+
+	lm_bytes := lmResponse(challenge, auth.Password)
+	lm := lm_bytes[:]
+	nt_bytes := ntResponse(challenge, auth.Password)
+	nt := nt_bytes[:]
+
+	return buildNTLMResponsePayload(lm, nt, flags, auth.Domain, auth.Workstation, auth.UserName)
+}
+
 func (auth *ntlmAuth) Free() {
 }
diff --git a/tds.go b/tds.go
index 9419836..832c4fd 100644
--- a/tds.go
+++ b/tds.go
@@ -100,13 +100,15 @@ const (
 // prelogin fields
 // http://msdn.microsoft.com/en-us/library/dd357559.aspx
 const (
-	preloginVERSION    = 0
-	preloginENCRYPTION = 1
-	preloginINSTOPT    = 2
-	preloginTHREADID   = 3
-	preloginMARS       = 4
-	preloginTRACEID    = 5
-	preloginTERMINATOR = 0xff
+	preloginVERSION         = 0
+	preloginENCRYPTION      = 1
+	preloginINSTOPT         = 2
+	preloginTHREADID        = 3
+	preloginMARS            = 4
+	preloginTRACEID         = 5
+	preloginFEDAUTHREQUIRED = 6
+	preloginNONCEOPT        = 7
+	preloginTERMINATOR      = 0xff
 )
 
 const (
@@ -245,6 +247,12 @@ const (
 	fReadOnlyIntent = 32
 )
 
+// OptionFlags3
+// https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/773a62b6-ee89-4c02-9e5e-344882630aac
+const (
+	fExtension = 0x10
+)
+
 type login struct {
 	TDSVersion     uint32
 	PacketSize     uint32
@@ -269,6 +277,89 @@ type login struct {
 	SSPI           []byte
 	AtchDBFile     string
 	ChangePassword string
+	FeatureExt     featureExts
+}
+
+type featureExts struct {
+	features map[byte]featureExt
+}
+
+type featureExt interface {
+	featureID() byte
+	toBytes() []byte
+}
+
+func (e *featureExts) Add(f featureExt) error {
+	if f == nil {
+		return nil
+	}
+	id := f.featureID()
+	if _, exists := e.features[id]; exists {
+		f := "Login error: Feature with ID '%v' is already present in FeatureExt block."
+		return fmt.Errorf(f, id)
+	}
+	if e.features == nil {
+		e.features = make(map[byte]featureExt)
+	}
+	e.features[id] = f
+	return nil
+}
+
+func (e featureExts) toBytes() []byte {
+	if len(e.features) == 0 {
+		return nil
+	}
+	var d []byte
+	for featureID, f := range e.features {
+		featureData := f.toBytes()
+
+		hdr := make([]byte, 5)
+		hdr[0] = featureID                                               // FedAuth feature extension BYTE
+		binary.LittleEndian.PutUint32(hdr[1:], uint32(len(featureData))) // FeatureDataLen DWORD
+		d = append(d, hdr...)
+
+		d = append(d, featureData...) // FeatureData *BYTE
+	}
+	if d != nil {
+		d = append(d, 0xff) // Terminator
+	}
+	return d
+}
+
+type featureExtFedAuthSTS struct {
+	FedAuthEcho  bool
+	FedAuthToken string
+	Nonce        []byte
+}
+
+func (e *featureExtFedAuthSTS) featureID() byte {
+	return 0x02
+}
+
+func (e *featureExtFedAuthSTS) toBytes() []byte {
+	if e == nil {
+		return nil
+	}
+
+	options := byte(0x01) << 1 // 0x01 => STS bFedAuthLibrary 7BIT
+	if e.FedAuthEcho {
+		options |= 1 // fFedAuthEcho
+	}
+
+	d := make([]byte, 5)
+	d[0] = options
+
+	// looks like string in
+	// https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/f88b63bb-b479-49e1-a87b-deda521da508
+	tokenBytes := str2ucs2(e.FedAuthToken)
+	binary.LittleEndian.PutUint32(d[1:], uint32(len(tokenBytes))) // Should be a signed int32, but since the length is relatively small, this should work
+	d = append(d, tokenBytes...)
+
+	if len(e.Nonce) == 32 {
+		d = append(d, e.Nonce...)
+	}
+
+	return d
 }
 
 type loginHeader struct {
@@ -295,7 +386,7 @@ type loginHeader struct {
 	ServerNameOffset     uint16
 	ServerNameLength     uint16
 	ExtensionOffset      uint16
-	ExtensionLenght      uint16
+	ExtensionLength      uint16
 	CtlIntNameOffset     uint16
 	CtlIntNameLength     uint16
 	LanguageOffset       uint16
@@ -357,6 +448,8 @@ func sendLogin(w *tdsBuffer, login login) error {
 	database := str2ucs2(login.Database)
 	atchdbfile := str2ucs2(login.AtchDBFile)
 	changepassword := str2ucs2(login.ChangePassword)
+	featureExt := login.FeatureExt.toBytes()
+
 	hdr := loginHeader{
 		TDSVersion:           login.TDSVersion,
 		PacketSize:           login.PacketSize,
@@ -405,7 +498,18 @@ func sendLogin(w *tdsBuffer, login login) error {
 	offset += uint16(len(atchdbfile))
 	hdr.ChangePasswordOffset = offset
 	offset += uint16(len(changepassword))
-	hdr.Length = uint32(offset)
+
+	featureExtOffset := uint32(0)
+	featureExtLen := len(featureExt)
+	if featureExtLen > 0 {
+		hdr.OptionFlags3 |= fExtension
+		hdr.ExtensionOffset = offset
+		hdr.ExtensionLength = 4
+		offset += hdr.ExtensionLength // DWORD
+		featureExtOffset = uint32(offset)
+	}
+	hdr.Length = uint32(offset) + uint32(featureExtLen)
+
 	var err error
 	err = binary.Write(w, binary.LittleEndian, &hdr)
 	if err != nil {
@@ -455,6 +559,16 @@ func sendLogin(w *tdsBuffer, login login) error {
 	if err != nil {
 		return err
 	}
+	if featureExtOffset > 0 {
+		err = binary.Write(w, binary.LittleEndian, featureExtOffset)
+		if err != nil {
+			return err
+		}
+		_, err = w.Write(featureExt)
+		if err != nil {
+			return err
+		}
+	}
 	return w.FinishPacket()
 }
 
@@ -844,15 +958,23 @@ initiate_connection:
 		AppName:      p.appname,
 		TypeFlags:    p.typeFlags,
 	}
-	auth, auth_ok := getAuth(p.user, p.password, p.serverSPN, p.workstation)
-	if auth_ok {
+	auth, authOk := getAuth(p.user, p.password, p.serverSPN, p.workstation)
+	switch {
+	case p.fedAuthAccessToken != "": // accesstoken ignores user/password
+		featurext := &featureExtFedAuthSTS{
+			FedAuthEcho:  len(fields[preloginFEDAUTHREQUIRED]) > 0 && fields[preloginFEDAUTHREQUIRED][0] == 1,
+			FedAuthToken: p.fedAuthAccessToken,
+			Nonce:        fields[preloginNONCEOPT],
+		}
+		login.FeatureExt.Add(featurext)
+	case authOk:
 		login.SSPI, err = auth.InitialBytes()
 		if err != nil {
 			return nil, err
 		}
 		login.OptionFlags2 |= fIntSecurity
 		defer auth.Free()
-	} else {
+	default:
 		login.UserName = p.user
 		login.Password = p.password
 	}
diff --git a/token.go b/token.go
index 1acac8a..25385e8 100644
--- a/token.go
+++ b/token.go
@@ -17,20 +17,21 @@ type token byte
 
 // token ids
 const (
-	tokenReturnStatus token = 121 // 0x79
-	tokenColMetadata  token = 129 // 0x81
-	tokenOrder        token = 169 // 0xA9
-	tokenError        token = 170 // 0xAA
-	tokenInfo         token = 171 // 0xAB
-	tokenReturnValue  token = 0xAC
-	tokenLoginAck     token = 173 // 0xad
-	tokenRow          token = 209 // 0xd1
-	tokenNbcRow       token = 210 // 0xd2
-	tokenEnvChange    token = 227 // 0xE3
-	tokenSSPI         token = 237 // 0xED
-	tokenDone         token = 253 // 0xFD
-	tokenDoneProc     token = 254
-	tokenDoneInProc   token = 255
+	tokenReturnStatus  token = 121 // 0x79
+	tokenColMetadata   token = 129 // 0x81
+	tokenOrder         token = 169 // 0xA9
+	tokenError         token = 170 // 0xAA
+	tokenInfo          token = 171 // 0xAB
+	tokenReturnValue   token = 0xAC
+	tokenLoginAck      token = 173 // 0xad
+	tokenFeatureExtAck token = 174 // 0xae
+	tokenRow           token = 209 // 0xd1
+	tokenNbcRow        token = 210 // 0xd2
+	tokenEnvChange     token = 227 // 0xE3
+	tokenSSPI          token = 237 // 0xED
+	tokenDone          token = 253 // 0xFD
+	tokenDoneProc      token = 254
+	tokenDoneInProc    token = 255
 )
 
 // done flags
@@ -447,6 +448,22 @@ func parseLoginAck(r *tdsBuffer) loginAckStruct {
 	return res
 }
 
+// https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/2eb82f8e-11f0-46dc-b42d-27302fa4701a
+func parseFeatureExtAck(r *tdsBuffer) {
+	// at most 1 featureAck per feature in featureExt
+	// go-mssqldb will add at most 1 feature, the spec defines 7 different features
+	for i := 0; i < 8; i++ {
+		featureID := r.byte() // FeatureID
+		if featureID == 0xff {
+			return
+		}
+		size := r.uint32() // FeatureAckDataLen
+		d := make([]byte, size)
+		r.ReadFull(d)
+	}
+	panic("parsed more than 7 featureAck's, protocol implementation error?")
+}
+
 // http://msdn.microsoft.com/en-us/library/dd357363.aspx
 func parseColMetadata72(r *tdsBuffer) (columns []columnStruct) {
 	count := r.uint16()
@@ -577,6 +594,8 @@ func processSingleResponse(sess *tdsSession, ch chan tokenStruct, outs map[strin
 		case tokenLoginAck:
 			loginAck := parseLoginAck(sess.buf)
 			ch <- loginAck
+		case tokenFeatureExtAck:
+			parseFeatureExtAck(sess.buf)
 		case tokenOrder:
 			order := parseOrder(sess.buf)
 			ch <- order


Test patch:
diff --git a/accesstokenconnector_test.go b/accesstokenconnector_test.go
new file mode 100644
index 0000000..826dedb
--- /dev/null
+++ b/accesstokenconnector_test.go
@@ -0,0 +1,92 @@
+// +build go1.10
+
+package mssql
+
+import (
+	"context"
+	"database/sql/driver"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+)
+
+func TestNewAccessTokenConnector(t *testing.T) {
+	dsn := "Server=server.database.windows.net;Database=db"
+	tp := func() (string, error) { return "token", nil }
+	type args struct {
+		dsn           string
+		tokenProvider func() (string, error)
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    func(driver.Connector) error
+		wantErr bool
+	}{
+		{
+			name: "Happy path",
+			args: args{
+				dsn:           dsn,
+				tokenProvider: tp},
+			want: func(c driver.Connector) error {
+				tc, ok := c.(*accessTokenConnector)
+				if !ok {
+					return fmt.Errorf("Expected driver to be of type *accessTokenConnector, but got %T", c)
+				}
+				p := tc.Connector.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")
+				}
+				t, err := tc.accessTokenProvider()
+				if t != "token" || err != nil {
+					return fmt.Errorf("Unexpected results from tokenProvider: %v, %v", t, err)
+				}
+				return nil
+			},
+			wantErr: false,
+		},
+		{
+			name: "Nil tokenProvider gives error",
+			args: args{
+				dsn:           dsn,
+				tokenProvider: nil},
+			want:    nil,
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := NewAccessTokenConnector(tt.args.dsn, tt.args.tokenProvider)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("NewAccessTokenConnector() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if tt.want != nil {
+				if err := tt.want(got); err != nil {
+					t.Error(err)
+				}
+			}
+		})
+	}
+}
+
+func TestAccessTokenConnectorFailsToConnectIfNoAccessToken(t *testing.T) {
+	errorText := "This is a test"
+	dsn := "Server=server.database.windows.net;Database=db"
+	tp := func() (string, error) { return "", errors.New(errorText) }
+	sut, err := NewAccessTokenConnector(dsn, tp)
+	if err != nil {
+		t.Fatalf("expected err==nil, but got %+v", err)
+	}
+	_, err = sut.Connect(context.TODO())
+	if err == nil || !strings.Contains(err.Error(), errorText) {
+		t.Fatalf("expected error to contain %q, but got %q", errorText, err)
+	}
+}
diff --git a/ntlm_test.go b/ntlm_test.go
index db56f4e..a85317e 100644
--- a/ntlm_test.go
+++ b/ntlm_test.go
@@ -3,8 +3,10 @@
 package mssql
 
 import (
+	"bytes"
 	"encoding/hex"
 	"testing"
+	"time"
 )
 
 func TestLMOWFv1(t *testing.T) {
@@ -74,3 +76,56 @@ func TestNTLMSessionResponse(t *testing.T) {
 		t.Errorf("got:\n%sexpected:\n%s", hex.Dump(nt[:]), hex.Dump(val[:]))
 	}
 }
+
+func TestNTLMV2Response(t *testing.T) {
+	target := "DOMAIN"
+	username := "user"
+	password := "SecREt01"
+	challenge := [8]byte{ 0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef}
+	targetInformationBlock, _ := hex.DecodeString("02000c0044004f004d00410049004e0001000c005300450052005600450052000400140064006f006d00610069006e002e0063006f006d00030022007300650072007600650072002e0064006f006d00610069006e002e0063006f006d0000000000")
+	nonceBytes, _ := hex.DecodeString("ffffff0011223344")
+	var nonce [8]byte
+	copy(nonce[:8], nonceBytes[:])
+	timestamp, err := time.Parse(time.RFC3339, "2006-01-02T15:04:05Z")
+	if err != nil {
+		panic(err)
+	}
+
+	expectedNTLMV2Response, _ := hex.DecodeString("5c788afec59c1fef3f90bf6ea419c02501010000000000000fc4a4d5fdf6b200ffffff00112233440000000002000c0044004f004d00410049004e0001000c005300450052005600450052000400140064006f006d00610069006e002e0063006f006d00030022007300650072007600650072002e0064006f006d00610069006e002e0063006f006d000000000000000000")
+	expectedLMV2Response, _ := hex.DecodeString("d6e6152ea25d03b7c6ba6629c2d6aaf0ffffff0011223344")
+	ntlmV2Response, lmV2Response := getNTLMv2AndLMv2ResponsePayloads(target, username, password, challenge, nonce, targetInformationBlock, timestamp)
+	if bytes.Compare(ntlmV2Response, expectedNTLMV2Response) != 0 {
+		t.Errorf("got:\n%s\nexpected:\n%s", hex.Dump(ntlmV2Response), hex.Dump(expectedNTLMV2Response))
+	}
+
+	if bytes.Compare(lmV2Response, expectedLMV2Response) != 0 {
+		t.Errorf("got:\n%s\nexpected:\n%s", hex.Dump(ntlmV2Response), hex.Dump(expectedNTLMV2Response))
+	}
+}
+
+func TestGetNTLMv2TargetInfoFields(t *testing.T) {
+	type2Message, _ := hex.DecodeString("4e544c4d53535000020000000600060038000000058289026999bc21067c77f40000000000000000ac00ac003e0000000a0039380000000f4600570042000200060046005700420001000c00590037004100410041003400040022006000700065002e00610058006e0071006e0070006e00650074002e0063006f006d00030030007900370041004100410034002e006000700065002e00610058006e0071006e0070006e00650074002e0063006f006d00050024006100610058006d002e00610058006e0071006e0070006e00650074002e0063006f006d00070008007d9647e8aed6d50100000000")
+	info, target, err := getNTLMv2TargetInfoFields(type2Message)
+	if err != nil {
+		t.Errorf("got:\n%e\nexpected:\nnil", err)
+	}
+
+	expectedTarget := "FWB"
+	if target != expectedTarget {
+		t.Errorf("got:\n%s\nexpected:\n%s", target, expectedTarget)
+	}
+
+	expectedResponseLength := 172
+	responseLength := len(info)
+	if responseLength != expectedResponseLength {
+		t.Errorf("got:\n%d\nexpected:\n%d", responseLength, expectedResponseLength)
+	}
+}
+
+func TestGetNTLMv2TargetInfoFieldsInvalidMessage(t *testing.T) {
+	type2Message, _ := hex.DecodeString("4e544c4d53535000020000000600060038000000058289026999bc21067c77f40000000000000000ac00ac003e0000000a0039380000000f4600570042000200060046005700420001000c00590037004100410041003400040022006000700065002e00610058006e0071006e0070006e00650074002e0063006f006d00030030007900370041004100410034002e006000700065002e00610058006e0071006e0070006e00650074002e0063006f006d00050024006100610058006d002e00610058006e0071006e0070006e00650074002e0063006f006d00070008007")
+	_, _, err := getNTLMv2TargetInfoFields(type2Message)
+	if err == nil {
+		t.Error("expected to get an error")
+	}
+}
diff --git a/tds_test.go b/tds_test.go
index 00360d8..e725a66 100644
--- a/tds_test.go
+++ b/tds_test.go
@@ -5,7 +5,6 @@ import (
 	"context"
 	"database/sql"
 	"encoding/hex"
-	"fmt"
 	"net/url"
 	"os"
 	"runtime"
@@ -62,11 +61,60 @@ func TestSendLogin(t *testing.T) {
 		116, 0, 104, 0}
 	out := memBuf.Bytes()
 	if !bytes.Equal(ref, out) {
-		fmt.Println("Expected:")
-		fmt.Print(hex.Dump(ref))
-		fmt.Println("Returned:")
-		fmt.Print(hex.Dump(out))
-		t.Error("input output don't match")
+		t.Log("Expected:")
+		t.Log(hex.Dump(ref))
+		t.Log("Returned:")
+		t.Log(hex.Dump(out))
+		t.Fatal("input output don't match")
+	}
+}
+
+func TestSendLoginWithFeatureExt(t *testing.T) {
+	memBuf := new(MockTransport)
+	buf := newTdsBuffer(1024, memBuf)
+	login := login{
+		TDSVersion:     verTDS74,
+		PacketSize:     0x1000,
+		ClientProgVer:  0x01060100,
+		ClientPID:      100,
+		ClientTimeZone: -4 * 60,
+		ClientID:       [6]byte{0x12, 0x34, 0x56, 0x78, 0x90, 0xab},
+		OptionFlags1:   0xe0,
+		OptionFlags3:   8,
+		HostName:       "subdev1",
+		AppName:        "appname",
+		ServerName:     "servername",
+		CtlIntName:     "library",
+		Language:       "en",
+		Database:       "database",
+		ClientLCID:     0x204,
+	}
+	login.FeatureExt.Add(&featureExtFedAuthSTS{
+		FedAuthToken: "fedauthtoken",
+	})
+	err := sendLogin(buf, login)
+	if err != nil {
+		t.Error("sendLogin should succeed")
+	}
+	ref := []byte{
+		16, 1, 0, 223, 0, 0, 1, 0, 215, 0, 0, 0, 4, 0, 0, 116, 0, 16, 0, 0, 0, 1,
+		6, 1, 100, 0, 0, 0, 0, 0, 0, 0, 224, 0, 0, 24, 16, 255, 255, 255, 4, 2, 0,
+		0, 94, 0, 7, 0, 108, 0, 0, 0, 108, 0, 0, 0, 108, 0, 7, 0, 122, 0, 10, 0, 176,
+		0, 4, 0, 142, 0, 7, 0, 156, 0, 2, 0, 160, 0, 8, 0, 18, 52, 86, 120, 144, 171,
+		176, 0, 0, 0, 176, 0, 0, 0, 176, 0, 0, 0, 0, 0, 0, 0, 115, 0, 117, 0, 98,
+		0, 100, 0, 101, 0, 118, 0, 49, 0, 97, 0, 112, 0, 112, 0, 110, 0, 97, 0,
+		109, 0, 101, 0, 115, 0, 101, 0, 114, 0, 118, 0, 101, 0, 114, 0, 110, 0, 97,
+		0, 109, 0, 101, 0, 108, 0, 105, 0, 98, 0, 114, 0, 97, 0, 114, 0, 121, 0, 101,
+		0, 110, 0, 100, 0, 97, 0, 116, 0, 97, 0, 98, 0, 97, 0, 115, 0, 101, 0, 180, 0,
+		0, 0, 2, 29, 0, 0, 0, 2, 24, 0, 0, 0, 102, 0, 101, 0, 100, 0, 97, 0, 117, 0,
+		116, 0, 104, 0, 116, 0, 111, 0, 107, 0, 101, 0, 110, 0, 255}
+	out := memBuf.Bytes()
+	if !bytes.Equal(ref, out) {
+		t.Log("Expected:")
+		t.Log(hex.Dump(ref))
+		t.Log("Returned:")
+		t.Log(hex.Dump(out))
+		t.Fatal("input output don't match")
 	}
 }
 
@@ -109,7 +157,7 @@ loop:
 		case []interface{}:
 			lastRow = token
 		default:
-			fmt.Println("unknown token", tok)
+			t.Log("unknown token", tok)
 		}
 	}
 
@@ -176,8 +224,6 @@ func (l testLogger) Println(v ...interface{}) {
 	l.t.Log(v...)
 }
 
-
-
 func TestConnect(t *testing.T) {
 	checkConnStr(t)
 	SetLogger(testLogger{t})
@@ -413,7 +459,7 @@ func TestSSPIAuth(t *testing.T) {
 
 func TestUcs22Str(t *testing.T) {
 	// Test valid input
-	s, err := ucs22str([]byte{0x31, 0, 0x32, 0, 0x33, 0})  // 123 in UCS2 encoding
+	s, err := ucs22str([]byte{0x31, 0, 0x32, 0, 0x33, 0}) // 123 in UCS2 encoding
 	if err != nil {
 		t.Errorf("ucs22str should not fail for valid ucs2 byte sequence: %s", err)
 	}
@@ -429,7 +475,7 @@ func TestUcs22Str(t *testing.T) {
 }
 
 func TestReadUcs2(t *testing.T) {
-	buf := bytes.NewBuffer([]byte{0x31, 0, 0x32, 0, 0x33, 0})  // 123 in UCS2 encoding
+	buf := bytes.NewBuffer([]byte{0x31, 0, 0x32, 0, 0x33, 0}) // 123 in UCS2 encoding
 	s, err := readUcs2(buf, 3)
 	if err != nil {
 		t.Errorf("readUcs2 should not fail for valid ucs2 byte sequence: %s", err)
@@ -447,7 +493,7 @@ func TestReadUcs2(t *testing.T) {
 
 func TestReadUsVarChar(t *testing.T) {
 	// should succeed for valid buffer
-	buf := bytes.NewBuffer([]byte{3, 0, 0x31, 0, 0x32, 0, 0x33, 0})  // 123 in UCS2 encoding with length prefix 3 uint16
+	buf := bytes.NewBuffer([]byte{3, 0, 0x31, 0, 0x32, 0, 0x33, 0}) // 123 in UCS2 encoding with length prefix 3 uint16
 	s, err := readUsVarChar(buf)
 	if err != nil {
 		t.Errorf("readUsVarChar should not fail for valid ucs2 byte sequence: %s", err)
@@ -487,4 +533,4 @@ func TestReadBVarByte(t *testing.T) {
 	if err == nil {
 		t.Error("readUsVarByte should fail on short buffer, but it didn't")
 	}
-}
\ No newline at end of file
+}