fix(okta): random per-install device token (#23)

pkg/provider/okta/okta.go

@@ -3,8 +3,10 @@ package okta
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
 	"crypto/tls"
 	"encoding/base64"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"html"
@@ -13,13 +15,16 @@ import (
 	"net/http"
 	"net/http/cookiejar"
 	"net/url"
+	"os"
+	"path/filepath"
 	"regexp"
 	"strconv"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/PuerkitoBio/goquery"
+	"github.com/google/uuid"
 	"github.com/marshallbrekka/go-u2fhost"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -360,10 +365,83 @@ func (oc *Client) getDeviceTokenFromOkta(loginDetails *creds.LoginDetails) (stri
 	return "", fmt.Errorf("unable to get a device token from okta")
 }
 
+// deviceTokenDir is where the device token file fallback lives. A variable so
+// tests can point it elsewhere.
+var deviceTokenDir = func() (string, error) {
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return "", err
+	}
+	return filepath.Join(home, ".aws", "saml2aws"), nil
+}
+
+// deviceToken returns the Okta device token for this install. A random UUID is
+// generated on first use and kept in the credential store under
+// <idp url>/deviceToken, so Okta keeps remembering the device across runs.
+// Without a usable credential store the token is kept in a 0600 file under the
+// saml2aws config directory instead. The token is never derived from the username.
+func deviceToken(loginDetails *creds.LoginDetails) string {
+	if !credentials.SupportsStorage() {
+		logger.Debug("no credential store, using okta device token file")
+		return fileDeviceToken(loginDetails.URL, uuid.NewString())
+	}
+
+	key := loginDetails.URL + "/deviceToken"
+	_, dt, err := credentials.CurrentHelper.Get(key)
+	switch {
+	case err == nil && dt != "":
+		return dt
+	case err != nil && !credentials.IsErrCredentialsNotFound(err):
+		// Don't rotate a token we could not read; it may still be in the store.
+		logger.Warnf("unable to read okta device token from the credential store, using the device token file | err: %v", err)
+		return fileDeviceToken(loginDetails.URL, uuid.NewString())
+	}
+
+	dt = uuid.NewString()
+	if err := credentials.SaveCredentials(key, loginDetails.Username, dt); err != nil {
+		logger.Warnf("unable to store okta device token in the credential store, using the device token file; Okta may ask for MFA again | err: %v", err)
+		return fileDeviceToken(loginDetails.URL, dt)
+	}
+	return dt
+}
+
+// fileDeviceToken returns the device token kept in the fallback file for idpURL,
+// creating the file with candidate if it does not exist yet. If the file can't
+// be used, candidate is returned for this run only.
+func fileDeviceToken(idpURL, candidate string) string {
+	dir, err := deviceTokenDir()
+	if err != nil {
+		logger.Warnf("unable to find the saml2aws config directory, okta device token will not be remembered | err: %v", err)
+		return candidate
+	}
+	sum := sha256.Sum256([]byte(idpURL))
+	path := filepath.Join(dir, "okta-device-token-"+hex.EncodeToString(sum[:8]))
+
+	if b, err := os.ReadFile(path); err == nil {
+		if dt := strings.TrimSpace(string(b)); dt != "" {
+			return dt
+		}
+	} else if !os.IsNotExist(err) {
+		logger.Warnf("unable to read okta device token file %s, device will not be remembered | err: %v", path, err)
+		return candidate
+	}
+
+	if err := os.MkdirAll(dir, 0700); err != nil {
+		logger.Warnf("unable to create %s, okta device token will not be remembered | err: %v", dir, err)
+		return candidate
+	}
+	if err := os.WriteFile(path, []byte(candidate), 0600); err != nil {
+		logger.Warnf("unable to write okta device token file %s, device will not be remembered | err: %v", path, err)
+		return candidate
+	}
+	// WriteFile keeps the mode of an existing file; make sure it is private.
+	_ = os.Chmod(path, 0600)
+	return candidate
+}
+
 // setDeviceTokenCookie sets the DT cookie in the HTTP Client cookie jar
-// using the okta_<loginDetails.Username>_saml2aws, we reduce making an extra api call
-// this func can be uplifted in the future to set custom device tokens or used with
-// getDeviceTokenFromOkta function
+// using a per-install device token (see deviceToken), we reduce making an extra api call
+// this func can be uplifted in the future to use getDeviceTokenFromOkta
 func (oc *Client) setDeviceTokenCookie(loginDetails *creds.LoginDetails) error {
 
 	// getDeviceTokenFromOkta is not used but doing this to keep the function code
@@ -385,8 +463,8 @@ func (oc *Client) setDeviceTokenCookie(loginDetails *creds.LoginDetails) error {
 	cookie := http.Cookie{
 		Name:    "DT",
 		Secure:  true,
-		Expires: time.Now().Add(time.Hour * 24 * 30),                    // 30 Days -> this time might not matter as this cookie is set on every saml2aws login request
-		Value:   fmt.Sprintf("okta_%s_saml2aws", loginDetails.Username), // Okta recommends using an UUID but this should be unique enough. Also, this is key to remembering Okta MFA device
+		Expires: time.Now().Add(time.Hour * 24 * 30), // 30 Days -> this time might not matter as this cookie is set on every saml2aws login request
+		Value:   deviceToken(loginDetails),           // key to remembering the Okta MFA device
 	}
 	cookies = append(cookies, &cookie)
 	oc.client.Jar.SetCookies(baseURL, cookies)

pkg/provider/okta/okta_devicetoken_test.go

@@ -0,0 +1,201 @@
+package okta
+
+import (
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"os"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"testing"
+
+	"github.com/sirupsen/logrus"
+	logtest "github.com/sirupsen/logrus/hooks/test"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/versent/saml2aws/v2/helper/credentials"
+)
+
+// memHelper is an in-memory credentials store for tests.
+type memHelper struct {
+	store  map[string]*credentials.Credentials
+	getErr error
+	addErr error
+}
+
+func (m *memHelper) Add(c *credentials.Credentials) error {
+	if m.addErr != nil {
+		return m.addErr
+	}
+	m.store[c.ServerURL] = c
+	return nil
+}
+
+func (m *memHelper) Delete(serverURL string) error {
+	delete(m.store, serverURL)
+	return nil
+}
+
+func (m *memHelper) Get(serverURL string) (string, string, error) {
+	if m.getErr != nil {
+		return "", "", m.getErr
+	}
+	c, ok := m.store[serverURL]
+	if !ok {
+		return "", "", credentials.ErrCredentialsNotFound
+	}
+	return c.Username, c.Secret, nil
+}
+
+func (m *memHelper) SupportsCredentialStorage() bool { return true }
+
+func useMemHelper(t *testing.T) *memHelper {
+	t.Helper()
+	old := credentials.CurrentHelper
+	h := &memHelper{store: map[string]*credentials.Credentials{}}
+	credentials.CurrentHelper = h
+	t.Cleanup(func() { credentials.CurrentHelper = old })
+	return h
+}
+
+func dtCookie(t *testing.T, oc *Client, rawURL string) string {
+	t.Helper()
+	u, err := url.Parse(rawURL)
+	require.NoError(t, err)
+	for _, c := range oc.client.Jar.Cookies(&url.URL{Scheme: u.Scheme, Host: u.Host, Path: "/"}) {
+		if c.Name == "DT" {
+			return c.Value
+		}
+	}
+	t.Fatal("DT cookie not set")
+	return ""
+}
+
+func TestDeviceTokenIsRandomAndStable(t *testing.T) {
+	ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	defer ts.Close()
+	h := useMemHelper(t)
+
+	oc, loginDetails := setupTestClient(t, ts, "PUSH")
+	require.NoError(t, oc.setDeviceTokenCookie(loginDetails))
+	first := dtCookie(t, oc, loginDetails.URL)
+
+	assert.NotContains(t, first, loginDetails.Username, "device token must not be derived from the username")
+	assert.NotEqual(t, "okta_"+loginDetails.Username+"_saml2aws", first)
+	_, stored, err := h.Get(loginDetails.URL + "/deviceToken")
+	require.NoError(t, err, "device token must be persisted in the credential store")
+	assert.Equal(t, first, stored)
+
+	// A second run (new client) reuses the stored token.
+	oc2, loginDetails2 := setupTestClient(t, ts, "PUSH")
+	require.NoError(t, oc2.setDeviceTokenCookie(loginDetails2))
+	assert.Equal(t, first, dtCookie(t, oc2, loginDetails2.URL))
+
+	// A different install (empty store) gets a different token.
+	useMemHelper(t)
+	oc3, loginDetails3 := setupTestClient(t, ts, "PUSH")
+	require.NoError(t, oc3.setDeviceTokenCookie(loginDetails3))
+	assert.NotEqual(t, first, dtCookie(t, oc3, loginDetails3.URL))
+	assert.False(t, strings.HasPrefix(dtCookie(t, oc3, loginDetails3.URL), "okta_"))
+}
+
+func TestDeviceTokenFileFallbackWithoutCredentialStore(t *testing.T) {
+	ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	defer ts.Close()
+	dir := t.TempDir()
+	useDeviceTokenDir(t, dir)
+
+	// The default helper has no storage, so the token goes to a private file.
+	oc, loginDetails := setupTestClient(t, ts, "PUSH")
+	require.NoError(t, oc.setDeviceTokenCookie(loginDetails))
+	first := dtCookie(t, oc, loginDetails.URL)
+	assert.NotContains(t, first, loginDetails.Username)
+
+	files, err := filepath.Glob(filepath.Join(dir, "okta-device-token-*"))
+	require.NoError(t, err)
+	require.Len(t, files, 1)
+	if runtime.GOOS != "windows" {
+		fi, err := os.Stat(files[0])
+		require.NoError(t, err)
+		assert.Equal(t, os.FileMode(0600), fi.Mode().Perm())
+	}
+
+	// A later run reads the same token back.
+	oc2, loginDetails2 := setupTestClient(t, ts, "PUSH")
+	require.NoError(t, oc2.setDeviceTokenCookie(loginDetails2))
+	assert.Equal(t, first, dtCookie(t, oc2, loginDetails2.URL))
+}
+
+func TestDeviceTokenTransientGetErrorDoesNotRotate(t *testing.T) {
+	ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	defer ts.Close()
+	h := useMemHelper(t)
+
+	oc, loginDetails := setupTestClient(t, ts, "PUSH")
+	require.NoError(t, oc.setDeviceTokenCookie(loginDetails))
+	stored := dtCookie(t, oc, loginDetails.URL)
+
+	// The keychain is locked or unavailable for this run.
+	h.getErr = errors.New("keychain locked")
+	hook := logtest.NewGlobal()
+	oc2, loginDetails2 := setupTestClient(t, ts, "PUSH")
+	require.NoError(t, oc2.setDeviceTokenCookie(loginDetails2))
+
+	_, still, err := (&memHelper{store: h.store}).Get(loginDetails.URL + "/deviceToken")
+	require.NoError(t, err)
+	assert.Equal(t, stored, still, "a transient Get error must not overwrite the stored token")
+	assert.True(t, hasWarning(hook), "a transient Get error must be reported as a warning")
+}
+
+func TestDeviceTokenSaveErrorWarnsAndUsesFile(t *testing.T) {
+	ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	defer ts.Close()
+	h := useMemHelper(t)
+	h.addErr = errors.New("keychain write denied")
+	dir := t.TempDir()
+	useDeviceTokenDir(t, dir)
+
+	hook := logtest.NewGlobal()
+	oc, loginDetails := setupTestClient(t, ts, "PUSH")
+	require.NoError(t, oc.setDeviceTokenCookie(loginDetails))
+	first := dtCookie(t, oc, loginDetails.URL)
+
+	assert.NotEqual(t, "okta_"+loginDetails.Username+"_saml2aws", first, "must not fall back to the predictable token")
+	assert.True(t, hasWarning(hook), "a failed save must be reported as a warning")
+
+	// The token kept in the file is reused while saves keep failing.
+	oc2, loginDetails2 := setupTestClient(t, ts, "PUSH")
+	require.NoError(t, oc2.setDeviceTokenCookie(loginDetails2))
+	assert.Equal(t, first, dtCookie(t, oc2, loginDetails2.URL))
+}
+
+func hasWarning(hook *logtest.Hook) bool {
+	for _, e := range hook.AllEntries() {
+		if e.Level == logrus.WarnLevel && strings.Contains(e.Message, "device token") {
+			return true
+		}
+	}
+	return false
+}
+
+func useDeviceTokenDir(t *testing.T, dir string) {
+	t.Helper()
+	old := deviceTokenDir
+	deviceTokenDir = func() (string, error) { return dir, nil }
+	t.Cleanup(func() { deviceTokenDir = old })
+}
+
+// TestMain keeps every test in this package away from the real
+// ~/.aws/saml2aws directory.
+func TestMain(m *testing.M) {
+	dir, err := os.MkdirTemp("", "saml2aws-okta-test")
+	if err != nil {
+		panic(err)
+	}
+	deviceTokenDir = func() (string, error) { return dir, nil }
+	code := m.Run()
+	_ = os.RemoveAll(dir)
+	os.Exit(code)
+}

pkg/provider/okta/okta_test.go

@@ -665,15 +665,15 @@ func TestSetDeviceTokenCookie(t *testing.T) {
 	err = oc.setDeviceTokenCookie(loginDetails)
 	assert.Nil(t, err)
 
-	expectedDT := fmt.Sprintf("okta_%s_saml2aws", loginDetails.Username)
+	legacyDT := fmt.Sprintf("okta_%s_saml2aws", loginDetails.Username)
 	actualDT := ""
 	for _, c := range oc.client.Jar.Cookies(&url.URL{Scheme: "https", Host: "idp.example.com", Path: "/abc"}) {
 		if c.Name == "DT" {
 			actualDT = c.Value
 		}
 	}
 	assert.NotEqual(t, actualDT, "")
-	assert.Equal(t, expectedDT, actualDT)
+	assert.NotEqual(t, legacyDT, actualDT)
 
 }