raft/auth.go

package raft

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base32"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"igit.com/xbase/raft/db"
)

// ==================== Auth Constants ====================

const (
	ActionRead  = "read"
	ActionWrite = "write" // Includes Set, Del
	ActionAdmin = "admin" // Cluster ops

	SystemKeyPrefix   = "system."
	AuthUserPrefix    = "system.user."
	AuthRolePrefix    = "system.role."
	AuthSessionPrefix = "system.session."
	AuthLockPrefix    = "system.lock."
	AuthConfigKey     = "system.config"
)

// ==================== Auth Data Structures ====================

// LoginState tracks login failures (Synced via Raft)
type LoginState struct {
	FailureCount int64 `json:"failure_count"`
	LastFailure  int64 `json:"last_failure"` // Unix timestamp
	LockUntil    int64 `json:"lock_until"`   // Unix timestamp
}

// Constraint defines numeric constraints on values
type Constraint struct {
	Min *float64 `json:"min,omitempty"`
	Max *float64 `json:"max,omitempty"`
}

// Permission defines access to a key pattern
type Permission struct {
	KeyPattern string      `json:"key"`     // Prefix match. "*" means all.
	Actions    []string    `json:"actions"` // "read", "write", "admin"
	Constraint *Constraint `json:"constraint,omitempty"`
}

// Role defines a set of permissions
type Role struct {
	Name        string       `json:"name"`
	Permissions []Permission `json:"permissions"`
	ParentRoles []string     `json:"parent_roles,omitempty"` // Inherit permissions from these roles
}

// User defines a system user
type User struct {
	Username     string `json:"username"`
	PasswordHash string `json:"password_hash"` // SHA256(salt + password)
	Salt         string `json:"salt"`
	Roles        []string `json:"roles"`
	
	// Specific overrides
	AllowPermissions []Permission `json:"allow_permissions"`
	DenyPermissions  []Permission `json:"deny_permissions"` // Higher priority

	Email      string `json:"email"`
	Icon       string `json:"icon"`
	IsOnline   bool   `json:"is_online"`
	MFASecret  string `json:"mfa_secret"`  // Base32 encoded TOTP secret
	MFAEnabled bool   `json:"mfa_enabled"`

	WhitelistIPs []string `json:"whitelist_ips"`
	BlacklistIPs []string `json:"blacklist_ips"`

	// Optimized In-Memory Structure (Not serialized)
	EffectivePermissions []Permission `json:"-"`
}

// Session represents an active login session (Local Only)
type Session struct {
	Token    string `json:"token"`
	Username string `json:"username"`
	LoginIP  string `json:"login_ip"`
	Expiry   int64  `json:"expiry"` // Unix timestamp
	
	// Permission Cache: key+action -> bool
	permCache sync.Map `json:"-"`
}

// AuthConfig defines global auth settings
type AuthConfig struct {
	Enabled bool `json:"enabled"`
}

// ==================== Auth Manager ====================

// AuthManager handles authentication and authorization
type AuthManager struct {
	server *KVServer
	mu     sync.RWMutex

	// In-memory cache
	users       map[string]*User
	roles       map[string]*Role
	sessions    map[string]*Session    // Local sessions
	loginStates map[string]*LoginState // Synced login states
	enabled     bool
}

// NewAuthManager creates a new AuthManager
func NewAuthManager(server *KVServer) *AuthManager {
	return &AuthManager{
		server:      server,
		users:       make(map[string]*User),
		roles:       make(map[string]*Role),
		sessions:    make(map[string]*Session),
		loginStates: make(map[string]*LoginState),
		enabled:     false,
	}
}

// HashPassword hashes a password with salt (Exported for tools)
func HashPassword(password, salt string) string {
	hash := sha256.Sum256([]byte(salt + password))
	return hex.EncodeToString(hash[:])
}

// Internal helper
func hashPassword(password, salt string) string {
	return HashPassword(password, salt)
}

