llmgw/internal/store/memberships.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"time"
)

// ErrNotMember means a person asked to use a project they are not in.
var ErrNotMember = errors.New("not a member of that project")

// Roles a membership can carry.
const (
	RoleMember = "member"
	// RoleManager is reserved for delegated administration of a scope.
	RoleManager = "manager"
)

// Membership is a person's place at one level of the tree.
type Membership struct {
	UserID  string
	Scope   Scope // org, team or project
	ScopeID string
	Role    string
	// Path is where it is: "Acme", "Acme / ml", "Acme / ml / training".
	Path    string
	AddedAt time.Time
	AddedBy string
}

func validMembershipScope(s Scope) bool {
	return s == ScopeOrg || s == ScopeTeam || s == ScopeProject
}

// AddMembership puts a person in an organization, team or project. Adding one
// that exists changes nothing but its role.
func (s *DB) AddMembership(ctx context.Context, userID string, scope Scope, scopeID, role, by string, at time.Time) error {
	if !validMembershipScope(scope) {
		return fmt.Errorf("a membership is in an org, team or project, not %q", scope)
	}
	if role == "" {
		role = RoleMember
	}
	table := map[Scope]string{ScopeOrg: "organizations", ScopeTeam: "teams", ScopeProject: "projects"}[scope]
	if err := s.exists(ctx, string(scope), `SELECT COUNT(*) FROM `+table+` WHERE id = ?`, scopeID); err != nil {
		return err
	}
	_, err := s.db.ExecContext(ctx, `
		INSERT INTO memberships (user_id, scope_type, scope_id, role, added_at, added_by)
		VALUES (?, ?, ?, ?, ?, ?)
		ON CONFLICT (user_id, scope_type, scope_id) DO UPDATE SET role = excluded.role`,
		userID, string(scope), scopeID, role, millis(at), nullString(by))
	return err
}

// RemoveMembership takes a person out of a scope and revokes the keys they
// made in any project they can no longer use, returning how many.
//
// Revoked rather than turned personal: a key made in a project spends that
// project's budget, and quietly moving it to the person would carry access
// out of the project the operator just removed them from.
func (s *DB) RemoveMembership(ctx context.Context, userID string, scope Scope, scopeID string, at time.Time) (int, error) {
	var revoked int
	err := s.inTx(ctx, func(c conn) error {
		res, err := c.ExecContext(ctx,
			`DELETE FROM memberships WHERE user_id = ? AND scope_type = ? AND scope_id = ?`,
			userID, string(scope), scopeID)
		if err != nil {
			return err
		}
		if n, _ := res.RowsAffected(); n == 0 {
			return ErrNotFound
		}
		res, err = c.ExecContext(ctx, `
			UPDATE api_keys SET revoked_at = ?
			WHERE user_id = ? AND project_id IS NOT NULL AND revoked_at IS NULL
			  AND project_id NOT IN (`+usableProjects+`)`,
			millis(at), userID, userID, userID, userID)
		if err != nil {
			return err
		}
		n, _ := res.RowsAffected()
		revoked = int(n)
		return nil
	})
	return revoked, err
}

// usableProjects selects every project a person may use: those they are a
// member of, and every project under a team or organization they are a member
// of. Takes the user id three times.
const usableProjects = `
	SELECT m.scope_id FROM memberships m
	  WHERE m.user_id = ? AND m.scope_type = 'project'
	UNION
	SELECT p.id FROM projects p JOIN memberships m ON m.scope_type = 'team' AND m.scope_id = p.team_id
	  WHERE m.user_id = ?
	UNION
	SELECT p.id FROM projects p JOIN teams t ON t.id = p.team_id
	  JOIN memberships m ON m.scope_type = 'org' AND m.scope_id = t.org_id
	  WHERE m.user_id = ?`

