github.com/moqsien/xraycore@v1.8.5/core/config.go (about)

     1  package core
     2  
     3  import (
     4  	"io"
     5  	"strings"
     6  
     7  	"github.com/moqsien/xraycore/common"
     8  	"github.com/moqsien/xraycore/common/buf"
     9  	"github.com/moqsien/xraycore/common/cmdarg"
    10  	"github.com/moqsien/xraycore/main/confloader"
    11  	"google.golang.org/protobuf/proto"
    12  )
    13  
    14  // ConfigFormat is a configurable format of Xray config file.
    15  type ConfigFormat struct {
    16  	Name      string
    17  	Extension []string
    18  	Loader    ConfigLoader
    19  }
    20  
    21  // ConfigLoader is a utility to load Xray config from external source.
    22  type ConfigLoader func(input interface{}) (*Config, error)
    23  
    24  // ConfigBuilder is a builder to build core.Config from filenames and formats
    25  type ConfigBuilder func(files []string, formats []string) (*Config, error)
    26  
    27  var (
    28  	configLoaderByName    = make(map[string]*ConfigFormat)
    29  	configLoaderByExt     = make(map[string]*ConfigFormat)
    30  	ConfigBuilderForFiles ConfigBuilder
    31  )
    32  
    33  // RegisterConfigLoader add a new ConfigLoader.
    34  func RegisterConfigLoader(format *ConfigFormat) error {
    35  	name := strings.ToLower(format.Name)
    36  	if _, found := configLoaderByName[name]; found {
    37  		return newError(format.Name, " already registered.")
    38  	}
    39  	configLoaderByName[name] = format
    40  
    41  	for _, ext := range format.Extension {
    42  		lext := strings.ToLower(ext)
    43  		if f, found := configLoaderByExt[lext]; found {
    44  			return newError(ext, " already registered to ", f.Name)
    45  		}
    46  		configLoaderByExt[lext] = format
    47  	}
    48  
    49  	return nil
    50  }
    51  
    52  func GetFormatByExtension(ext string) string {
    53  	switch strings.ToLower(ext) {
    54  	case "pb", "protobuf":
    55  		return "protobuf"
    56  	case "yaml", "yml":
    57  		return "yaml"
    58  	case "toml":
    59  		return "toml"
    60  	case "json", "jsonc":
    61  		return "json"
    62  	default:
    63  		return ""
    64  	}
    65  }
    66  
    67  func getExtension(filename string) string {
    68  	idx := strings.LastIndexByte(filename, '.')
    69  	if idx == -1 {
    70  		return ""
    71  	}
    72  	return filename[idx+1:]
    73  }
    74  
    75  func getFormat(filename string) string {
    76  	return GetFormatByExtension(getExtension(filename))
    77  }
    78  
    79  func LoadConfig(formatName string, input interface{}) (*Config, error) {
    80  	switch v := input.(type) {
    81  	case cmdarg.Arg:
    82  		formats := make([]string, len(v))
    83  		hasProtobuf := false
    84  		for i, file := range v {
    85  			var f string
    86  
    87  			if formatName == "auto" {
    88  				if file != "stdin:" {
    89  					f = getFormat(file)
    90  				} else {
    91  					f = "json"
    92  				}
    93  			} else {
    94  				f = formatName
    95  			}
    96  
    97  			if f == "" {
    98  				return nil, newError("Failed to get format of ", file).AtWarning()
    99  			}
   100  
   101  			if f == "protobuf" {
   102  				hasProtobuf = true
   103  			}
   104  			formats[i] = f
   105  		}
   106  
   107  		// only one protobuf config file is allowed
   108  		if hasProtobuf {
   109  			if len(v) == 1 {
   110  				return configLoaderByName["protobuf"].Loader(v)
   111  			} else {
   112  				return nil, newError("Only one protobuf config file is allowed").AtWarning()
   113  			}
   114  		}
   115  
   116  		// to avoid import cycle
   117  		return ConfigBuilderForFiles(v, formats)
   118  
   119  	case io.Reader:
   120  		if f, found := configLoaderByName[formatName]; found {
   121  			return f.Loader(v)
   122  		} else {
   123  			return nil, newError("Unable to load config in", formatName).AtWarning()
   124  		}
   125  	}
   126  
   127  	return nil, newError("Unable to load config").AtWarning()
   128  }
   129  
   130  func loadProtobufConfig(data []byte) (*Config, error) {
   131  	config := new(Config)
   132  	if err := proto.Unmarshal(data, config); err != nil {
   133  		return nil, err
   134  	}
   135  	return config, nil
   136  }
   137  
   138  func init() {
   139  	common.Must(RegisterConfigLoader(&ConfigFormat{
   140  		Name:      "Protobuf",
   141  		Extension: []string{"pb"},
   142  		Loader: func(input interface{}) (*Config, error) {
   143  			switch v := input.(type) {
   144  			case cmdarg.Arg:
   145  				r, err := confloader.LoadConfig(v[0])
   146  				common.Must(err)
   147  				data, err := buf.ReadAllToBytes(r)
   148  				common.Must(err)
   149  				return loadProtobufConfig(data)
   150  			case io.Reader:
   151  				data, err := buf.ReadAllToBytes(v)
   152  				common.Must(err)
   153  				return loadProtobufConfig(data)
   154  			default:
   155  				return nil, newError("unknow type")
   156  			}
   157  		},
   158  	}))
   159  }