// Helper to generate token
func generateToken() string {
	// In a real system use crypto/rand, but here we use math/rand seeded in main or just simple timestamp mix for example
	// Since we can't easily access crypto/rand without import and "no external lib" usually implies standard lib only,
	// crypto/rand IS standard lib.
	h := sha256.Sum256([]byte(fmt.Sprintf("%d", time.Now().UnixNano())))
	return hex.EncodeToString(h[:])
}

// ValidateTOTP validates a TOTP code against a secret
func ValidateTOTP(secret string, code string) bool {
	// Decode base32 secret
	key, err := base32.StdEncoding.DecodeString(strings.ToUpper(secret))
	if err != nil {
		// Try with padding if failed? Or assume standard encoding.
		// Some libs produce non-padded.
		if padding := len(secret) % 8; padding != 0 {
			secret += strings.Repeat("=", 8-padding)
			key, err = base32.StdEncoding.DecodeString(strings.ToUpper(secret))
		}
		if err != nil {
			return false
		}
	}

	// Current time steps (30s window)
	// Check current, prev, and next window to allow for slight skew
	now := time.Now().Unix()
	return checkTOTP(key, now, code) || checkTOTP(key, now-30, code) || checkTOTP(key, now+30, code)
}

// GenerateMFASecret generates a new random Base32 secret for TOTP
func GenerateMFASecret() (string, error) {
	b := make([]byte, 20) // 160 bits is recommended
	_, err := rand.Read(b)
	if err != nil {
		return "", err
	}
	return base32.StdEncoding.EncodeToString(b), nil
}

// GenerateOTPAuthURL generates a URL for QR code generation
func GenerateOTPAuthURL(username, secret, issuer string) string {
	// otpauth://totp/Issuer:User?secret=SECRET&issuer=Issuer
	return fmt.Sprintf("otpauth://totp/%s:%s?secret=%s&issuer=%s", 
		issuer, username, secret, issuer)
}

// validateTOTP kept for backward compatibility if used internally, forwarding to exported
func validateTOTP(secret string, code string) bool {
	return ValidateTOTP(secret, code)
}

func checkTOTP(key []byte, timestamp int64, code string) bool {
	step := timestamp / 30
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, uint64(step))

	h := hmac.New(sha1.New, key)
	h.Write(buf)
	sum := h.Sum(nil)

	offset := sum[len(sum)-1] & 0xf
	val := int64(((int(sum[offset]) & 0x7f) << 24) |
		((int(sum[offset+1]) & 0xff) << 16) |
		((int(sum[offset+2]) & 0xff) << 8) |
		(int(sum[offset+3]) & 0xff))

	otp := val % 1000000
	expected := fmt.Sprintf("%06d", otp)
	return expected == code
}

// UpdateCache updates the internal cache based on key-value changes
// Called by the state machine when applying logs
func (am *AuthManager) UpdateCache(key, value string, isDelete bool) {
	am.mu.Lock()
	defer am.mu.Unlock()

	if key == AuthConfigKey {
		if isDelete {
			am.enabled = false
		} else {
			var config AuthConfig
			if err := json.Unmarshal([]byte(value), &config); err == nil {
				am.enabled = config.Enabled
			}
		}
		return
	}

	if strings.HasPrefix(key, AuthUserPrefix) {
		username := strings.TrimPrefix(key, AuthUserPrefix)
		if isDelete {
			delete(am.users, username)
		} else {
			var user User
			if err := json.Unmarshal([]byte(value), &user); err == nil {
				// Rebuild effective permissions immediately
				am.rebuildEffectivePermissions(&user)
				am.users[username] = &user
				am.server.Raft.config.Logger.Info("Updated user cache for %s", username)
			} else {
				am.server.Raft.config.Logger.Error("Failed to unmarshal user %s: %v", username, err)
			}
		}
		return
	}

	if strings.HasPrefix(key, AuthRolePrefix) {
		rolename := strings.TrimPrefix(key, AuthRolePrefix)
		if isDelete {
			delete(am.roles, rolename)
		} else {
			var role Role
			if err := json.Unmarshal([]byte(value), &role); err == nil {
				am.roles[rolename] = &role
			}
		}
		// When roles change, we must rebuild ALL users' effective permissions
		// This might be expensive if there are many users, but for "backend system" it's fine.
		for _, u := range am.users {
			am.rebuildEffectivePermissions(u)
		}
		return
	}

	if strings.HasPrefix(key, AuthLockPrefix) {
		username := strings.TrimPrefix(key, AuthLockPrefix)
		if isDelete {
			delete(am.loginStates, username)
		} else {
			var state LoginState
			if err := json.Unmarshal([]byte(value), &state); err == nil {
				am.loginStates[username] = &state
			}
		}
		return
	}

	// Session is local only, ignore AuthSessionPrefix from Raft (unless we decide to sync sessions later)
	if strings.HasPrefix(key, AuthSessionPrefix) {
		return
	}
	
	// Invalidate permission cache on any auth change
	// For simplicity, we clear all local session caches when Users/Roles change
	// (We could be more granular but this guarantees safety)
	for _, s := range am.sessions {
		s.permCache.Range(func(key, value interface{}) bool {
			s.permCache.Delete(key)
			return true
		})
	}
}

