config.go

// Copyright 2026 Columnar Technologies Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package config

import (
	"archive/tar"
	"compress/gzip"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"maps"
	"os"
	"path/filepath"
	"runtime"
	"slices"
	"strings"

	"github.com/pelletier/go-toml/v2"
)

const adbcEnvVar = "ADBC_DRIVER_PATH"

var platformTuple string

var ErrInvalidManifest = errors.New("invalid manifest")

func init() {
	os := runtime.GOOS
	switch os {
	case "darwin":
		os = "macos"
	case "windows", "freebsd", "linux", "openbsd":
	default:
		os = "unknown"
	}

	arch := runtime.GOARCH
	switch arch {
	case "386":
		arch = "x86"
	case "ppc":
		arch = "powerpc"
	case "ppc64":
		arch = "powerpc64"
	case "ppc64le":
		arch = "powerpc64le"
	case "wasm":
		arch = "wasm64"
	default:
	}

	platformTuple = os + "_" + arch
}

func PlatformTuple() string {
	return platformTuple
}

type Config struct {
	Level    ConfigLevel
	Location string
	Drivers  map[string]DriverInfo
	Exists   bool
	Err      error
}

type ConfigLevel int

const (
	ConfigUnknown ConfigLevel = iota
	ConfigSystem
	ConfigUser
	ConfigEnv
)

func (c ConfigLevel) String() string {
	switch c {
	case ConfigSystem:
		return "system"
	case ConfigUser:
		return "user"
	case ConfigEnv:
		return "env"
	default:
		return "unknown"
	}
}

var validLevelArgConfigValues = []ConfigLevel{ConfigUser, ConfigSystem}

func (c *ConfigLevel) UnmarshalText(b []byte) error {
	switch strings.ToLower(strings.TrimSpace(string(b))) {
	case "system":
		*c = ConfigSystem
	case "user":
		*c = ConfigUser
	default:
		names := make([]string, len(validLevelArgConfigValues))
		for i, lvl := range validLevelArgConfigValues {
			names[i] = lvl.String()
		}
		return fmt.Errorf("unknown config level %q, valid values are: %s", string(b), strings.Join(names, ", "))
	}
	return nil
}

func EnsureLocation(cfg Config) (string, error) {
	loc := cfg.Location
	if cfg.Level == ConfigEnv {
		list := splitConfigList(loc)
		if len(list) == 0 {
			return "", errors.New("ADBC_DRIVER_PATH is empty, must be set to valid path to use")
		}
		loc = list[0]
	}

	if _, err := os.Stat(loc); err != nil {
		if errors.Is(err, fs.ErrNotExist) {
			if err := os.MkdirAll(loc, 0o755); err != nil {
				return "", fmt.Errorf("failed to create config directory %s: %w", loc, err)
			}
			// Create a .gitignore with "*" in it.
			//
			// This depends on the if block it's in: We only want to create this file
			// if we also had to create `loc` in the same call.
			if cfg.Level == ConfigEnv {
				gitignorePath := filepath.Join(loc, ".gitignore")
				_ = os.WriteFile(gitignorePath, []byte("*\n"), 0o644)
			}
		} else {
			return "", fmt.Errorf("failed to stat config directory %s: %w", loc, err)
		}
	}

	return loc, nil
}

func loadConfig(lvl ConfigLevel) Config {
	cfg := Config{Level: lvl, Location: lvl.ConfigLocation()}
	if cfg.Location == "" {
		return cfg
	}

	if lvl == ConfigEnv {
		pathList := filepath.SplitList(cfg.Location)
		slices.Reverse(pathList)
		finalDrivers := make(map[string]DriverInfo)
		for _, p := range pathList {
			drivers, err := loadDir(p)
			if err != nil && !errors.Is(err, fs.ErrNotExist) {
				cfg.Err = fmt.Errorf("error loading drivers from %s: %w", p, err)
				return cfg
			}
			maps.Copy(finalDrivers, drivers)
		}
		cfg.Exists, cfg.Drivers = len(finalDrivers) > 0, finalDrivers
		return cfg
	}

	drivers, err := loadDir(cfg.Location)
	if err != nil {
		if !errors.Is(err, fs.ErrNotExist) {
			cfg.Err = fmt.Errorf("error loading drivers from %s: %w", cfg.Location, err)
		}
		return cfg
	}

	cfg.Exists, cfg.Drivers = true, drivers
	return cfg
}

