// Copyright © 2023 Ory Corp // SPDX-License-Identifier: Apache-2.0 package testhelpers import ( "bytes" "context" "net/http" "net/http/httptest" "net/url" "strings" "testing" "time" "github.com/ory/kratos/identity" "github.com/ory/x/urlx" "github.com/tidwall/gjson" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/ory/kratos/driver" "github.com/ory/kratos/driver/config" kratos "github.com/ory/kratos/pkg/httpclient" "github.com/ory/kratos/selfservice/flow/login" "github.com/ory/kratos/x" "github.com/ory/x/ioutilx" ) func NewLoginUIFlowEchoServer(t *testing.T, reg driver.Registry) *httptest.Server { ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { e, err := reg.LoginFlowPersister().GetLoginFlow(r.Context(), x.ParseUUID(r.URL.Query().Get("flow"))) require.NoError(t, err) reg.Writer().Write(w, r, e) })) ts.URL = strings.ReplaceAll(ts.URL, "127.0.0.1", "localhost") reg.Config().MustSet(t.Context(), config.ViperKeySelfServiceLoginUI, ts.URL+"/login-ts") t.Cleanup(ts.Close) return ts } func NewLoginUIWith401Response(t *testing.T, c *config.Config) *httptest.Server { ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusUnauthorized) })) ts.URL = strings.ReplaceAll(ts.URL, "127.0.0.1", "localhost") c.MustSet(t.Context(), config.ViperKeySelfServiceLoginUI, ts.URL+"/login-ts") t.Cleanup(ts.Close) return ts } type initFlowOptions struct { aal identity.AuthenticatorAssuranceLevel returnTo string refresh bool oauth2LoginChallenge string via string identitySchema string ctx context.Context expectActive *string } func newInitFlowOptions(opts []InitFlowWithOption) *initFlowOptions { return new(initFlowOptions).apply(opts) } func (o *initFlowOptions) apply(opts []InitFlowWithOption) *initFlowOptions { for _, opt := range opts { opt(o) } return o } func getURLFromInitOptions(ts *httptest.Server, path string, forced bool, opts ...InitFlowWithOption) string { o := new(initFlowOptions).apply(opts) q := url.Values{} if forced || o.refresh { q.Set("refresh", "true") } if o.aal != "" { q.Set("aal", string(o.aal)) } if o.returnTo != "" { q.Set("return_to", string(o.returnTo)) } if o.oauth2LoginChallenge != "" { q.Set("login_challenge", o.oauth2LoginChallenge) } if o.via != "" { q.Set("via", o.via) } if o.identitySchema != "" { q.Set("identity_schema", o.identitySchema) } u := urlx.ParseOrPanic(ts.URL + path) u.RawQuery = q.Encode() return u.String() } type InitFlowWithOption func(*initFlowOptions) func InitFlowWithAAL(aal identity.AuthenticatorAssuranceLevel) InitFlowWithOption { return func(o *initFlowOptions) { o.aal = aal } } func InitFlowWithReturnTo(returnTo string) InitFlowWithOption { return func(o *initFlowOptions) { o.returnTo = returnTo } } func InitFlowWithRefresh() InitFlowWithOption { return func(o *initFlowOptions) { o.refresh = true } } func InitFlowWithContext(ctx context.Context) InitFlowWithOption { return func(o *initFlowOptions) { o.ctx = ctx } } func InitFlowWithOAuth2LoginChallenge(hlc string) InitFlowWithOption { return func(o *initFlowOptions) { o.oauth2LoginChallenge = hlc } } func InitFlowWithIdentitySchema(schema string) InitFlowWithOption { return func(o *initFlowOptions) { o.identitySchema = schema } } func ExpectActive(active string) InitFlowWithOption { return func(o *initFlowOptions) { o.expectActive = &active } } // InitFlowWithVia sets the `via` query parameter which is used by the code MFA flows to determine the trait to use to send the code to the user func InitFlowWithVia(via string) InitFlowWithOption { return func(o *initFlowOptions) { o.via = via } } // Deprecated: use InitializeLoginFlowViaBrowserCtx instead. func InitializeLoginFlowViaBrowser(t *testing.T, client *http.Client, ts *httptest.Server, forced, isSPA, expectInitError, expectGetError bool, opts ...InitFlowWithOption) *kratos.LoginFlow { return InitializeLoginFlowViaBrowserCtx(t.Context(), t, client, ts, forced, isSPA, expectInitError, expectGetError, opts...) } func InitializeLoginFlowViaBrowserCtx(ctx context.Context, t *testing.T, client *http.Client, ts *httptest.Server, forced, isSPA, expectInitError, expectGetError bool, opts ...InitFlowWithOption) *kratos.LoginFlow { publicClient := NewSDKCustomClient(ts, client) req, err := http.NewRequest("GET", getURLFromInitOptions(ts, login.RouteInitBrowserFlow, forced, opts...), nil) require.NoError(t, err) o := newInitFlowOptions(opts) if isSPA { req.Header.Set("Accept", "application/json") } res, err := client.Do(req.WithContext(ctx)) require.NoError(t, err) body := x.MustReadAll(res.Body) if isSPA { require.Truef(t, gjson.ValidBytes(body), "body is not valid JSON: %s", body) } require.NoError(t, res.Body.Close()) require.Equal(t, 200, res.StatusCode, "%s", body) if expectInitError { require.NotNil(t, res.Request.URL) require.Contains(t, res.Request.URL.String(), "error-ts") } flowID := res.Request.URL.Query().Get("flow") if isSPA { flowID = gjson.GetBytes(body, "id").String() } if !expectGetError { require.NotEmpty(t, flowID) } rs, r, err := publicClient.FrontendAPI.GetLoginFlow(ctx).Id(flowID).Execute() if expectGetError { require.Error(t, err) require.Nil(t, rs) } else { require.NoError(t, err) require.NoErrorf(t, err, "%s", ioutilx.MustReadAll(r.Body)) assert.Equal(t, o.expectActive, rs.Active) } return rs } func InitializeLoginFlowViaAPICtx(ctx context.Context, t *testing.T, client *http.Client, ts *httptest.Server, forced bool, opts ...InitFlowWithOption) *kratos.LoginFlow { return initializeLoginFlowViaAPI(ctx, t, client, ts, forced, false, opts...) } func initializeLoginFlowViaAPI(ctx context.Context, t *testing.T, client *http.Client, ts *httptest.Server, forced bool, expectError bool, opts ...InitFlowWithOption) *kratos.LoginFlow { publicClient := NewSDKCustomClient(ts, client) o := new(initFlowOptions).apply(opts) req := publicClient.FrontendAPI.CreateNativeLoginFlow(ctx).Refresh(forced) if o.aal != "" { req = req.Aal(string(o.aal)) } if o.via != "" { req = req.Via(o.via) } if o.identitySchema != "" { req = req.IdentitySchema(o.identitySchema) } rs, res, err := req.Execute() if expectError { require.Error(t, err) require.Nil(t, rs) } else { require.NoError(t, err, "%s", ioutilx.MustReadAll(res.Body)) assert.Equal(t, o.expectActive, rs.Active) } return rs } // Deprecated: use InitializeLoginFlowViaAPICtx instead. func InitializeLoginFlowViaAPI(t *testing.T, client *http.Client, ts *httptest.Server, forced bool, opts ...InitFlowWithOption) *kratos.LoginFlow { return InitializeLoginFlowViaAPICtx(t.Context(), t, client, ts, forced, opts...) } func InitializeLoginFlowViaAPIExpectError(ctx context.Context, t *testing.T, client *http.Client, ts *httptest.Server, forced bool, opts ...InitFlowWithOption) *kratos.LoginFlow { return initializeLoginFlowViaAPI(ctx, t, client, ts, forced, true, opts...) } // Deprecated: use LoginMakeRequestCtx instead. func LoginMakeRequest( t *testing.T, isAPI bool, isSPA bool, f *kratos.LoginFlow, hc *http.Client, values string, opt ...RequestOption, ) (string, *http.Response) { return LoginMakeRequestCtx(t.Context(), t, isAPI, isSPA, f, hc, values, opt...) } type RequestOption func(*http.Request) func WithHeader(key, value string) RequestOption { return func(r *http.Request) { r.Header.Add(key, value) } } func LoginMakeRequestCtx( ctx context.Context, t *testing.T, isAPI bool, isSPA bool, f *kratos.LoginFlow, hc *http.Client, values string, opt ...RequestOption, ) (string, *http.Response) { require.NotEmpty(t, f.Ui.Action) req := NewPostRequest(t, isAPI, f.Ui.Action, bytes.NewBufferString(values)) if isSPA && !isAPI { req.Header.Set("Accept", "application/json") req.Header.Set("Sec-Fetch-Mode", "cors") } for _, o := range opt { o(req) } res, err := hc.Do(req.WithContext(ctx)) require.NoError(t, err, "action: %s", f.Ui.Action) defer func() { _ = res.Body.Close() }() return string(ioutilx.MustReadAll(res.Body)), res } func GetLoginFlow(t *testing.T, client *http.Client, ts *httptest.Server, flowID string) *kratos.LoginFlow { publicClient := NewSDKCustomClient(ts, client) rs, _, err := publicClient.FrontendAPI.GetLoginFlow(context.Background()).Id(flowID).Execute() require.NoError(t, err) return rs } // Deprecated: use SubmitLoginFormCtx instead. func SubmitLoginForm( t *testing.T, isAPI bool, hc *http.Client, publicTS *httptest.Server, withValues func(v url.Values), isSPA bool, forced bool, expectedStatusCode int, expectedURL string, opts ...InitFlowWithOption, ) string { return SubmitLoginFormCtx(context.Background(), t, isAPI, hc, publicTS, withValues, isSPA, forced, expectedStatusCode, expectedURL, opts...) } // SubmitLoginFormCtx initiates a login flow (for Browser and API!), fills out the form and modifies // the form values with `withValues`, and submits the form. Returns the body and checks for expectedStatusCode and // expectedURL on completion func SubmitLoginFormCtx( ctx context.Context, t *testing.T, isAPI bool, hc *http.Client, publicTS *httptest.Server, withValues func(v url.Values), isSPA bool, forced bool, expectedStatusCode int, expectedURL string, opts ...InitFlowWithOption, ) string { if hc == nil { hc = NewTestClient(t) if !isAPI { hc = NewClientWithCookies(t) } } hc.Transport = NewTransportWithLogger(hc.Transport, t) var f *kratos.LoginFlow if isAPI { f = InitializeLoginFlowViaAPICtx(ctx, t, hc, publicTS, forced, opts...) } else { f = InitializeLoginFlowViaBrowserCtx(ctx, t, hc, publicTS, forced, isSPA, false, false, opts...) } time.Sleep(time.Millisecond) // add a bit of delay to allow `1ns` to time out. payload := SDKFormFieldsToURLValues(f.Ui.Nodes) withValues(payload) b, res := LoginMakeRequestCtx(ctx, t, isAPI, isSPA, f, hc, EncodeFormAsJSON(t, isAPI, payload)) assert.EqualValuesf(t, expectedStatusCode, res.StatusCode, "%s", b) assert.Containsf(t, res.Request.URL.String(), expectedURL, "%+v\n\t%s", res.Request, b) t.Logf("%+v", res.Header) return b }