ThothII/tools/tht/internal/authconfig/store.go

package authconfig

import (
	"bytes"
	"context"
	"errors"
	"io"
	"os"
	"path/filepath"
	"time"

	"github.com/aritmolab/thothii/tools/tht/internal/authprojection"
	"github.com/aritmolab/thothii/tools/tht/internal/safeio"
	"github.com/gofrs/flock"
	"gopkg.in/yaml.v3"
)

const (
	authFileName  = "auth.yaml"
	usersFileName = "users.yaml"
	lockFileName  = ".auth.lock"
)

// Load reads auth.yaml and, for local authentication, the paired users.yaml registry. All reads
// are bounded and reject non-private, linked, or symlinked files.
func Load(directory string) (Config, Registry, error) {
	return load(directory)
}

func load(directory string) (Config, Registry, error) {
	if err := requirePrivateDirectory(directory); err != nil {
		return Config{}, Registry{}, err
	}
	authContents, err := readPrivateFile(filepath.Join(directory, authFileName))
	if err != nil {
		return Config{}, Registry{}, err
	}
	var config Config
	if err := decodeStrictYAML(authContents, &config); err != nil || !validConfig(config) {
		return Config{}, Registry{}, errInvalidAuthenticationConfig
	}
	if config.Mode != "local" {
		return config, Registry{}, nil
	}
	usersContents, err := readPrivateFile(filepath.Join(directory, config.Local.UsersFile))
	if err != nil {
		return Config{}, Registry{}, err
	}
	var registry Registry
	if err := decodeStrictYAML(usersContents, &registry); err != nil {
		return Config{}, Registry{}, errInvalidAuthenticationConfig
	}
	if err := validateRegistry(registry); err != nil {
		return Config{}, Registry{}, errInvalidAuthenticationConfig
	}
	return config, registry, nil
}

// loadSnapshotBytes reads only canonical, validated authentication files and returns a detached,
// bounded projection snapshot. Callers that coordinate a mutation hold .auth.lock around it.
func loadSnapshotBytes(directory string) (authprojection.Snapshot, error) {
	if err := requirePrivateDirectory(directory); err != nil {
		return authprojection.Snapshot{}, err
	}
	authContents, err := readPrivateFile(filepath.Join(directory, authFileName))
	if err != nil {
		return authprojection.Snapshot{}, err
	}
	var config Config
	if err := decodeStrictYAML(authContents, &config); err != nil || !validConfig(config) {
		return authprojection.Snapshot{}, errInvalidAuthenticationConfig
	}
	var usersContents []byte
	if config.Mode == "local" {
		usersContents, err = readPrivateFile(filepath.Join(directory, config.Local.UsersFile))
		if err != nil {
			return authprojection.Snapshot{}, err
		}
		var registry Registry
		if err := decodeStrictYAML(usersContents, &registry); err != nil || validateRegistry(registry) != nil {
			return authprojection.Snapshot{}, errInvalidAuthenticationConfig
		}
	}
	return authprojection.NewSnapshot(config.Mode, append([]byte(nil), authContents...), append([]byte(nil), usersContents...))
}

// MutateUsers serializes the entire read-check-write transaction under the configuration lock.
// The resulting registry is revalidated and atomically replaced only after all invariants hold.
func MutateUsers(directory string, mutate func(*Registry) error) error {
	if mutate == nil {
		return errInvalidAuthenticationConfig
	}
	if err := requirePrivateDirectory(directory); err != nil {
		return err
	}
	lock, err := acquireLock(directory)
	if err != nil {
		return err
	}
	defer func() { _ = lock.Unlock() }()

	config, registry, err := load(directory)
	if err != nil || config.Mode != "local" {
		return errInvalidAuthenticationConfig
	}
	before := cloneRegistry(registry)
	if err := mutate(&registry); err != nil {
		return errInvalidAuthenticationConfig
	}
	if err := applyMutationInvariants(before, &registry); err != nil {
		return errInvalidAuthenticationConfig
	}
	contents, err := yaml.Marshal(registry)
	if err != nil {
		return errInvalidAuthenticationConfig
	}
	contents = append(contents, '\n')
	if err := safeio.ReplaceCanonicalRegular(filepath.Join(directory, config.Local.UsersFile), contents, 0o600); err != nil {
		return err
	}
	return nil
}

func validConfig(config Config) bool {
	if config.Version != 1 || (config.Mode != "local" && config.Mode != "oidc") {
		return false
	}
	return config.Mode != "local" || config.Local.UsersFile == usersFileName
}

func decodeStrictYAML(contents []byte, destination any) error {
	decoder := yaml.NewDecoder(bytes.NewReader(contents))
	decoder.KnownFields(true)
	if err := decoder.Decode(destination); err != nil {
		return errInvalidAuthenticationConfig
	}
	var extra any
	if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) {
		return errInvalidAuthenticationConfig
	}
	return nil
}

func requirePrivateDirectory(directory string) error {
	return safeio.ValidatePrivateDirectory(directory)
}

func readPrivateFile(path string) ([]byte, error) {
	if err := safeio.ValidatePrivateRegular(path); err != nil {
		return nil, err
	}
	before, err := os.Lstat(path)
	if err != nil {
		return nil, safeio.ErrUnsafeFile
	}
	contents, err := safeio.ReadCanonicalRegular(path, maxYAMLBytes)
	if err != nil {
		return nil, err
	}
	after, err := os.Lstat(path)
	if err != nil || !os.SameFile(before, after) {
		return nil, safeio.ErrUnsafeFile
	}
	if err := safeio.ValidatePrivateRegular(path); err != nil {
		return nil, safeio.ErrUnsafeFile
	}
	return contents, nil
}

func acquireLock(directory string) (*flock.Flock, error) {
	return acquireLockContext(context.Background(), directory)
}

func acquireLockContext(ctx context.Context, directory string) (*flock.Flock, error) {
	if ctx == nil || ctx.Err() != nil {
		return nil, errInvalidAuthenticationConfig
	}
	path := filepath.Join(directory, lockFileName)
	if _, err := os.Lstat(path); errors.Is(err, os.ErrNotExist) {
		if err := safeio.WriteCanonicalNewFile(path, nil, 0o600); err != nil {
			// A competing mutation may have created the lock after our Lstat. Accept only
			// that exact safe/private lock; every other creation failure remains unsafe.
			if validateErr := safeio.ValidatePrivateRegular(path); validateErr != nil {
				return nil, safeio.ErrUnsafeFile
			}
		}
	} else if err != nil {
		return nil, safeio.ErrUnsafeFile
	}
	if err := safeio.ValidatePrivateRegular(path); err != nil {
		return nil, err
	}
	lock := flock.New(path, flock.SetPermissions(0o600))
	locked, err := lock.TryLock()
	if err != nil {
		return nil, errInvalidAuthenticationConfig
	}
	if !locked {
		if hook := coordinatorHooks().onCanonicalLockContention; hook != nil {
			hook()
		}
		locked, err = lock.TryLockContext(ctx, 10*time.Millisecond)
		if err != nil || !locked {
			return nil, errInvalidAuthenticationConfig
		}
	}
	if _, err := readPrivateFile(path); err != nil {
		_ = lock.Unlock()
		return nil, err
	}
	return lock, nil
}