relspecgo/pkg/assetloader/embed.go

package assetloader

import (
	"encoding/base64"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"unicode/utf8"

	"git.warky.dev/wdevs/relspecgo/pkg/safefile"
	"git.warky.dev/wdevs/relspecgo/pkg/sqllex"
)

const ScriptSourcePathMetadataKey = "source_path"

const DefaultMaxAssetBytes int64 = 64 << 20 // 64 MiB per manifest/embedded asset

// EmbedOptions are trusted caller configuration, not script/model metadata.
// An empty AssetRoot defaults to the SQL file's directory; external reads are rejected.
type EmbedOptions struct {
	AssetRoot    string
	MaxFileBytes int64
}

var (
	embedDirectivePattern = regexp.MustCompile(`(?m)^\s*--\s*@embed:\s*(.+?)\s*$`)
	embedAttrPattern      = regexp.MustCompile(`([a-zA-Z_][a-zA-Z0-9_]*)=("[^"]*"|'[^']*'|\S+)`)
	embedVarPattern       = regexp.MustCompile(`^:[a-zA-Z_][a-zA-Z0-9_]*$`)
)

// ProcessEmbedDirectives expands SQL comments in the form:
//
//	-- @embed: path=... var=:... mode=text|base64
//
// Paths are resolved relative to sqlPath. Text mode embeds a quoted UTF-8 SQL
// string literal. Base64 mode embeds a quoted base64 literal suitable for
// decode(:var, 'base64').
func ProcessEmbedDirectives(sqlPath, sql string) (string, error) {
	return ProcessEmbedDirectivesWithOptions(sqlPath, sql, EmbedOptions{})
}

func embedDirectives(sql string) (result string, directives []string) {
	// Only standalone line comments are directives, and only in SQL code.
	result = sqllex.Walk(sql, nil, func(comment string) string {
		if match := embedDirectivePattern.FindStringSubmatch(comment); match != nil {
			directives = append(directives, match[1])
			return ""
		}
		return comment
	})
	return result, directives
}

func HasEmbedDirectives(sql string) bool {
	_, directives := embedDirectives(sql)
	return len(directives) > 0
}

func ProcessEmbedDirectivesWithOptions(sqlPath, sql string, options EmbedOptions) (string, error) {
	result, directives := embedDirectives(sql)
	if len(directives) == 0 {
		return sql, nil
	}

	if sqlPath == "" {
		return "", fmt.Errorf("sql path is required for embed directives")
	}

	assetRoot := options.AssetRoot
	if assetRoot == "" {
		assetRoot = filepath.Dir(sqlPath)
	}
	absRoot, err := filepath.Abs(assetRoot)
	if err != nil {
		return "", err
	}
	absSQL, err := filepath.Abs(sqlPath)
	if err != nil {
		return "", err
	}
	scriptRel, err := filepath.Rel(absRoot, absSQL)
	if err != nil {
		return "", err
	}
	if _, err := safefile.Relative(filepath.ToSlash(scriptRel)); err != nil {
		return "", fmt.Errorf("SQL source path is outside asset root: %w", err)
	}
	root, err := openAssetRoot(absRoot)
	if err != nil {
		return "", err
	}
	defer root.Close()
	limit := options.MaxFileBytes
	if limit == 0 {
		limit = DefaultMaxAssetBytes
	}
	literals := map[string]string{}
	used := map[string]bool{}
	for i, directive := range directives {
		literal, placeholder, err := embedDirectiveLiteral(root, filepath.Dir(scriptRel), limit, sqlPath, directive, i+1)
		if err != nil {
			return "", err
		}
		if _, exists := literals[placeholder]; exists {
			return "", fmt.Errorf("duplicate embed placeholder %s", placeholder)
		}
		literals[placeholder] = literal
	}
	result = sqllex.Walk(result, func(name string) string {
		if literal, ok := literals[name]; ok {
			used[name] = true
			return literal
		}
		return name
	}, nil)
	for name := range literals {
		if !used[name] {
			return "", fmt.Errorf("%s: embed placeholder %s not found in executable SQL", sqlPath, name)
		}
	}

	return result, nil
}

