core/auth/memberships/store.go

package main

import (
	"database/sql"
	_ "embed"
	"fmt"
	"strings"

	"github.com/ncruces/go-sqlite3"
)

//go:embed schema.sql
var schema string

const (
	ErrorUniqueConstraintViolation     = 2067
	ErrorConstraintPrimaryKeyViolation = 1555
)

type scanner interface {
	Scan(dest ...any) error
}

type Store interface {
	// Initializes store with admin user and their groups.
	Init(user User, groups []Group) error
	CreateUser(tx *sql.Tx, user User) error
	GetAllUsers(tx *sql.Tx) ([]User, error)
	GetUser(tx *sql.Tx, id string) (User, error)
	GetUsers(tx *sql.Tx, ids []string) ([]User, error)
	CreateGroup(tx *sql.Tx, userId string, group Group) error
	GetGroup(tx *sql.Tx, id string) (Group, error)
	GetAllGroups(tx *sql.Tx) ([]Group, error)
	AddMemberUser(tx *sql.Tx, groupId, userId string) error
	AddOwnerUser(tx *sql.Tx, groupId, userId string) error
	AddMemberGroup(tx *sql.Tx, groupId, otherId string) error
	AddOwnerGroup(tx *sql.Tx, groupId, otherId string) error
	GetMemberUsers(tx *sql.Tx, groupId string) ([]User, error)
	GetOwnerUsers(tx *sql.Tx, groupId string) ([]User, error)
	GetMemberGroups(tx *sql.Tx, groupId string) ([]Group, error)
	GetOwnerGroups(tx *sql.Tx, groupId string) ([]Group, error)
	RemoveMemberUser(tx *sql.Tx, groupId, userId string) error
	RemoveOwnerUser(tx *sql.Tx, groupId, userId string) error
	RemoveMemberGroup(tx *sql.Tx, groupId, otherId string) error
	RemoveOwnerGroup(tx *sql.Tx, groupId, otherId string) error
	GetGroupsUserCanActAs(tx *sql.Tx, userId string) ([]Group, error)
	GetGroupsGroupCanActAs(tx *sql.Tx, groupId string) ([]Group, error)
	GetGroupsUserOwns(tx *sql.Tx, userId string) ([]Group, error)
	GetGroupsUserIsMemberOf(tx *sql.Tx, userId string) ([]Group, error)
	GetUserPublicKeys(tx *sql.Tx, userId string) ([]string, error)
	AddUserPublicKey(tx *sql.Tx, userId, publicKey string) error
	RemoveUserPublicKey(tx *sql.Tx, userId, publicKey string) error
}

type SQLiteStore struct {
	db *sql.DB
}

func NewSQLiteStore(db *sql.DB) (*SQLiteStore, error) {
	_, err := db.Exec(schema)
	if err != nil {
		return nil, err
	}
	return &SQLiteStore{db: db}, nil
}

func (s *SQLiteStore) Init(user User, groups []Group) error {
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	row := tx.QueryRow("SELECT COUNT(*) FROM groups")
	var count int
	if err := row.Scan(&count); err != nil {
		return err
	}
	if count != 0 {
		return fmt.Errorf("Store already initialised")
	}
	if err := s.CreateUser(tx, user); err != nil {
		return err
	}
	for _, g := range groups {
		if err := s.CreateGroup(tx, user.Id, g); err != nil {
			return err
		}
		if err := s.AddMemberUser(tx, g.Id, user.Id); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func (s *SQLiteStore) CreateUser(tx *sql.Tx, user User) (err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			} else {
				tx.Rollback()
			}
		}()
	}
	query := "INSERT INTO users (id, username, email) VALUES (?, ?, ?)"
	_, err = tx.Exec(query, user.Id, user.Username, user.Email)
	return
}

func (s *SQLiteStore) scanUser(row scanner) (User, error) {
	var ret User
	if err := row.Scan(&ret.Id, &ret.Username, &ret.Email); err != nil {
		return User{}, err
	}
	return ret, nil
}

