From fb9eb31be9e5f1e27d47dec5a602a16b0d248e6e Mon Sep 17 00:00:00 2001 From: Kenneth Jenkins <51246568+kenjenkins@users.noreply.github.com> Date: Wed, 3 Jan 2024 09:48:11 -0800 Subject: [PATCH 1/2] authenticateflow: change how sessions are deleted (#4893) The identity manager expects to be able to read session ID and user ID from any deleted databroker session records. The session.Delete() wrapper method is not compatible with this expectation, as it calls Put() with a record containing an empty session. The stateful authentication flow currently calls session.Delete() from its RevokeSession() method. The result is that the identity manager will not correctly track sessions deleted by the the stateful authentication flow, and will still try to use them during session refresh and user info refresh. Instead, let's change the stateful authentication flow RevokeSession() method to perform deletions in a way that is compatible with the current identity manager code. That is, include the existing session data in the Put() call to delete the revoked session. --- internal/authenticateflow/stateful.go | 36 ++++++- internal/authenticateflow/stateful_test.go | 117 +++++++++++++++++++++ 2 files changed, 150 insertions(+), 3 deletions(-) diff --git a/internal/authenticateflow/stateful.go b/internal/authenticateflow/stateful.go index 51e0fa512..64030e4f7 100644 --- a/internal/authenticateflow/stateful.go +++ b/internal/authenticateflow/stateful.go @@ -27,6 +27,7 @@ import ( "github.com/pomerium/pomerium/pkg/grpc/databroker" "github.com/pomerium/pomerium/pkg/grpc/session" "github.com/pomerium/pomerium/pkg/grpc/user" + "github.com/pomerium/pomerium/pkg/grpcutil" ) // Stateful implements the stateful authentication flow. In this flow, the @@ -261,15 +262,44 @@ func (s *Stateful) RevokeSession( return "" } + // Note: session.Delete() cannot be used safely, because the identity + // manager expects to be able to read both session ID and user ID from + // deleted session records. Instead, we match the behavior used in the + // identity manager itself: fetch the existing databroker session record, + // explicitly set the DeletedAt timestamp, and Put() that record back. + + res, err := s.dataBrokerClient.Get(ctx, &databroker.GetRequest{ + Type: grpcutil.GetTypeURL(new(session.Session)), + Id: sessionState.ID, + }) + if err != nil { + err = fmt.Errorf("couldn't get session to be revoked: %w", err) + log.Ctx(ctx).Warn().Err(err).Msg("authenticate: failed to revoke access token") + return "" + } + + record := res.GetRecord() + + var sess session.Session + if err := record.GetData().UnmarshalTo(&sess); err != nil { + err = fmt.Errorf("couldn't unmarshal data of session to be revoked: %w", err) + log.Ctx(ctx).Warn().Err(err).Msg("authenticate: failed to revoke access token") + return "" + } + var rawIDToken string - sess, _ := session.Get(ctx, s.dataBrokerClient, sessionState.ID) - if sess != nil && sess.OauthToken != nil { + if sess.OauthToken != nil { rawIDToken = sess.GetIdToken().GetRaw() if err := authenticator.Revoke(ctx, manager.FromOAuthToken(sess.OauthToken)); err != nil { log.Ctx(ctx).Warn().Err(err).Msg("authenticate: failed to revoke access token") } } - if err := session.Delete(ctx, s.dataBrokerClient, sessionState.ID); err != nil { + + record.DeletedAt = timestamppb.Now() + _, err = s.dataBrokerClient.Put(ctx, &databroker.PutRequest{ + Records: []*databroker.Record{record}, + }) + if err != nil { log.Ctx(ctx).Warn().Err(err). Msg("authenticate: failed to delete session from session store") } diff --git a/internal/authenticateflow/stateful_test.go b/internal/authenticateflow/stateful_test.go index 51f1a62a3..a219677be 100644 --- a/internal/authenticateflow/stateful_test.go +++ b/internal/authenticateflow/stateful_test.go @@ -1,24 +1,37 @@ package authenticateflow import ( + "context" "encoding/base64" "errors" + "fmt" "net/http" "net/http/httptest" "net/url" "strings" "testing" + "time" + "github.com/golang/mock/gomock" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/oauth2" + "google.golang.org/grpc" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/timestamppb" "github.com/pomerium/pomerium/config" "github.com/pomerium/pomerium/internal/encoding" "github.com/pomerium/pomerium/internal/encoding/mock" + "github.com/pomerium/pomerium/internal/identity" "github.com/pomerium/pomerium/internal/sessions" mstore "github.com/pomerium/pomerium/internal/sessions/mock" "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/databroker/mock_databroker" + "github.com/pomerium/pomerium/pkg/grpc/session" + "github.com/pomerium/pomerium/pkg/protoutil" ) func TestStatefulSignIn(t *testing.T) { @@ -269,3 +282,107 @@ func TestStatefulCallback(t *testing.T) { }) } } + +func TestStatefulRevokeSession(t *testing.T) { + opts := config.NewDefaultOptions() + flow, err := NewStateful(&config.Config{Options: opts}, nil) + require.NoError(t, err) + + ctrl := gomock.NewController(t) + client := mock_databroker.NewMockDataBrokerServiceClient(ctrl) + flow.dataBrokerClient = client + + // Exercise the happy path (no errors): calling RevokeSession() should + // fetch and delete a session record from the databroker and make a request + // to the identity provider to revoke the corresponding OAuth2 token. + + ctx := context.Background() + authenticator := &mockAuthenticator{} + sessionState := &sessions.State{ID: "session-id"} + tokenExpiry := time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC) + + client.EXPECT().Get(ctx, protoEqualMatcher{ + &databroker.GetRequest{ + Type: "type.googleapis.com/session.Session", + Id: "session-id", + }, + }).Return(&databroker.GetResponse{ + Record: &databroker.Record{ + Version: 123456, + Type: "type.googleapis.com/session.Session", + Id: "session-id", + Data: protoutil.NewAny(&session.Session{ + Id: "session-id", + UserId: "user-id", + IdToken: &session.IDToken{ + Raw: "[raw-id-token]", + }, + OauthToken: &session.OAuthToken{ + AccessToken: "[oauth-access-token]", + TokenType: "Bearer", + RefreshToken: "[oauth-refresh-token]", + ExpiresAt: timestamppb.New(tokenExpiry), + }, + }), + }, + }, nil) + + client.EXPECT().Put(ctx, gomock.Any()).DoAndReturn( + func(_ context.Context, r *databroker.PutRequest, _ ...grpc.CallOption) (*databroker.PutResponse, error) { + require.Len(t, r.Records, 1) + record := r.GetRecord() + assert.Equal(t, "type.googleapis.com/session.Session", record.Type) + assert.Equal(t, "session-id", record.Id) + assert.Equal(t, uint64(123456), record.Version) + + // The session record received in this PutRequest should have a + // DeletedAt timestamp, as well as the same session ID and user ID + // as was returned in the previous GetResponse. + assert.NotNil(t, record.DeletedAt) + var s session.Session + record.GetData().UnmarshalTo(&s) + assert.Equal(t, "session-id", s.Id) + assert.Equal(t, "user-id", s.UserId) + return nil, nil + }) + + idToken := flow.RevokeSession(ctx, nil, authenticator, sessionState) + + assert.Equal(t, "[raw-id-token]", idToken) + assert.Equal(t, &oauth2.Token{ + AccessToken: "[oauth-access-token]", + TokenType: "Bearer", + RefreshToken: "[oauth-refresh-token]", + Expiry: tokenExpiry, + }, authenticator.revokedToken) +} + +// protoEqualMatcher implements gomock.Matcher using proto.Equal. +// TODO: move this to a testutil package? +type protoEqualMatcher struct { + expected proto.Message +} + +func (m protoEqualMatcher) Matches(x interface{}) bool { + p, ok := x.(proto.Message) + if !ok { + return false + } + return proto.Equal(m.expected, p) +} + +func (m protoEqualMatcher) String() string { + return fmt.Sprintf("is equal to %v (%T)", m.expected, m.expected) +} + +type mockAuthenticator struct { + identity.Authenticator + + revokedToken *oauth2.Token + revokeError error +} + +func (a *mockAuthenticator) Revoke(_ context.Context, token *oauth2.Token) error { + a.revokedToken = token + return a.revokeError +} From 915c230d8d4e34f0fd5c1e284c8803c99d69b140 Mon Sep 17 00:00:00 2001 From: Kenneth Jenkins <51246568+kenjenkins@users.noreply.github.com> Date: Wed, 3 Jan 2024 17:18:18 -0800 Subject: [PATCH 2/2] authenticateflow: verify redirect in Callback test (#4894) Add assertions for the redirect URL in the unit test for the stateful authentication flow Callback() method. Remove a commented-out line of debug logging. --- internal/authenticateflow/stateful_test.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/internal/authenticateflow/stateful_test.go b/internal/authenticateflow/stateful_test.go index a219677be..8d4aae270 100644 --- a/internal/authenticateflow/stateful_test.go +++ b/internal/authenticateflow/stateful_test.go @@ -261,7 +261,6 @@ func TestStatefulCallback(t *testing.T) { } r := httptest.NewRequest(http.MethodGet, uri.String(), nil) - // fmt.Println(uri.String()) r.Host = r.URL.Host r.Header.Set("Accept", "application/json") @@ -272,13 +271,14 @@ func TestStatefulCallback(t *testing.T) { if err != nil { t.Errorf("unexpected error: %v", err) } + location, _ := url.Parse(w.Result().Header.Get("Location")) + assert.Equal(t, "example.com", location.Host) + assert.Equal(t, "ok", location.Query().Get("pomerium_callback_uri")) } else { if err == nil || !strings.Contains(err.Error(), tt.wantErrorMsg) { t.Errorf("expected error containing %q; got %v", tt.wantErrorMsg, err) } } - - // XXX: assert redirect URL }) } }