// rebuildEffectivePermissions flattens roles into User.EffectivePermissions
// This is the "Special Memory Design" to ensure high performance
func (am *AuthManager) rebuildEffectivePermissions(user *User) {
	// Optimization: Pre-allocate slice?
	var effective []Permission

	// 1. User Deny Permissions (Keep separate or handle in check? Logic says Deny > Allow)
	// We'll handle Deny explicitly in CheckPermission, but for Allow, we merge Roles.
	// Actually, the prompt says "efficient... prevent traversing role keys for every access".
	// So we should collect all ALLOW rules.
	
	// Collect all allow permissions
	// Order: User Allow -> Role Allow -> Parent Role Allow
	
	// User Allow
	effective = append(effective, user.AllowPermissions...)

	// Roles (Recursive collection with loop detection)
	visited := make(map[string]bool)
	var collectRoles func(roles []string)
	collectRoles = func(roles []string) {
		for _, rName := range roles {
			if visited[rName] {
				continue
			}
			visited[rName] = true
			if role, ok := am.roles[rName]; ok {
				effective = append(effective, role.Permissions...)
				collectRoles(role.ParentRoles)
			}
		}
	}
	collectRoles(user.Roles)
	
	user.EffectivePermissions = effective
	// Debug logging for root or specific users could go here
	if user.Username == "root" {
		am.server.Raft.config.Logger.Info("Rebuilt effective permissions for root: %d permissions", len(effective))
		for _, p := range effective {
			am.server.Raft.config.Logger.Debug("Root Perm: %s %v", p.KeyPattern, p.Actions)
		}
	}
}

