From 8b783cbe56de7f52dc1fc99162cc83b0c93c499c Mon Sep 17 00:00:00 2001
From: flymin <905370712@qq.com>
Date: Tue, 24 May 2022 23:33:18 +0800
Subject: [PATCH] add access control with md5sum

---
 README.md            | 11 +++++++++++
 cli/server/server.go |  5 +++++
 config/config.go     |  5 +++++
 handle/handle.go     | 40 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 61 insertions(+)

diff --git a/README.md b/README.md
index f373092..5e1e50e 100644
--- a/README.md
+++ b/README.md
@@ -64,6 +64,16 @@ TLS_MIN_VERS=
 #   To accept missing referrer header, add a blank entry (start comma):
 #   'REFERRERS=,http://localhost,https://another.name'
 REFERRERS=
+
+# Use key / code parameter in the request URL for access control. The code is
+# computed by requested PATH and your key.
+# Example:
+#   ACCESS_KEY=username
+#   To access your file, either access:
+#   http://$HOST:$PORT/my/place/my.file?key=username
+#   or access (md5sum of "/my/place/my.fileusername"):
+#   http://$HOST:$PORT/my/place/my.file?code=44356A355E89D9EE7B2D5687E48024B0
+ACCESS_KEY=
 ```
 
 ### YAML Configuration File
@@ -85,6 +95,7 @@ tls-cert: ""
 tls-key: ""
 tls-min-vers: ""
 url-prefix: ""
+access_key: ""
 ```
 
 Example configuration with possible alternative values:
diff --git a/cli/server/server.go b/cli/server/server.go
index 19d55cb..ea754ad 100644
--- a/cli/server/server.go
+++ b/cli/server/server.go
@@ -67,6 +67,11 @@ func handlerSelector() (handler http.HandlerFunc) {
 		handler = handle.AddCorsWildcardHeaders(handler)
 	}
 
+	// If configured, apply key code access control.
+	if "" != config.Get.AccessKey {
+		handler = handle.AddAccessKey(handler, config.Get.AccessKey)
+	}
+
 	return
 }
 
diff --git a/config/config.go b/config/config.go
index 9b8ef00..216e174 100644
--- a/config/config.go
+++ b/config/config.go
@@ -28,6 +28,7 @@ var (
 		TLSMinVersStr string   `yaml:"tls-min-vers"`
 		URLPrefix     string   `yaml:"url-prefix"`
 		Referrers     []string `yaml:"referrers"`
+		AccessKey     string   `yaml:"accessKey"`
 	}
 )
 
@@ -43,6 +44,7 @@ const (
 	tlsKeyKey      = "TLS_KEY"
 	tlsMinVersKey  = "TLS_MIN_VERS"
 	urlPrefixKey   = "URL_PREFIX"
+	accessKeyKey   = "ACCESS_KEY"
 )
 
 var (
@@ -57,6 +59,7 @@ var (
 	defaultTLSMinVers  = ""
 	defaultURLPrefix   = ""
 	defaultCors        = false
+	defaultAccessKey   = ""
 )
 
 func init() {
@@ -76,6 +79,7 @@ func setDefaults() {
 	Get.TLSMinVersStr = defaultTLSMinVers
 	Get.URLPrefix = defaultURLPrefix
 	Get.Cors = defaultCors
+	Get.AccessKey = defaultAccessKey
 }
 
 // Load the configuration file.
@@ -126,6 +130,7 @@ func overrideWithEnvVars() {
 	Get.TLSMinVersStr = envAsStr(tlsMinVersKey, Get.TLSMinVersStr)
 	Get.URLPrefix = envAsStr(urlPrefixKey, Get.URLPrefix)
 	Get.Referrers = envAsStrSlice(referrersKey, Get.Referrers)
+	Get.AccessKey = envAsStr(accessKeyKey, Get.AccessKey)
 }
 
 // validate the configuration.
diff --git a/handle/handle.go b/handle/handle.go
index e347b55..fb59298 100644
--- a/handle/handle.go
+++ b/handle/handle.go
@@ -2,6 +2,7 @@ package handle
 
 import (
 	"crypto/tls"
+	"crypto/md5"
 	"fmt"
 	"log"
 	"net/http"
@@ -152,6 +153,45 @@ func AddCorsWildcardHeaders(serve http.HandlerFunc) http.HandlerFunc {
 	}
 }
 
+// Access Control through url parameters. The access key is set by ACCESS_KEY.
+// md5sum is computed by queried path + access key
+// (e.g. "/my/file" + ACCESS_KEY)
+func AddAccessKey(serve http.HandlerFunc, accessKey string) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		// Get key or md5sum from this access.
+		keys, key_ok := r.URL.Query()["key"]
+		var code string
+    	if !key_ok || len(keys[0]) < 1 {
+			// In case a code is provided
+			codes, code_ok := r.URL.Query()["code"]
+    		if !code_ok || len(codes[0]) < 1 {
+				http.NotFound(w, r)
+				return
+			}
+			code = codes[0]
+		} else {
+			// In case a key is provided, convert to code.
+			data := []byte(r.URL.Path + keys[0])
+			hash := md5.Sum(data)
+			code = fmt.Sprintf("%x", hash)
+		}
+		code = strings.ToUpper(code)
+
+		// Compute the correct md5sum of this access.
+		local_data := []byte(r.URL.Path + accessKey)
+		hash := md5.Sum(local_data)
+		local_code := fmt.Sprintf("%x", hash)
+		local_code = strings.ToUpper(local_code)
+		
+		// Compare the two.
+		if code != local_code {
+			http.NotFound(w, r)
+			return
+		}
+		serve(w, r)
+	}
+}
+
 // Listening function for serving the handler function.
 func Listening() ListenerFunc {
 	return func(binding string, handler http.HandlerFunc) error {