Optimize code

Signed-off-by: Jianhui Zhao <jianhuizhao329@gmail.com>
This commit is contained in:
Jianhui Zhao
2019-04-24 23:41:54 +08:00
parent baed5351a2
commit afe0957a29
2 changed files with 45 additions and 34 deletions
+21 -6
View File
@@ -7,6 +7,7 @@ import (
log "github.com/sirupsen/logrus"
_ "github.com/zhaojh329/rttys/statik"
"net/http"
"os"
"strconv"
"sync"
"time"
@@ -60,7 +61,7 @@ func httpAuth(w http.ResponseWriter, r *http.Request) bool {
return true
}
func httpLogin(cfg *rttysConfig, username, password string) bool {
func httpLogin(cfg *RttysConfig, username, password string) bool {
if cfg.username != "" {
if cfg.username != username {
return false
@@ -76,7 +77,7 @@ func httpLogin(cfg *rttysConfig, username, password string) bool {
return login(username, password)
}
func httpStart(br *Broker, cfg *rttysConfig) {
func httpStart(br *Broker, cfg *RttysConfig) {
statikFS, err := fs.New()
if err != nil {
log.Fatal(err)
@@ -150,11 +151,25 @@ func httpStart(br *Broker, cfg *rttysConfig) {
time.AfterFunc(5*time.Second, cleanHttpSession)
if cfg.cert != "" && cfg.key != "" {
log.Info("Listen on: ", cfg.addr, "SSL on")
log.Fatal(http.ListenAndServeTLS(cfg.addr, cfg.cert, cfg.key, nil))
if cfg.sslCert != "" && cfg.sslKey != "" {
_, err := os.Lstat(cfg.sslCert)
if err != nil {
log.Error(err)
cfg.sslCert = ""
}
_, err = os.Lstat(cfg.sslKey)
if err != nil {
log.Error(err)
cfg.sslKey = ""
}
}
if cfg.sslCert != "" && cfg.sslKey != "" {
log.Info("Listen on: ", cfg.addr, " SSL on")
log.Fatal(http.ListenAndServeTLS(cfg.addr, cfg.sslCert, cfg.sslKey, nil))
} else {
log.Info("Listen on: ", cfg.addr, "SSL off")
log.Info("Listen on: ", cfg.addr, " SSL off")
log.Fatal(http.ListenAndServe(cfg.addr, nil))
}
}
+24 -28
View File
@@ -36,10 +36,10 @@ import (
log "github.com/sirupsen/logrus"
)
type rttysConfig struct {
type RttysConfig struct {
addr string
cert string
key string
sslCert string
sslKey string
username string
password string
}
@@ -78,35 +78,31 @@ func genUniqueID(extra string) string {
return hex.EncodeToString(h.Sum(nil))
}
func parseConfig() *rttysConfig {
addr := flag.String("addr", ":5912", "address to listen")
cert := flag.String("ssl-cert", "", "certFile Path")
key := flag.String("ssl-key", "", "keyFile Path")
func setConfigOpt(yamlCfg *yaml.File, name string, opt *string) {
val, err := yamlCfg.Get("addr")
if err != nil {
return
}
*opt = val
}
func parseConfig() *RttysConfig {
cfg := &RttysConfig{}
flag.StringVar(&cfg.addr, "addr", ":5912", "address to listen")
flag.StringVar(&cfg.sslCert, "ssl-cert", "./rttys.crt", "certFile Path")
flag.StringVar(&cfg.sslKey, "ssl-key", "./rttys.key", "keyFile Path")
conf := flag.String("conf", "./rttys.conf", "config file to load")
flag.Parse()
cfg := &rttysConfig{}
config, _ := yaml.ReadFile(*conf)
if config != nil {
cfg.addr, _ = config.Get("addr")
cfg.cert, _ = config.Get("ssl-cert")
cfg.key, _ = config.Get("ssl-key")
cfg.username, _ = config.Get("username")
cfg.password, _ = config.Get("password")
}
if cfg.addr == "" {
cfg.addr = *addr
}
if cfg.cert == "" {
cfg.cert = *cert
}
if cfg.key == "" {
cfg.key = *key
yamlCfg, err := yaml.ReadFile(*conf)
if err == nil {
setConfigOpt(yamlCfg, "addr", &cfg.addr)
setConfigOpt(yamlCfg, "ssl-cert", &cfg.sslCert)
setConfigOpt(yamlCfg, "ssl-key", &cfg.sslKey)
setConfigOpt(yamlCfg, "username", &cfg.username)
setConfigOpt(yamlCfg, "password", &cfg.password)
}
return cfg