Improve DSN env vairable expansion

Author: pquerna

pkg/database/database.go

@@ -8,6 +8,7 @@ import (
 	"net/url"
 	"os"
 	"regexp"
+	"strings"
 
 	"github.com/conductorone/baton-sql/pkg/database/hdb"
 	"github.com/conductorone/baton-sql/pkg/database/mysql"
@@ -50,8 +51,153 @@ func updateFromEnv(dsn string) (string, error) {
 	return result, nil
 }
 
+// extractPlaceholders replaces ${...} placeholders with unique numeric sentinels
+// and looks up the environment variable values immediately.
+// Numeric sentinels (999000, 999001, etc.) are valid in all URL components:
+// ports (must be numeric), hostnames, userinfo, paths, and query strings.
+// This allows us to parse the URL structure before expanding environment variables.
+// Returns: the string with sentinels, a mapping of sentinel->value, and any error.
+func extractPlaceholders(s string) (string, map[string]string, error) {
+	mapping := make(map[string]string)
+	counter := 0
+	var err error
+
+	result := DSNREnvRegex.ReplaceAllStringFunc(s, func(match string) string {
+		sentinel := fmt.Sprintf("999%03d", counter)
+		varName := match[2 : len(match)-1] // Extract VAR from ${VAR}
+
+		// Look up the environment variable immediately
+		value, exists := os.LookupEnv(varName)
+		if !exists {
+			err = errors.Join(err, fmt.Errorf("environment variable %s is not set", varName))
+			return sentinel // Return sentinel anyway to allow URL parsing for better error messages
+		}
+
+		mapping[sentinel] = value
+		counter++
+		return sentinel
+	})
+
+	if err != nil {
+		return "", nil, err
+	}
+
+	return result, mapping, nil
+}
+
+// expandWithMapping expands sentinels in a string by replacing them with their
+// corresponding values from the mapping.
+func expandWithMapping(s string, mapping map[string]string) string {
+	result := s
+	for sentinel, value := range mapping {
+		result = strings.ReplaceAll(result, sentinel, value)
+	}
+	return result
+}
+
+// expandUserInfo expands environment variable placeholders in the URL's user info component.
+// It handles the special case where an entire "user:password" string might be in a single variable.
+// The url.UserPassword function automatically handles URL encoding of special characters.
+func expandUserInfo(parsedUrl *url.URL, mapping map[string]string) {
+	if parsedUrl.User == nil {
+		return
+	}
+
+	username := parsedUrl.User.Username()
+	password, hasPass := parsedUrl.User.Password()
+
+	// Expand sentinels in username
+	expandedUser := expandWithMapping(username, mapping)
+
+	// Expand sentinels in password
+	var expandedPass string
+	if hasPass {
+		expandedPass = expandWithMapping(password, mapping)
+	}
+
+	// Handle the case where the entire userinfo (user:password) is in a single variable.
+	// For example: ${CREDENTIALS} where CREDENTIALS="admin:p@ss#word"
+	if strings.Contains(expandedUser, ":") && !hasPass {
+		parts := strings.SplitN(expandedUser, ":", 2)
+		expandedUser = parts[0]
+		expandedPass = parts[1]
+		hasPass = true
+	}
+
+	// url.UserPassword automatically handles URL encoding of special characters
+	if hasPass {
+		parsedUrl.User = url.UserPassword(expandedUser, expandedPass)
+	} else {
+		parsedUrl.User = url.User(expandedUser)
+	}
+}
+
+// expandHost expands environment variable placeholders in the URL's host component.
+func expandHost(parsedUrl *url.URL, mapping map[string]string) {
+	if parsedUrl.Host != "" {
+		parsedUrl.Host = expandWithMapping(parsedUrl.Host, mapping)
+	}
+}
+
+// expandPath expands environment variable placeholders in the URL's path component.
+func expandPath(parsedUrl *url.URL, mapping map[string]string) {
+	if parsedUrl.Path != "" {
+		parsedUrl.Path = expandWithMapping(parsedUrl.Path, mapping)
+	}
+}
+
+// expandQuery expands environment variable placeholders in the URL's query string.
+func expandQuery(parsedUrl *url.URL, mapping map[string]string) {
+	if parsedUrl.RawQuery != "" {
+		parsedUrl.RawQuery = expandWithMapping(parsedUrl.RawQuery, mapping)
+	}
+}
+
+// expandFragment expands environment variable placeholders in the URL's fragment.
+func expandFragment(parsedUrl *url.URL, mapping map[string]string) {
+	if parsedUrl.Fragment != "" {
+		parsedUrl.Fragment = expandWithMapping(parsedUrl.Fragment, mapping)
+	}
+}
+
+// expandDSN expands environment variable placeholders in a DSN using a three-phase approach:
+// 1. Replace ${...} with safe sentinel values and lookup env vars
+// 2. Parse the URL structure with sentinels
+// 3. Expand sentinels component-by-component with appropriate encoding
+//
+// This approach ensures that special characters in environment variables (like #, @, :)
+// don't break URL parsing, since they are expanded after the URL structure is established.
+func expandDSN(dsn string) (string, error) {
+	// Phase 1: Replace ${...} with sentinels and lookup env vars
+	sentinelDSN, mapping, err := extractPlaceholders(dsn)
+	if err != nil {
+		return "", err
+	}
+
+	// If there are no placeholders, return as-is
+	if len(mapping) == 0 {
+		return dsn, nil
+	}
+
+	// Phase 2: Parse with sentinels
+	parsedUrl, err := url.Parse(sentinelDSN)
+	if err != nil {
+		return "", fmt.Errorf("invalid DSN structure: %w", err)
+	}
+
+	// Phase 3: Expand by component with appropriate encoding
+	expandUserInfo(parsedUrl, mapping)
+	expandHost(parsedUrl, mapping)
+	expandPath(parsedUrl, mapping)
+	expandQuery(parsedUrl, mapping)
+	expandFragment(parsedUrl, mapping)
+
+	return parsedUrl.String(), nil
+}
+
 func Connect(ctx context.Context, dsn string, user string, password string) (*sql.DB, DbEngine, error) {
-	populatedDSN, err := updateFromEnv(dsn)
+	// Use the new expandDSN function which handles special characters correctly
+	populatedDSN, err := expandDSN(dsn)
 	if err != nil {
 		return nil, Unknown, err
 	}