feat: Implement OAuth2 PKCE in SSORoleCredentialsProvider by patrickod · Pull Request #359 · ByteNess/aws-vault

patrickod · 2026-05-05T20:42:59Z

Ported from tailscale/aws-vault@949040e

Inspired by AWS CLI this PR adds support for OAuth2 "authorization code" grant flow accompanied by Proof Key for Code Exchange (PKCE) by OAuth Public Clients and makes it the default.

In the "device code" flow the user is entirely responsible for matching the one time password displayed by their terminal and browser. This is error prone at best and a one click penalty for every login at worst.

In the "authorization code" flow we create an ephemeral HTTP server to host an OAuth2 callback on 127.0.0.1. Users are redirected here after confirming their login attempt and the code comparison step is performed for them. AWS enforces the use of 127.0.0.1 as a callback destination host so a remote attacker who successfully convinces a user to confirm their malicious login prompt would not be able to access the resulting credential as the user would be redirected to their own localhost and not an attacker controlled destination.

If we determine that we are not able to open a browser (either from the user specifying --stdout or observing that we are executing in an SSH session) we revert to the previous "device code" flow and print the URLs for copy/pasting into browsers.

mbevc1

Thanks @patrickod and looks great! Few minor suggestions and couple of questions:

I wonder if it would make sense to add tests for new code
We should probably capture changes in the USAGE.md.

Seems defaults are changing here as well, would it make sense to revert the logic to preserve current behaviour as default? 🤔

mbevc1 · 2026-05-06T18:33:08Z

 		BoolVar(&input.UseStdout)

+	cmd.Flag("device-code", "Use the device-code OAuth2 flow for SSO instead of the default PKCE browser flow").
+		BoolVar(&input.UseDeviceCode)


Could we add ENV handling as well?
OverrideDefaultFromEnvar("AWS_VAULT_DEVICE_CODE").

mbevc1 · 2026-05-06T18:33:20Z

 		BoolVar(&input.UseStdout)

+	cmd.Flag("device-code", "Use the device-code OAuth2 flow for SSO instead of the default PKCE browser flow").
+		BoolVar(&input.UseDeviceCode)


Could we add ENV handling as well?
OverrideDefaultFromEnvar("AWS_VAULT_DEVICE_CODE").

mbevc1 · 2026-05-06T18:48:58Z

