From 5f458dd1a6d631f324e4af9a4f5429ffdf199342 Mon Sep 17 00:00:00 2001
From: Endless Paradox <129645532+EndlessParadox1@users.noreply.github.com>
Date: Mon, 11 Mar 2024 22:22:58 +0800
Subject: [PATCH] feat(auth): add proxy-server authentication (#3877)

---
 auth.go      | 21 +++++++++++++++++++++
 auth_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/auth.go b/auth.go
index 2503c51..cc6c5a7 100644
--- a/auth.go
+++ b/auth.go
@@ -15,6 +15,7 @@ import (
 
 // AuthUserKey is the cookie name for user credential in basic auth.
 const AuthUserKey = "user"
+const AuthProxyUserKey = "proxy_user"
 
 // Accounts defines a key/value for user/pass list of authorized logins.
 type Accounts map[string]string
@@ -89,3 +90,23 @@ func authorizationHeader(user, password string) string {
 	base := user + ":" + password
 	return "Basic " + base64.StdEncoding.EncodeToString(bytesconv.StringToBytes(base))
 }
+
+func BasicAuthForProxy(accounts Accounts, realm string) HandlerFunc {
+	if realm == "" {
+		realm = "Proxy Authorization Required"
+	}
+	realm = "Basic realm=" + strconv.Quote(realm)
+	pairs := processAccounts(accounts)
+	return func(c *Context) {
+		proxyUser, found := pairs.searchCredential(c.requestHeader("Proxy-Authorization"))
+		if !found {
+			// Credentials doesn't match, we return 407 and abort handlers chain.
+			c.Header("Proxy-Authenticate", realm)
+			c.AbortWithStatus(http.StatusProxyAuthRequired)
+			return
+		}
+		// The proxy_user credentials was found, set proxy_user's id to key AuthProxyUserKey in this context, the proxy_user's id can be read later using
+		// c.MustGet(gin.AuthProxyUserKey).
+		c.Set(AuthProxyUserKey, proxyUser)
+	}
+}
diff --git a/auth_test.go b/auth_test.go
index 42b6f8f..f717592 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -137,3 +137,40 @@ func TestBasicAuth401WithCustomRealm(t *testing.T) {
 	assert.Equal(t, http.StatusUnauthorized, w.Code)
 	assert.Equal(t, "Basic realm=\"My Custom \\\"Realm\\\"\"", w.Header().Get("WWW-Authenticate"))
 }
+
+func TestBasicAuthForProxySucceed(t *testing.T) {
+	accounts := Accounts{"admin": "password"}
+	router := New()
+	router.Use(BasicAuthForProxy(accounts, ""))
+	router.Any("/*proxyPath", func(c *Context) {
+		c.String(http.StatusOK, c.MustGet(AuthProxyUserKey).(string))
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest("GET", "/test", nil)
+	req.Header.Set("Proxy-Authorization", authorizationHeader("admin", "password"))
+	router.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+	assert.Equal(t, "admin", w.Body.String())
+}
+
+func TestBasicAuthForProxy407(t *testing.T) {
+	called := false
+	accounts := Accounts{"foo": "bar"}
+	router := New()
+	router.Use(BasicAuthForProxy(accounts, ""))
+	router.Any("/*proxyPath", func(c *Context) {
+		called = true
+		c.String(http.StatusOK, c.MustGet(AuthProxyUserKey).(string))
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest("GET", "/test", nil)
+	req.Header.Set("Proxy-Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte("admin:password")))
+	router.ServeHTTP(w, req)
+
+	assert.False(t, called)
+	assert.Equal(t, http.StatusProxyAuthRequired, w.Code)
+	assert.Equal(t, "Basic realm=\"Proxy Authorization Required\"", w.Header().Get("Proxy-Authenticate"))
+}