warpbox-dev/backend/libs/services/auth.go

package services

import (
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/mail"
	"os"
	"sort"
	"strings"
	"time"

	"go.etcd.io/bbolt"
	"golang.org/x/crypto/argon2"
)

var (
	usersBucket       = []byte("users")
	userEmailsBucket  = []byte("user_emails")
	sessionsBucket    = []byte("sessions")
	invitesBucket     = []byte("invites")
	collectionsBucket = []byte("collections")
)

const (
	UserRoleAdmin = "admin"
	UserRoleUser  = "user"

	UserStatusActive   = "active"
	UserStatusDisabled = "disabled"
)

var (
	ErrInvalidCredentials = errors.New("invalid credentials")
	ErrRegistrationClosed = errors.New("registration is closed")
	ErrInviteInvalid      = errors.New("invite is invalid")
	ErrUserDisabled       = errors.New("user is disabled")
)

type AuthService struct {
	db      *bbolt.DB
	baseURL string
}

type User struct {
	ID           string    `json:"id"`
	Username     string    `json:"username"`
	Email        string    `json:"email"`
	PasswordHash string    `json:"passwordHash"`
	Role         string    `json:"role"`
	Status       string    `json:"status"`
	CreatedAt    time.Time `json:"createdAt"`
	UpdatedAt    time.Time `json:"updatedAt"`
}

type PublicUser struct {
	ID        string
	Username  string
	Email     string
	Role      string
	Status    string
	CreatedAt time.Time
}

type Session struct {
	ID        string    `json:"id"`
	UserID    string    `json:"userId"`
	TokenHash string    `json:"tokenHash"`
	ExpiresAt time.Time `json:"expiresAt"`
	CreatedAt time.Time `json:"createdAt"`
}

type Invite struct {
	ID           string     `json:"id"`
	UserID       string     `json:"userId,omitempty"`
	Email        string     `json:"email"`
	Role         string     `json:"role"`
	TokenHash    string     `json:"tokenHash"`
	CreatedBy    string     `json:"createdBy"`
	CreatedAt    time.Time  `json:"createdAt"`
	ExpiresAt    time.Time  `json:"expiresAt"`
	UsedAt       *time.Time `json:"usedAt,omitempty"`
	UsedByUserID string     `json:"usedByUserId,omitempty"`
}

type Collection struct {
	ID        string    `json:"id"`
	UserID    string    `json:"userId"`
	Name      string    `json:"name"`
	CreatedAt time.Time `json:"createdAt"`
	UpdatedAt time.Time `json:"updatedAt"`
}

type InviteResult struct {
	Invite Invite
	URL    string
	Token  string
}