func (s *SQLiteStore) scanUsers(rows *sql.Rows) ([]User, error) {
	var ret []User
	for rows.Next() {
		if u, err := s.scanUser(rows); err != nil {
			return nil, err
		} else {
			ret = append(ret, u)
		}
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return ret, nil
}

func (s *SQLiteStore) scanGroup(row scanner) (Group, error) {
	var ret Group
	if err := row.Scan(&ret.Id, &ret.Title, &ret.Description); err != nil {
		return Group{}, err
	}
	return ret, nil
}

func (s *SQLiteStore) scanGroups(rows *sql.Rows) ([]Group, error) {
	var ret []Group
	for rows.Next() {
		if g, err := s.scanGroup(rows); err != nil {
			return nil, err
		} else {
			ret = append(ret, g)
		}
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return ret, nil
}

func (s *SQLiteStore) GetAllUsers(tx *sql.Tx) (ret []User, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	query := "SELECT id, username, email FROM users"
	var rows *sql.Rows
	if rows, err = tx.Query(query); err == nil {
		ret, err = s.scanUsers(rows)
	}
	return
}

func (s *SQLiteStore) GetUser(tx *sql.Tx, id string) (ret User, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			} else {
				tx.Rollback()
			}
		}()
	}
	query := "SELECT id, username, email FROM users WHERE id = ?"
	row := tx.QueryRow(query, id)
	ret, err = s.scanUser(row)
	return
}

func (s *SQLiteStore) GetUsers(tx *sql.Tx, ids []string) (ret []User, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	idPlaceholders := make([]string, len(ids))
	idParams := make([]any, len(ids))
	for i, id := range ids {
		idPlaceholders[i] = "?"
		idParams[i] = id
	}
	query := fmt.Sprintf("SELECT id, username, email FROM users WHERE id IN (%s)", strings.Join(idPlaceholders, ", "))
	var rows *sql.Rows
	if rows, err = tx.Query(query, idParams...); err == nil {
		ret, err = s.scanUsers(rows)
	}
	return
}

func (s *SQLiteStore) CreateGroup(tx *sql.Tx, userId string, group Group) (err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			} else {
				tx.Rollback()
			}
		}()
	}
	query := "INSERT INTO groups (id, title, description) VALUES (?, ?, ?)"
	if _, err = tx.Exec(query, group.Id, group.Title, group.Description); err != nil {
		sqliteErr, ok := err.(*sqlite3.Error)
		if ok && sqliteErr.ExtendedCode() == ErrorConstraintPrimaryKeyViolation {
			err = fmt.Errorf("Group with id %s already exists", group.Id)
		}
		return
	}
	err = s.addUser(tx, group.Id, userId, "owner")
	return
}

func (s *SQLiteStore) GetGroup(tx *sql.Tx, id string) (ret Group, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			} else {
				tx.Rollback()
			}
		}()
	}
	query := "SELECT id, title, description FROM groups WHERE id = ?"
	row := tx.QueryRow(query, id)
	ret, err = s.scanGroup(row)
	return
}
func (s *SQLiteStore) GetAllGroups(tx *sql.Tx) (ret []Group, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	query := "SELECT id, title, description FROM groups"
	var rows *sql.Rows
	if rows, err = tx.Query(query); err == nil {
		ret, err = s.scanGroups(rows)
	}
	return
}

func (s *SQLiteStore) addUser(tx *sql.Tx, groupId, userId, membershipType string) (err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	query := "INSERT INTO memberships (id, membership_type, member_type, user_id) VALUES (?, ?, ?, ?)"
	_, err = tx.Exec(query, groupId, membershipType, "user", userId)
	return
}

func (s *SQLiteStore) addGroup(tx *sql.Tx, groupId, otherId, membershipType string) (err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	query := "INSERT INTO memberships (id, membership_type, member_type, group_id) VALUES (?, ?, ?, ?)"
	_, err = tx.Exec(query, groupId, membershipType, "group", otherId)
	return
}

