From 9345fe39c9e585340173e53efcce3577a214e3da Mon Sep 17 00:00:00 2001
From: Wang Yan <wangyan@vmware.com>
Date: Fri, 15 Nov 2024 21:40:27 +0800
Subject: [PATCH] update csrf key generation (#21154)

* update csrf key generation

Fixes #21060

Do not generate a random key if the provided key has an invalid length.

Signed-off-by: wang yan <wangyan@vmware.com>

* fix ut check

Signed-off-by: wang yan <wangyan@vmware.com>

---------

Signed-off-by: wang yan <wangyan@vmware.com>
---
 src/server/middleware/csrf/csrf.go      | 11 +++++++--
 src/server/middleware/csrf/csrf_test.go | 32 ++++++++++++++++++-------
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/src/server/middleware/csrf/csrf.go b/src/server/middleware/csrf/csrf.go
index 459b99603..82b1755b4 100644
--- a/src/server/middleware/csrf/csrf.go
+++ b/src/server/middleware/csrf/csrf.go
@@ -67,9 +67,16 @@ func attach(handler http.Handler) http.Handler {
 func Middleware() func(handler http.Handler) http.Handler {
 	once.Do(func() {
 		key := os.Getenv(csrfKeyEnv)
-		if len(key) != 32 {
-			log.Warningf("Invalid CSRF key from environment: %s, generating random key...", key)
+		if len(key) == 0 {
 			key = utils.GenerateRandomString()
+		} else if len(key) != 32 {
+			log.Errorf("Invalid CSRF key length from the environment: %s. Please ensure the key length is 32 characters.", key)
+			protect = func(_ http.Handler) http.Handler {
+				return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+					lib_http.SendError(w, errors.New("invalid CSRF key length from the environment. Please ensure the key length is 32 characters"))
+				})
+			}
+			return
 		}
 		secureFlag = secureCookie()
 		protect = csrf.Protect([]byte(key), csrf.RequestHeader(tokenHeader),
diff --git a/src/server/middleware/csrf/csrf_test.go b/src/server/middleware/csrf/csrf_test.go
index a4975ae58..1b9e75512 100644
--- a/src/server/middleware/csrf/csrf_test.go
+++ b/src/server/middleware/csrf/csrf_test.go
@@ -4,6 +4,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -14,6 +15,10 @@ import (
 	_ "github.com/goharbor/harbor/src/pkg/config/inmemory"
 )
 
+func resetMiddleware() {
+	once = sync.Once{}
+}
+
 func TestMain(m *testing.M) {
 	test.InitDatabaseFromEnv()
 	conf := map[string]interface{}{}
@@ -32,7 +37,6 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 }
 
 func TestMiddleware(t *testing.T) {
-	srv := Middleware()(&handler{})
 	cases := []struct {
 		req         *http.Request
 		statusCode  int
@@ -60,6 +64,7 @@ func TestMiddleware(t *testing.T) {
 		},
 	}
 	for _, c := range cases {
+		srv := Middleware()(&handler{})
 		rec := httptest.NewRecorder()
 		srv.ServeHTTP(rec, c.req)
 		assert.Equal(t, c.statusCode, rec.Result().StatusCode)
@@ -67,13 +72,24 @@ func TestMiddleware(t *testing.T) {
 	}
 }
 
-func hasCookie(resp *http.Response, name string) bool {
-	for _, c := range resp.Cookies() {
-		if c != nil && c.Name == name {
-			return true
-		}
-	}
-	return false
+func TestMiddlewareInvalidKey(t *testing.T) {
+	originalEnv := os.Getenv(csrfKeyEnv)
+	defer os.Setenv(csrfKeyEnv, originalEnv)
+
+	t.Run("invalid CSRF key", func(t *testing.T) {
+		os.Setenv(csrfKeyEnv, "invalidkey")
+		resetMiddleware()
+		middleware := Middleware()
+		testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			t.Error("handler should not be reached when CSRF key is invalid")
+		})
+
+		handler := middleware(testHandler)
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+		rec := httptest.NewRecorder()
+		handler.ServeHTTP(rec, req)
+		assert.Equal(t, http.StatusInternalServerError, rec.Code)
+	})
 }
 
 func TestSecureCookie(t *testing.T) {