kalzu-value-store/auth/auth.go

package auth

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	badger "github.com/dgraph-io/badger/v4"
	"github.com/sirupsen/logrus"

	"kvs/types"
	"kvs/utils"
)

// AuthContext holds authentication information for a request
type AuthContext struct {
	UserUUID string   `json:"user_uuid"`
	Scopes   []string `json:"scopes"`
	Groups   []string `json:"groups"`
}

// AuthService handles authentication operations
type AuthService struct {
	db     *badger.DB
	logger *logrus.Logger
	config *types.Config
}

// NewAuthService creates a new authentication service
func NewAuthService(db *badger.DB, logger *logrus.Logger, config *types.Config) *AuthService {
	return &AuthService{
		db:     db,
		logger: logger,
		config: config,
	}
}

// StoreAPIToken stores an API token in BadgerDB with TTL
func (s *AuthService) StoreAPIToken(tokenString string, userUUID string, scopes []string, expiresAt int64) error {
	tokenHash := utils.HashToken(tokenString)
	
	apiToken := types.APIToken{
		TokenHash: tokenHash,
		UserUUID:  userUUID,
		Scopes:    scopes,
		IssuedAt:  time.Now().Unix(),
		ExpiresAt: expiresAt,
	}

	tokenData, err := json.Marshal(apiToken)
	if err != nil {
		return err
	}

	return s.db.Update(func(txn *badger.Txn) error {
		entry := badger.NewEntry([]byte(TokenStorageKey(tokenHash)), tokenData)
		
		// Set TTL to the token expiration time
		ttl := time.Until(time.Unix(expiresAt, 0))
		if ttl > 0 {
			entry = entry.WithTTL(ttl)
		}
		
		return txn.SetEntry(entry)
	})
}

// GetAPIToken retrieves an API token from BadgerDB by hash
func (s *AuthService) GetAPIToken(tokenHash string) (*types.APIToken, error) {
	var apiToken types.APIToken
	
	err := s.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte(TokenStorageKey(tokenHash)))
		if err != nil {
			return err
		}

		return item.Value(func(val []byte) error {
			return json.Unmarshal(val, &apiToken)
		})
	})

	if err != nil {
		return nil, err
	}

	return &apiToken, nil
}

// ExtractTokenFromHeader extracts the Bearer token from the Authorization header
func ExtractTokenFromHeader(r *http.Request) (string, error) {
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {
		return "", fmt.Errorf("missing authorization header")
	}

	parts := strings.Split(authHeader, " ")
	if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" {
		return "", fmt.Errorf("invalid authorization header format")
	}

	return parts[1], nil
}

// GetUserGroups retrieves all groups that a user belongs to
func (s *AuthService) GetUserGroups(userUUID string) ([]string, error) {
	var user types.User
	err := s.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte(UserStorageKey(userUUID)))
		if err != nil {
			return err
		}

		return item.Value(func(val []byte) error {
			return json.Unmarshal(val, &user)
		})
	})

	if err != nil {
		return nil, err
	}

	return user.Groups, nil
}

// AuthenticateRequest validates the JWT token and returns authentication context
func (s *AuthService) AuthenticateRequest(r *http.Request) (*AuthContext, error) {
	// Extract token from header
	tokenString, err := ExtractTokenFromHeader(r)
	if err != nil {
		return nil, err
	}

	// Validate JWT token
	claims, err := ValidateJWT(tokenString)
	if err != nil {
		return nil, fmt.Errorf("invalid token: %v", err)
	}

	// Verify token exists in our database (not revoked)
	tokenHash := utils.HashToken(tokenString)
	_, err = s.GetAPIToken(tokenHash)
	if err == badger.ErrKeyNotFound {
		return nil, fmt.Errorf("token not found or revoked")
	}
	if err != nil {
		return nil, fmt.Errorf("failed to verify token: %v", err)
	}

	// Get user's groups
	groups, err := s.GetUserGroups(claims.UserUUID)
	if err != nil {
		s.logger.WithError(err).WithField("user_uuid", claims.UserUUID).Warn("Failed to get user groups")
		groups = []string{} // Continue with empty groups on error
	}

	return &AuthContext{
		UserUUID: claims.UserUUID,
		Scopes:   claims.Scopes,
		Groups:   groups,
	}, nil
}

// CheckResourcePermission checks if a user has permission to perform an operation on a resource
func (s *AuthService) CheckResourcePermission(authCtx *AuthContext, resourceKey string, operation string) bool {
	// Get resource metadata
	var metadata types.ResourceMetadata
	err := s.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte(ResourceMetadataKey(resourceKey)))
		if err != nil {
			return err
		}

		return item.Value(func(val []byte) error {
			return json.Unmarshal(val, &metadata)
		})
	})

	// If no metadata exists, use default permissions
	if err == badger.ErrKeyNotFound {
		metadata = types.ResourceMetadata{
			OwnerUUID:   authCtx.UserUUID, // Treat requester as owner for new resources
			GroupUUID:   "",
			Permissions: types.DefaultPermissions,
		}
	} else if err != nil {
		s.logger.WithError(err).WithField("resource_key", resourceKey).Warn("Failed to get resource metadata")
		return false
	}

	// Check user relationship to resource
	isOwner, isGroupMember := CheckUserResourceRelationship(authCtx.UserUUID, &metadata, authCtx.Groups)

	// Check permission
	return CheckPermission(metadata.Permissions, operation, isOwner, isGroupMember)
}

// GetAuthContext retrieves auth context from request context
func GetAuthContext(ctx context.Context) *AuthContext {
	if authCtx, ok := ctx.Value("auth").(*AuthContext); ok {
		return authCtx
	}
	return nil
}

// HasUsers checks if any users exist in the database
func (s *AuthService) HasUsers() (bool, error) {
	var hasUsers bool
	
	err := s.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.PrefetchValues = false // We only need to check if keys exist
		iterator := txn.NewIterator(opts)
		defer iterator.Close()
		
		// Look for any key starting with "user:"
		prefix := []byte("user:")
		for iterator.Seek(prefix); iterator.ValidForPrefix(prefix); iterator.Next() {
			hasUsers = true
			return nil // Found at least one user, can exit early
		}
		
		return nil
	})
	
	return hasUsers, err
}

// StoreResourceMetadata stores or updates resource metadata in BadgerDB
func (s *AuthService) StoreResourceMetadata(path string, metadata *types.ResourceMetadata) error {
	now := time.Now().Unix()
	if metadata.CreatedAt == 0 {
		metadata.CreatedAt = now
	}
	metadata.UpdatedAt = now

	metadataData, err := json.Marshal(metadata)
	if err != nil {
		return err
	}

	return s.db.Update(func(txn *badger.Txn) error {
		return txn.Set([]byte(ResourceMetadataKey(path)), metadataData)
	})
}

// GetResourceMetadata retrieves resource metadata from BadgerDB
func (s *AuthService) GetResourceMetadata(path string) (*types.ResourceMetadata, error) {
	var metadata types.ResourceMetadata
	
	err := s.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte(ResourceMetadataKey(path)))
		if err != nil {
			return err
		}

		return item.Value(func(val []byte) error {
			return json.Unmarshal(val, &metadata)
		})
	})

	if err != nil {
		return nil, err
	}

	return &metadata, nil
}