func (s *SQLiteStore) removeUser(tx *sql.Tx, groupId, userId, membershipType string) (err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	query := "DELETE FROM memberships WHERE id = ? AND membership_type = ? AND member_type = ? AND user_id = ?"
	_, err = tx.Exec(query, groupId, membershipType, "user", userId)
	return
}

func (s *SQLiteStore) removeGroup(tx *sql.Tx, groupId, otherId, membershipType string) (err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	query := "DELETE FROM memberships WHERE id = ? AND membership_type = ? AND member_type = ? AND group_id = ?"
	_, err = tx.Exec(query, groupId, membershipType, "group", otherId)
	return
}

func (s *SQLiteStore) getUsers(tx *sql.Tx, groupId, membershipType string) (ret []User, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	query := `
SELECT u.id, u.username, u.email
FROM users u
JOIN memberships m
ON u.id = m.user_id
WHERE m.id = ? AND membership_type = ? AND member_type = ?
`
	var rows *sql.Rows
	rows, err = tx.Query(query, groupId, membershipType, "user")
	if err != nil {
		return
	}
	ret, err = s.scanUsers(rows)
	return
}

func (s *SQLiteStore) getGroups(tx *sql.Tx, groupId, membershipType string) (ret []Group, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	query := `
SELECT g.id, g.title, g.description
FROM groups AS g
JOIN memberships AS m
ON g.id = m.group_id
WHERE m.id = ? AND m.membership_type = ? AND m.member_type = ?`
	var rows *sql.Rows
	if rows, err = tx.Query(query, groupId, membershipType, "group"); err == nil {
		ret, err = s.scanGroups(rows)
	}
	return
}

func (s *SQLiteStore) AddMemberUser(tx *sql.Tx, groupId, userId string) error {
	return s.addUser(tx, groupId, userId, "member")
}

func (s *SQLiteStore) AddOwnerUser(tx *sql.Tx, groupId, userId string) error {
	return s.addUser(tx, groupId, userId, "owner")
}

func (s *SQLiteStore) AddMemberGroup(tx *sql.Tx, groupId, otherId string) error {
	return s.addGroup(tx, groupId, otherId, "member")
}

func (s *SQLiteStore) AddOwnerGroup(tx *sql.Tx, groupId, otherId string) error {
	return s.addGroup(tx, groupId, otherId, "owner")
}

func (s *SQLiteStore) GetMemberUsers(tx *sql.Tx, groupId string) ([]User, error) {
	return s.getUsers(tx, groupId, "member")
}

func (s *SQLiteStore) GetOwnerUsers(tx *sql.Tx, groupId string) ([]User, error) {
	return s.getUsers(tx, groupId, "owner")
}

func (s *SQLiteStore) GetMemberGroups(tx *sql.Tx, groupId string) ([]Group, error) {
	return s.getGroups(tx, groupId, "member")
}

func (s *SQLiteStore) GetOwnerGroups(tx *sql.Tx, groupId string) ([]Group, error) {
	return s.getGroups(tx, groupId, "owner")
}

func (s *SQLiteStore) RemoveMemberUser(tx *sql.Tx, groupId, userId string) error {
	return s.removeUser(tx, groupId, userId, "member")
}

func (s *SQLiteStore) RemoveOwnerUser(tx *sql.Tx, groupId, userId string) error {
	return s.removeUser(tx, groupId, userId, "owner")
}

func (s *SQLiteStore) RemoveMemberGroup(tx *sql.Tx, groupId, otherId string) error {
	return s.removeGroup(tx, groupId, otherId, "member")
}

func (s *SQLiteStore) RemoveOwnerGroup(tx *sql.Tx, groupId, otherId string) error {
	return s.removeGroup(tx, groupId, otherId, "owner")
}