// IsSubset checks if the requested permissions are a subset of the grantor's permissions.
// This is used for fine-grained delegation: you can only grant what you have.
func (am *AuthManager) IsSubset(grantorToken string, requestedPerms []Permission) bool {
	am.mu.RLock()
	grantorSession, ok := am.sessions[grantorToken]
	am.mu.RUnlock()

	// 1. Internal System Bypass
	if grantorToken == "SYSTEM_INTERNAL" {
		return true
	}

	if !ok {
		return false
	}

	am.mu.RLock()
	grantor, ok := am.users[grantorSession.Username]
	am.mu.RUnlock()

	if !ok {
		return false
	}

	// 2. Check each requested permission against grantor's effective permissions
	for _, req := range requestedPerms {
		allowed := false
		
		// If grantor has Deny for this, they definitely can't grant it (even if they have Allow)
		// Logic: Deny > Allow. If I am denied "secret.*", I cannot grant "secret.doc".
		// Check Grantor Deny
		for _, deny := range grantor.DenyPermissions {
			// If deny covers req, then grantor effectively doesn't have it.
			// deny covers req if deny pattern matches req pattern (or is broader) AND deny actions cover req actions.
			// Actually, Deny check is usually done at access time.
			// But here we check "Capability to Grant".
			// Simplified: If Grantor allows it, we assume they can grant it.
			// But we must respect Deny.
			if matchKey(deny.KeyPattern, req.KeyPattern) || (req.KeyPattern != "*" && matchKey(deny.KeyPattern, strings.TrimSuffix(req.KeyPattern, "*"))) {
				// Potential overlap. If actions overlap, fail.
				if hasCommonAction(deny.Actions, req.Actions) {
					return false
				}
			}
		}

		// Check Grantor Allow (Effective)
		for _, grant := range grantor.EffectivePermissions {
			// Grantor must have a permission that COVERS the requested permission
			// 1. Key Pattern: grant.Key must cover req.Key
			//    e.g. grant="user.*" covers req="user.alice"
			//    e.g. grant="*" covers everything
			if matchKey(grant.KeyPattern, req.KeyPattern) || (req.KeyPattern != "*" && matchKey(grant.KeyPattern, strings.TrimSuffix(req.KeyPattern, "*"))) {
				// 2. Actions: grant.Actions must be superset of req.Actions
				if isActionSubset(req.Actions, grant.Actions) {
					// 3. Constraints: Grant constraints must be looser or equal to Req constraints
					// (If Grant allows 0-100, Req can be 0-50. If Grant allows 0-50, Req cannot be 0-100)
					if isConstraintSubset(req.Constraint, grant.Constraint) {
						allowed = true
						break
					}
				}
			}
		}
		if !allowed {
			return false
		}
	}
	return true
}

// Helpers for Subset Check

func hasCommonAction(a1, a2 []string) bool {
	for _, x := range a1 {
		for _, y := range a2 {
			if x == y || x == "*" || y == "*" {
				return true
			}
		}
	}
	return false
}

func isActionSubset(sub, super []string) bool {
	// Check if every action in 'sub' is present in 'super' (or super has "*")
	for _, s := range sub {
		found := false
		for _, S := range super {
			if S == "*" || S == s {
				found = true
				break
			}
		}
		if !found {
			return false
		}
	}
	return true
}

func isConstraintSubset(sub, super *Constraint) bool {
	// If super has no constraint, it allows everything (subset is valid whatever it is)
	if super == nil {
		return true
	}
	// If super has constraint but sub doesn't, sub is broader -> Fail
	if sub == nil {
		return false
	}
	
	// Check Min: Super.Min <= Sub.Min
	if super.Min != nil {
		if sub.Min == nil || *sub.Min < *super.Min {
			return false
		}
	}
	// Check Max: Super.Max >= Sub.Max
	if super.Max != nil {
		if sub.Max == nil || *sub.Max > *super.Max {
			return false
		}
	}
	return true
}

// IsEnabled checks if auth is enabled
func (am *AuthManager) IsEnabled() bool {
	am.mu.RLock()
	defer am.mu.RUnlock()
	return am.enabled
}

// CheckPermission checks if a user has permission to perform an action on a key
func (am *AuthManager) CheckPermission(token string, key string, action string, value string) error {
	am.mu.RLock()
	enabled := am.enabled
	am.mu.RUnlock()

	if !enabled {
		return nil
	}
	
	// Internal system bypass
	if token == "SYSTEM_INTERNAL" {
		return nil
	}

	am.mu.RLock()
	session, ok := am.sessions[token]
	am.mu.RUnlock()

	if !ok {
		return fmt.Errorf("invalid or expired session")
	}

	if time.Now().Unix() > session.Expiry {
		return fmt.Errorf("session expired")
	}

	// Check Session Local Cache
	// Cache key: "action:key" (value is harder to cache due to constraints, but if value is empty/read, we can cache)
	// For writes with value constraints, caching is risky unless we include value in cache key or don't cache constraint checks.
	// For READS (action="read"), value is empty, so we can cache.
	// For WRITES without value (Del), or simple sets, we can cache if no constraints.
	// Let's safe cache only for Read or if we verify constraint logic is stable.
	// To be safe and simple: Cache the result of (key, action). If constraints exist, the checkPerms function will fail if value invalid.
	// But if we cache "true", we skip checkPerms.
	// So we can only cache if the permission granted did NOT depend on value constraints.
	// That's complex.
	// Simplified: Cache only "read" and "admin" actions?
	cacheKey := action + ":" + key
	if action == ActionRead {
		if val, hit := session.permCache.Load(cacheKey); hit {
			if allowed, ok := val.(bool); ok {
				if allowed {
					return nil
				} else {
					return fmt.Errorf("permission denied (cached)")
				}
			}
		}
	}

	am.mu.RLock()
	user, userOk := am.users[session.Username]
	am.mu.RUnlock()

	if !userOk {
		return fmt.Errorf("user not found")
	}

	err := am.checkUserPermission(user, key, action, value)
	
	// Update Cache for Read
	if action == ActionRead {
		session.permCache.Store(cacheKey, err == nil)
	}
	
	return err
}

