配置项提取出来

LubyRuffy · Oct 17, 2022 · 99637c0 · 99637c0
1 parent 01166dd
commit 99637c0
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 19 deletions.
diff --git a/api/api.go b/api/api.go
@@ -11,7 +11,6 @@ import (
 	"github.com/golang-jwt/jwt/v4"
 	"github.com/kabukky/httpscerts"
 	"github.com/mitchellh/mapstructure"
-	"github.com/spf13/viper"
 	"gorm.io/gorm"
 	"log"
 	"net/http"
@@ -20,6 +19,9 @@ import (
 )
 
 var (
+	EnableTls   bool // 启动tls开关
+	EnableDebug bool // 启动debug开关
+
 	srv         *http.Server // http服务器
 	Version     = "v0.1.6"
 	Prefix      = "/api"
@@ -291,7 +293,7 @@ func defaultHandler(c *gin.Context) {
 }
 
 func Start(addr string) error {
-	if viper.GetBool("debug.gin") {
+	if EnableDebug {
 		gin.SetMode(gin.DebugMode)
 	} else {
 		gin.SetMode(gin.ReleaseMode)
@@ -318,7 +320,7 @@ func Start(addr string) error {
 	log.Println("api server listened at:", addr)
 
 	var err error
-	if viper.GetBool("tls") {
+	if EnableTls {
 		err = router.RunTLS(addr, "cert.pem", "key.pem")
 		if err != nil {
 			if err = httpscerts.Generate("cert.pem", "key.pem", ""); err != nil {

diff --git a/api/check.go b/api/check.go
@@ -35,6 +35,8 @@ func headerString(r *http.Response) string {
 }
 
 var (
+	EnableErrorCheckLog bool // 是否启用错误日志：在检查失败的情况下也记录日志
+
 	once               sync.Once
 	myPublicIP         string // 公网ip，用于检查代理是否匿名
 	defaultCheckUrl    = "http://ip.bmh.im/h"
@@ -150,17 +152,20 @@ func checkProtocolHost(protocol string, host string) *proxyResult {
 	}
 
 	var err error
-	//defer func() {
-	//	errStr := ""
-	//	if err != nil {
-	//		errStr = err.Error()
-	//	}
-	//	models.GetDB().Create(&models.CheckLog{
-	//		ProxyType: protocol,
-	//		Host:      host,
-	//		Error:     errStr,
-	//	})
-	//}()
+	defer func() {
+		if EnableErrorCheckLog {
+			errStr := ""
+			if err != nil {
+				errStr = err.Error()
+			}
+			models.GetDB().Create(&models.CheckLog{
+				ProxyType: protocol,
+				Host:      host,
+				Error:     errStr,
+			})
+		}
+
+	}()
 
 	if transportFunc, ok := Transports[protocol]; ok {
 		client := defaultHttpClient(transportFunc(host))

diff --git a/go.sum b/go.sum
@@ -38,8 +38,6 @@ cloud.google.com/go/storage v1.14.0/go.mod h1:GrKmX003DSIwi9o29oFT7YDnHYwZoctc3f
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
-github.com/LubyRuffy/gorestful v0.0.0-20221014061004-7e3441c69ae5 h1:GNWzrYPtEsD5/qZyTt4MjqZGRF1ychnR1fkFhdl6P5k=
-github.com/LubyRuffy/gorestful v0.0.0-20221014061004-7e3441c69ae5/go.mod h1:4h7u9GJv8xLmhi97PtpyvT8wyfnSunC1alLXERK9CrU=
 github.com/LubyRuffy/gorestful v0.0.0-20221016124020-540a7f24139f h1:lnEIxeFrznUTp+RAr4r/X5bfbP3f2Rs2GDUhGdWDwnI=
 github.com/LubyRuffy/gorestful v0.0.0-20221016124020-540a7f24139f/go.mod h1:4h7u9GJv8xLmhi97PtpyvT8wyfnSunC1alLXERK9CrU=
 github.com/LubyRuffy/myip v0.0.0-20220808131427-cd1832923c09 h1:JJP4NZjc+n9D1ShFjaW9RpmWhGm96IhKTslfWaG8rXA=
@@ -274,7 +272,6 @@ github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PK
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/subosito/gotenv v1.3.0 h1:mjC+YW8QpAdXibNi+vNWgzmgBH4+5l5dCXv8cNysBLI=
 github.com/subosito/gotenv v1.3.0/go.mod h1:YzJjq/33h7nrwdY+iHMhEOEEbW0ovIz0tB6t6PwAXzs=
-github.com/ugorji/go v1.2.7 h1:qYhyWUUd6WbiM+C6JZAUkIJt/1WrjzNHY9+KCIjVqTo=
 github.com/ugorji/go v1.2.7/go.mod h1:nF9osbDWLy6bDVv/Rtoh6QgnvNDpmCalQV5urGCCS6M=
 github.com/ugorji/go/codec v1.2.7 h1:YPXUKf7fYbp/y8xloBqZOw2qaVggbfwMlI8WM3wZUJ0=
 github.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY=
@@ -368,8 +365,6 @@ golang.org/x/net v0.0.0-20201209123823-ac852fbbde11/go.mod h1:m0MpNAwzfU5UDzcl9v
 golang.org/x/net v0.0.0-20201224014010-6772e930b67b/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
-golang.org/x/net v0.0.0-20221012135044-0b7e1fb9d458 h1:MgJ6t2zo8v0tbmLCueaCbF1RM+TtB0rs3Lv8DGtOIpY=
-golang.org/x/net v0.0.0-20221012135044-0b7e1fb9d458/go.mod h1:YDH+HFinaLZZlnHAfSS6ZXJJ9M9t4Dl22yv3iI2vPwk=
 golang.org/x/net v0.0.0-20221014081412-f15817d10f9b h1:tvrvnPFcdzp294diPnrdZZZ8XUt2Tyj7svb7X52iDuU=
 golang.org/x/net v0.0.0-20221014081412-f15817d10f9b/go.mod h1:YDH+HFinaLZZlnHAfSS6ZXJJ9M9t4Dl22yv3iI2vPwk=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=

diff --git a/main.go b/main.go
@@ -25,12 +25,14 @@ func main() {
 	viper.SetDefault("dbfile", "rproxy.sqlite")
 	viper.SetDefault("debug.dbsql", false)
 	viper.SetDefault("tls", false)
+	viper.SetDefault("logerror", false)
 
 	viper.AddConfigPath(filepath.Dir(os.Args[0]))
 	viper.SetConfigType("yaml")
 	viper.SetConfigName("config")
 
 	pflag.Bool("tls", false, "enable tls")
+	pflag.Bool("logerror", false, "enable error check log")
 	pflag.String("addr", ":8088", "bind addr")
 	pflag.String("dbfile", "rproxy.sqlite", "sqlite database file")
 	pflag.Parse()
@@ -50,6 +52,16 @@ func main() {
 	// 检查数据库
 	go ipdb.UpdateIpDatabase()
 
+	if viper.GetBool("logerror") {
+		api.EnableErrorCheckLog = true // 打开proxy检查错误的日志记录
+	}
+	if viper.GetBool("tls") {
+		api.EnableTls = true
+	}
+	if viper.GetBool("debug.gin") {
+		api.EnableDebug = true
+	}
+
 	// 启动web
 	if err = api.Start(viper.GetString("addr")); err != nil {
 		panic(err)