Config using viper

configure/channel.go

@@ -22,14 +22,14 @@ var RoomKeys = &RoomKeysType{
 var saveInLocal = true
 
 func Init() {
-	saveInLocal = GetRedisAddr() == nil
+	saveInLocal = len(Config.GetString("redis_addr")) == 0
 	if saveInLocal {
 		return
 	}
 
 	RoomKeys.redisCli = redis.NewClient(&redis.Options{
-		Addr:     *GetRedisAddr(),
-		Password: *GetRedisPwd(),
+		Addr:     Config.GetString("redis_addr"),
+		Password: Config.GetString("redis_pwd"),
 		DB:       0,
 	})
 

configure/liveconfig.go

@@ -1,11 +1,15 @@
 package configure
 
 import (
+	"bytes"
 	"encoding/json"
-	"flag"
-	"io/ioutil"
+	"fmt"
+	"strings"
 
+	"github.com/kr/pretty"
 	log "github.com/sirupsen/logrus"
+	"github.com/spf13/pflag"
+	"github.com/spf13/viper"
 )
 
 /*
@@ -20,30 +24,47 @@ import (
   ]
 }
 */
-var (
-	redisAddr = flag.String("redis_addr", "", "redis addr to save room keys ex. localhost:6379")
-	redisPwd  = flag.String("redis_pwd", "", "redis password")
-)
 
 type Application struct {
-	Appname    string   `json:"appname"`
-	Live       bool     `json:"liveon"`
-	Hls        bool     `json:"hls"`
-	StaticPush []string `json:"static_push"`
+	Appname    string   `mapstructure:"appname"`
+	Live       bool     `mapstructure:"live"`
+	Hls        bool     `mapstructure:"hls"`
+	StaticPush []string `mapstructure:"static_push"`
 }
-type JWTCfg struct {
-	Secret    string `json:"secret"`
-	Algorithm string `json:"algorithm"`
+
+type Applications []Application
+
+type JWT struct {
+	Secret    string `mapstructure:"secret"`
+	Algorithm string `mapstructure:"algorithm"`
 }
 type ServerCfg struct {
-	RedisAddr string `json:"redis_addr"`
-	RedisPwd  string `json:"redis_pwd"`
-	JWTCfg    `json:"jwt"`
-	Server    []Application `json:"server"`
+	Level        string        `mapstructure:"level"`
+	ConfigFile   string        `mapstructure:"config_file"`
+	FLVDir       string        `mapstructure:"flv_dir"`
+	RTMPAddr     string        `mapstructure:"rtmp_addr"`
+	HTTPFLVAddr  string        `mapstructure:"httpflv_addr"`
+	HLSAddr      string        `mapstructure:"hls_addr"`
+	APIAddr      string        `mapstructure:"api_addr"`
+	RedisAddr    string        `mapstructure:"redis_addr"`
+	RedisPwd     string        `mapstructure:"redis_pwd"`
+	ReadTimeout  int           `mapstructure:"read_timeout"`
+	WriteTimeout int           `mapstructure:"write_timeout"`
+	GopNum       int           `mapstructure:"gop_num"`
+	JWT          JWT           `mapstructure:"jwt"`
+	Server       []Application `mapstructure:"server"`
 }
 
 // default config
-var RtmpServercfg = ServerCfg{
+var defaultConf = ServerCfg{
+	ConfigFile:   "livego.yaml",
+	RTMPAddr:     ":1935",
+	HTTPFLVAddr:  ":7001",
+	HLSAddr:      ":7002",
+	APIAddr:      ":8090",
+	WriteTimeout: 10,
+	ReadTimeout:  10,
+	GopNum:       1,
 	Server: []Application{{
 		Appname:    "livego",
 		Live:       true,
@@ -52,47 +73,65 @@ var RtmpServercfg = ServerCfg{
 	}},
 }
 
-func LoadConfig(configfilename string) {
-	defer Init()
+var Config = viper.New()
 
-	log.Infof("starting load configure file %s", configfilename)
-	data, err := ioutil.ReadFile(configfilename)
-	if err != nil {
-		log.Warningf("ReadFile %s error:%v", configfilename, err)
-		log.Info("Using default config")
-		return
+func initLog() {
+	if l, err := log.ParseLevel(Config.GetString("level")); err == nil {
+		log.SetLevel(l)
+		log.SetReportCaller(l == log.DebugLevel)
 	}
-
-	err = json.Unmarshal(data, &RtmpServercfg)
-	if err != nil {
-		log.Errorf("json.Unmarshal error:%v", err)
-		log.Info("Using default config")
-	}
-	log.Debugf("get config json data:%v", RtmpServercfg)
 }
 
-func GetRedisAddr() *string {
-	if len(RtmpServercfg.RedisAddr) > 0 {
-		*redisAddr = RtmpServercfg.RedisAddr
-	}
+func LoadConfig() {
+	defer Init()
 
-	if len(*redisAddr) == 0 {
-		return nil
+	// Default config
+	b, _ := json.Marshal(defaultConf)
+	defaultConfig := bytes.NewReader(b)
+	Config.MergeConfig(defaultConfig)
+
+	// Flags
+	pflag.String("rtmp_addr", ":1935", "RTMP server listen address")
+	pflag.String("httpflv_addr", ":7001", "HTTP-FLV server listen address")
+	pflag.String("hls_addr", ":7002", "HLS server listen address")
+	pflag.String("api_addr", ":8090", "HTTP manage interface server listen address")
+	pflag.String("config_file", "livego.yaml", "configure filename")
+	pflag.String("level", "info", "Log level")
+	pflag.String("flv_dir", "tmp", "output flv file at flvDir/APP/KEY_TIME.flv")
+	pflag.Int("read_timeout", 10, "read time out")
+	pflag.Int("write_timeout", 10, "write time out")
+	pflag.Int("gop_num", 1, "gop num")
+	pflag.Parse()
+	Config.BindPFlags(pflag.CommandLine)
+
+	// File
+	Config.SetConfigFile(Config.GetString("config_file"))
+	Config.AddConfigPath(".")
+	err := Config.ReadInConfig()
+	if err != nil {
+		log.Error(err)
+		log.Info("Using default config")
 	}
 
-	return redisAddr
-}
+	// Environment
+	replacer := strings.NewReplacer(".", "_")
+	Config.SetEnvKeyReplacer(replacer)
+	Config.AllowEmptyEnv(true)
+	Config.AutomaticEnv()
 
-func GetRedisPwd() *string {
-	if len(RtmpServercfg.RedisPwd) > 0 {
-		*redisPwd = RtmpServercfg.RedisPwd
-	}
+	// Log
+	initLog()
 
-	return redisPwd
+	c := ServerCfg{}
+	Config.Unmarshal(&c)
+	log.Debugf("Current configurations: \n%# v", pretty.Formatter(c))
 }
 
 func CheckAppName(appname string) bool {
-	for _, app := range RtmpServercfg.Server {
+	apps := Applications{}
+	Config.UnmarshalKey("server", &apps)
+	fmt.Println(apps)
+	for _, app := range apps {
 		if app.Appname == appname {
 			return app.Live
 		}
@@ -101,15 +140,16 @@ func CheckAppName(appname string) bool {
 }
 
 func GetStaticPushUrlList(appname string) ([]string, bool) {
-	for _, app := range RtmpServercfg.Server {
+	apps := Applications{}
+	Config.UnmarshalKey("server", &apps)
+	for _, app := range apps {
 		if (app.Appname == appname) && app.Live {
 			if len(app.StaticPush) > 0 {
 				return app.StaticPush, true
 			} else {
 				return nil, false
 			}
 		}
-
 	}
 	return nil, false
 }

container/flv/muxer.go

@@ -1,14 +1,14 @@
 package flv
 
 import (
-	"flag"
 	"fmt"
 	"os"
 	"path"
 	"strings"
 	"time"
 
 	"livego/av"
+	"livego/configure"
 	"livego/protocol/amf"
 	"livego/utils/pio"
 	"livego/utils/uid"
@@ -18,7 +18,6 @@ import (
 
 var (
 	flvHeader = []byte{0x46, 0x4c, 0x56, 0x01, 0x05, 0x00, 0x00, 0x00, 0x09}
-	flvDir    = flag.String("flvDir", "tmp", "output flv file at flvDir/APP/KEY_TIME.flv")
 )
 
 /*
@@ -152,13 +151,15 @@ func (f *FlvDvr) GetWriter(info av.Info) av.WriteCloser {
 		return nil
 	}
 
-	err := os.MkdirAll(path.Join(*flvDir, paths[0]), 0755)
+	flvDir := configure.Config.GetString("flv_dir")
+
+	err := os.MkdirAll(path.Join(flvDir, paths[0]), 0755)
 	if err != nil {
 		log.Error("mkdir error: ", err)
 		return nil
 	}
 
-	fileName := fmt.Sprintf("%s_%d.%s", path.Join(*flvDir, info.Key), time.Now().Unix(), "flv")
+	fileName := fmt.Sprintf("%s_%d.%s", path.Join(flvDir, info.Key), time.Now().Unix(), "flv")
 	log.Debug("flv dvr save stream to: ", fileName)
 	w, err := os.OpenFile(fileName, os.O_CREATE|os.O_RDWR, 0755)
 	if err != nil {

go.mod

@@ -7,11 +7,13 @@ require (
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/go-redis/redis/v7 v7.2.0
 	github.com/gorilla/mux v1.7.4 // indirect
+	github.com/kr/pretty v0.1.0
 	github.com/orcaman/concurrent-map v0.0.0-20190826125027-8c72a8bb44f6
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/satori/go.uuid v1.2.0
 	github.com/sirupsen/logrus v1.5.0
-	github.com/smartystreets/goconvey v1.6.4 // indirect
+	github.com/spf13/pflag v1.0.3
+	github.com/spf13/viper v1.6.3
 	github.com/stretchr/testify v1.4.0
 	github.com/urfave/negroni v1.0.0 // indirect
 )