proxy-lite/internal/noodle/database.go

package noodle

import (
	"errors"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/google/uuid"
	"go.mills.io/bitcask/v2"
)

var ErrUserNotFound = errors.New("user not found")
var ErrLastAdminRemoval = errors.New("cannot remove the last admin")

type Database struct {
	connection *bitcask.Bitcask
	Handle     *bitcask.Collection
	Users      *bitcask.Collection
	mu         sync.Mutex
}

func NewDatabase(path string) *Database {
	db, err := bitcask.Open(path)
	if err != nil {
		log.Fatal(err)
	}
	return &Database{
		connection: db,
		Handle:     db.Collection("noodles"),
		Users:      db.Collection("users"),
	}
}

func (db *Database) MakeID() string {
	id, err := uuid.NewUUID()
	if err != nil {
		log.Print(err)
		return ""
	}
	return id.String()
}

func (db *Database) GetAll() []Noodle {
	db.mu.Lock()
	defer db.mu.Unlock()

	var data []Noodle
	if err := db.Handle.List(&data); err != nil {
		log.Print(err)
		return nil
	}
	return data
}

func (db *Database) GetAllGeneric() []interface{} {
	db.mu.Lock()
	defer db.mu.Unlock()

	var data []interface{}
	if err := db.Handle.List(&data); err != nil {
		log.Print(err)
		return nil
	}
	return data
}

func (db *Database) Get(id string) Noodle {
	db.mu.Lock()
	defer db.mu.Unlock()

	var item Noodle
	log.Printf("Looking up noodle key='%s'", id)
	log.Printf("key='%s' exists=%t", id, db.Handle.Has(id))
	if err := db.Handle.Get(id, &item); err != nil {
		log.Print(err)
		return Noodle{}
	}
	return item
}

func (db *Database) Add(item Noodle) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if err := db.Handle.Add(item.Id, item); err != nil {
		log.Print(err)
		return err
	}
	return nil
}

func (db *Database) Update(item Noodle) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if err := db.Handle.Add(item.Id, item); err != nil {
		log.Print(err)
		return err
	}
	return nil
}

func (db *Database) Delete(id string) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if err := db.Handle.Delete(id); err != nil {
		log.Print(err)
		return err
	}
	return nil
}

func (db *Database) Close() error {
	return db.connection.Close()
}

func (db *Database) GetAllUsers() []User {
	db.mu.Lock()
	defer db.mu.Unlock()

	var data []User
	if err := db.Users.List(&data); err != nil {
		log.Print(err)
		return nil
	}
	return data
}

func (db *Database) GetUserByUsername(username string) (User, error) {
	db.mu.Lock()
	defer db.mu.Unlock()

	var data []User
	if err := db.Users.List(&data); err != nil {
		log.Print(err)
		return User{}, err
	}

	for _, user := range data {
		if user.Username == username {
			return user, nil
		}
	}

	return User{}, ErrUserNotFound
}

func (db *Database) AddUser(user User) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	var data []User
	if err := db.Users.List(&data); err != nil {
		log.Print(err)
		return err
	}
	for _, existing := range data {
		if existing.Username == user.Username {
			return fmt.Errorf("user %q already exists", user.Username)
		}
	}

	if err := db.Users.Add(user.Id, user); err != nil {
		log.Print(err)
		return err
	}
	return nil
}

func (db *Database) UpdateUser(user User) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if err := db.Users.Add(user.Id, user); err != nil {
		log.Print(err)
		return err
	}
	return nil
}

func (db *Database) SetUserRole(username, role string) (User, error) {
	db.mu.Lock()
	defer db.mu.Unlock()

	var data []User
	if err := db.Users.List(&data); err != nil {
		log.Print(err)
		return User{}, err
	}

	var target User
	found := false
	adminCount := 0
	for _, user := range data {
		if user.Role == UserRoleAdmin {
			adminCount++
		}
		if user.Username == username {
			target = user
			found = true
		}
	}

	if !found {
		return User{}, ErrUserNotFound
	}
	if target.Role == UserRoleAdmin && role != UserRoleAdmin && adminCount <= 1 {
		return User{}, ErrLastAdminRemoval
	}

	target.Role = role
	if err := db.Users.Add(target.Id, target); err != nil {
		log.Print(err)
		return User{}, err
	}
	return target, nil
}

func (db *Database) SetUserPassword(username, passwordHash string) (User, error) {
	db.mu.Lock()
	defer db.mu.Unlock()

	var data []User
	if err := db.Users.List(&data); err != nil {
		log.Print(err)
		return User{}, err
	}

	for _, user := range data {
		if user.Username != username {
			continue
		}
		user.PasswordHash = passwordHash
		if err := db.Users.Add(user.Id, user); err != nil {
			log.Print(err)
			return User{}, err
		}
		return user, nil
	}

	return User{}, ErrUserNotFound
}

func (db *Database) DeleteUser(username string) (User, error) {
	db.mu.Lock()
	defer db.mu.Unlock()

	var data []User
	if err := db.Users.List(&data); err != nil {
		log.Print(err)
		return User{}, err
	}

	var target User
	found := false
	adminCount := 0
	for _, user := range data {
		if user.Role == UserRoleAdmin {
			adminCount++
		}
		if user.Username == username {
			target = user
			found = true
		}
	}

	if !found {
		return User{}, ErrUserNotFound
	}
	if target.Role == UserRoleAdmin && adminCount <= 1 {
		return User{}, ErrLastAdminRemoval
	}

	if err := db.Users.Delete(target.Id); err != nil {
		log.Print(err)
		return User{}, err
	}
	return target, nil
}

func (db *Database) SetIsUp(id string, isUp bool) (Noodle, error) {
	db.mu.Lock()
	defer db.mu.Unlock()

	var item Noodle
	if err := db.Handle.Get(id, &item); err != nil {
		log.Print(err)
		return Noodle{}, err
	}

	item.IsUp = isUp
	if err := db.Handle.Add(item.Id, item); err != nil {
		log.Print(err)
		return Noodle{}, err
	}

	return item, nil
}

func (db *Database) DeleteByID(id string) (Noodle, error) {
	db.mu.Lock()
	defer db.mu.Unlock()

	var item Noodle
	if err := db.Handle.Get(id, &item); err != nil {
		log.Print(err)
		return Noodle{}, err
	}

	if err := db.Handle.Delete(id); err != nil {
		log.Print(err)
		return Noodle{}, err
	}

	return item, nil
}

func (db *Database) TickExpirations() []Noodle {
	db.mu.Lock()
	defer db.mu.Unlock()

	var noodles []Noodle
	if err := db.Handle.List(&noodles); err != nil {
		log.Print(err)
		return nil
	}

	var stopped []Noodle
	for _, item := range noodles {
		if !item.IsUp {
			continue
		}

		if item.Expiration <= 0 {
			item.IsUp = false
			stopped = append(stopped, item)
			if err := db.Handle.Delete(item.Id); err != nil {
				log.Print(err)
			}
			continue
		}

		item.Expiration -= time.Second
		if item.Expiration <= 0 {
			item.Expiration = 0
			item.IsUp = false
			stopped = append(stopped, item)
			if err := db.Handle.Delete(item.Id); err != nil {
				log.Print(err)
			}
			continue
		}

		if err := db.Handle.Add(item.Id, item); err != nil {
			log.Print(err)
		}
	}

	return stopped
}