Refactor and enhance .env loading functionality

Refactored .env loading to support custom configurations via `LoadWithConfig` and implement new utility functions for type-safe environment variable retrieval (`GetBool`, `GetFloat`, `MustGetString`). Introduced depth limit, validation for env variable names, handling of comments, quoted values, and error reporting for invalid formats.

Author: raiesbo

goenv.go

@@ -1,44 +1,90 @@
 package goenv
 
 import (
+	"bufio"
+	"fmt"
 	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
 )
 
 const (
-	EnvFile = ".env"
+	DefaultEnvFile = ".env"
+	MaxDepth       = 10 // Prevent infinite recursion
 )
 
-// getPaths converts an array of directory entries into an array of file paths by concatenating each
-// entry's name with a base path.
-func getPaths(newDirs []os.DirEntry, basePath string) []string {
-	paths := make([]string, len(newDirs))
+type Config struct {
+	EnvFiles    []string
+	MaxDepth    int
+	StopOnFirst bool
+	Prefix      string
+}
 
-	for i, dir := range newDirs {
-		paths[i] = filepath.Join(basePath, dir.Name())
+func DefaultConfig() *Config {
+	return &Config{
+		EnvFiles:    []string{".env"},
+		MaxDepth:    MaxDepth,
+		StopOnFirst: true,
+		Prefix:      "",
 	}
+}
 
-	return paths
+// LoadWithConfig loads environment variables with custom configuration
+func LoadWithConfig(config *Config) error {
+	if config == nil {
+		config = DefaultConfig()
+	}
+
+	return loadFromDirectory(".", config, 0, make(map[string]bool))
 }
 
-// loadVarsFromFile parses an .env file at the given path and loads its variables into the environment.
-func loadVarsFromFile(path string) error {
-	fileData, err := os.ReadFile(path)
+// Load provides backward compatibility with default behavior
+func Load() error {
+	return LoadWithConfig(nil)
+}
+
+// loadFromDirectory recursively searches for .env files with proper error handling and cycle detection
+func loadFromDirectory(dir string, config *Config, depth int, visited map[string]bool) error {
+	if depth > config.MaxDepth {
+		return nil
+	}
+
+	// Get absolute path to detect cycles
+	absPath, err := filepath.Abs(dir)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get absolute path for %s: %w", dir, err)
 	}
 
-	for _, line := range strings.Split(string(fileData), "\n") {
-		if line == "" || strings.HasPrefix(line, "#") {
-			continue
+	if visited[absPath] {
+		return nil // Skip already visited directories
+	}
+	visited[absPath] = true
+
+	// Check for .env files in current directory
+	for _, envFile := range config.EnvFiles {
+		envPath := filepath.Join(dir, envFile)
+		if fileExists(envPath) {
+			if err := loadVarsFromFile(envPath); err != nil {
+				return fmt.Errorf("failed to load %s: %w", envPath, err)
+			}
+			if config.StopOnFirst {
+				return nil
+			}
 		}
+	}
+
+	// Read directory entries
+	entries, err := os.ReadDir(dir)
+	if err != nil {
+		return fmt.Errorf("failed to read directory %s: %w", dir, err)
+	}
 
-		parts := strings.Split(line, "=")
-		if len(parts) == 2 {
-			key, value := strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1])
-			if err := os.Setenv(key, value); err != nil {
+	// Recursively search subdirectories
+	for _, entry := range entries {
+		if entry.IsDir() && !strings.HasPrefix(entry.Name(), ".") {
+			subDir := filepath.Join(dir, entry.Name())
+			if err := loadFromDirectory(subDir, config, depth+1, visited); err != nil {
 				return err
 			}
 		}
@@ -47,56 +93,159 @@ func loadVarsFromFile(path string) error {
 	return nil
 }
 
-// Load recursively scans all directories of a project until it finds a .env file. Once found, it reads
-// the file and loads its values as environment variables.
-func Load() error {
-	dirsQueue := []string{"./"}
+// fileExists checks if a file exists and is not a directory
+func fileExists(path string) bool {
+	info, err := os.Stat(path)
+	return err == nil && !info.IsDir()
+}
+
+// loadVarsFromFile parses an .env file with improved error handling and format support
+func loadVarsFromFile(path string) error {
+	file, err := os.Open(path)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	lineNum := 0
 
-	for len(dirsQueue) > 0 {
-		path := dirsQueue[0]
+	for scanner.Scan() {
+		lineNum++
+		line := strings.TrimSpace(scanner.Text())
 
-		file, err := os.Stat(path)
-		if err != nil {
+		// Skip empty lines and comments
+		if line == "" || strings.HasPrefix(line, "#") {
+			continue
+		}
+
+		// Parse key=value pairs
+		if err := parseAndSetEnvVar(line, path, lineNum); err != nil {
 			return err
 		}
+	}
 
-		if file.IsDir() {
-			children, err := os.ReadDir(path)
-			if err != nil {
-				return err
-			}
-			dirsQueue = append(dirsQueue, getPaths(children, path)...)
-		} else if file.Name() == EnvFile {
-			return loadVarsFromFile(path)
+	return scanner.Err()
+}
+
+// parseAndSetEnvVar parses a single environment variable line
+func parseAndSetEnvVar(line, filePath string, lineNum int) error {
+	parts := strings.SplitN(line, "=", 2)
+	if len(parts) != 2 {
+		return fmt.Errorf("invalid format in %s at line %d: %s", filePath, lineNum, line)
+	}
+
+	key := strings.TrimSpace(parts[0])
+	value := strings.TrimSpace(parts[1])
+
+	// Handle quoted values
+	if len(value) >= 2 {
+		if (strings.HasPrefix(value, "\"") && strings.HasSuffix(value, "\"")) ||
+			(strings.HasPrefix(value, "'") && strings.HasSuffix(value, "'")) {
+			value = value[1 : len(value)-1]
 		}
+	}
 
-		dirsQueue = dirsQueue[1:]
+	// Validate key format
+	if !isValidEnvKey(key) {
+		return fmt.Errorf("invalid environment variable name in %s at line %d: %s", filePath, lineNum, key)
 	}
 
-	return nil
+	return os.Setenv(key, value)
+}
+
+// isValidEnvKey validates environment variable key format
+func isValidEnvKey(key string) bool {
+	if key == "" {
+		return false
+	}
+
+	for i, r := range key {
+		if i == 0 {
+			if !((r >= 'A' && r <= 'Z') || (r >= 'a' && r <= 'z') || r == '_') {
+				return false
+			}
+		} else {
+			if !((r >= 'A' && r <= 'Z') || (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '_') {
+				return false
+			}
+		}
+	}
+
+	return true
 }
 
-// GetString returns the value of an environment variable if it exists; otherwise, it returns the fallback value.
+// GetString returns the value of an environment variable with fallback
 func GetString(key, fallback string) string {
-	val, ok := os.LookupEnv(key)
-	if !ok {
+	if val, exists := os.LookupEnv(key); exists {
+		return val
+	}
+	return fallback
+}
+
+// GetInt returns the integer value of an environment variable with fallback
+func GetInt(key string, fallback int) int {
+	val, exists := os.LookupEnv(key)
+	if !exists {
 		return fallback
 	}
 
-	return val
+	if intVal, err := strconv.Atoi(val); err == nil {
+		return intVal
+	}
+
+	return fallback
 }
 
-// GetInt returns the integer value of an environment variable if it exists; otherwise, it returns the fallback value.
-func GetInt(key string, fallback int) int {
-	val, ok := os.LookupEnv(key)
-	if !ok {
+// GetBool returns the boolean value of an environment variable with fallback
+func GetBool(key string, fallback bool) bool {
+	val, exists := os.LookupEnv(key)
+	if !exists {
 		return fallback
 	}
 
-	intVal, err := strconv.Atoi(val)
-	if err != nil {
+	if boolVal, err := strconv.ParseBool(val); err == nil {
+		return boolVal
+	}
+
+	return fallback
+}
+
+// GetFloat returns the float64 value of an environment variable with fallback
+func GetFloat(key string, fallback float64) float64 {
+	val, exists := os.LookupEnv(key)
+	if !exists {
 		return fallback
 	}
 
-	return intVal
+	if floatVal, err := strconv.ParseFloat(val, 64); err == nil {
+		return floatVal
+	}
+
+	return fallback
+}
+
+// MustGetString returns the value of an environment variable or panics if not found
+func MustGetString(key string) string {
+	val, exists := os.LookupEnv(key)
+	if !exists {
+		panic(fmt.Sprintf("required environment variable %s not found", key))
+	}
+	return val
+}
+
+// LoadFile loads a specific .env file
+func LoadFile(path string) error {
+	return loadVarsFromFile(path)
+}
+
+// Unload removes all environment variables loaded from .env files
+// Note: This is a simplified implementation - tracking loaded vars would be better
+func Unload(keys []string) error {
+	for _, key := range keys {
+		if err := os.Unsetenv(key); err != nil {
+			return err
+		}
+	}
+	return nil
 }