From afe0957a29f1fdd7d7ee43c770b4bee41cfdee46 Mon Sep 17 00:00:00 2001
From: Jianhui Zhao <jianhuizhao329@gmail.com>
Date: Wed, 24 Apr 2019 23:41:54 +0800
Subject: [PATCH] Optimize code

Signed-off-by: Jianhui Zhao <jianhuizhao329@gmail.com>
---
 http.go | 27 +++++++++++++++++++++------
 main.go | 52 ++++++++++++++++++++++++----------------------------
 2 files changed, 45 insertions(+), 34 deletions(-)

diff --git a/http.go b/http.go
index 62c06f6..cc8da27 100644
--- a/http.go
+++ b/http.go
@@ -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))
 	}
 }
diff --git a/main.go b/main.go
index 7246e96..5a0d1d8 100644
--- a/main.go
+++ b/main.go
@@ -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