func (am *AuthManager) checkUserPermission(user *User, key, action, value string) error {
	// 1. Check Deny Permissions (User level - Highest Priority)
	for _, perm := range user.DenyPermissions {
		if matchKey(perm.KeyPattern, key) && hasAction(perm.Actions, action) {
			return fmt.Errorf("permission denied (explicit deny)")
		}
	}

	// 2. Check Effective Permissions (Flattened Allow Rules)
	// This uses the pre-calculated list, avoiding role traversal
	if checkPerms(user.EffectivePermissions, key, action, value) {
		return nil
	}

	return fmt.Errorf("permission denied")
}

// checkRole and checkRoleFullAccess are deprecated by effective permissions model
// but we keep checkRoleFullAccess logic optimized or use effective permissions?
// Effective permissions might be large, but "HasFullAccess" usually looks for "*".
// We can scan EffectivePermissions for "*".

func matchKey(pattern, key string) bool {
	if pattern == "*" {
		return true
	}
	if pattern == key {
		return true
	}
	// Prefix match: "user.asin" matches "user.asin.china"
	if strings.HasSuffix(pattern, "*") {
		prefix := strings.TrimSuffix(pattern, "*")
		return strings.HasPrefix(key, prefix)
	}
	// Hierarchy match: pattern "a.b" matches "a.b.c"
	if strings.HasPrefix(key, pattern+".") {
		return true
	}
	return false
}

func hasAction(actions []string, target string) bool {
	for _, a := range actions {
		if a == target || a == "*" {
			return true
		}
	}
	return false
}

func checkPerms(perms []Permission, key, action, value string) bool {
	for _, perm := range perms {
		if matchKey(perm.KeyPattern, key) && hasAction(perm.Actions, action) {
			// Check Value Constraints if action is write
			if action == ActionWrite && perm.Constraint != nil && value != "" {
				// Try to parse value as float
				val, err := strconv.ParseFloat(value, 64)
				if err == nil {
					if perm.Constraint.Min != nil && val < *perm.Constraint.Min {
						return false
					}
					if perm.Constraint.Max != nil && val > *perm.Constraint.Max {
						return false
					}
				} else {
					// If value is not a number but constraint exists, what to do?
					// Strict mode: fail.
					return false
				}
			}
			return true
		}
	}
	return false
}

