multi-domain login redirects (#5564)

Add a new 'depends_on' route configuration option taking a list of 
additional hosts to redirect through on login. Update the authorize 
service and proxy service to support a chain of /.pomerium/callback
redirects. Add an integration test for this feature.

internal/authenticateflow/authenticateflow_int_test.go

@@ -0,0 +1,98 @@
+package authenticateflow_test
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/cookiejar"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/pomerium/pomerium/config"
+	"github.com/pomerium/pomerium/internal/testenv"
+	"github.com/pomerium/pomerium/internal/testenv/scenarios"
+	"github.com/pomerium/pomerium/internal/testenv/snippets"
+	"github.com/pomerium/pomerium/internal/testenv/upstreams"
+	"github.com/pomerium/pomerium/internal/testenv/values"
+)
+
+func newHTTPUpstream(
+	env testenv.Environment, subdomain string,
+) (upstreams.HTTPUpstream, testenv.Route) {
+	up := upstreams.HTTP(nil)
+	up.Handle("/", func(w http.ResponseWriter, _ *http.Request) { fmt.Fprintln(w, "hello world") })
+	route := up.Route().
+		From(env.SubdomainURL(subdomain)).
+		To(values.Bind(up.Addr(), func(addr string) string {
+			// override the target protocol to use http://
+			return fmt.Sprintf("http://%s", addr)
+		})).
+		Policy(func(p *config.Policy) { p.AllowAnyAuthenticatedUser = true })
+	env.AddUpstream(up)
+	return up, route
+}
+
+func TestMultiDomainLogin(t *testing.T) {
+	env := testenv.New(t)
+
+	env.Add(scenarios.NewIDP([]*scenarios.User{{Email: "test@example.com"}}))
+
+	// Create three routes to be linked via multi-domain login...
+	upstreamA, routeA := newHTTPUpstream(env, "a")
+	upstreamB, routeB := newHTTPUpstream(env, "b")
+	upstreamC, routeC := newHTTPUpstream(env, "c")
+	// ...and one route that will not be involved.
+	upstreamD, routeD := newHTTPUpstream(env, "d")
+
+	// Configure route A to redirect through routes B and C on login.
+	routeA.Policy(func(p *config.Policy) {
+		p.DependsOn = []string{
+			strings.TrimPrefix(routeB.URL().Value(), "https://"),
+			strings.TrimPrefix(routeC.URL().Value(), "https://"),
+		}
+	})
+
+	env.Start()
+	snippets.WaitStartupComplete(env)
+
+	// By default the testenv code will use a separate http.Client for each
+	// separate route. Instead we specifically want to test the cross-route
+	// behavior for a single client.
+	cj, err := cookiejar.New(nil)
+	require.NoError(t, err)
+	sharedClient := http.Client{Jar: cj}
+	withSharedClient := upstreams.ClientHook(
+		func(_ *http.Client) *http.Client { return &sharedClient })
+
+	assertResponseOK := func(resp *http.Response, err error) {
+		t.Helper()
+		require.NoError(t, err)
+		assert.Equal(t, http.StatusOK, resp.StatusCode)
+		io.ReadAll(resp.Body)
+		resp.Body.Close()
+	}
+	assertRedirect := func(resp *http.Response, err error) {
+		t.Helper()
+		require.NoError(t, err)
+		assert.Equal(t, http.StatusFound, resp.StatusCode)
+		io.ReadAll(resp.Body)
+		resp.Body.Close()
+	}
+
+	// Log in to the first route.
+	assertResponseOK(upstreamA.Get(routeA, withSharedClient, upstreams.AuthenticateAs("test@example.com")))
+
+	// The client should also be authenticated for routes B and C without any
+	// additional login redirects.
+	sharedClient.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
+		return http.ErrUseLastResponse
+	}
+	assertResponseOK(upstreamB.Get(routeB, withSharedClient))
+	assertResponseOK(upstreamC.Get(routeC, withSharedClient))
+
+	// The client should not be authenticated for route D.
+	assertRedirect(upstreamD.Get(routeD, withSharedClient))
+}