// CanUseProject reports whether a person may make and use keys in a project.
func (s *DB) CanUseProject(ctx context.Context, userID, projectID string) (bool, error) {
	var n int
	err := s.db.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM (`+usableProjects+`) u WHERE u.scope_id = ?`,
		userID, userID, userID, projectID).Scan(&n)
	return n > 0, err
}

// ProjectsFor lists the projects a person may use, with their paths.
func (s *DB) ProjectsFor(ctx context.Context, userID string) ([]ProjectPath, error) {
	return queryAll(ctx, s.db, `
		SELECT p.id, o.name || ' / ' || t.name || ' / ' || p.name
		FROM projects p
		JOIN teams t ON t.id = p.team_id
		JOIN organizations o ON o.id = t.org_id
		WHERE p.id IN (`+usableProjects+`)
		ORDER BY 2`, []any{userID, userID, userID}, func(rows *sql.Rows) (ProjectPath, error) {
		var p ProjectPath
		return p, rows.Scan(&p.ID, &p.Path)
	})
}

// membershipPath names the scope of a membership row as a path.
const membershipPath = `
	CASE m.scope_type
	  WHEN 'org' THEN (SELECT o.name FROM organizations o WHERE o.id = m.scope_id)
	  WHEN 'team' THEN (SELECT o.name || ' / ' || t.name FROM teams t
	                    JOIN organizations o ON o.id = t.org_id WHERE t.id = m.scope_id)
	  ELSE (SELECT o.name || ' / ' || t.name || ' / ' || p.name FROM projects p
	        JOIN teams t ON t.id = p.team_id JOIN organizations o ON o.id = t.org_id
	        WHERE p.id = m.scope_id)
	END`

func scanMembership(rows *sql.Rows) (Membership, error) {
	var m Membership
	var scope string
	var path, by sql.NullString
	var added int64
	err := rows.Scan(&m.UserID, &scope, &m.ScopeID, &m.Role, &added, &by, &path)
	m.Scope, m.Path, m.AddedBy, m.AddedAt = Scope(scope), path.String, by.String, fromMillis(added)
	return m, err
}

// MembershipsOf lists where a person is a member, by path.
func (s *DB) MembershipsOf(ctx context.Context, userID string) ([]Membership, error) {
	return queryAll(ctx, s.db, `
		SELECT m.user_id, m.scope_type, m.scope_id, m.role, m.added_at, m.added_by, `+membershipPath+`
		FROM memberships m WHERE m.user_id = ? ORDER BY 7`, []any{userID}, scanMembership)
}

// AllMemberships lists every membership, for a page that shows the tree with
// its people.
func (s *DB) AllMemberships(ctx context.Context) ([]Membership, error) {
	return queryAll(ctx, s.db, `
		SELECT m.user_id, m.scope_type, m.scope_id, m.role, m.added_at, m.added_by, `+membershipPath+`
		FROM memberships m ORDER BY 7`, nil, scanMembership)
}

// SetPersonalDefaults sets the limits a person's own use falls under when
// nobody wrote a limit for that person: what a personal key, with no project
// to be bounded by, gets. Nil clears them.
func (s *DB) SetPersonalDefaults(l *Limits) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if l == nil {
		s.personal = nil
		return
	}
	c := *l
	s.personal = &c
}

// Resolve walks a key's chain and returns the most restrictive of each limit.
func (s *DB) Resolve(ctx context.Context, keyID string) (Effective, error) {
	var e Effective
	var band, profile, project, user sql.NullString
	err := s.db.QueryRowContext(ctx,
		`SELECT id, band, routing_profile, project_id, user_id FROM api_keys WHERE id = ?`, keyID).
		Scan(&e.KeyID, &band, &profile, &project, &user)
	if errors.Is(err, sql.ErrNoRows) {
		return Effective{}, ErrNotFound
	}
	if err != nil {
		return Effective{}, err
	}
	e.Band, e.RoutingProfile = band.String, profile.String
	e.ProjectID, e.UserID = project.String, user.String
	// A person's key works while the person's account does: suspending an
	// account closed its sessions and left every key it had made working.
	if e.UserID != "" {
		var state string
		err := s.db.QueryRowContext(ctx, `SELECT state FROM users WHERE id = ?`, e.UserID).Scan(&state)
		if errors.Is(err, sql.ErrNoRows) || (err == nil && UserState(state) != UserActive) {
			return Effective{}, ErrOwnerInactive
		}
		if err != nil {
			return Effective{}, err
		}
	}
	return e, s.resolveChain(ctx, &e, true)
}

// ErrOwnerInactive is a key whose owner's account is suspended, pending, or
// gone.
var ErrOwnerInactive = errors.New("the key's owner is not an active account")

// ResolveFor resolves a caller that has no key row: the chat page, spending
// as a person in a project or, with no project, personally.
func (s *DB) ResolveFor(ctx context.Context, callerID, userID, projectID string) (Effective, error) {
	e := Effective{KeyID: callerID, UserID: userID, ProjectID: projectID}
	return e, s.resolveChain(ctx, &e, false)
}

// resolveChain fills in the tree above a caller and folds every limit that
// applies: the project's chain, each per-member limit on it for this person,
// the person's own, and the key's.
//
// One pool per level. A budget on a team is spent by everything beneath it; a
// person's limit by every key they own; a per-member limit by one person
// within one scope.
func (s *DB) resolveChain(ctx context.Context, e *Effective, keyLevel bool) error {
	if e.ProjectID != "" {
		err := s.db.QueryRowContext(ctx,
			`SELECT t.id, t.org_id FROM projects p JOIN teams t ON t.id = p.team_id WHERE p.id = ?`,
			e.ProjectID).Scan(&e.TeamID, &e.OrgID)
		if errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		}
		if err != nil {
			return err
		}
	}

	type step struct {
		scope Scope
		id    string
		// counter is where use is counted, when it is not the scope's own id.
		counter string
	}
	var steps []step
	if e.ProjectID != "" {
		steps = append(steps, step{ScopeOrg, e.OrgID, ""}, step{ScopeTeam, e.TeamID, ""},
			step{ScopeProject, e.ProjectID, ""})
		if e.UserID != "" {
			for _, lvl := range []step{{ScopeOrg, e.OrgID, ""}, {ScopeTeam, e.TeamID, ""}, {ScopeProject, e.ProjectID, ""}} {
				ms := MemberScope(lvl.scope)
				steps = append(steps, step{ms, lvl.id, MemberOwner(ms, lvl.id, e.UserID).ID})
			}
		}
	}
	if e.UserID != "" {
		steps = append(steps, step{ScopeUser, e.UserID, ""})
	}
	if keyLevel {
		steps = append(steps, step{ScopeKey, e.KeyID, ""})
	}

	sawUser := false
	for _, st := range steps {
		l, err := s.LimitsFor(ctx, st.scope, st.id)
		if errors.Is(err, ErrNotFound) {
			continue
		}
		if err != nil {
			return err
		}
		if st.counter != "" {
			l.ScopeID = st.counter
		}
		if st.scope == ScopeUser {
			sawUser = true
		}
		e.tighten(l)
		// Kept whole as well as folded in. Enforcement walks this; the folded
		// values are for showing somebody their ceiling.
		e.Chain = append(e.Chain, l)
	}

	// A personal caller has no project to be bounded by. With no limit written
	// for the person, the personal defaults stand in for one — counted in the
	// person's own pool, so every personal key they hold shares it.
	if e.ProjectID == "" && e.UserID != "" && !sawUser {
		s.mu.RLock()
		p := s.personal
		s.mu.RUnlock()
		if p != nil {
			l := *p
			l.Scope, l.ScopeID = ScopeUser, e.UserID
			e.tighten(l)
			e.Chain = append(e.Chain, l)
		}
	}
	s.applyDefaults(e)
	return nil
}

// ManagedOrgs lists the organizations a person manages.
func (s *DB) ManagedOrgs(ctx context.Context, userID string) ([]Org, error) {
	return queryAll(ctx, s.db, `
		SELECT o.id, o.name FROM memberships m JOIN organizations o ON o.id = m.scope_id
		WHERE m.user_id = ? AND m.scope_type = 'org' AND m.role = 'manager'
		ORDER BY o.name`, []any{userID}, func(rows *sql.Rows) (Org, error) {
		var o Org
		return o, rows.Scan(&o.ID, &o.Name)
	})
}

// ManagesOrg reports whether a person manages an organization.
func (s *DB) ManagesOrg(ctx context.Context, userID, orgID string) (bool, error) {
	var n int
	err := s.db.QueryRowContext(ctx, `
		SELECT COUNT(*) FROM memberships
		WHERE user_id = ? AND scope_type = 'org' AND scope_id = ? AND role = 'manager'`,
		userID, orgID).Scan(&n)
	return n > 0, err
}