package colabauth
import (
"context"
"crypto/subtle"
"encoding/json"
"errors"
"fmt"
"html/template"
"io"
"net/http"
"net/url"
"os"
"path/filepath"
"regexp"
"strings"
"sync"
"time"
"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
"github.com/sigbit/mcp-auth-proxy/v2/pkg/auth"
"github.com/sigbit/mcp-auth-proxy/v2/pkg/utils"
"golang.org/x/oauth2"
)
const (
PageEndpoint = "/.auth/colab"
LoginEndpoint = "/.auth/colab/login"
CallbackEndpoint = "/.auth/colab/callback"
authURL = "https://accounts.google.com/o/oauth2/v2/auth"
tokenURL = "https://oauth2.googleapis.com/token"
userinfoURL = "https://openidconnect.googleapis.com/v1/userinfo"
probeURL = "https://colab.research.google.com/tun/m/assignments?authuser=0"
pendingTTL = 10 * time.Minute
maxPending = 32
maxProfileBytes = 8 << 10
maxProbeBytes = 1 << 20
maxFormBytes = 4 << 10
requestTimeout = 15 * time.Second
sessionCSRF = "colab_csrf"
sessionNonce = "colab_nonce"
)
var xssiPrefix = regexp.MustCompile(`^\)\]\}'[^\n]*\n`)
var defaultScopes = []string{"openid", "https://www.googleapis.com/auth/userinfo.email"}
type pending struct {
nonce string
verifier string
expires time.Time
}
type Handler struct {
oauth oauth2.Config
subject string
account string
managedPath string
origin string
userinfoURL string
probeURL string
client *http.Client
now func() time.Time
storeMu sync.Mutex
mu sync.Mutex
pending map[string]pending
}
func validOrigin(u *url.URL) bool {
if u.Host == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || (u.Path != "" && u.Path != "/") {
return false
}
return u.Scheme == "https" || (u.Scheme == "http" && (u.Hostname() == "127.0.0.1" || u.Hostname() == "::1"))
}
// FromEnv returns nil when linking is not configured.
func FromEnv(externalURL string) (*Handler, error) {
id, secret := os.Getenv("COLAB_GOOGLE_CLIENT_ID"), os.Getenv("COLAB_GOOGLE_CLIENT_SECRET")
if id == "" && secret == "" {
return nil, nil
}
if id == "" || secret == "" {
return nil, errors.New("COLAB_GOOGLE_CLIENT_ID and COLAB_GOOGLE_CLIENT_SECRET must be set together")
}
subject, managed := os.Getenv("MCP_ALLOWED_SUBJECT"), os.Getenv("COLAB_MANAGED_CREDENTIALS")
if subject == "" {
return nil, errors.New("MCP_ALLOWED_SUBJECT is required for colab account linking")
}
if !filepath.IsAbs(managed) {
return nil, errors.New("COLAB_MANAGED_CREDENTIALS must be an absolute path for colab account linking")
}
scopes := strings.Fields(os.Getenv("COLAB_GOOGLE_SCOPES"))
if len(scopes) == 0 {
scopes = defaultScopes
}
h, err := New(externalURL, id, secret, subject, managed, strings.TrimSpace(os.Getenv("COLAB_GOOGLE_ACCOUNT")), scopes)
if err != nil {
return nil, err
}
if err := h.initializeAccounts(); err != nil {
return nil, errors.New("could not initialize colab account store")
}
return h, nil
}
func New(externalURL, clientID, clientSecret, subject, managedPath, account string, scopes []string) (*Handler, error) {
u, err := url.Parse(externalURL)
if err != nil || !validOrigin(u) {
return nil, errors.New("invalid external URL for colab account linking")
}
origin := u.Scheme + "://" + u.Host
return &Handler{
oauth: oauth2.Config{
ClientID: clientID,
ClientSecret: clientSecret,
Endpoint: oauth2.Endpoint{AuthURL: authURL, TokenURL: tokenURL, AuthStyle: oauth2.AuthStyleInParams},
RedirectURL: origin + CallbackEndpoint,
Scopes: scopes,
},
subject: subject,
account: account,
managedPath: managedPath,
origin: origin,
userinfoURL: userinfoURL,
probeURL: probeURL,
client: &http.Client{
Timeout: requestTimeout,
CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
},
now: time.Now,
pending: map[string]pending{},
}, nil
}
func (h *Handler) SetupRoutes(router gin.IRouter, authRouter *auth.AuthRouter) {
g := router.Group(PageEndpoint, privateHeaders, authRouter.RequireAuth(), h.requireOwner)
g.GET("", h.handlePage)
g.POST("/login", h.handleStart)
g.POST("/remove", h.handleRemove)
g.GET("/callback", h.handleCallback)
}
func privateHeaders(c *gin.Context) {
c.Header("Cache-Control", "no-store")
// no-referrer turns the Origin of native form posts into null.
c.Header("Referrer-Policy", "strict-origin")
c.Header("X-Content-Type-Options", "nosniff")
}
func (h *Handler) requireOwner(c *gin.Context) {
user, _ := sessions.Default(c).Get(auth.SessionKeyUserID).(string)
if user == "" || subtle.ConstantTimeCompare([]byte(user), []byte(h.subject)) != 1 {
h.fail(c, http.StatusForbidden, "forbidden")
return
}
c.Next()
}
var pageTemplate = template.Must(template.New("page").Parse(`<!doctype html>
<meta charset="utf-8"><title>colab accounts</title>
<h1>colab accounts</h1>
{{if .Accounts}}<ul>{{range .Accounts}}<li>connected as {{.Email}} <code>{{.AccountID}}</code>
<form method="post" action="/.auth/colab/remove">
<input type="hidden" name="csrf" value="{{$.CSRF}}">
<input type="hidden" name="account_id" value="{{.AccountID}}">
<button>remove from pool</button></form></li>{{end}}</ul>
<p>removing an account does not stop its runtimes or revoke google access.</p>
{{else}}<p>not connected.</p>{{end}}
<form method="post" action="/.auth/colab/login">
<input type="hidden" name="csrf" value="{{.CSRF}}">
<button>connect google account</button>
</form>
<p>connect another account to add it; reconnecting the same account keeps its id. up to 16 accounts.</p>
`))
func (h *Handler) handlePage(c *gin.Context) {
session := sessions.Default(c)
token, err := utils.GenerateState()
if err != nil {
h.fail(c, http.StatusInternalServerError, "something went wrong")
return
}
session.Set(sessionCSRF, token)
if err := session.Save(); err != nil {
h.fail(c, http.StatusInternalServerError, "something went wrong")
return
}
store, err := h.readAccounts()
if err != nil {
h.fail(c, http.StatusInternalServerError, "could not read accounts")
return
}
c.Header("Content-Type", "text/html; charset=utf-8")
pageTemplate.Execute(c.Writer, map[string]any{"Accounts": store.Accounts, "CSRF": token})
}
func (h *Handler) validForm(c *gin.Context) bool {
session := sessions.Default(c)
want, _ := session.Get(sessionCSRF).(string)
if origin := c.GetHeader("Origin"); origin != "" && origin != h.origin {
h.fail(c, http.StatusForbidden, "forbidden")
return false
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxFormBytes)
if err := c.Request.ParseForm(); err != nil {
h.fail(c, http.StatusBadRequest, "invalid form")
return false
}
got := c.Request.PostForm.Get("csrf")
if want == "" || subtle.ConstantTimeCompare([]byte(got), []byte(want)) != 1 {
h.fail(c, http.StatusForbidden, "forbidden")
return false
}
return true
}
func (h *Handler) handleRemove(c *gin.Context) {
if !h.validForm(c) {
return
}
removed, err := h.removeAccount(c.Request.PostForm.Get("account_id"))
if err != nil {
h.fail(c, http.StatusInternalServerError, "could not remove account")
return
}
if !removed {
h.fail(c, http.StatusNotFound, "account not found")
return
}
c.Redirect(http.StatusSeeOther, PageEndpoint)
}
func (h *Handler) handleStart(c *gin.Context) {
if !h.validForm(c) {
return
}
session := sessions.Default(c)
state, err := utils.GenerateState()
nonce, err2 := utils.GenerateState()
if err != nil || err2 != nil {
h.fail(c, http.StatusInternalServerError, "something went wrong")
return
}
verifier := oauth2.GenerateVerifier()
now := h.now()
h.mu.Lock()
h.prune(now)
if len(h.pending) >= maxPending {
h.mu.Unlock()
h.fail(c, http.StatusTooManyRequests, "too many pending links, try again later")
return
}
h.pending[state] = pending{nonce: nonce, verifier: verifier, expires: now.Add(pendingTTL)}
h.mu.Unlock()
session.Set(sessionNonce, nonce)
if err := session.Save(); err != nil {
h.take(state, nonce)
h.fail(c, http.StatusInternalServerError, "something went wrong")
return
}
c.Redirect(http.StatusFound, h.oauth.AuthCodeURL(state,
oauth2.AccessTypeOffline,
oauth2.SetAuthURLParam("prompt", "consent select_account"),
oauth2.S256ChallengeOption(verifier)))
}
func (h *Handler) prune(now time.Time) {
for k, p := range h.pending {
if !now.Before(p.expires) {
delete(h.pending, k)
}
}
}
// take consumes the pending entry only when it is live and bound to nonce.
func (h *Handler) take(state, nonce string) (pending, bool) {
h.mu.Lock()
defer h.mu.Unlock()
h.prune(h.now())
p, ok := h.pending[state]
if !ok || nonce == "" || subtle.ConstantTimeCompare([]byte(p.nonce), []byte(nonce)) != 1 {
return pending{}, false
}
delete(h.pending, state)
return p, true
}
func (h *Handler) handleCallback(c *gin.Context) {
session := sessions.Default(c)
nonce, _ := session.Get(sessionNonce).(string)
state := c.Query("state")
if state == "" {
h.fail(c, http.StatusBadRequest, "invalid link request")
return
}
p, ok := h.take(state, nonce)
if !ok {
h.fail(c, http.StatusBadRequest, "invalid link request")
return
}
session.Delete(sessionNonce)
if err := session.Save(); err != nil {
h.fail(c, http.StatusInternalServerError, "something went wrong")
return
}
code := c.Query("code")
if c.Query("error") != "" || code == "" {
h.fail(c, http.StatusBadRequest, "google did not grant access")
return
}
status, err := h.link(c.Request.Context(), code, p.verifier)
if err != nil {
h.fail(c, status, err.Error())
return
}
c.Redirect(http.StatusFound, PageEndpoint)
}
type credentials struct {
AccountID string `json:"account_id,omitempty"`
GoogleSubject string `json:"google_subject,omitempty"`
Type string `json:"type"`
ClientID string `json:"client_id"`
ClientSecret string `json:"client_secret"`
RefreshToken string `json:"refresh_token"`
Token string `json:"token"`
Expiry string `json:"expiry"`
Email string `json:"email,omitempty"`
Scopes []string `json:"scopes,omitempty"`
}
// link exchanges the code and persists credentials only after every check passes.
func (h *Handler) link(ctx context.Context, code, verifier string) (int, error) {
ctx, cancel := context.WithTimeout(ctx, requestTimeout)
defer cancel()
ctx = context.WithValue(ctx, oauth2.HTTPClient, h.client)
token, err := h.oauth.Exchange(ctx, code, oauth2.VerifierOption(verifier))
if err != nil {
return http.StatusBadGateway, errors.New("google token exchange failed")
}
if token.RefreshToken == "" {
return http.StatusBadGateway, errors.New("google returned no refresh token")
}
email, googleSubject, err := h.profileIdentity(ctx, token.AccessToken)
if err != nil {
return http.StatusBadGateway, errors.New("could not verify google account")
}
if h.account != "" && !strings.EqualFold(email, h.account) {
return http.StatusForbidden, errors.New("google account is not the configured account")
}
if err := h.probeColab(ctx, token.AccessToken); err != nil {
return http.StatusBadGateway, errors.New("could not verify colab access for this google client or account")
}
creds := credentials{
Type: "authorized_user",
ClientID: h.oauth.ClientID,
ClientSecret: h.oauth.ClientSecret,
RefreshToken: token.RefreshToken,
Token: token.AccessToken,
Expiry: token.Expiry.UTC().Format(time.RFC3339),
Email: email,
GoogleSubject: googleSubject,
Scopes: h.oauth.Scopes,
}
if err := h.save(creds); err != nil {
if errors.Is(err, errPoolFull) {
return http.StatusConflict, errPoolFull
}
return http.StatusInternalServerError, errors.New("could not save credentials")
}
return http.StatusOK, nil
}
func (h *Handler) profileIdentity(ctx context.Context, accessToken string) (string, string, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, h.userinfoURL, nil)
if err != nil {
return "", "", err
}
req.Header.Set("Authorization", "Bearer "+accessToken)
res, err := h.client.Do(req)
if err != nil {
return "", "", err
}
defer res.Body.Close()
body, err := io.ReadAll(io.LimitReader(res.Body, maxProfileBytes+1))
if err != nil || res.StatusCode != http.StatusOK || len(body) > maxProfileBytes {
return "", "", errors.New("bad profile response")
}
var profile struct {
Sub string `json:"sub"`
Email string `json:"email"`
EmailVerified bool `json:"email_verified"`
}
if json.Unmarshal(body, &profile) != nil || profile.Sub == "" || profile.Email == "" || !profile.EmailVerified {
return "", "", errors.New("unverified profile")
}
return profile.Email, profile.Sub, nil
}
func (h *Handler) probeColab(ctx context.Context, accessToken string) error {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, h.probeURL, nil)
if err != nil {
return err
}
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Accept", "application/json")
req.Header.Set("X-Colab-Client-Agent", "colab-cli")
req.Header.Set("X-Colab-Tunnel", "Google")
res, err := h.client.Do(req)
if err != nil {
return err
}
defer res.Body.Close()
body, err := io.ReadAll(io.LimitReader(res.Body, maxProbeBytes+1))
if err != nil || res.StatusCode != http.StatusOK || len(body) > maxProbeBytes {
return errors.New("bad colab response")
}
body = xssiPrefix.ReplaceAll(body, nil)
var data any
if err := json.Unmarshal(body, &data); err != nil {
return err
}
if obj, ok := data.(map[string]any); ok {
data = obj["assignments"]
}
if _, ok := data.([]any); !ok {
return errors.New("invalid assignments list")
}
return nil
}
func ensurePrivateDir(dir string) error {
info, err := os.Stat(dir)
if errors.Is(err, os.ErrNotExist) {
if err := os.MkdirAll(dir, 0o700); err != nil {
return err
}
return os.Chmod(dir, 0o700)
}
if err != nil {
return err
}
if !info.IsDir() || info.Mode().Perm()&0o077 != 0 {
return errors.New("credential directory is not private")
}
return nil
}
func (h *Handler) fail(c *gin.Context, status int, message string) {
page := fmt.Sprintf("<!doctype html><title>colab account</title><p>%s</p>", template.HTMLEscapeString(message))
c.Data(status, "text/html; charset=utf-8", []byte(page))
c.Abort()
}