// FindDriverConfigsIn lists installed drivers from an explicit location without
// consulting environment variables, so it is safe for request-scoped concurrent
// callers. A location containing the OS list separator is treated as a path list
// (matching the env config level).
func FindDriverConfigsIn(location string) []DriverInfo {
	if location == "" {
		return nil
	}
	paths := splitConfigList(location)
	slices.Reverse(paths)
	merged := make(map[string]DriverInfo)
	for _, p := range paths {
		if p == "" {
			continue
		}
		drivers, err := loadDir(p)
		if err != nil {
			continue
		}
		maps.Copy(merged, drivers)
	}
	return slices.Collect(maps.Values(merged))
}

func getEnvConfigDir() string {
	envConfigLoc := filepath.SplitList(os.Getenv(adbcEnvVar))
	if venv := os.Getenv("VIRTUAL_ENV"); venv != "" {
		envConfigLoc = append(envConfigLoc, filepath.Join(venv, "etc", "adbc", "drivers"))
	}

	if conda := os.Getenv("CONDA_PREFIX"); conda != "" {
		envConfigLoc = append(envConfigLoc, filepath.Join(conda, "etc", "adbc", "drivers"))
	}

	envConfigLoc = slices.DeleteFunc(envConfigLoc, func(s string) bool {
		return s == ""
	})

	return strings.Join(envConfigLoc, string(filepath.ListSeparator))
}

func InstallDriver(cfg Config, shortName string, downloaded *os.File, platform string) (Manifest, error) {
	var (
		loc string
		err error
	)
	if loc, err = EnsureLocation(cfg); err != nil {
		return Manifest{}, fmt.Errorf("could not ensure config location: %w", err)
	}
	base := strings.TrimSuffix(strings.TrimSuffix(filepath.Base(downloaded.Name()), ".tar.gz"), ".tgz")
	finalDir := filepath.Join(loc, base)

	if err := os.MkdirAll(finalDir, 0o755); err != nil {
		return Manifest{}, fmt.Errorf("failed to create driver directory %s: %w", finalDir, err)
	}

	manifest, err := InflateTarball(downloaded, finalDir)
	if err != nil {
		return Manifest{}, fmt.Errorf("failed to extract tarball: %w", err)
	}

	driverPath := filepath.Join(finalDir, manifest.Files.Driver)

	manifest.DriverInfo.ID = shortName
	manifest.DriverInfo.Source = "dbc"
	manifest.DriverInfo.Driver.Shared.Set(platform, driverPath)

	return manifest, nil
}

// TODO: Unexport once we refactor sync.go. sync.go has it's own separate
// installation routine which it probably shouldn't.
func InflateTarball(f *os.File, outDir string) (Manifest, error) {
	defer f.Close()
	var m Manifest

	if _, err := f.Seek(0, io.SeekStart); err != nil {
		return m, fmt.Errorf("could not seek to start: %w", err)
	}
	rdr, err := gzip.NewReader(f)
	if err != nil {
		return m, fmt.Errorf("could not create gzip reader: %w", err)
	}
	defer rdr.Close()

	t := tar.NewReader(rdr)
	for {
		hdr, err := t.Next()
		if errors.Is(err, io.EOF) {
			break
		}

		if err != nil {
			return m, fmt.Errorf("error reading tarball: %w", err)
		}

		// Return a helpful error if an entry is a directory. dbc doesn't support
		// installing driver tarballs that contain directories.
		if hdr.Typeflag == tar.TypeDir {
			return m, fmt.Errorf("found a directory entry when trying to extract %s which isn't supported. driver archives shouldn't contain subdirectories", f.Name())
		}

		if hdr.Name != "MANIFEST" {
			next, err := os.Create(filepath.Join(outDir, hdr.Name))
			if err != nil {
				return m, fmt.Errorf("could not create file %s: %w", hdr.Name, err)
			}

			if _, err = io.Copy(next, t); err != nil {
				next.Close()
				return m, fmt.Errorf("could not write file from tarball %s: %w", hdr.Name, err)
			}
			next.Close()
		} else {
			m, err = decodeManifest(t, "", false)
			if err != nil {
				return m, fmt.Errorf("could not decode manifest: %w", err)
			}

		}
	}

	return m, nil
}