func NewAuthService(db *bbolt.DB, baseURL string) (*AuthService, error) {
	service := &AuthService{db: db, baseURL: strings.TrimRight(baseURL, "/")}
	err := db.Update(func(tx *bbolt.Tx) error {
		for _, bucket := range [][]byte{usersBucket, userEmailsBucket, sessionsBucket, invitesBucket, collectionsBucket} {
			if _, err := tx.CreateBucketIfNotExists(bucket); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return service, nil
}

func (s *AuthService) BootstrapAvailable() (bool, error) {
	count := 0
	err := s.db.View(func(tx *bbolt.Tx) error {
		return tx.Bucket(usersBucket).ForEach(func(_, _ []byte) error {
			count++
			return nil
		})
	})
	return count == 0, err
}

func (s *AuthService) CreateBootstrapUser(username, email, password string) (User, error) {
	available, err := s.BootstrapAvailable()
	if err != nil {
		return User{}, err
	}
	if !available {
		return User{}, ErrRegistrationClosed
	}
	return s.createUser(username, email, password, UserRoleAdmin)
}

func (s *AuthService) Login(email, password string) (User, string, error) {
	user, err := s.UserByEmail(email)
	if err != nil {
		return User{}, "", ErrInvalidCredentials
	}
	if user.Status != UserStatusActive {
		return User{}, "", ErrUserDisabled
	}
	if !VerifyPasswordHash(user.PasswordHash, password) {
		return User{}, "", ErrInvalidCredentials
	}

	token := randomID(32)
	session := Session{
		ID:        randomID(12),
		UserID:    user.ID,
		TokenHash: tokenHash(token),
		CreatedAt: time.Now().UTC(),
		ExpiresAt: time.Now().UTC().Add(30 * 24 * time.Hour),
	}
	err = s.db.Update(func(tx *bbolt.Tx) error {
		data, err := json.Marshal(session)
		if err != nil {
			return err
		}
		return tx.Bucket(sessionsBucket).Put([]byte(session.ID), data)
	})
	return user, session.ID + "." + token, err
}

func (s *AuthService) UserForSession(raw string) (User, Session, error) {
	sessionID, token, ok := strings.Cut(raw, ".")
	if !ok || sessionID == "" || token == "" {
		return User{}, Session{}, os.ErrNotExist
	}

	var session Session
	err := s.db.View(func(tx *bbolt.Tx) error {
		data := tx.Bucket(sessionsBucket).Get([]byte(sessionID))
		if data == nil {
			return os.ErrNotExist
		}
		return json.Unmarshal(data, &session)
	})
	if err != nil {
		return User{}, Session{}, err
	}
	if time.Now().UTC().After(session.ExpiresAt) || subtle.ConstantTimeCompare([]byte(tokenHash(token)), []byte(session.TokenHash)) != 1 {
		return User{}, Session{}, os.ErrPermission
	}
	user, err := s.UserByID(session.UserID)
	if err != nil {
		return User{}, Session{}, err
	}
	if user.Status != UserStatusActive {
		return User{}, Session{}, ErrUserDisabled
	}
	return user, session, nil
}

func (s *AuthService) Logout(raw string) error {
	sessionID, _, ok := strings.Cut(raw, ".")
	if !ok || sessionID == "" {
		return nil
	}
	return s.db.Update(func(tx *bbolt.Tx) error {
		return tx.Bucket(sessionsBucket).Delete([]byte(sessionID))
	})
}

func (s *AuthService) CreateInvite(email, role, createdBy string, expiresIn time.Duration) (InviteResult, error) {
	email, err := normalizeEmail(email)
	if err != nil {
		return InviteResult{}, err
	}
	if role == "" {
		role = UserRoleUser
	}
	if role != UserRoleAdmin && role != UserRoleUser {
		role = UserRoleUser
	}
	if expiresIn <= 0 {
		expiresIn = 7 * 24 * time.Hour
	}

	token := randomID(32)
	invite := Invite{
		ID:        randomID(12),
		Email:     email,
		Role:      role,
		TokenHash: tokenHash(token),
		CreatedBy: createdBy,
		CreatedAt: time.Now().UTC(),
		ExpiresAt: time.Now().UTC().Add(expiresIn),
	}
	err = s.saveInvite(invite)
	if err != nil {
		return InviteResult{}, err
	}
	return InviteResult{
		Invite: invite,
		Token:  token,
		URL:    fmt.Sprintf("%s/invite/%s", s.baseURL, token),
	}, nil
}

func (s *AuthService) AcceptInvite(token, username, password string) (User, error) {
	invite, err := s.InviteByToken(token)
	if err != nil {
		return User{}, err
	}
	if invite.UsedAt != nil || time.Now().UTC().After(invite.ExpiresAt) {
		return User{}, ErrInviteInvalid
	}

	var user User
	if invite.UserID != "" {
		user, err = s.UserByID(invite.UserID)
		if err != nil {
			return User{}, err
		}
		if err := s.SetPassword(user.ID, password); err != nil {
			return User{}, err
		}
		user, _ = s.UserByID(user.ID)
	} else {
		user, err = s.createUser(username, invite.Email, password, invite.Role)
		if err != nil {
			return User{}, err
		}
	}

	now := time.Now().UTC()
	invite.UsedAt = &now
	invite.UsedByUserID = user.ID
	if err := s.saveInvite(invite); err != nil {
		return User{}, err
	}
	return user, nil
}

func (s *AuthService) InviteByToken(token string) (Invite, error) {
	hash := tokenHash(token)
	var match Invite
	err := s.db.View(func(tx *bbolt.Tx) error {
		return tx.Bucket(invitesBucket).ForEach(func(_, value []byte) error {
			var invite Invite
			if err := json.Unmarshal(value, &invite); err != nil {
				return err
			}
			if subtle.ConstantTimeCompare([]byte(hash), []byte(invite.TokenHash)) == 1 {
				match = invite
			}
			return nil
		})
	})
	if err != nil {
		return Invite{}, err
	}
	if match.ID == "" {
		return Invite{}, ErrInviteInvalid
	}
	return match, nil
}

func (s *AuthService) CreatePasswordResetInvite(userID, createdBy string) (InviteResult, error) {
	user, err := s.UserByID(userID)
	if err != nil {
		return InviteResult{}, err
	}
	result, err := s.CreateInvite(user.Email, user.Role, createdBy, 24*time.Hour)
	if err != nil {
		return InviteResult{}, err
	}
	result.Invite.UserID = user.ID
	if err := s.saveInvite(result.Invite); err != nil {
		return InviteResult{}, err
	}
	return result, nil
}

func (s *AuthService) ListUsers() ([]User, error) {
	users := make([]User, 0)
	err := s.db.View(func(tx *bbolt.Tx) error {
		return tx.Bucket(usersBucket).ForEach(func(_, value []byte) error {
			var user User
			if err := json.Unmarshal(value, &user); err != nil {
				return err
			}
			users = append(users, user)
			return nil
		})
	})
	sort.Slice(users, func(i, j int) bool {
		return users[i].CreatedAt.After(users[j].CreatedAt)
	})
	return users, err
}

func (s *AuthService) DisableUser(userID string, disabled bool) error {
	user, err := s.UserByID(userID)
	if err != nil {
		return err
	}
	if disabled {
		user.Status = UserStatusDisabled
	} else {
		user.Status = UserStatusActive
	}
	user.UpdatedAt = time.Now().UTC()
	return s.saveUser(user)
}

func (s *AuthService) SetPassword(userID, password string) error {
	if len(password) < 8 {
		return fmt.Errorf("password must be at least 8 characters")
	}
	user, err := s.UserByID(userID)
	if err != nil {
		return err
	}
	user.PasswordHash = HashPassword(password)
	user.UpdatedAt = time.Now().UTC()
	return s.saveUser(user)
}

func (s *AuthService) UserByID(id string) (User, error) {
	var user User
	err := s.db.View(func(tx *bbolt.Tx) error {
		data := tx.Bucket(usersBucket).Get([]byte(id))
		if data == nil {
			return os.ErrNotExist
		}
		return json.Unmarshal(data, &user)
	})
	return user, err
}

func (s *AuthService) UserByEmail(email string) (User, error) {
	email, err := normalizeEmail(email)
	if err != nil {
		return User{}, err
	}
	var userID string
	err = s.db.View(func(tx *bbolt.Tx) error {
		data := tx.Bucket(userEmailsBucket).Get([]byte(email))
		if data == nil {
			return os.ErrNotExist
		}
		userID = string(data)
		return nil
	})
	if err != nil {
		return User{}, err
	}
	return s.UserByID(userID)
}

func (s *AuthService) CreateCollection(userID, name string) (Collection, error) {
	name = strings.TrimSpace(name)
	if name == "" {
		return Collection{}, fmt.Errorf("collection name is required")
	}
	collection := Collection{
		ID:        randomID(10),
		UserID:    userID,
		Name:      name,
		CreatedAt: time.Now().UTC(),
		UpdatedAt: time.Now().UTC(),
	}
	return collection, s.saveCollection(collection)
}

func (s *AuthService) ListCollections(userID string) ([]Collection, error) {
	collections := make([]Collection, 0)
	err := s.db.View(func(tx *bbolt.Tx) error {
		return tx.Bucket(collectionsBucket).ForEach(func(_, value []byte) error {
			var collection Collection
			if err := json.Unmarshal(value, &collection); err != nil {
				return err
			}
			if collection.UserID == userID {
				collections = append(collections, collection)
			}
			return nil
		})
	})
	sort.Slice(collections, func(i, j int) bool {
		return strings.ToLower(collections[i].Name) < strings.ToLower(collections[j].Name)
	})
	return collections, err
}

func (s *AuthService) CollectionOwnedBy(collectionID, userID string) bool {
	if collectionID == "" {
		return true
	}
	collection, err := s.CollectionByID(collectionID)
	return err == nil && collection.UserID == userID
}

func (s *AuthService) CollectionByID(id string) (Collection, error) {
	var collection Collection
	err := s.db.View(func(tx *bbolt.Tx) error {
		data := tx.Bucket(collectionsBucket).Get([]byte(id))
		if data == nil {
			return os.ErrNotExist
		}
		return json.Unmarshal(data, &collection)
	})
	return collection, err
}

func (s *AuthService) PublicUser(user User) PublicUser {
	return PublicUser{
		ID:        user.ID,
		Username:  user.Username,
		Email:     user.Email,
		Role:      user.Role,
		Status:    user.Status,
		CreatedAt: user.CreatedAt,
	}
}

func (s *AuthService) createUser(username, email, password, role string) (User, error) {
	username = strings.TrimSpace(username)
	if username == "" {
		return User{}, fmt.Errorf("username is required")
	}
	email, err := normalizeEmail(email)
	if err != nil {
		return User{}, err
	}
	if len(password) < 8 {
		return User{}, fmt.Errorf("password must be at least 8 characters")
	}
	if role != UserRoleAdmin && role != UserRoleUser {
		role = UserRoleUser
	}

	now := time.Now().UTC()
	user := User{
		ID:           randomID(12),
		Username:     username,
		Email:        email,
		PasswordHash: HashPassword(password),
		Role:         role,
		Status:       UserStatusActive,
		CreatedAt:    now,
		UpdatedAt:    now,
	}
	return user, s.db.Update(func(tx *bbolt.Tx) error {
		if existing := tx.Bucket(userEmailsBucket).Get([]byte(email)); existing != nil {
			return fmt.Errorf("email is already registered")
		}
		data, err := json.Marshal(user)
		if err != nil {
			return err
		}
		if err := tx.Bucket(usersBucket).Put([]byte(user.ID), data); err != nil {
			return err
		}
		return tx.Bucket(userEmailsBucket).Put([]byte(email), []byte(user.ID))
	})
}

func (s *AuthService) saveUser(user User) error {
	data, err := json.Marshal(user)
	if err != nil {
		return err
	}
	return s.db.Update(func(tx *bbolt.Tx) error {
		return tx.Bucket(usersBucket).Put([]byte(user.ID), data)
	})
}

func (s *AuthService) saveInvite(invite Invite) error {
	data, err := json.Marshal(invite)
	if err != nil {
		return err
	}
	return s.db.Update(func(tx *bbolt.Tx) error {
		return tx.Bucket(invitesBucket).Put([]byte(invite.ID), data)
	})
}

func (s *AuthService) saveCollection(collection Collection) error {
	data, err := json.Marshal(collection)
	if err != nil {
		return err
	}
	return s.db.Update(func(tx *bbolt.Tx) error {
		return tx.Bucket(collectionsBucket).Put([]byte(collection.ID), data)
	})
}

func normalizeEmail(email string) (string, error) {
	email = strings.ToLower(strings.TrimSpace(email))
	if email == "" {
		return "", fmt.Errorf("email is required")
	}
	if _, err := mail.ParseAddress(email); err != nil {
		return "", fmt.Errorf("email is invalid")
	}
	return email, nil
}

func tokenHash(token string) string {
	sum := sha256.Sum256([]byte("warpbox-session:" + token))
	return hex.EncodeToString(sum[:])
}

func HashPassword(password string) string {
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		salt = []byte(randomID(16))[:16]
	}
	hash := argon2.IDKey([]byte(password), salt, 1, 64*1024, 4, 32)
	return "argon2id$v=19$m=65536,t=1,p=4$" + base64.RawStdEncoding.EncodeToString(salt) + "$" + base64.RawStdEncoding.EncodeToString(hash)
}

func VerifyPasswordHash(encoded, password string) bool {
	parts := strings.Split(encoded, "$")
	if len(parts) != 5 || parts[0] != "argon2id" {
		return false
	}
	salt, err := base64.RawStdEncoding.DecodeString(parts[3])
	if err != nil {
		return false
	}
	expected, err := base64.RawStdEncoding.DecodeString(parts[4])
	if err != nil {
		return false
	}
	actual := argon2.IDKey([]byte(password), salt, 1, 64*1024, 4, uint32(len(expected)))
	return subtle.ConstantTimeCompare(actual, expected) == 1
}