authenticate: add databroker versions to session cookie (#2709)

* authenticate: add databroker versions to session cookie
authorize: wait for databroker synchronization on updated sessions

* fix test

authenticate/handlers.go

@@ -186,19 +186,13 @@ func (a *Authenticate) SignIn(w http.ResponseWriter, r *http.Request) error {
 
 	jwtAudience := []string{state.redirectURL.Host, redirectURL.Host}
 
-	var callbackURL *url.URL
 	// if the callback is explicitly set, set it and add an additional audience
 	if callbackStr := r.FormValue(urlutil.QueryCallbackURI); callbackStr != "" {
-		callbackURL, err = urlutil.ParseAndValidateURL(callbackStr)
+		callbackURL, err := urlutil.ParseAndValidateURL(callbackStr)
 		if err != nil {
 			return httputil.NewError(http.StatusBadRequest, err)
 		}
 		jwtAudience = append(jwtAudience, callbackURL.Host)
-	} else {
-		// otherwise, assume callback is the same host as redirect
-		callbackURL, _ = urlutil.DeepCopy(redirectURL)
-		callbackURL.Path = "/.pomerium/callback/"
-		callbackURL.RawQuery = ""
 	}
 
 	// add an additional claim for the forward-auth host, if set
@@ -219,11 +213,8 @@ func (a *Authenticate) SignIn(w http.ResponseWriter, r *http.Request) error {
 		return httputil.NewError(http.StatusBadRequest, err)
 	}
 
-	callbackParams := callbackURL.Query()
-
 	if r.FormValue(urlutil.QueryIsProgrammatic) == "true" {
 		newSession.Programmatic = true
-		callbackParams.Set(urlutil.QueryIsProgrammatic, "true")
 	}
 
 	// sign the route session, as a JWT
@@ -237,10 +228,10 @@ func (a *Authenticate) SignIn(w http.ResponseWriter, r *http.Request) error {
 	// base64 our encrypted payload for URL-friendlyness
 	encodedJWT := base64.URLEncoding.EncodeToString(encryptedJWT)
 
-	// add our encoded and encrypted route-session JWT to a query param
-	callbackParams.Set(urlutil.QuerySessionEncrypted, encodedJWT)
-	callbackParams.Set(urlutil.QueryRedirectURI, redirectURL.String())
-	callbackURL.RawQuery = callbackParams.Encode()
+	callbackURL, err := urlutil.GetCallbackURL(r, encodedJWT)
+	if err != nil {
+		return httputil.NewError(http.StatusBadRequest, err)
+	}
 
 	// build our hmac-d redirect URL with our session, pointing back to the
 	// proxy's callback URL which is responsible for setting our new route-session
@@ -665,10 +656,17 @@ func (a *Authenticate) getWebauthnState(ctx context.Context) (*webauthn.State, e
 		return nil, err
 	}
 
+	ss, err := a.getSessionFromCtx(ctx)
+	if err != nil {
+		return nil, err
+	}
+
 	return &webauthn.State{
 		SharedKey:    state.sharedKey,
 		Client:       state.dataBrokerClient,
 		Session:      s,
+		SessionState: ss,
+		SessionStore: state.sessionStore,
 		RelyingParty: state.webauthnRelyingParty,
 	}, nil
 }

authenticate/handlers/webauthn/webauthn.go

@@ -4,6 +4,7 @@ package webauthn
 import (
 	"bytes"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -19,10 +20,13 @@ import (
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
+	"github.com/pomerium/pomerium/internal/encoding/jws"
 	"github.com/pomerium/pomerium/internal/frontend"
 	"github.com/pomerium/pomerium/internal/httputil"
 	"github.com/pomerium/pomerium/internal/middleware"
+	"github.com/pomerium/pomerium/internal/sessions"
 	"github.com/pomerium/pomerium/internal/urlutil"
+	"github.com/pomerium/pomerium/pkg/cryptutil"
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
 	"github.com/pomerium/pomerium/pkg/grpc/device"
 	"github.com/pomerium/pomerium/pkg/grpc/session"
@@ -42,6 +46,8 @@ type State struct {
 	SharedKey    []byte
 	Client       databroker.DataBrokerServiceClient
 	Session      *session.Session
+	SessionState *sessions.State
+	SessionStore sessions.SessionStore
 	RelyingParty *webauthn.RelyingParty
 }
 
@@ -177,21 +183,14 @@ func (h *Handler) handleAuthenticate(w http.ResponseWriter, r *http.Request, sta
 		}
 	}
 
-	// save the session
+	// update the session
 	state.Session.DeviceCredentials = append(state.Session.DeviceCredentials, &session.Session_DeviceCredential{
 		TypeId: deviceType.GetId(),
 		Credential: &session.Session_DeviceCredential_Id{
 			Id: webauthnutil.GetDeviceCredentialID(serverCredential.ID),
 		},
 	})
-	_, err = session.Put(ctx, state.Client, state.Session)
-	if err != nil {
-		return err
-	}
-
-	// redirect
-	httputil.Redirect(w, r, redirectURIParam, http.StatusFound)
-	return nil
+	return h.saveSessionAndRedirect(w, r, state, redirectURIParam)
 }
 
 func (h *Handler) handleRegister(w http.ResponseWriter, r *http.Request, state *State) error {
@@ -286,21 +285,14 @@ func (h *Handler) handleRegister(w http.ResponseWriter, r *http.Request, state *
 		return err
 	}
 
-	// save the session
+	// update the session
 	state.Session.DeviceCredentials = append(state.Session.DeviceCredentials, &session.Session_DeviceCredential{
 		TypeId: deviceType.GetId(),
 		Credential: &session.Session_DeviceCredential_Id{
 			Id: webauthnutil.GetDeviceCredentialID(serverCredential.ID),
 		},
 	})
-	_, err = session.Put(ctx, state.Client, state.Session)
-	if err != nil {
-		return err
-	}
-
-	// redirect
-	httputil.Redirect(w, r, redirectURIParam, http.StatusFound)
-	return nil
+	return h.saveSessionAndRedirect(w, r, state, redirectURIParam)
 }
 
 func (h *Handler) handleView(w http.ResponseWriter, r *http.Request, state *State) error {
@@ -351,6 +343,52 @@ func (h *Handler) handleView(w http.ResponseWriter, r *http.Request, state *Stat
 	return err
 }
 
+func (h *Handler) saveSessionAndRedirect(w http.ResponseWriter, r *http.Request, state *State, rawRedirectURI string) error {
+	// save the session to the databroker
+	res, err := session.Put(r.Context(), state.Client, state.Session)
+	if err != nil {
+		return err
+	}
+
+	// add databroker versions to the session cookie and save
+	state.SessionState.Version = sessions.Version(fmt.Sprint(res.GetServerVersion()))
+	state.SessionState.DatabrokerServerVersion = res.GetServerVersion()
+	state.SessionState.DatabrokerRecordVersion = res.GetRecord().GetVersion()
+	err = state.SessionStore.SaveSession(w, r, state.SessionState)
+	if err != nil {
+		return err
+	}
+
+	// sign+encrypt the session JWT
+	encoder, err := jws.NewHS256Signer(state.SharedKey)
+	if err != nil {
+		return err
+	}
+
+	signedJWT, err := encoder.Marshal(state.SessionState)
+	if err != nil {
+		return err
+	}
+
+	cipher, err := cryptutil.NewAEADCipher(state.SharedKey)
+	if err != nil {
+		return err
+	}
+
+	encryptedJWT := cryptutil.Encrypt(cipher, signedJWT, nil)
+	encodedJWT := base64.URLEncoding.EncodeToString(encryptedJWT)
+
+	// redirect to the proxy callback URL with the session
+	callbackURL, err := urlutil.GetCallbackURL(r, encodedJWT)
+	if err != nil {
+		return err
+	}
+
+	signedCallbackURL := urlutil.NewSignedURL(state.SharedKey, callbackURL)
+	httputil.Redirect(w, r, signedCallbackURL.String(), http.StatusFound)
+	return nil
+}
+
 func getKnownDeviceCredentials(
 	ctx context.Context,
 	client databroker.DataBrokerServiceClient,