vaultik/internal/config/size.go

package config

import (
	"errors"
	"fmt"
	"math"
	"reflect"
	"slices"
	"strconv"
	"strings"

	"github.com/dustin/go-humanize"
	"gopkg.in/yaml.v3"
)

var (
	errSizeType     = errors.New("size must be a number or string")
	errSizeTooLarge = errors.New("size exceeds maximum supported value")
	errNotASize     = errors.New(
		"is not a size; use a number of bytes or a unit such as 10MB")
)

// Size represents a byte size that can be specified in configuration files.
// It can unmarshal from both numeric values (interpreted as bytes) and
// human-readable strings like "10MB", "2.5GB", or "1TB".
type Size int64

// UnmarshalYAML implements yaml.Unmarshaler for Size, allowing it to be
// parsed from YAML configuration files. It accepts both numeric values
// (interpreted as bytes) and string values with units (e.g., "10MB").
func (s *Size) UnmarshalYAML(unmarshal func(any) error) error {
	// Try to unmarshal as int64 first
	var intVal int64

	err := unmarshal(&intVal)
	if err == nil {
		*s = Size(intVal)

		return nil
	}

	// A number too large for an int64 decodes as a float64. Its text can be
	// 1e+20, which is not a size string.
	var floatVal float64

	err = unmarshal(&floatVal)
	if err == nil && floatVal > math.MaxInt64 {
		return errSizeTooLarge
	}

	// Try to unmarshal as string
	var strVal string

	err = unmarshal(&strVal)
	if err != nil {
		return errSizeType
	}

	size, err := ParseSize(strVal)
	if err != nil {
		return err
	}

	*s = size

	return nil
}

// checkSizes returns an error naming the first size setting in node whose
// value is not a size. node is the config, or the part of it at the dotted
// keys, after smartconfig has filled in its ${...} references.
func checkSizes(node *yaml.Node, keys []string) error {
	typ, err := SettingType(keys)
	if err != nil {
		// Unknown keys are refused before this check, so keys is in one of
		// two places. The env section holds no settings. A value of the
		// wrong shape, such as a map given for hostname or inside exclude,
		// is reported later by the YAML decoder.
		return nil //nolint:nilerr // not an error here
	}

	if typ == reflect.TypeFor[Size]() {
		var size Size

		err = node.Decode(&size)
		if err != nil {
			return fmt.Errorf("%s: %w", strings.Join(keys, "."), err)
		}

		return nil
	}

	switch node.Kind {
	case yaml.DocumentNode:
		for _, child := range node.Content {
			err = checkSizes(child, keys)
			if err != nil {
				return err
			}
		}
	case yaml.MappingNode:
		for i := 0; i+1 < len(node.Content); i += 2 {
			err = checkSizes(node.Content[i+1],
				slices.Concat(keys, []string{node.Content[i].Value}))
			if err != nil {
				return err
			}
		}
	case yaml.SequenceNode:
		for i, child := range node.Content {
			err = checkSizes(child, slices.Concat(keys, []string{strconv.Itoa(i)}))
			if err != nil {
				return err
			}
		}
	case yaml.ScalarNode, yaml.AliasNode:
		// A scalar holds no settings, and the filled config has no aliases.
	}

	return nil
}

// Int64 returns the size as int64 bytes.
// This is useful when the size needs to be passed to APIs that expect
// a numeric byte count.
func (s Size) Int64() int64 {
	return int64(s)
}

// String returns the size as a human-readable string.
// For example, 1048576 bytes would be formatted as "1.0 MB".
// This implements the fmt.Stringer interface.
func (s Size) String() string {
	//nolint:gosec // G115: sizes are non-negative by construction
	return humanize.Bytes(uint64(s))
}

// ParseSize parses a size string, such as "10MB", into a Size value. An
// error quotes s.
func ParseSize(s string) (Size, error) {
	bytes, err := humanize.ParseBigBytes(s)
	if err != nil {
		return 0, fmt.Errorf("%q %w", s, errNotASize)
	}

	if !bytes.IsInt64() {
		return 0, fmt.Errorf("%w: %s", errSizeTooLarge, s)
	}

	return Size(bytes.Int64()), nil
}