// Login authenticates a user and returns a token
// Modified: Sessions are LOCAL, Lock State is CLUSTER-WIDE
func (am *AuthManager) Login(username, password, totpCode, ip string) (string, error) {
	am.mu.RLock()
	if !am.enabled {
		am.mu.RUnlock()
		return "", fmt.Errorf("auth not enabled")
	}
	user, ok := am.users[username]
	lockState := am.loginStates[username]
	am.mu.RUnlock()

	if !ok {
		return "", fmt.Errorf("invalid credentials")
	}

	// 1. Check Lock State
	if lockState != nil && time.Now().Unix() < lockState.LockUntil {
		return "", fmt.Errorf("account locked, please try again in %d seconds", lockState.LockUntil-time.Now().Unix())
	}

	// 2. Check IP (Allow/Deny)
	if len(user.WhitelistIPs) > 0 {
		allowed := false
		for _, w := range user.WhitelistIPs {
			if w == ip {
				allowed = true
				break
			}
		}
		if !allowed {
			return "", fmt.Errorf("ip not allowed")
		}
	}
	if len(user.BlacklistIPs) > 0 {
		for _, b := range user.BlacklistIPs {
			if b == ip {
				return "", fmt.Errorf("ip blacklisted")
			}
		}
	}

	// 3. Verify Password
	pwdValid := hashPassword(password, user.Salt) == user.PasswordHash
	if !pwdValid {
		// Handle Failure Logic (Sync to Cluster)
		go am.handleLoginFailure(username, lockState)
		return "", fmt.Errorf("invalid credentials")
	}

	// 4. Verify TOTP
	if user.MFAEnabled {
		if totpCode == "" {
			return "", fmt.Errorf("mfa code required")
		}
		if !validateTOTP(user.MFASecret, totpCode) {
			// TOTP failure counts as login failure
			go am.handleLoginFailure(username, lockState)
			return "", fmt.Errorf("invalid mfa code")
		}
	}

	// 5. Login Success - Clear Lock if needed
	if lockState != nil && lockState.FailureCount > 0 {
		go am.handleLoginSuccess(username)
	}

	// 6. Create Session (Local Only)
	token := generateToken()
	session := &Session{
		Token:    token,
		Username: username,
		LoginIP:  ip,
		Expiry:   time.Now().Add(24 * time.Hour).Unix(),
	}

	am.mu.Lock()
	am.sessions[token] = session
	am.mu.Unlock()

	return token, nil
}

func (am *AuthManager) handleLoginFailure(username string, state *LoginState) {
	// Calculate new state
	now := time.Now().Unix()
	newState := LoginState{
		LastFailure: now,
	}
	if state != nil {
		newState.FailureCount = state.FailureCount + 1
	} else {
		newState.FailureCount = 1
	}

	// Logic: 3 consecutive errors -> lock 1 min. Each subsequent error -> lock 1 min.
	if newState.FailureCount >= 3 {
		newState.LockUntil = now + 60 // 1 minute lock
	}

	data, _ := json.Marshal(newState)
	// We use Set (Async) because we don't want to block login failure response too long,
	// but SetSync is safer to ensure next attempt sees the lock.
	// Given the requirement "lock immediately", we should probably use SetSync or just accept slight delay.
	// Using Set for now to avoid blocking the user response too much, but for strictness SetSync is better.
	am.server.Set(AuthLockPrefix+username, string(data))
}

func (am *AuthManager) handleLoginSuccess(username string) {
	// Clear failure count
	// We could delete the key, or set failure count to 0
	am.server.Del(AuthLockPrefix + username)
}

// Logout invalidates a session (Local Only)
func (am *AuthManager) Logout(token string) error {
	am.mu.Lock()
	defer am.mu.Unlock()
	delete(am.sessions, token)
	return nil
}

// GetSession returns session info
func (am *AuthManager) GetSession(token string) (*Session, error) {
	am.mu.RLock()
	defer am.mu.RUnlock()

	session, ok := am.sessions[token]
	if !ok {
		return nil, fmt.Errorf("invalid session")
	}
	if time.Now().Unix() > session.Expiry {
		return nil, fmt.Errorf("session expired")
	}
	return session, nil
}

// GetRoleEffectivePermissions returns all permissions a role grants (including inherited)
func (am *AuthManager) GetRoleEffectivePermissions(roleName string) ([]Permission, error) {
	am.mu.RLock()
	defer am.mu.RUnlock()
	
	role, ok := am.roles[roleName]
	if !ok {
		return nil, fmt.Errorf("role not found")
	}
	
	var effective []Permission
	visited := make(map[string]bool)
	
	var collect func(r *Role)
	collect = func(r *Role) {
		if visited[r.Name] {
			return
		}
		visited[r.Name] = true
		effective = append(effective, r.Permissions...)
		
		for _, pName := range r.ParentRoles {
			if parent, ok := am.roles[pName]; ok {
				collect(parent)
			}
		}
	}
	
	collect(role)
	return effective, nil
}