func embedDirectiveLiteral(root *os.Root, scriptDir string, maxBytes int64, sqlPath, raw string, directiveNumber int) (literal, placeholder string, err error) {
	attrs, err := parseEmbedAttrs(raw)
	if err != nil {
		return "", "", fmt.Errorf("%s embed directive %d: %w", sqlPath, directiveNumber, err)
	}

	pathValue := attrs["path"]
	varValue := attrs["var"]
	modeValue := attrs["mode"]
	if pathValue == "" {
		return "", "", fmt.Errorf("%s embed directive %d: missing path", sqlPath, directiveNumber)
	}
	if !embedVarPattern.MatchString(varValue) {
		return "", "", fmt.Errorf("%s embed directive %d: var must be a named placeholder like :asset", sqlPath, directiveNumber)
	}
	if modeValue != "text" && modeValue != "base64" {
		return "", "", fmt.Errorf("%s embed directive %d: mode must be text or base64", sqlPath, directiveNumber)
	}

	relative, err := safefile.Relative(pathValue)
	if err != nil {
		return "", "", fmt.Errorf("%s embed directive %d: %w", sqlPath, directiveNumber, err)
	}
	resolved := filepath.Join(scriptDir, relative)
	data, err := safefile.Read(root, resolved, maxBytes)
	if err != nil {
		return "", "", fmt.Errorf("%s embed directive %d: reading %s: %w", sqlPath, directiveNumber, resolved, err)
	}

	switch modeValue {
	case "text":
		if !utf8.Valid(data) {
			return "", "", fmt.Errorf("%s embed directive %d: %s is not valid UTF-8", sqlPath, directiveNumber, resolved)
		}
		literal, err := sqlStringLiteral(string(data))
		if err != nil {
			return "", "", fmt.Errorf("%s embed directive %d: %w", sqlPath, directiveNumber, err)
		}
		return literal, varValue, nil
	case "base64":
		literal, err := sqlStringLiteral(base64.StdEncoding.EncodeToString(data))
		if err != nil {
			return "", "", fmt.Errorf("%s embed directive %d: %w", sqlPath, directiveNumber, err)
		}
		return literal, varValue, nil
	default:
		return "", "", fmt.Errorf("%s embed directive %d: mode must be text or base64", sqlPath, directiveNumber)
	}
}

func parseEmbedAttrs(raw string) (map[string]string, error) {
	attrs := map[string]string{}
	matches := embedAttrPattern.FindAllStringSubmatchIndex(raw, -1)
	if len(matches) == 0 {
		return nil, fmt.Errorf("expected path, var, and mode attributes")
	}

	lastEnd := 0
	for _, match := range matches {
		gap := strings.TrimSpace(raw[lastEnd:match[0]])
		if gap != "" {
			return nil, fmt.Errorf("invalid attribute syntax near %q", gap)
		}

		key := raw[match[2]:match[3]]
		value := raw[match[4]:match[5]]
		if _, exists := attrs[key]; exists {
			return nil, fmt.Errorf("duplicate attribute %q", key)
		}
		unquoted, err := unquoteEmbedValue(value)
		if err != nil {
			return nil, fmt.Errorf("invalid %s value: %w", key, err)
		}
		attrs[key] = unquoted
		lastEnd = match[1]
	}
	if tail := strings.TrimSpace(raw[lastEnd:]); tail != "" {
		return nil, fmt.Errorf("invalid attribute syntax near %q", tail)
	}

	for key := range attrs {
		if key != "path" && key != "var" && key != "mode" {
			return nil, fmt.Errorf("unknown attribute %q", key)
		}
	}

	return attrs, nil
}

func unquoteEmbedValue(value string) (string, error) {
	if len(value) < 2 {
		return value, nil
	}
	if value[0] == '"' {
		return strconv.Unquote(value)
	}
	if value[0] == '\'' && value[len(value)-1] == '\'' {
		return value[1 : len(value)-1], nil
	}
	return value, nil
}

func sqlStringLiteral(value string) (string, error) {
	if strings.ContainsRune(value, '\x00') {
		return "", fmt.Errorf("embedded text contains NUL byte")
	}
	prefix := ""
	if strings.Contains(value, "\\") {
		prefix = "E"
		value = strings.ReplaceAll(value, "\\", "\\\\")
	}
	return prefix + "'" + strings.ReplaceAll(value, "'", "''") + "'", nil
}