+	go func() {
+		if err := cbServer.Serve(); err != nil && !errors.Is(err, http.ErrServerClosed) {


select on ctx.Done() and tear down via defer:

Suggested change

go func() {

if err := cbServer.Serve(); err != nil && !errors.Is(err, http.ErrServerClosed) {

// ensure the callback server is torn down on any exit path

defer func() {

if err := cbServer.h.Close(); err != nil {

log.Printf("Failed to close oauthCallbackServer: %s", err)

}

}()

go func() {

if err := cbServer.Serve(); err != nil && !errors.Is(err, http.ErrServerClosed) {

log.Printf("Failed to run oauthCallbackServer: %s", err)

}

}()

go func() {

if err := cbServer.Serve(); err != nil && !errors.Is(err, http.ErrServerClosed) {

mbevc1 · 2026-05-06T18:50:59Z

+	r := <-cbServer.resultChan
+	// tear down the callback server
+	if err := cbServer.h.Close(); err != nil {
+		log.Printf("Failed to close oauthCallbackServer: %s", err)
+	}


Suggested change

r := <-cbServer.resultChan

// tear down the callback server

if err := cbServer.h.Close(); err != nil {

log.Printf("Failed to close oauthCallbackServer: %s", err)

}

// await the authorization code, or context cancellation

var r oauthCallbackResult

select {

case r = <-cbServer.resultChan:

case <-ctx.Done():

return nil, fmt.Errorf("aborted waiting for OAuth callback: %w", ctx.Err())

}

if r.err != nil {

return nil, r.err

}

if r.code == "" {

return nil, errors.New("no authorization code received")

}

mbevc1 · 2026-05-06T19:02:07Z

+			log.Printf("Failed to run oauthCallbackServer: %s", err)
+		}
+	}()
+	// keep a copy of the redirectURI for the CreateToken call (as the server will be closed after the code is received)


Minor nit on improving the comment here

Suggested change

// keep a copy of the redirectURI for the CreateToken call (as the server will be closed after the code is received)

// redirectURI is used in both the authorize URL and the later CreateToken call;

// bind it once so both sides agree on the same value.

mbevc1 · 2026-05-06T19:06:05Z

+
+	// respond with a success message
+	io.WriteString(w, "Authorization code received, you can close this tab now.")
+}


Just ordering when to emit the message:

Suggested change

}

// respond with a success message before signalling completion, so the

// response is fully written before the main goroutine tears the server down

io.WriteString(w, "Authorization code received, you can close this tab now.")

s.resultChan <- oauthCallbackResult{code: code}

mbevc1 · 2026-05-18T21:14:09Z

Hi @patrickod, did you manage to look at the feedback?

mbevc1 · 2026-06-08T13:49:08Z

@patrickod are you still working on this feature? I might need a rebase to trigger new Windows tests

patrickod · 2026-06-09T18:29:53Z

@mbevc1 sorry I entered into a hectic period right as I cut this PR. I am happy to polish this up but as it might be next week until that happens.

For context this fork has been successfully deployed at Tailscale for a few months now, my hope was to encourage wider adoption of the PKCE flow as it is strictly more secure and closes a phishing vulnerability in the traditional flow. RE changing the default YMMV but I would argue this is an acceptable upgrade and is likely to be a net positive (one fewer click) for the majority of users. For those operating on remote systems where no browser is available to complete the flow the traditional flow remains a viable alternative.

Ported from tailscale/aws-vault@949040e Inspired by AWS CLI this PR adds support for OAuth2 "authorization code" grant flow accompanied by Proof Key for Code Exchange (PKCE) by OAuth Public Clients and makes it the default. In the "device code" flow the user is entirely responsible for matching the one time password displayed by their terminal and browser. This is error prone at best and a one click penalty for every login at worst. In the "authorization code" flow we create an ephemeral HTTP server to host an OAuth2 callback on 127.0.0.1. Users are redirected here after confirming their login attempt and the code comparison step is automated for them. As AWS enforces the use of 127.0.0.1 as a callback destination host a remote attacker who successfully convinces a user to confirm their malicious login prompt would not be able to access the resulting credential as the user would be redirected to their own localhost and not an attacker controlled destination. If we determine that we are not able to open a browser (either from the user specifying --stdout or observing that we are executing in an SSH session) we revert to the previous "device code" flow and print the URLs for copy/pasting into browsers.

mbevc1 · 2026-06-10T14:10:34Z

+// and sends the authorization code received via a channel.
+func newOauthCallbackServer() (*oauthCallbackServer, error) {
+	// select a random port for the callback server
+	ln, err := net.Listen("tcp", ":0")


We should probably scope this to localhost only:

Suggested change

ln, err := net.Listen("tcp", ":0")

ln, err := net.Listen("tcp", "127.0.0.1:0")

mbevc1 · 2026-06-10T14:14:51Z

+	state := r.URL.Query().Get("state")
+	if subtle.ConstantTimeCompare([]byte(state), []byte(s.state)) != 1 {
+		http.Error(w, "Invalid state", http.StatusBadRequest)
+		s.resultChan <- oauthCallbackResult{err: errors.New("invalid state")}


A note and probably intentional?

A wrong/missing state (e.g., a port probe) pushes an error into resultChan, which the main goroutine treats as terminal. AWS CLI just ignores mismatched callbacks and keeps waiting.

mbevc1 · 2026-06-10T14:18:41Z

Thanks @patrickod . It's a bit harder to review on a force-push, but I've left few comments, if you can have a look at those please?

Also I wonder and noticed the following:

login has no opt-out: cli/login.go never sets SSOUseDeviceCode and has no --device-code flag, so SSO profiles via aws-vault login are forced onto PKCE.
Missing error param handling - AWS CLI handles error/error_description in its callback handler, is this something we want to handle as well?

patrickod requested a review from mbevc1 as a code owner May 5, 2026 20:43

mbevc1 changed the title ~~vault: Implement OAuth2 PKCE in SSORoleCredentialsProvider~~ feat: Implement OAuth2 PKCE in SSORoleCredentialsProvider May 5, 2026

github-actions Bot added the feat label May 5, 2026

patrickod force-pushed the patrickod/pkce-byteness-7.10.5 branch from 11c7540 to bf36047 Compare May 5, 2026 21:27

mbevc1 requested changes May 6, 2026

View reviewed changes

patrickod force-pushed the patrickod/pkce-byteness-7.10.5 branch from bf36047 to c86138a Compare June 10, 2026 01:03

github-actions Bot added the documentation Improvements or additions to documentation label Jun 10, 2026

patrickod requested a review from mbevc1 June 10, 2026 01:20

mbevc1 reviewed Jun 10, 2026

View reviewed changes

		go func() {
		if err := cbServer.Serve(); err != nil && !errors.Is(err, http.ErrServerClosed) {

-	r := <-cbServer.resultChan
-	// tear down the callback server
-	if err := cbServer.h.Close(); err != nil {
-		log.Printf("Failed to close oauthCallbackServer: %s", err)
-	}
+    // await the authorization code, or context cancellation
+    var r oauthCallbackResult
+    select {
+    case r = <-cbServer.resultChan:
+    case <-ctx.Done():
+        return nil, fmt.Errorf("aborted waiting for OAuth callback: %w", ctx.Err())
+    }
+    if r.err != nil {
+        return nil, r.err
+    }
+    if r.code == "" {
+        return nil, errors.New("no authorization code received")
+    }

	// keep a copy of the redirectURI for the CreateToken call (as the server will be closed after the code is received)
	// redirectURI is used in both the authorize URL and the later CreateToken call;
	// bind it once so both sides agree on the same value.

-}
+    // respond with a success message before signalling completion, so the
+    // response is fully written before the main goroutine tears the server down
+    io.WriteString(w, "Authorization code received, you can close this tab now.")
+    s.resultChan <- oauthCallbackResult{code: code}

	ln, err := net.Listen("tcp", ":0")
	ln, err := net.Listen("tcp", "127.0.0.1:0")

Uh oh!

Conversation

patrickod commented May 5, 2026

Uh oh!

mbevc1 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mbevc1 commented May 18, 2026

Uh oh!

mbevc1 commented Jun 8, 2026

Uh oh!

patrickod commented Jun 9, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mbevc1 commented Jun 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants