From 3f3a29573d91d2226a01135c99c9b3335eaa1c37 Mon Sep 17 00:00:00 2001
From: weizhichen <weizhichen@microsoft.com>
Date: Tue, 26 Mar 2024 12:14:29 +0000
Subject: [PATCH] fix: strip service account token

---
 pkg/csi-common/utils.go      | 48 +++++++++++++++++++++++++++++++++++-
 pkg/csi-common/utils_test.go | 38 ++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/pkg/csi-common/utils.go b/pkg/csi-common/utils.go
index b46146dabd..e25af77cd0 100644
--- a/pkg/csi-common/utils.go
+++ b/pkg/csi-common/utils.go
@@ -17,6 +17,7 @@ limitations under the License.
 package csicommon
 
 import (
+	"encoding/json"
 	"fmt"
 	"net"
 	"os"
@@ -98,7 +99,7 @@ func getLogLevel(method string) int32 {
 func LogGRPC(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
 	level := klog.Level(getLogLevel(info.FullMethod))
 	klog.V(level).Infof("GRPC call: %s", info.FullMethod)
-	klog.V(level).Infof("GRPC request: %s", protosanitizer.StripSecrets(req))
+	klog.V(level).Infof("GRPC request: %s", StripSensitiveValue(protosanitizer.StripSecrets(req), "csi.storage.k8s.io/serviceAccount.tokens"))
 
 	resp, err := handler(ctx, req)
 	if err != nil {
@@ -108,3 +109,48 @@ func LogGRPC(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, h
 	}
 	return resp, err
 }
+
+type stripSensitiveValue struct {
+	// volume_context[key] is the value to be stripped.
+	key string
+	// req is the csi grpc request stripped by `protosanitizer.StripSecrets`
+	req fmt.Stringer
+}
+
+func StripSensitiveValue(req fmt.Stringer, key string) fmt.Stringer {
+	return &stripSensitiveValue{
+		key: key,
+		req: req,
+	}
+}
+
+func (s *stripSensitiveValue) String() string {
+	return stripSensitiveValueByKey(s.req, s.key)
+}
+
+func stripSensitiveValueByKey(req fmt.Stringer, key string) string {
+	var parsed map[string]interface{}
+
+	err := json.Unmarshal([]byte(req.String()), &parsed)
+	if err != nil || parsed == nil {
+		return req.String()
+	}
+
+	volumeContext, ok := parsed["volume_context"].(map[string]interface{})
+	if !ok {
+		return req.String()
+	}
+
+	if _, ok := volumeContext[key]; !ok {
+		return req.String()
+	}
+
+	volumeContext[key] = "***stripped***"
+
+	b, err := json.Marshal(parsed)
+	if err != nil {
+		return req.String()
+	}
+
+	return string(b)
+}
diff --git a/pkg/csi-common/utils_test.go b/pkg/csi-common/utils_test.go
index 0bbde2de8e..7ab9c7c2c0 100644
--- a/pkg/csi-common/utils_test.go
+++ b/pkg/csi-common/utils_test.go
@@ -127,6 +127,44 @@ func TestLogGRPC(t *testing.T) {
 			},
 			`GRPC request: {"starting_token":"testtoken"}`,
 		},
+		{
+			"NodeStageVolumeRequest with service account token",
+			&csi.NodeStageVolumeRequest{
+				VolumeContext: map[string]string{
+					"csi.storage.k8s.io/serviceAccount.tokens": "testtoken",
+					"csi.storage.k8s.io/testfield":             "testvalue",
+				},
+				XXX_sizecache: 100,
+			},
+			`GRPC request: {"volume_context":{"csi.storage.k8s.io/serviceAccount.tokens":"***stripped***","csi.storage.k8s.io/testfield":"testvalue"}}`,
+		},
+		{
+			"NodePublishVolumeRequest with service account token",
+			&csi.NodePublishVolumeRequest{
+				VolumeContext: map[string]string{
+					"csi.storage.k8s.io/serviceAccount.tokens": "testtoken",
+					"csi.storage.k8s.io/testfield":             "testvalue",
+				},
+				XXX_sizecache: 100,
+			},
+			`GRPC request: {"volume_context":{"csi.storage.k8s.io/serviceAccount.tokens":"***stripped***","csi.storage.k8s.io/testfield":"testvalue"}}`,
+		},
+		{
+			"with secrets and service account token",
+			&csi.NodeStageVolumeRequest{
+				VolumeId: "vol_1",
+				Secrets: map[string]string{
+					"account_name": "k8s",
+					"account_key":  "testkey",
+				},
+				VolumeContext: map[string]string{
+					"csi.storage.k8s.io/serviceAccount.tokens": "testtoken",
+					"csi.storage.k8s.io/testfield":             "testvalue",
+				},
+				XXX_sizecache: 100,
+			},
+			`GRPC request: {"secrets":"***stripped***","volume_context":{"csi.storage.k8s.io/serviceAccount.tokens":"***stripped***","csi.storage.k8s.io/testfield":"testvalue"},"volume_id":"vol_1"}`,
+		},
 	}
 
 	for _, test := range tests {