// GetUser returns user info (Public for internal use)
func (am *AuthManager) GetUser(username string) (*User, error) {
	am.mu.RLock()
	defer am.mu.RUnlock()
	user, ok := am.users[username]
	if !ok {
		return nil, fmt.Errorf("user not found")
	}
	return user, nil
}

// RegisterUser creates a new user
func (am *AuthManager) RegisterUser(username, password string, roles []string) error {
	am.mu.RLock()
	if _, exists := am.users[username]; exists {
		am.mu.RUnlock()
		return fmt.Errorf("user already exists")
	}
	am.mu.RUnlock()

	salt := fmt.Sprintf("%d", time.Now().UnixNano()) // Simple salt
	user := User{
		Username:     username,
		Salt:         salt,
		PasswordHash: hashPassword(password, salt),
		Roles:        roles,
	}
	data, _ := json.Marshal(user)
	// Use SetSync to ensure user is created before returning
	return am.server.SetSync(AuthUserPrefix+username, string(data))
}

// ChangePassword updates a user's password
func (am *AuthManager) ChangePassword(username, newPassword string) error {
	am.mu.RLock()
	user, exists := am.users[username]
	am.mu.RUnlock()
	
	if !exists {
		return fmt.Errorf("user not found")
	}
	
	newUser := *user // Shallow copy to modify
	newUser.Salt = fmt.Sprintf("%d", time.Now().UnixNano())
	newUser.PasswordHash = hashPassword(newPassword, newUser.Salt)
	
	data, _ := json.Marshal(newUser)
	// Use SetSync to ensure password is updated before returning
	return am.server.SetSync(AuthUserPrefix+username, string(data))
}

// UpdateUser updates generic fields of a user (including roles)
func (am *AuthManager) UpdateUser(user User) error {
	am.mu.RLock()
	_, exists := am.users[user.Username]
	am.mu.RUnlock()
	
	if !exists {
		return fmt.Errorf("user not found")
	}
	
	// Ensure we don't overwrite passwordhash/salt blindly if they are empty
	// But usually UpdateUser is called with a fully populated user object from GetUser + mods
	// So we just save it.
	
	data, _ := json.Marshal(user)
	return am.server.SetSync(AuthUserPrefix+user.Username, string(data))
}

// CreateRole creates a new role
func (am *AuthManager) CreateRole(name string) error {
	am.mu.RLock()
	if _, exists := am.roles[name]; exists {
		am.mu.RUnlock()
		return fmt.Errorf("role already exists")
	}
	am.mu.RUnlock()

	role := Role{Name: name}
	data, _ := json.Marshal(role)
	return am.server.SetSync(AuthRolePrefix+name, string(data))
}

// DeleteRole deletes a role
func (am *AuthManager) DeleteRole(name string) error {
	am.mu.RLock()
	if _, exists := am.roles[name]; !exists {
		am.mu.RUnlock()
		return fmt.Errorf("role not found")
	}
	am.mu.RUnlock()
	return am.server.DelSync(AuthRolePrefix + name)
}

// UpdateRole updates an existing role
func (am *AuthManager) UpdateRole(role Role) error {
	am.mu.RLock()
	if _, exists := am.roles[role.Name]; !exists {
		am.mu.RUnlock()
		return fmt.Errorf("role not found")
	}
	am.mu.RUnlock()
	
	data, _ := json.Marshal(role)
	return am.server.SetSync(AuthRolePrefix+role.Name, string(data))
}

// GetRole returns a role definition
func (am *AuthManager) GetRole(name string) (*Role, error) {
	am.mu.RLock()
	defer am.mu.RUnlock()
	role, ok := am.roles[name]
	if !ok {
		return nil, fmt.Errorf("role not found")
	}
	return role, nil
}

