package colabauth
import (
"crypto/rand"
"encoding/base64"
"encoding/json"
"errors"
"io"
"os"
"path/filepath"
"regexp"
"strings"
)
const maxAccounts = 16
const maxAccountBytes = 1 << 20
var errPoolFull = errors.New("account pool is full; remove an account first")
var accountIDPattern = regexp.MustCompile(`^acct_[A-Za-z0-9_-]{43}$`)
type accountStore struct {
Version int `json:"version"`
Accounts []credentials `json:"accounts"`
}
func (h *Handler) accountPath() string { return h.managedPath + ".accounts.json" }
func readPrivateJSON(path string, value any) error {
f, err := os.Open(path)
if err != nil {
return err
}
defer f.Close()
info, err := f.Stat()
if err != nil {
return err
}
if !info.Mode().IsRegular() || info.Mode().Perm()&0o077 != 0 {
return errors.New("account file is not private")
}
data, err := io.ReadAll(io.LimitReader(f, maxAccountBytes+1))
if err != nil {
return err
}
if len(data) > maxAccountBytes {
return errors.New("account file too large")
}
return json.Unmarshal(data, value)
}
func (h *Handler) readAccounts() (accountStore, error) {
var store accountStore
err := readPrivateJSON(h.accountPath(), &store)
if errors.Is(err, os.ErrNotExist) {
return accountStore{Version: 1, Accounts: []credentials{}}, nil
}
if err != nil {
return store, err
}
if store.Version != 1 || store.Accounts == nil || len(store.Accounts) > maxAccounts {
return store, errors.New("invalid account store")
}
seen := map[string]bool{}
for _, c := range store.Accounts {
if !accountIDPattern.MatchString(c.AccountID) || seen[c.AccountID] || !validCredentials(c) {
return store, errors.New("invalid account record")
}
seen[c.AccountID] = true
}
return store, nil
}
func validCredentials(c credentials) bool {
return c.RefreshToken != "" && c.ClientID != "" && c.ClientSecret != ""
}
func newAccountID() (string, error) {
var data [32]byte
if _, err := rand.Read(data[:]); err != nil {
return "", err
}
return "acct_" + base64.RawURLEncoding.EncodeToString(data[:]), nil
}
// import once, without changing the credential an older gateway would use.
func (h *Handler) initializeAccounts() error {
h.storeMu.Lock()
defer h.storeMu.Unlock()
return h.initializeAccountsLocked()
}
func (h *Handler) initializeAccountsLocked() error {
if _, err := os.Stat(h.accountPath()); err == nil {
_, err = h.readAccounts()
return err
} else if !errors.Is(err, os.ErrNotExist) {
return err
}
var old credentials
if err := readPrivateJSON(h.managedPath, &old); errors.Is(err, os.ErrNotExist) {
return nil
} else if err != nil {
return err
}
if !validCredentials(old) {
return errors.New("invalid legacy credential")
}
id, err := newAccountID()
if err != nil {
return err
}
old.AccountID = id
return writePrivateJSON(h.accountPath(), accountStore{Version: 1, Accounts: []credentials{old}})
}
func (h *Handler) save(creds credentials) error {
if !validCredentials(creds) || creds.GoogleSubject == "" || creds.Email == "" {
return errors.New("invalid account credential")
}
h.storeMu.Lock()
defer h.storeMu.Unlock()
if err := h.initializeAccountsLocked(); err != nil {
return err
}
store, err := h.readAccounts()
if err != nil {
return err
}
for i, old := range store.Accounts {
// older credentials have no Google subject. Upgrade that one verified email match.
if old.GoogleSubject == creds.GoogleSubject || (old.GoogleSubject == "" && strings.EqualFold(old.Email, creds.Email)) {
creds.AccountID = old.AccountID
store.Accounts[i] = creds
return writePrivateJSON(h.accountPath(), store)
}
}
if len(store.Accounts) >= maxAccounts {
return errPoolFull
}
creds.AccountID, err = newAccountID()
if err != nil {
return err
}
store.Accounts = append(store.Accounts, creds)
return writePrivateJSON(h.accountPath(), store)
}
func (h *Handler) removeAccount(id string) (bool, error) {
h.storeMu.Lock()
defer h.storeMu.Unlock()
store, err := h.readAccounts()
if err != nil {
return false, err
}
for i, account := range store.Accounts {
if account.AccountID == id {
store.Accounts = append(store.Accounts[:i], store.Accounts[i+1:]...)
return true, writePrivateJSON(h.accountPath(), store)
}
}
return false, nil
}
func writePrivateJSON(path string, value any) error {
data, err := json.MarshalIndent(value, "", " ")
if err != nil {
return err
}
if len(data)+1 > maxAccountBytes {
return errors.New("account file too large")
}
dir := filepath.Dir(path)
if err := ensurePrivateDir(dir); err != nil {
return err
}
tmp, err := os.CreateTemp(dir, ".colab-credentials-*")
if err != nil {
return err
}
defer os.Remove(tmp.Name())
defer tmp.Close()
if err := tmp.Chmod(0o600); err != nil {
return err
}
if _, err := tmp.Write(append(data, '\n')); err != nil {
return err
}
if err := tmp.Sync(); err != nil {
return err
}
if err := tmp.Close(); err != nil {
return err
}
return os.Rename(tmp.Name(), path)
}