AUTH-5682 Org token flow in Access logins should pass CF_AppSession cookie

- Refactor HandleRedirects function and add unit tests
- Move signal test to its own file because of OS specific instructions

token/signal_test.go

@@ -0,0 +1,54 @@
+//go:build linux || darwin
+
+package token
+
+import (
+	"os"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSignalHandler(t *testing.T) {
+	sigHandler := signalHandler{signals: []os.Signal{syscall.SIGUSR1}}
+	handlerRan := false
+	done := make(chan struct{})
+	timer := time.NewTimer(time.Second)
+	sigHandler.register(func() {
+		handlerRan = true
+		done <- struct{}{}
+	})
+
+	p, err := os.FindProcess(os.Getpid())
+	require.Nil(t, err)
+	p.Signal(syscall.SIGUSR1)
+
+	// Blocks for up to one second to make sure the handler callback runs before the assert.
+	select {
+	case <-done:
+		assert.True(t, handlerRan)
+	case <-timer.C:
+		t.Fail()
+	}
+	sigHandler.deregister()
+}
+
+func TestSignalHandlerClose(t *testing.T) {
+	sigHandler := signalHandler{signals: []os.Signal{syscall.SIGUSR1}}
+	done := make(chan struct{})
+	timer := time.NewTimer(time.Second)
+	sigHandler.register(func() { done <- struct{}{} })
+	sigHandler.deregister()
+
+	p, err := os.FindProcess(os.Getpid())
+	require.Nil(t, err)
+	p.Signal(syscall.SIGUSR1)
+	select {
+	case <-done:
+		t.Fail()
+	case <-timer.C:
+	}
+}

token/token.go

@@ -23,9 +23,11 @@ import (
 const (
 	keyName                    = "token"
 	tokenCookie                = "CF_Authorization"
+	appSessionCookie           = "CF_AppSession"
 	appDomainHeader            = "CF-Access-Domain"
 	appAUDHeader               = "CF-Access-Aud"
 	AccessLoginWorkerPath      = "/cdn-cgi/access/login"
+	AccessAuthorizedWorkerPath = "/cdn-cgi/access/authorized"
 )
 
 var (
@@ -297,20 +299,41 @@ func GetAppInfo(reqURL *url.URL) (*AppInfo, error) {
 	return &AppInfo{location.Hostname(), aud, domain}, nil
 }
 
+func handleRedirects(req *http.Request, via []*http.Request, orgToken string) error {
+	// attach org token to login request
+	if strings.Contains(req.URL.Path, AccessLoginWorkerPath) {
+		req.AddCookie(&http.Cookie{Name: tokenCookie, Value: orgToken})
+	}
+
+	// attach app session cookie to authorized request
+	if strings.Contains(req.URL.Path, AccessAuthorizedWorkerPath) {
+		// We need to check and see if the CF_APP_SESSION cookie was set
+		for _, prevReq := range via {
+			if prevReq != nil && prevReq.Response != nil {
+				for _, c := range prevReq.Response.Cookies() {
+					if c.Name == appSessionCookie {
+						req.AddCookie(&http.Cookie{Name: appSessionCookie, Value: c.Value})
+						return nil
+					}
+				}
+			}
+		}
+
+	}
+
+	// stop after hitting authorized endpoint since it will contain the app token
+	if len(via) > 0 && strings.Contains(via[len(via)-1].URL.Path, AccessAuthorizedWorkerPath) {
+		return http.ErrUseLastResponse
+	}
+	return nil
+}
+
 // exchangeOrgToken attaches an org token to a request to the appURL and returns an app token. This uses the Access SSO
 // flow to automatically generate and return an app token without the login page.
 func exchangeOrgToken(appURL *url.URL, orgToken string) (string, error) {
 	client := &http.Client{
 		CheckRedirect: func(req *http.Request, via []*http.Request) error {
-			// attach org token to login request
-			if strings.Contains(req.URL.Path, AccessLoginWorkerPath) {
-				req.AddCookie(&http.Cookie{Name: tokenCookie, Value: orgToken})
-			}
-			// stop after hitting authorized endpoint since it will contain the app token
-			if strings.Contains(via[len(via)-1].URL.Path, "cdn-cgi/access/authorized") {
-				return http.ErrUseLastResponse
-			}
-			return nil
+			return handleRedirects(req, via, orgToken)
 		},
 		Timeout: time.Second * 7,
 	}

token/token_test.go

@@ -1,54 +1,82 @@
-//go:build linux
-
 package token
 
 import (
-	"os"
-	"syscall"
+	"net/http"
+	"net/url"
 	"testing"
-	"time"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
-func TestSignalHandler(t *testing.T) {
-	sigHandler := signalHandler{signals: []os.Signal{syscall.SIGUSR1}}
-	handlerRan := false
-	done := make(chan struct{})
-	timer := time.NewTimer(time.Second)
-	sigHandler.register(func() {
-		handlerRan = true
-		done <- struct{}{}
-	})
+func TestHandleRedirects_AttachOrgToken(t *testing.T) {
+	req, _ := http.NewRequest("GET", "http://example.com/cdn-cgi/access/login", nil)
+	via := []*http.Request{}
+	orgToken := "orgTokenValue"
 
-	p, err := os.FindProcess(os.Getpid())
-	require.Nil(t, err)
-	p.Signal(syscall.SIGUSR1)
+	handleRedirects(req, via, orgToken)
 
-	// Blocks for up to one second to make sure the handler callback runs before the assert.
-	select {
-	case <-done:
-		assert.True(t, handlerRan)
-	case <-timer.C:
-		t.Fail()
-	}
-	sigHandler.deregister()
-}
-
-func TestSignalHandlerClose(t *testing.T) {
-	sigHandler := signalHandler{signals: []os.Signal{syscall.SIGUSR1}}
-	done := make(chan struct{})
-	timer := time.NewTimer(time.Second)
-	sigHandler.register(func() { done <- struct{}{} })
-	sigHandler.deregister()
-
-	p, err := os.FindProcess(os.Getpid())
-	require.Nil(t, err)
-	p.Signal(syscall.SIGUSR1)
-	select {
-	case <-done:
-		t.Fail()
-	case <-timer.C:
+	// Check if the orgToken cookie is attached
+	cookies := req.Cookies()
+	found := false
+	for _, cookie := range cookies {
+		if cookie.Name == tokenCookie && cookie.Value == orgToken {
+			found = true
+			break
+		}
+	}
+
+	if !found {
+		t.Errorf("OrgToken cookie not attached to the request.")
+	}
+}
+
+func TestHandleRedirects_AttachAppSessionCookie(t *testing.T) {
+	req, _ := http.NewRequest("GET", "http://example.com/cdn-cgi/access/authorized", nil)
+	via := []*http.Request{
+		{
+			URL: &url.URL{Path: "/cdn-cgi/access/login"},
+			Response: &http.Response{
+				Header: http.Header{"Set-Cookie": {"CF_AppSession=appSessionValue"}},
+			},
+		},
+	}
+	orgToken := "orgTokenValue"
+
+	err := handleRedirects(req, via, orgToken)
+
+	// Check if the appSessionCookie is attached to the request
+	cookies := req.Cookies()
+	found := false
+	for _, cookie := range cookies {
+		if cookie.Name == appSessionCookie && cookie.Value == "appSessionValue" {
+			found = true
+			break
+		}
+	}
+
+	if !found {
+		t.Errorf("AppSessionCookie not attached to the request.")
+	}
+
+	if err != nil {
+		t.Errorf("Expected no error, got %v", err)
+	}
+}
+
+func TestHandleRedirects_StopAtAuthorizedEndpoint(t *testing.T) {
+	req, _ := http.NewRequest("GET", "http://example.com/cdn-cgi/access/authorized", nil)
+	via := []*http.Request{
+		{
+			URL: &url.URL{Path: "other"},
+		},
+		{
+			URL: &url.URL{Path: AccessAuthorizedWorkerPath},
+		},
+	}
+	orgToken := "orgTokenValue"
+
+	err := handleRedirects(req, via, orgToken)
+
+	// Check if ErrUseLastResponse is returned
+	if err != http.ErrUseLastResponse {
+		t.Errorf("Expected ErrUseLastResponse, got %v", err)
 	}
 }