func decodeManifest(r io.Reader, driverName string, requireShared bool) (Manifest, error) {
	var di tomlDriverInfo
	if err := toml.NewDecoder(r).Decode(&di); err != nil {
		return Manifest{}, fmt.Errorf("error decoding manifest: %w", err)
	}

	if di.ManifestVersion > currentManifestVersion {
		return Manifest{}, fmt.Errorf("manifest version %d is unsupported, only %d and lower are supported by this version of dbc",
			di.ManifestVersion, currentManifestVersion)
	}

	// Callers can assume these fields are set so return an error if they aren't
	if di.Name == "" {
		return Manifest{}, fmt.Errorf("%w: name is required", ErrInvalidManifest)
	}
	if di.Version == nil {
		return Manifest{}, fmt.Errorf("%w: version is required", ErrInvalidManifest)
	}

	result := Manifest{
		DriverInfo: DriverInfo{
			ID:        driverName,
			Name:      di.Name,
			Publisher: di.Publisher,
			License:   di.License,
			Version:   di.Version,
			Source:    di.Source,
			AdbcInfo:  di.AdbcInfo,
		},
		Files:       di.Files,
		PostInstall: di.PostInstall,
	}

	result.Driver.Entrypoint = di.Driver.Entrypoint
	switch s := di.Driver.Shared.(type) {
	case string:
		result.Driver.Shared.defaultPath = s
	case map[string]any:
		result.Driver.Shared.platformMap = make(map[string]string)
		for k, v := range s {
			if strVal, ok := v.(string); ok {
				result.Driver.Shared.platformMap[k] = strVal
			} else {
				return Manifest{}, fmt.Errorf("%w: invalid type for platform %s, expected string", ErrInvalidManifest, k)
			}
		}
	default:
		if requireShared {
			return Manifest{}, fmt.Errorf("%w: invalid type for 'Driver.shared' in manifest, expected string or table", ErrInvalidManifest)
		}
	}

	return result, nil
}

// Common, non-platform-specific code for uninstalling a driver. Called by
// platform-specific UninstallDriver function.
func UninstallDriverShared(info DriverInfo) error {
	// For the User and System config levels, info.FilePath is set to the
	// appropriate registry key instead of the filesystem on windows so we
	// handle that here first.
	filesystemLocation := info.FilePath
	if strings.Contains(info.FilePath, "HKCU\\") {
		filesystemLocation = ConfigUser.ConfigLocation()
	} else if strings.Contains(info.FilePath, "HKLM\\") {
		filesystemLocation = ConfigSystem.ConfigLocation()
	}

	root, err := os.OpenRoot(filesystemLocation)
	if err != nil {
		return fmt.Errorf("error opening driver path %s: %w", info.FilePath, err)
	}
	defer root.Close()

	for sharedPath := range info.Driver.Shared.Paths() {
		// Make sharedPath relative to info.FilePath and use it within root
		// to ensure that nothing can escape the intended directory.
		// (i.e. avoid malicious driver manifests)
		sharedPath, err = filepath.Rel(filesystemLocation, sharedPath)
		if err != nil {
			// If we can't make it relative, something is wrong, skip
			continue
		}

		// dbc installs drivers in a folder, other tools may not so we handle each
		// differently.
		if info.Source == "dbc" {
			sharedDir := filepath.Dir(sharedPath)
			// Edge case when manifest is ill-formed: if sharedPath is set to the
			// folder containing the shared library instead of the shared library
			// itself, sharedDir is info.FilePath and we definitely don't want to
			// remove that
			if sharedDir == "." {
				continue
			}

			if err := root.RemoveAll(sharedDir); err != nil {
				// Ignore only when not found. This supports manifest-only drivers.
				// TODO: Come up with a better mechanism to handle manifest-only drivers
				// and remove this continue when we do
				if errors.Is(err, fs.ErrNotExist) {
					continue
				}
				return fmt.Errorf("error removing driver %s: %w", info.ID, err)
			}
		} else {
			if err := root.Remove(sharedPath); err != nil {
				// Ignore only when not found. This supports manifest-only drivers.
				// TODO: Come up with a better mechanism to handle manifest-only drivers
				// and remove this continue when we do
				if errors.Is(err, fs.ErrNotExist) {
					continue
				}
				return fmt.Errorf("error removing driver %s: %w", info.ID, err)
			}
		}
	}

	// Special handling to clean up manifest-only drivers
	//
	// Manifest only drivers can come with extra files such as a LICENSE and we
	// create a folder next to the driver manifest to store them, same as we'd
	// store the actual driver shared library. Above, we find the path of this
	// folder by looking at the Driver.shared path. For manifest-only drivers,
	// Driver.shared is not a valid path (it's just a name), so this trick doesn't
	// work. We do want to clean this folder up so here we guess what it is and
	// try to remove it e.g., "somedriver_macos_arm64_v1.2.3."
	extraFolder := fmt.Sprintf("%s_%s_v%s", info.ID, platformTuple, info.Version)
	extraFolder = filepath.Clean(extraFolder)
	finfo, err := root.Stat(extraFolder)
	if err == nil && finfo.IsDir() && extraFolder != "." {
		_ = root.RemoveAll(extraFolder)
		// ignore errors
	}

	return nil
}