internal/authenticateflow/stateful.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"
 
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
@@ -163,7 +164,9 @@ func (s *Stateful) SignIn(
 	// base64 our encrypted payload for URL-friendlyness
 	encodedJWT := base64.URLEncoding.EncodeToString(encryptedJWT)
 
-	callbackURL, err := urlutil.GetCallbackURL(r, encodedJWT)
+	additionalHosts := strings.Split(r.FormValue(urlutil.QueryAdditionalHosts), ",")
+
+	callbackURL, err := urlutil.GetCallbackURL(r, encodedJWT, additionalHosts)
 	if err != nil {
 		return httputil.NewError(http.StatusBadRequest, err)
 	}
@@ -324,7 +327,11 @@ func (s *Stateful) LogAuthenticateEvent(*http.Request) {}
 // AuthenticateSignInURL returns a URL to redirect the user to the authenticate
 // domain.
 func (s *Stateful) AuthenticateSignInURL(
-	ctx context.Context, queryParams url.Values, redirectURL *url.URL, idpID string,
+	ctx context.Context,
+	queryParams url.Values,
+	redirectURL *url.URL,
+	idpID string,
+	additionalHosts []string,
 ) (string, error) {
 	signinURL := s.authenticateURL.ResolveReference(&url.URL{
 		Path: "/.pomerium/sign_in",
@@ -335,6 +342,9 @@ func (s *Stateful) AuthenticateSignInURL(
 	}
 	queryParams.Set(urlutil.QueryRedirectURI, redirectURL.String())
 	queryParams.Set(urlutil.QueryIdentityProviderID, idpID)
+	if len(additionalHosts) > 0 {
+		queryParams.Set(urlutil.QueryAdditionalHosts, strings.Join(additionalHosts, ","))
+	}
 	otel.GetTextMapPropagator().Inject(ctx, trace.PomeriumURLQueryCarrier(queryParams))
 	signinURL.RawQuery = queryParams.Encode()
 	redirectTo := urlutil.NewSignedURL(s.sharedKey, signinURL).String()
@@ -387,6 +397,23 @@ func (s *Stateful) Callback(w http.ResponseWriter, r *http.Request) error {
 		redirectURL.RawQuery = q.Encode()
 	}
 
+	// Redirect chaining for multi-domain login.
+	additionalHosts := r.URL.Query().Get(urlutil.QueryAdditionalHosts)
+	if additionalHosts != "" {
+		nextHops := strings.Split(additionalHosts, ",")
+		log.Ctx(r.Context()).Debug().Strs("next-hops", nextHops).Msg("multi-domain login callback")
+
+		callbackURL, err := urlutil.GetCallbackURL(r, encryptedSession, nextHops[1:])
+		if err != nil {
+			return httputil.NewError(http.StatusInternalServerError,
+				fmt.Errorf("proxy: couldn't get next hop callback URL: %w", err))
+		}
+		callbackURL.Host = nextHops[0]
+		signedCallbackURL := urlutil.NewSignedURL(s.sharedKey, callbackURL)
+		httputil.Redirect(w, r, signedCallbackURL.String(), http.StatusFound)
+		return nil
+	}
+
 	// redirect
 	httputil.Redirect(w, r, redirectURL.String(), http.StatusFound)
 	return nil

internal/authenticateflow/stateful_test.go

@@ -128,7 +128,7 @@ func TestStatefulAuthenticateSignInURL(t *testing.T) {
 
 	t.Run("NilQueryParams", func(t *testing.T) {
 		redirectURL := &url.URL{Scheme: "https", Host: "example.com"}
-		u, err := flow.AuthenticateSignInURL(context.Background(), nil, redirectURL, "fake-idp-id")
+		u, err := flow.AuthenticateSignInURL(context.Background(), nil, redirectURL, "fake-idp-id", nil)
 		assert.NoError(t, err)
 		parsed, _ := url.Parse(u)
 		assert.NoError(t, urlutil.NewSignedURL(key, parsed).Validate())
@@ -143,7 +143,7 @@ func TestStatefulAuthenticateSignInURL(t *testing.T) {
 		redirectURL := &url.URL{Scheme: "https", Host: "example.com"}
 		q := url.Values{}
 		q.Set("foo", "bar")
-		u, err := flow.AuthenticateSignInURL(context.Background(), q, redirectURL, "fake-idp-id")
+		u, err := flow.AuthenticateSignInURL(context.Background(), q, redirectURL, "fake-idp-id", nil)
 		assert.NoError(t, err)
 		parsed, _ := url.Parse(u)
 		assert.NoError(t, urlutil.NewSignedURL(key, parsed).Validate())
@@ -155,6 +155,21 @@ func TestStatefulAuthenticateSignInURL(t *testing.T) {
 		assert.Equal(t, "fake-idp-id", q.Get("pomerium_idp_id"))
 		assert.Equal(t, "bar", q.Get("foo"))
 	})
+	t.Run("AdditionalHosts", func(t *testing.T) {
+		redirectURL := &url.URL{Scheme: "https", Host: "example.com"}
+		additionalHosts := []string{"foo.example.com", "bar.example.com:1234"}
+		u, err := flow.AuthenticateSignInURL(context.Background(), nil, redirectURL, "fake-idp-id", additionalHosts)
+		assert.NoError(t, err)
+		parsed, _ := url.Parse(u)
+		assert.NoError(t, urlutil.NewSignedURL(key, parsed).Validate())
+		assert.Equal(t, "https", parsed.Scheme)
+		assert.Equal(t, "authenticate.example.com", parsed.Host)
+		assert.Equal(t, "/.pomerium/sign_in", parsed.Path)
+		q := parsed.Query()
+		assert.Equal(t, "https://example.com", parsed.Query().Get("pomerium_redirect_uri"))
+		assert.Equal(t, "fake-idp-id", q.Get("pomerium_idp_id"))
+		assert.Equal(t, "foo.example.com,bar.example.com:1234", q.Get("pomerium_additional_hosts"))
+	})
 }
 
 func TestStatefulGetIdentityProviderIDForURLValues(t *testing.T) {
@@ -277,6 +292,7 @@ func TestStatefulCallback(t *testing.T) {
 				}
 				location, _ := url.Parse(w.Result().Header.Get("Location"))
 				assert.Equal(t, "example.com", location.Host)
+				assert.Equal(t, "/", location.Path)
 				assert.Equal(t, "ok", location.Query().Get("pomerium_callback_uri"))
 			} else {
 				if err == nil || !strings.Contains(err.Error(), tt.wantErrorMsg) {
@@ -287,6 +303,60 @@ func TestStatefulCallback(t *testing.T) {
 	}
 }
 
+func TestStatefulCallback_AdditionalHosts(t *testing.T) {
+	opts := config.NewDefaultOptions()
+	opts.SharedKey = "80ldlrU2d7w+wVpKNfevk6fmb8otEx6CqOfshj2LwhQ="
+	sharedKey, _ := opts.GetSharedKey()
+
+	flow, err := NewStateful(
+		context.Background(),
+		trace.NewNoopTracerProvider(),
+		&config.Config{Options: opts},
+		&mstore.Store{Session: &sessions.State{}},
+	)
+	require.NoError(t, err)
+
+	redirectURI := "https://route.example.com/"
+	callbackURI := &url.URL{
+		Scheme: "https",
+		Host:   "route.example.com",
+		Path:   "/.pomerium/callback",
+		RawQuery: url.Values{
+			urlutil.QuerySessionEncrypted: []string{goodEncryptionString},
+			urlutil.QueryRedirectURI:      []string{redirectURI},
+			urlutil.QueryAdditionalHosts:  []string{"foo.example.com,bar.example.com"},
+		}.Encode(),
+	}
+	signedCallbackURI := urlutil.NewSignedURL(sharedKey, callbackURI)
+
+	doCallback := func(uri string) *http.Response {
+		t.Helper()
+		r := httptest.NewRequest(http.MethodGet, uri, nil)
+		r.Host = r.URL.Host
+
+		w := httptest.NewRecorder()
+		err = flow.Callback(w, r)
+		require.NoError(t, err)
+		return w.Result()
+	}
+
+	// Callback() should serve redirects to the additional hosts before the final redirect URI.
+	res := doCallback(signedCallbackURI.String())
+	location, _ := url.Parse(res.Header.Get("Location"))
+	assert.Equal(t, "foo.example.com", location.Host)
+	assert.Equal(t, "/.pomerium/callback/", location.Path)
+
+	res = doCallback(location.String())
+	location, _ = url.Parse(res.Header.Get("Location"))
+	assert.Equal(t, "bar.example.com", location.Host)
+	assert.Equal(t, "/.pomerium/callback/", location.Path)
+
+	res = doCallback(location.String())
+	location, _ = url.Parse(res.Header.Get("Location"))
+	assert.Equal(t, "route.example.com", location.Host)
+	assert.Equal(t, "/", location.Path)
+}
+
 func TestStatefulRevokeSession(t *testing.T) {
 	opts := config.NewDefaultOptions()
 	flow, err := NewStateful(context.Background(), trace.NewNoopTracerProvider(), &config.Config{Options: opts}, nil)

internal/authenticateflow/stateless.go

@@ -355,7 +355,11 @@ func getUserClaim(profile *identitypb.Profile, field string) *string {
 // AuthenticateSignInURL returns a URL to redirect the user to the authenticate
 // domain.
 func (s *Stateless) AuthenticateSignInURL(
-	ctx context.Context, queryParams url.Values, redirectURL *url.URL, idpID string,
+	ctx context.Context,
+	queryParams url.Values,
+	redirectURL *url.URL,
+	idpID string,
+	_ []string,
 ) (string, error) {
 	authenticateHPKEPublicKey, err := s.authenticateKeyFetcher.FetchPublicKey(ctx)
 	if err != nil {