func (s *SQLiteStore) traverseGroups(tx *sql.Tx, groups ...Group) ([]Group, error) {
	seen := map[string]struct{}{}
	for _, g := range groups {
		seen[g.Id] = struct{}{}
	}
	for i := 0; i < len(groups); i++ {
		g := groups[i]
		parents, err := s.getGroupGroups(tx, g.Id, "member")
		if err != nil {
			return nil, err
		}
		for _, p := range parents {
			if _, ok := seen[p.Id]; !ok {
				groups = append(groups, p)
				seen[p.Id] = struct{}{}
			}
		}
	}
	return groups, nil
}

func (s *SQLiteStore) GetGroupsUserCanActAs(tx *sql.Tx, userId string) (ret []Group, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	var memberOf []Group
	memberOf, err = s.GetGroupsUserIsMemberOf(tx, userId)
	if err != nil {
		return
	}
	ret, err = s.traverseGroups(tx, memberOf...)
	return
}

func (s *SQLiteStore) GetGroupsGroupCanActAs(tx *sql.Tx, groupId string) (ret []Group, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	var g Group
	g, err = s.GetGroup(tx, groupId)
	if err != nil {
		return
	}
	var groups []Group
	if groups, err = s.traverseGroups(tx, g); err == nil {
		ret = groups[1:]
	}
	return
}

func (s *SQLiteStore) getUserGroups(tx *sql.Tx, userId, membershipType string) (ret []Group, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			}
			if err != nil {
				tx.Rollback()
			}
		}()
	}
	query := `
SELECT g.id, g.title, g.description
FROM groups AS g
JOIN memberships AS m
ON g.id = m.id
WHERE m.user_id = ? AND m.membership_type = ? AND m.member_type = ?`
	var rows *sql.Rows
	if rows, err = tx.Query(query, userId, membershipType, "user"); err == nil {
		ret, err = s.scanGroups(rows)
	}
	return
}

func (s *SQLiteStore) getGroupGroups(tx *sql.Tx, groupId, membershipType string) (ret []Group, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			} else {
				tx.Rollback()
			}
		}()
	}
	query := `
SELECT g.id, g.title, g.description
FROM groups AS g
JOIN memberships AS m
ON g.id = m.id
WHERE m.group_id = ? AND m.membership_type = ? AND m.member_type = ?`
	var rows *sql.Rows
	if rows, err = tx.Query(query, groupId, membershipType, "group"); err == nil {
		ret, err = s.scanGroups(rows)
	}
	return
}

func (s *SQLiteStore) GetGroupsUserOwns(tx *sql.Tx, userId string) ([]Group, error) {
	return s.getUserGroups(tx, userId, "owner")
}

func (s *SQLiteStore) GetGroupsUserIsMemberOf(tx *sql.Tx, userId string) ([]Group, error) {
	return s.getUserGroups(tx, userId, "member")
}

func (s *SQLiteStore) AddUserPublicKey(tx *sql.Tx, userId, publicKey string) (err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			} else {
				tx.Rollback()
			}
		}()
	}
	query := "INSERT INTO keys (user_id, public_key) VALUES (?, ?)"
	_, err = tx.Exec(query, userId, publicKey)
	return
}

func (s *SQLiteStore) GetUserPublicKeys(tx *sql.Tx, userId string) (ret []string, err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			} else {
				tx.Rollback()
			}
		}()
	}
	query := "SELECT public_key FROM keys WHERE user_id = ?"
	var rows *sql.Rows
	if rows, err = tx.Query(query, userId); err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var publicKey string
		if err = rows.Scan(&publicKey); err != nil {
			return
		}
		ret = append(ret, publicKey)
	}
	err = rows.Err()
	return
}

func (s *SQLiteStore) RemoveUserPublicKey(tx *sql.Tx, userId, publicKey string) (err error) {
	if tx == nil {
		tx, err = s.db.Begin()
		if err != nil {
			return
		}
		defer func() {
			if err == nil {
				err = tx.Commit()
			} else {
				tx.Rollback()
			}
		}()
	}
	query := "DELETE FROM keys WHERE user_id = ? AND  public_key = ?"
	_, err = tx.Exec(query, userId, publicKey)
	return
}