dodo/pkg/config/file.go

package config

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"time"

	"github.com/aykhans/dodo/pkg/types"
	"github.com/aykhans/dodo/pkg/utils"
	"go.yaml.in/yaml/v4"
)

var _ IParser = ConfigFileParser{}

type ConfigFileParser struct {
	configFile types.ConfigFile
}

func NewConfigFileParser(configFile types.ConfigFile) *ConfigFileParser {
	return &ConfigFileParser{configFile}
}

// Parse parses config file arguments into a Config object.
// It can return the following errors:
// - types.ConfigFileReadError
// - types.UnmarshalError
// - types.FieldParseErrors
func (parser ConfigFileParser) Parse() (*Config, error) {
	var err error
	var configFileData []byte

	switch parser.configFile.LocationType() {
	case types.ConfigFileLocationLocal:
		configFileData, err = os.ReadFile(parser.configFile.Path())
		if err != nil {
			return nil, types.NewConfigFileReadError(err)
		}
	case types.ConfigFileLocationRemote:
		resp, err := http.Get(parser.configFile.Path())
		if err != nil {
			return nil, types.NewConfigFileReadError(err)
		}
		defer resp.Body.Close() //nolint:errcheck

		if resp.StatusCode != http.StatusOK {
			return nil, types.NewConfigFileReadError(errors.New("failed to retrieve remote config file: " + resp.Status))
		}

		configFileData, err = io.ReadAll(resp.Body)
		if err != nil {
			return nil, types.NewConfigFileReadError(err)
		}
	default:
		panic("unhandled config file location type")
	}

	switch parser.configFile.Type() {
	case types.ConfigFileTypeYAML, types.ConfigFileTypeUnknown:
		return parser.ParseYAML(configFileData)
	default:
		panic("unhandled config file type")
	}
}

type stringOrSliceField []string

func (ss *stringOrSliceField) UnmarshalYAML(node *yaml.Node) error {
	switch node.Kind {
	case yaml.ScalarNode:
		// Handle single string value
		*ss = []string{node.Value}
		return nil
	case yaml.SequenceNode:
		// Handle array of strings
		var slice []string
		if err := node.Decode(&slice); err != nil {
			return err //nolint:wrapcheck
		}
		*ss = slice
		return nil
	default:
		return fmt.Errorf("expected a string or a sequence of strings, but got %v", node.Kind)
	}
}

type configYAML struct {
	Files        stringOrSliceField `yaml:"files"`
	Method       *string            `yaml:"method"`
	URL          *string            `yaml:"url"`
	Timeout      *time.Duration     `yaml:"timeout"`
	DodosCount   *uint              `yaml:"dodos"`
	RequestCount *uint              `yaml:"requests"`
	Duration     *time.Duration     `yaml:"duration"`
	Yes          *bool              `yaml:"yes"`
	SkipVerify   *bool              `yaml:"skipVerify"`
	Params       stringOrSliceField `yaml:"params"`
	Headers      stringOrSliceField `yaml:"headers"`
	Cookies      stringOrSliceField `yaml:"cookies"`
	Bodies       stringOrSliceField `yaml:"body"`
	Proxies      stringOrSliceField `yaml:"proxy"`
}

// ParseYAML parses YAML config file arguments into a Config object.
// It can return the following errors:
// - types.UnmarshalError
// - types.FieldParseErrors
func (parser ConfigFileParser) ParseYAML(data []byte) (*Config, error) {
	var (
		config     = &Config{}
		parsedData = &configYAML{}
	)

	err := yaml.Unmarshal(data, &parsedData)
	if err != nil {
		return nil, types.NewUnmarshalError(err)
	}

	var fieldParseErrors []types.FieldParseError

	config.Method = parsedData.Method
	config.Timeout = parsedData.Timeout
	config.DodosCount = parsedData.DodosCount
	config.RequestCount = parsedData.RequestCount
	config.Duration = parsedData.Duration
	config.Yes = parsedData.Yes
	config.SkipVerify = parsedData.SkipVerify
	config.Params.Parse(parsedData.Params...)
	config.Headers.Parse(parsedData.Headers...)
	config.Cookies.Parse(parsedData.Cookies...)
	config.Bodies.Parse(parsedData.Bodies...)

	if len(parsedData.Files) > 0 {
		for i, configFile := range parsedData.Files {
			configFileParsed, err := types.ParseConfigFile(configFile)

			_ = utils.HandleErrorOrDie(err,
				utils.OnCustomError(func(err types.RemoteConfigFileParseError) error {
					fieldParseErrors = append(
						fieldParseErrors,
						types.NewFieldParseError(
							fmt.Sprintf("config-file[%d]", i),
							configFile,
							fmt.Errorf("parse error: %w", err),
						),
					)
					return nil
				}),
			)

			if err == nil {
				config.Files = append(config.Files, *configFileParsed)
			}
		}
	}

	if parsedData.URL != nil {
		urlParsed, err := url.Parse(*parsedData.URL)
		if err != nil {
			fieldParseErrors = append(fieldParseErrors, types.NewFieldParseError("url", *parsedData.URL, err))
		} else {
			config.URL = urlParsed
		}
	}

	for i, proxy := range parsedData.Proxies {
		err := config.Proxies.Parse(proxy)
		if err != nil {
			fieldParseErrors = append(
				fieldParseErrors,
				types.NewFieldParseError(fmt.Sprintf("proxy[%d]", i), proxy, err),
			)
		}
	}

	if len(fieldParseErrors) > 0 {
		return nil, types.NewFieldParseErrors(fieldParseErrors)
	}

	return config, nil
}