// ListUsers lists all users
func (am *AuthManager) ListUsers() []*User {
	am.mu.RLock()
	defer am.mu.RUnlock()
	users := make([]*User, 0, len(am.users))
	for _, u := range am.users {
		// Calculate IsOnline based on active sessions
		isOnline := false
		for _, s := range am.sessions {
			if s.Username == u.Username {
				if time.Now().Unix() <= s.Expiry {
					isOnline = true
					break
				}
			}
		}
		
		// Create a copy to avoid race conditions and not modify cache
		userCopy := *u
		userCopy.IsOnline = isOnline
		users = append(users, &userCopy)
	}
	return users
}

// ListRoles lists all roles
func (am *AuthManager) ListRoles() []*Role {
	am.mu.RLock()
	defer am.mu.RUnlock()
	roles := make([]*Role, 0, len(am.roles))
	for _, r := range am.roles {
		roles = append(roles, r)
	}
	return roles
}

// HasFullAccess checks if the user has "*" permission on all keys (superuser)
// allowing optimizations like pushing limits to DB engine.
func (am *AuthManager) HasFullAccess(token string) bool {
	am.mu.RLock()
	defer am.mu.RUnlock()

	if !am.enabled {
		return true
	}
	if token == "SYSTEM_INTERNAL" {
		return true
	}

	session, ok := am.sessions[token]
	if !ok || time.Now().Unix() > session.Expiry {
		return false
	}

	user, ok := am.users[session.Username]
	if !ok {
		return false
	}

	// Must not have any deny permissions that could block "*"
	// Actually if Deny is empty, and Allow has "*", it's full access.
	// If Deny has something, we can't assume full access blindly.
	if len(user.DenyPermissions) > 0 {
		return false
	}

	// Check Allow
	// Use EffectivePermissions instead of looping user.AllowPermissions
	for _, perm := range user.EffectivePermissions {
		if perm.KeyPattern == "*" && hasAction(perm.Actions, "*") {
			return true
		}
	}
	
	// Check Roles - No longer needed as EffectivePermissions includes them
	// BUT, if user has many roles, flattening helps.
	
	return false
}

func (am *AuthManager) checkRoleFullAccess(roleName string, visited map[string]bool) bool {
	// Deprecated / Unused in new model
	return false
}

// Admin helpers to bootstrap
func (am *AuthManager) CreateRootUser(password string) error {
	salt := "somesalt" // In production use random
	user := User{
		Username:     "root",
		Salt:         salt,
		PasswordHash: hashPassword(password, salt),
		AllowPermissions: []Permission{
			{KeyPattern: "*", Actions: []string{"*"}},
		},
	}
	data, _ := json.Marshal(user)
	return am.server.Set(AuthUserPrefix+"root", string(data))
}

func (am *AuthManager) EnableAuth() error {
	config := AuthConfig{Enabled: true}
	data, _ := json.Marshal(config)
	return am.server.Set(AuthConfigKey, string(data))
}

func (am *AuthManager) DisableAuth() error {
	// Only admin can do this via normal SetAuthenticated
	config := AuthConfig{Enabled: false}
	data, _ := json.Marshal(config)
	return am.server.Set(AuthConfigKey, string(data))
}

// LoadFromDB loads auth data from existing DB
func (am *AuthManager) LoadFromDB() error {
	prefix := SystemKeyPrefix
	
	// We use the DB engine's index to find all system keys
	am.server.DB.Index.WalkPrefix(prefix, func(key string, entry db.IndexEntry) bool {
		// Read value from storage
		// We need to access Storage via Engine.
		// Engine.Storage is exported.
		val, err := am.server.DB.Storage.ReadValue(entry.ValueOffset)
		if err != nil {
			// Skip corrupted or missing values
			return true
		}
		
		// Update Cache
		// Note: We are essentially replaying the state into the cache
		am.UpdateCache(key, val, false)
		return true
	})
	
	return nil
}