package colabauth
import (
"crypto/sha256"
"encoding/base64"
"encoding/json"
"io"
"net/http"
"net/http/cookiejar"
"net/http/httptest"
"net/url"
"os"
"path/filepath"
"regexp"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/gin-contrib/sessions"
"github.com/gin-contrib/sessions/cookie"
"github.com/gin-gonic/gin"
"github.com/sigbit/mcp-auth-proxy/v2/pkg/auth"
"github.com/stretchr/testify/require"
)
const owner = "did:plc:owner"
type env struct {
t *testing.T
h *Handler
srv *httptest.Server
path string
exchange atomic.Int32
// provider behaviour
refresh string
profile string
tokenCode int
challenge string
verifierOK atomic.Bool
assign func(http.ResponseWriter)
probes atomic.Int32
probeHdr http.Header
probeQuery string
}
func newEnv(t *testing.T, account string) *env {
t.Helper()
e := &env{t: t, refresh: "rt", tokenCode: 200,
profile: `{"sub":"1","email":"Me@Example.com","email_verified":true}`}
e.assign = func(w http.ResponseWriter) { w.Write([]byte(`{"assignments":[]}`)) }
provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/token":
e.exchange.Add(1)
require.NoError(t, r.ParseForm())
sum := sha256.Sum256([]byte(r.PostForm.Get("code_verifier")))
e.verifierOK.Store(base64.RawURLEncoding.EncodeToString(sum[:]) == e.challenge)
if e.tokenCode != 200 {
w.WriteHeader(e.tokenCode)
return
}
w.Header().Set("Content-Type", "application/json")
body := map[string]any{"access_token": "at", "token_type": "Bearer", "expires_in": 3600}
if e.refresh != "" {
body["refresh_token"] = e.refresh
}
json.NewEncoder(w).Encode(body)
case "/userinfo":
require.Equal(t, "Bearer at", r.Header.Get("Authorization"))
w.Write([]byte(e.profile))
case "/assignments":
e.probes.Add(1)
e.probeHdr, e.probeQuery = r.Header.Clone(), r.URL.RawQuery
e.assign(w)
}
}))
t.Cleanup(provider.Close)
e.path = filepath.Join(t.TempDir(), "creds", "colab.json")
h, err := New("https://gw.example/", "cid", "csecret", owner, e.path, account, defaultScopes)
require.NoError(t, err)
h.oauth.Endpoint.TokenURL = provider.URL + "/token"
h.userinfoURL = provider.URL + "/userinfo"
h.probeURL = provider.URL + "/assignments?authuser=0"
e.h = h
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(sessions.Sessions("session", cookie.NewStore([]byte("0123456789abcdef0123456789abcdef"))))
ar, err := auth.NewAuthRouter(nil, false, nil)
require.NoError(t, err)
r.GET("/test-login", func(c *gin.Context) {
s := sessions.Default(c)
s.Set(auth.SessionKeyAuthorized, true)
s.Set(auth.SessionKeyUserID, c.Query("u"))
s.Save()
})
h.SetupRoutes(r, ar)
e.srv = httptest.NewServer(r)
t.Cleanup(e.srv.Close)
return e
}
func (e *env) browser(user string) *http.Client {
jar, _ := cookiejar.New(nil)
c := &http.Client{Jar: jar, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
if user != "" {
res, err := c.Get(e.srv.URL + "/test-login?u=" + url.QueryEscape(user))
require.NoError(e.t, err)
res.Body.Close()
}
return c
}
func body(t *testing.T, res *http.Response) string {
defer res.Body.Close()
raw, err := io.ReadAll(res.Body)
require.NoError(t, err)
return string(raw)
}
var csrfRe = regexp.MustCompile(`name="csrf" value="([^"]+)"`)
func (e *env) csrf(c *http.Client) string {
res, err := c.Get(e.srv.URL + PageEndpoint)
require.NoError(e.t, err)
m := csrfRe.FindStringSubmatch(body(e.t, res))
require.NotNil(e.t, m)
return m[1]
}
func (e *env) post(c *http.Client, token, origin string) *http.Response {
req, _ := http.NewRequest(http.MethodPost, e.srv.URL+LoginEndpoint, strings.NewReader(url.Values{"csrf": {token}}.Encode()))
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
if origin != "" {
req.Header.Set("Origin", origin)
}
res, err := c.Do(req)
require.NoError(e.t, err)
return res
}
// start runs the login form and returns the google redirect.
func (e *env) start(c *http.Client) *url.URL {
res := e.post(c, e.csrf(c), "https://gw.example")
res.Body.Close()
require.Equal(e.t, http.StatusFound, res.StatusCode)
loc, err := url.Parse(res.Header.Get("Location"))
require.NoError(e.t, err)
e.challenge = loc.Query().Get("code_challenge")
return loc
}
func (e *env) callback(c *http.Client, query string) *http.Response {
res, err := c.Get(e.srv.URL + CallbackEndpoint + "?" + query)
require.NoError(e.t, err)
res.Body.Close()
return res
}
func (e *env) old() {
require.NoError(e.t, os.MkdirAll(filepath.Dir(e.path), 0o700))
require.NoError(e.t, os.WriteFile(e.path, []byte("old"), 0o600))
}
func (e *env) kept() {
got, err := os.ReadFile(e.path)
require.NoError(e.t, err)
require.Equal(e.t, "old", string(got))
}
func TestFlowWritesCompatibleCredentials(t *testing.T) {
e := newEnv(t, "me@example.COM")
c := e.browser(owner)
loc := e.start(c)
q := loc.Query()
require.Equal(t, "accounts.google.com", loc.Host)
require.Equal(t, "S256", q.Get("code_challenge_method"))
require.Equal(t, "offline", q.Get("access_type"))
require.Equal(t, "consent select_account", q.Get("prompt"))
require.Equal(t, "https://gw.example"+CallbackEndpoint, q.Get("redirect_uri"))
require.Equal(t, "openid https://www.googleapis.com/auth/userinfo.email", q.Get("scope"))
res := e.callback(c, "code=abc&state="+q.Get("state")+"&scope=x")
require.Equal(t, http.StatusFound, res.StatusCode)
require.Equal(t, PageEndpoint, res.Header.Get("Location"))
require.True(t, e.verifierOK.Load())
info, err := os.Stat(e.h.accountPath())
require.NoError(t, err)
require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
dir, _ := os.Stat(filepath.Dir(e.path))
require.Equal(t, os.FileMode(0o700), dir.Mode().Perm())
store, err := e.h.readAccounts()
require.NoError(t, err)
require.Len(t, store.Accounts, 1)
raw, err := json.Marshal(store.Accounts[0])
require.NoError(t, err)
var got map[string]any
require.NoError(t, json.Unmarshal(raw, &got))
require.Regexp(t, accountIDPattern, got["account_id"])
require.Equal(t, "1", got["google_subject"])
require.Equal(t, "authorized_user", got["type"])
require.Equal(t, "cid", got["client_id"])
require.Equal(t, "csecret", got["client_secret"])
require.Equal(t, "rt", got["refresh_token"])
require.Equal(t, "at", got["token"])
require.Equal(t, "Me@Example.com", got["email"])
exp, err := time.Parse(time.RFC3339, got["expiry"].(string))
require.NoError(t, err)
require.WithinDuration(t, time.Now().Add(time.Hour), exp, time.Minute)
// replay is rejected without another exchange
res = e.callback(c, "code=abc&state="+q.Get("state"))
require.Equal(t, http.StatusBadRequest, res.StatusCode)
require.EqualValues(t, 1, e.exchange.Load())
page, _ := c.Get(e.srv.URL + PageEndpoint)
require.Contains(t, body(t, page), "connected as Me@Example.com")
}
func TestOwnerGate(t *testing.T) {
e := newEnv(t, "")
anon := e.browser("")
res, _ := anon.Get(e.srv.URL + PageEndpoint)
res.Body.Close()
require.Equal(t, http.StatusFound, res.StatusCode)
require.Equal(t, auth.LoginEndpoint, res.Header.Get("Location"))
other := e.browser("did:plc:other")
for _, path := range []string{PageEndpoint, CallbackEndpoint + "?state=x&code=y"} {
res, _ := other.Get(e.srv.URL + path)
res.Body.Close()
require.Equal(t, http.StatusForbidden, res.StatusCode, path)
}
res = e.post(other, "x", "")
res.Body.Close()
require.Equal(t, http.StatusForbidden, res.StatusCode)
require.Zero(t, e.exchange.Load())
require.NoFileExists(t, e.h.accountPath())
}
func TestCSRFAndOrigin(t *testing.T) {
e := newEnv(t, "")
c := e.browser(owner)
token := e.csrf(c)
for name, res := range map[string]*http.Response{
"missing token": e.post(c, "", ""),
"wrong token": e.post(c, "nope", ""),
"bad origin": e.post(c, token, "https://evil.example"),
"null origin": e.post(c, token, "null"),
} {
res.Body.Close()
require.Equal(t, http.StatusForbidden, res.StatusCode, name)
}
res, _ := c.Get(e.srv.URL + LoginEndpoint)
res.Body.Close()
require.Equal(t, http.StatusNotFound, res.StatusCode)
res = e.post(c, token, "")
res.Body.Close()
require.Equal(t, http.StatusFound, res.StatusCode)
}
func TestStateBinding(t *testing.T) {
e := newEnv(t, "")
c := e.browser(owner)
state := e.start(c).Query().Get("state")
// another owner session cannot use the state, and does not burn it
thief := e.browser(owner)
require.Equal(t, http.StatusBadRequest, e.callback(thief, "code=c&state="+state).StatusCode)
require.Equal(t, http.StatusBadRequest, e.callback(c, "code=c&state=bogus").StatusCode)
require.Equal(t, http.StatusBadRequest, e.callback(c, "code=c").StatusCode)
require.Zero(t, e.exchange.Load())
require.Equal(t, http.StatusFound, e.callback(c, "code=c&state="+state).StatusCode)
}
func TestExpiryAndCap(t *testing.T) {
e := newEnv(t, "")
now := time.Now()
e.h.now = func() time.Time { return now }
c := e.browser(owner)
state := e.start(c).Query().Get("state")
now = now.Add(pendingTTL)
require.Equal(t, http.StatusBadRequest, e.callback(c, "code=c&state="+state).StatusCode)
require.Zero(t, e.exchange.Load())
for i := 0; i < maxPending; i++ {
e.start(c)
}
res := e.post(c, e.csrf(c), "")
res.Body.Close()
require.Equal(t, http.StatusTooManyRequests, res.StatusCode)
now = now.Add(pendingTTL)
e.start(c)
require.Len(t, e.h.pending, 1)
}
func TestFailuresKeepOldCredentials(t *testing.T) {
cases := map[string]struct {
account string
mutate func(*env)
query string
}{
"denied": {query: "error=access_denied&state=%s"},
"no code": {query: "state=%s"},
"token error": {mutate: func(e *env) { e.tokenCode = 400 }, query: "code=c&state=%s"},
"no refresh": {mutate: func(e *env) { e.refresh = "" }, query: "code=c&state=%s"},
"unverified": {mutate: func(e *env) { e.profile = `{"sub":"1","email":"a@b.c","email_verified":false}` }, query: "code=c&state=%s"},
"no sub": {mutate: func(e *env) { e.profile = `{"email":"a@b.c","email_verified":true}` }, query: "code=c&state=%s"},
"oversized": {mutate: func(e *env) {
e.profile = `{"sub":"1","email":"a@b.c","email_verified":true,"x":"` + strings.Repeat("a", maxProfileBytes) + `"}`
}, query: "code=c&state=%s"},
"wrong account": {account: "other@example.com", query: "code=c&state=%s"},
"colab denied": {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) { w.WriteHeader(403) } }, query: "code=c&state=%s"},
"colab error": {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) { w.WriteHeader(500) } }, query: "code=c&state=%s"},
"colab redirect": {mutate: func(e *env) {
e.assign = func(w http.ResponseWriter) {
w.Header().Set("Location", "/elsewhere")
w.WriteHeader(302)
}
}, query: "code=c&state=%s"},
"colab empty": {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) {} }, query: "code=c&state=%s"},
"colab malformed": {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) { w.Write([]byte(`{"assign`)) } }, query: "code=c&state=%s"},
"colab null": {mutate: func(e *env) { e.assign = func(w http.ResponseWriter) { w.Write([]byte(`null`)) } }, query: "code=c&state=%s"},
"colab wrong shape": {mutate: func(e *env) {
e.assign = func(w http.ResponseWriter) { w.Write([]byte(`{"assignments":{"endpoint":"x"}}`)) }
}, query: "code=c&state=%s"},
"colab oversized": {mutate: func(e *env) {
e.assign = func(w http.ResponseWriter) {
w.Write([]byte(`{"assignments":[],"x":"` + strings.Repeat("a", maxProbeBytes) + `"}`))
}
}, query: "code=c&state=%s"},
}
for name, tc := range cases {
t.Run(name, func(t *testing.T) {
e := newEnv(t, tc.account)
if tc.mutate != nil {
tc.mutate(e)
}
e.old()
c := e.browser(owner)
state := e.start(c).Query().Get("state")
res := e.callback(c, strings.Replace(tc.query, "%s", state, 1))
require.GreaterOrEqual(t, res.StatusCode, 400)
e.kept()
})
}
}
func TestFromEnv(t *testing.T) {
set := func(id, secret, subject, path string) {
t.Setenv("COLAB_GOOGLE_CLIENT_ID", id)
t.Setenv("COLAB_GOOGLE_CLIENT_SECRET", secret)
t.Setenv("MCP_ALLOWED_SUBJECT", subject)
t.Setenv("COLAB_MANAGED_CREDENTIALS", path)
}
set("", "", "", "")
h, err := FromEnv("https://gw.example/")
require.NoError(t, err)
require.Nil(t, h)
for _, args := range [][4]string{
{"id", "", owner, "/tmp/x"},
{"", "secret", owner, "/tmp/x"},
{"id", "secret", "", "/tmp/x"},
{"id", "secret", owner, ""},
{"id", "secret", owner, "relative.json"},
} {
set(args[0], args[1], args[2], args[3])
_, err := FromEnv("https://gw.example/")
require.Error(t, err, args)
}
set("id", "secret", owner, "/tmp/x.json")
t.Setenv("COLAB_GOOGLE_SCOPES", "openid email")
t.Setenv("COLAB_GOOGLE_ACCOUNT", " Me@Example.com ")
h, err = FromEnv("https://gw.example/")
require.NoError(t, err)
require.Equal(t, []string{"openid", "email"}, h.oauth.Scopes)
require.Equal(t, "Me@Example.com", h.account)
require.Equal(t, "https://gw.example/.auth/colab/callback", h.oauth.RedirectURL)
}
func TestNewRejectsBadOrigins(t *testing.T) {
for _, u := range []string{"http://gw.example", "https://u@gw.example", "https://gw.example/x", "https://gw.example/?a=1", "https://gw.example/#f", "https://gw.example/?", "gw.example", ""} {
_, err := New(u, "id", "secret", owner, "/tmp/x", "", defaultScopes)
require.Error(t, err, u)
}
for _, u := range []string{"https://gw.example", "https://gw.example/", "http://127.0.0.1:8080", "http://[::1]:8080/"} {
_, err := New(u, "id", "secret", owner, "/tmp/x", "", defaultScopes)
require.NoError(t, err, u)
}
}
func TestOversizeFormLeavesNoPending(t *testing.T) {
e := newEnv(t, "")
c := e.browser(owner)
token := e.csrf(c)
res := e.post(c, token+strings.Repeat("a", maxFormBytes), "")
res.Body.Close()
require.Equal(t, http.StatusBadRequest, res.StatusCode)
require.Empty(t, e.h.pending)
}
func TestPrivateHeadersEverywhere(t *testing.T) {
e := newEnv(t, "")
check := func(res *http.Response) {
res.Body.Close()
require.Equal(t, "no-store", res.Header.Get("Cache-Control"))
require.Equal(t, "strict-origin", res.Header.Get("Referrer-Policy"))
require.Equal(t, "nosniff", res.Header.Get("X-Content-Type-Options"))
}
anon := e.browser("")
res, _ := anon.Get(e.srv.URL + PageEndpoint)
check(res)
c := e.browser(owner)
res, _ = c.Get(e.srv.URL + PageEndpoint)
check(res)
res, _ = c.Get(e.srv.URL + CallbackEndpoint)
check(res)
check(e.post(c, "nope", ""))
state := e.start(c).Query().Get("state")
res, _ = c.Get(e.srv.URL + CallbackEndpoint + "?code=c&state=" + state)
check(res)
}
func TestRejectsNonPrivateExistingParent(t *testing.T) {
e := newEnv(t, "")
e.old()
require.NoError(t, os.Chmod(filepath.Dir(e.path), 0o755))
c := e.browser(owner)
state := e.start(c).Query().Get("state")
res := e.callback(c, "code=c&state="+state)
require.Equal(t, http.StatusInternalServerError, res.StatusCode)
e.kept()
}
func TestColabProbeRequestAndShapes(t *testing.T) {
for name, reply := range map[string]string{
"object": `{"assignments":[{"endpoint":"e"}]}`,
"array": `[{"endpoint":"e"}]`,
"xssi": ")]}'\n{\"assignments\":[]}",
"xssi suffix": ")]}'junk\n[]",
} {
t.Run(name, func(t *testing.T) {
e := newEnv(t, "")
e.assign = func(w http.ResponseWriter) { w.Write([]byte(reply)) }
c := e.browser(owner)
state := e.start(c).Query().Get("state")
require.Equal(t, http.StatusFound, e.callback(c, "code=c&state="+state).StatusCode)
require.EqualValues(t, 1, e.probes.Load())
require.Equal(t, "authuser=0", e.probeQuery)
require.Equal(t, "Bearer at", e.probeHdr.Get("Authorization"))
require.Equal(t, "application/json", e.probeHdr.Get("Accept"))
require.Equal(t, "colab-cli", e.probeHdr.Get("X-Colab-Client-Agent"))
require.Equal(t, "Google", e.probeHdr.Get("X-Colab-Tunnel"))
require.FileExists(t, e.h.accountPath())
})
}
}
func TestNoColabProbeBeforeProfileChecks(t *testing.T) {
for name, mutate := range map[string]func(*env){
"wrong account": func(e *env) { e.h.account = "other@example.com" },
"unverified": func(e *env) { e.profile = `{"sub":"1","email":"a@b.c","email_verified":false}` },
"no refresh": func(e *env) { e.refresh = "" },
} {
t.Run(name, func(t *testing.T) {
e := newEnv(t, "")
mutate(e)
c := e.browser(owner)
state := e.start(c).Query().Get("state")
require.GreaterOrEqual(t, e.callback(c, "code=c&state="+state).StatusCode, 400)
require.Zero(t, e.probes.Load())
require.NoFileExists(t, e.h.accountPath())
})
}
}
func TestProbeFailureHidesProviderDetail(t *testing.T) {
e := newEnv(t, "")
e.assign = func(w http.ResponseWriter) {
w.WriteHeader(403)
w.Write([]byte("secret-provider-detail"))
}
c := e.browser(owner)
state := e.start(c).Query().Get("state")
res, err := c.Get(e.srv.URL + CallbackEndpoint + "?code=c&state=" + state)
require.NoError(t, err)
require.Equal(t, http.StatusBadGateway, res.StatusCode)
require.NotContains(t, body(t, res), "secret-provider-detail")
require.NoFileExists(t, e.h.accountPath())
}
func TestProbeNetworkFailureKeepsOldCredential(t *testing.T) {
e := newEnv(t, "")
unavailable := httptest.NewServer(http.NotFoundHandler())
unavailable.Close()
e.h.probeURL = unavailable.URL
e.old()
c := e.browser(owner)
state := e.start(c).Query().Get("state")
require.Equal(t, http.StatusBadGateway, e.callback(c, "code=c&state="+state).StatusCode)
e.kept()
}
func TestProbeDoesNotFollowRedirects(t *testing.T) {
var followed atomic.Int32
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
followed.Add(1)
w.Write([]byte(`{"assignments":[]}`))
}))
defer target.Close()
e := newEnv(t, "")
e.assign = func(w http.ResponseWriter) {
w.Header().Set("Location", target.URL)
w.WriteHeader(http.StatusFound)
}
e.old()
c := e.browser(owner)
state := e.start(c).Query().Get("state")
require.Equal(t, http.StatusBadGateway, e.callback(c, "code=c&state="+state).StatusCode)
require.Zero(t, followed.Load())
e.kept()
}