package idp

import (
	"bytes"
	"context"
	"io"
	"net/http"
	"net/url"
	"strings"
	"testing"

	"github.com/ory/fosite"
	"github.com/stretchr/testify/require"
)

func TestRedirectPrefixes(t *testing.T) {
	const prefix = "https://oauth-redirect.googleusercontent.com/r/,http://127.0.0.1/"
	for _, tc := range []struct {
		uri string
		ok  bool
	}{
		{"https://oauth-redirect.googleusercontent.com/r/client", true},
		{"http://127.0.0.1/callback", true},
		{"http://127.0.0.1:1234/callback", false},
		{"https://oauth-redirect.googleusercontent.com.evil/r/client", false},
		{"https://oauth-redirect.googleusercontent.com/rubbish", false},
		{"https://oauth-redirect.googleusercontent.com/r/../evil", false},
		{"https://oauth-redirect.googleusercontent.com/r/%2e%2e/evil", false},
		{"https://oauth-redirect.googleusercontent.com/r/client?x=1", false},
		{"https://oauth-redirect.googleusercontent.com/r/client#x", false},
		{"https://evil@oauth-redirect.googleusercontent.com/r/client", false},
		{"http://localhost/callback", false},
		{"http://example.com/callback", false},
	} {
		t.Run(tc.uri, func(t *testing.T) { require.Equal(t, tc.ok, validRedirectURIs([]string{tc.uri}, prefix)) })
	}
	for _, config := range []string{"garbage", "https://good/r/,", "http://example.com/", "https://good/r/?x=1"} {
		require.False(t, validRedirectURIs([]string{"https://good/r/client"}, config))
	}
	require.False(t, validRedirectURIs(make([]string, 9), ""))
	require.False(t, validRedirectURIs([]string{strings.Repeat("x", 2049)}, ""))
}

func TestDCRRejectsRestrictedRedirects(t *testing.T) {
	for _, config := range []string{"", "broken", "https://good/r/"} {
		t.Run(config, func(t *testing.T) {
			t.Setenv("MCP_REDIRECT_PREFIXES", config)
			server, _, _ := setupTestServer(t)
			resp, err := http.Post(server.URL+RegistrationEndpoint, "application/json", bytes.NewBufferString(`{"redirect_uris":["https://evil/callback"],"token_endpoint_auth_method":"none"}`))
			require.NoError(t, err)
			defer resp.Body.Close()
			require.Equal(t, http.StatusBadRequest, resp.StatusCode)
		})
	}
}

func TestConsentEscapesClientAndRedirect(t *testing.T) {
	server, repo, _ := setupTestServer(t)
	const id = "<script>client</script>"
	const redirect = "https://client/callback?value=<script>"
	require.NoError(t, repo.RegisterClient(context.Background(), &fosite.DefaultClient{ID: id, RedirectURIs: []string{redirect}, GrantTypes: []string{"authorization_code"}, ResponseTypes: []string{"code"}, Public: true}))
	client := &http.Client{Jar: mustCookieJar(t), CheckRedirect: func(r *http.Request, via []*http.Request) error { return http.ErrUseLastResponse }}
	q := url.Values{"client_id": {id}, "redirect_uri": {redirect}, "response_type": {"code"}, "state": {"test-state"}}
	resp, err := client.Get(server.URL + AuthorizationEndpoint + "?" + q.Encode())
	require.NoError(t, err)
	resp.Body.Close()
	form, err := client.Get(server.URL + resp.Header.Get("Location"))
	require.NoError(t, err)
	defer form.Body.Close()
	require.Equal(t, 200, form.StatusCode)
	body, err := io.ReadAll(form.Body)
	require.NoError(t, err)
	require.Contains(t, string(body), "&lt;script&gt;client&lt;/script&gt;")
	require.Contains(t, string(body), "https://client/callback?value=&lt;script&gt;")
	require.NotContains(t, string(body), "<script>")
}