138 lines
3.1 KiB
Go
138 lines
3.1 KiB
Go
// Copyright (C) 2018 Storj Labs, Inc.
|
|
// See LICENSE for copying information.
|
|
|
|
package process
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
|
|
homedir "github.com/mitchellh/go-homedir"
|
|
"github.com/spf13/cobra"
|
|
"github.com/spf13/pflag"
|
|
"github.com/spf13/viper"
|
|
)
|
|
|
|
// ReadFlags will read in and bind flags for viper and pflag
|
|
func readFlags() {
|
|
pflag.CommandLine.AddGoFlagSet(flag.CommandLine)
|
|
pflag.Parse()
|
|
err := viper.BindPFlags(pflag.CommandLine)
|
|
if err != nil {
|
|
log.Print("error parsing command line flags into viper:", err)
|
|
}
|
|
}
|
|
|
|
// get default config folder
|
|
func configPath() string {
|
|
home, _ := homedir.Dir()
|
|
return filepath.Join(home, ".storj")
|
|
}
|
|
|
|
// get default config file
|
|
func defaultConfigFile(name string) string {
|
|
return filepath.Join(configPath(), name)
|
|
}
|
|
|
|
func generateConfig() error {
|
|
err := viper.WriteConfigAs(defaultConfigFile("main.json"))
|
|
return err
|
|
}
|
|
|
|
// ConfigEnvironment will read in command line flags, set the name of the config file,
|
|
// then look for configs in the current working directory and in $HOME/.storj
|
|
func ConfigEnvironment() (*viper.Viper, error) {
|
|
viper.SetEnvPrefix("storj")
|
|
viper.AutomaticEnv()
|
|
viper.SetConfigName("main")
|
|
viper.AddConfigPath(".")
|
|
viper.AddConfigPath(configPath())
|
|
|
|
// Override default config with a specific config
|
|
cfgFile := flag.String("config", "", "config file")
|
|
generate := flag.Bool("generate", false, "generate a default config in ~/.storj")
|
|
|
|
// if that file exists, set it as the config instead of reading in from default locations
|
|
if *cfgFile != "" && fileExists(*cfgFile) {
|
|
viper.SetConfigFile(*cfgFile)
|
|
}
|
|
|
|
err := viper.ReadInConfig()
|
|
|
|
if err != nil {
|
|
log.Print("could not read config file; defaulting to command line flags for configuration")
|
|
}
|
|
|
|
readFlags()
|
|
|
|
if *generate == true {
|
|
err := generateConfig()
|
|
if err != nil {
|
|
log.Print("unable to generate config file.", err)
|
|
}
|
|
}
|
|
|
|
v := viper.GetViper()
|
|
return v, nil
|
|
}
|
|
|
|
// check if file exists, handle error correctly if it doesn't
|
|
func fileExists(path string) bool {
|
|
_, err := os.Stat(path)
|
|
if err != nil {
|
|
if os.IsNotExist(err) {
|
|
return false
|
|
}
|
|
log.Fatalf("failed to check for file existence: %v", err)
|
|
}
|
|
return true
|
|
}
|
|
|
|
// Execute runs a *cobra.Command and sets up Storj-wide process configuration
|
|
// like a configuration file and logging.
|
|
func Execute(cmd *cobra.Command) {
|
|
cobra.OnInitialize(func() {
|
|
_, err := ConfigEnvironment()
|
|
if err != nil {
|
|
log.Fatal("error configuring environment", err)
|
|
}
|
|
})
|
|
|
|
Must(cmd.Execute())
|
|
}
|
|
|
|
// Main runs a Service
|
|
func Main(configFn func() (*viper.Viper, error), s ...Service) error {
|
|
if _, err := configFn(); err != nil {
|
|
return err
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
errors := make(chan error, len(s))
|
|
|
|
for _, service := range s {
|
|
go func(ctx context.Context, s Service, ch <-chan error) {
|
|
errors <- CtxService(s)(&cobra.Command{}, pflag.Args())
|
|
}(ctx, service, errors)
|
|
}
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil
|
|
case err := <-errors:
|
|
return err
|
|
}
|
|
}
|
|
|
|
// Must checks for errors
|
|
func Must(err error) {
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
}
|