Add referer feature to config and update request handling

arloor · Dec 27, 2023 · 37fe016 · 37fe016
1 parent 1e47a61
commit 37fe016
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 4 deletions.
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -13,6 +13,7 @@
             "program": "${workspaceFolder}/cmd/${workspaceFolderBasename}",
             "args": [
                 "--addr=localhost:7788",
+                "--refer=arloor"
                 // "--tls=true"
             ]
         }

diff --git a/internal/config/config.go b/internal/config/config.go
@@ -17,6 +17,7 @@ type Config struct {
 	LogPath   string `yaml:"log"`
 	WebPath   string `yaml:"content"`
 	BasicAuth string `yaml:"auth"`
+	Refer     string `yaml:"refer"`
 }
 
 var Instance Config
@@ -29,6 +30,7 @@ func init() {
 	flag.StringVar(&Instance.LogPath, "log", "/tmp/proxy.log", "日志文件路径")
 	flag.StringVar(&Instance.WebPath, "content", ".", "文件服务器目录")
 	flag.StringVar(&Instance.BasicAuth, "auth", "", "Basic Auth Header")
+	flag.StringVar(&Instance.Refer, "refer", "", "本站的referer特征")
 	flag.Parse()
 	initLog()
 	out, err := yaml.Marshal(Instance)

diff --git a/internal/server/httpHandlerFunc.go b/internal/server/httpHandlerFunc.go
@@ -24,6 +24,10 @@ func writeIp(w http.ResponseWriter, r *http.Request) {
 func fileHandlerFunc() http.HandlerFunc {
 
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if config.Instance.Refer != "" && r.Header.Get("referer") != "" && !strings.Contains(r.Header.Get("referer"),config.Instance.Refer){
+			HttpRequst.WithLabelValues(r.Header.Get("referer"), r.URL.Path).Inc()
+			HttpRequst.WithLabelValues("all", "all").Inc()
+		}
 
 		if containsDotDot(r.URL.Path) {
 			// Too many programs use r.URL.Path to construct the argument to

diff --git a/internal/server/server.go b/internal/server/server.go
@@ -18,10 +18,10 @@ var ssl_last_cert_update time.Time = time.Now()
 const ssl_cert_update_interval = 5 * time.Hour
 
 var (
-	// HttpRequst = promauto.NewCounterVec(prometheus.CounterOpts{
-	// 	Name: "req_from_out",
-	// 	Help: "Number of HTTP requests received",
-	// }, []string{"referer", "path"})
+	HttpRequst = promauto.NewCounterVec(prometheus.CounterOpts{
+		Name: "req_from_out",
+		Help: "Number of HTTP requests received",
+	}, []string{"referer", "path"})
 	ProxyAccess = promauto.NewCounterVec(prometheus.CounterOpts{
 		Name: "proxy_access",
 		Help: "num proxy_access",