test(totp): ensure context is cleaned up after use

ory · Oct 19, 2021 · 1905883 · 1905883
1 parent 8a210c4
commit 1905883
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/selfservice/strategy/totp/settings.go b/selfservice/strategy/totp/settings.go
@@ -30,7 +30,7 @@ import (
 	"github.com/ory/x/decoderx"
 )
 
-const internalContextKeyURL = "url"
+const InternalContextKeyURL = "url"
 
 func (s *Strategy) RegisterSettingsRoutes(_ *x.RouterPublic) {
 }
@@ -158,7 +158,7 @@ func (s *Strategy) continueSettingsFlow(
 }
 
 func (s *Strategy) continueSettingsFlowAddTOTP(w http.ResponseWriter, r *http.Request, ctxUpdate *settings.UpdateContext, p *submitSelfServiceSettingsFlowWithTotpMethodBody) (*identity.Identity, error) {
-	keyURL := gjson.GetBytes(ctxUpdate.Flow.InternalContext, flow.PrefixInternalContextKey(s.ID(), internalContextKeyURL)).String()
+	keyURL := gjson.GetBytes(ctxUpdate.Flow.InternalContext, flow.PrefixInternalContextKey(s.ID(), InternalContextKeyURL)).String()
 	if len(keyURL) == 0 {
 		return nil, errors.WithStack(herodot.ErrInternalServerError.WithReasonf("Could not find they TOTP key in the internal context. This is a code bug and should be reported to https://github.com/ory/kratos/."))
 	}
@@ -196,7 +196,7 @@ func (s *Strategy) continueSettingsFlowAddTOTP(w http.ResponseWriter, r *http.Re
 	i.SetCredentials(s.ID(), *c)
 
 	// Remove the TOTP URL from the internal context now that it is set!
-	ctxUpdate.Flow.InternalContext, err = sjson.DeleteBytes(ctxUpdate.Flow.InternalContext, flow.PrefixInternalContextKey(s.ID(), internalContextKeyURL))
+	ctxUpdate.Flow.InternalContext, err = sjson.DeleteBytes(ctxUpdate.Flow.InternalContext, flow.PrefixInternalContextKey(s.ID(), InternalContextKeyURL))
 	if err != nil {
 		return nil, err
 	}
@@ -254,7 +254,7 @@ func (s *Strategy) PopulateSettingsMethod(r *http.Request, id *identity.Identity
 			return err
 		}
 
-		f.InternalContext, err = sjson.SetBytes(f.InternalContext, flow.PrefixInternalContextKey(s.ID(), internalContextKeyURL), key.URL())
+		f.InternalContext, err = sjson.SetBytes(f.InternalContext, flow.PrefixInternalContextKey(s.ID(), InternalContextKeyURL), key.URL())
 		if err != nil {
 			return err
 		}

diff --git a/selfservice/strategy/totp/settings_test.go b/selfservice/strategy/totp/settings_test.go
@@ -3,6 +3,8 @@ package totp_test
 import (
 	"context"
 	"encoding/json"
+	"github.com/gofrs/uuid"
+	"github.com/ory/kratos/selfservice/flow"
 	"net/http"
 	"net/url"
 	"testing"
@@ -328,6 +330,10 @@ func TestCompleteSettings(t *testing.T) {
 				assert.EqualValues(t, settings.StateSuccess, gjson.Get(actual, "state").String(), actual)
 			}
 
+			actualFlow, err:= reg.SettingsFlowPersister().GetSettingsFlow(context.Background(),uuid.FromStringOrNil(f.Id))
+			require.NoError(t, err)
+			assert.Empty(t, gjson.GetBytes(actualFlow.InternalContext,flow.PrefixInternalContextKey(identity.CredentialsTypeTOTP, totp.InternalContextKeyURL)))
+
 			checkIdentity(t, id, key)
 		}