feat: bcs-common新增lane公用方法

bcs-common/pkg/header/header.go

@@ -0,0 +1,170 @@
+/*
+ * Tencent is pleased to support the open source community by making Blueking Container Service available.
+ * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
+ * Licensed under the MIT License (the "License"); you may not use this file except
+ * in compliance with the License. You may obtain a copy of the License at
+ * http://opensource.org/licenses/MIT
+ * Unless required by applicable law or agreed to in writing, software distributed under
+ * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Package header contains header related constants and functions
+package header
+
+import (
+	"context"
+	"net/http"
+	"net/textproto"
+	"strings"
+
+	"github.com/grpc-ecosystem/grpc-gateway/runtime"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+	"k8s.io/klog/v2"
+)
+
+// ContextValueKey is the key for context value
+type ContextValueKey string
+
+const (
+	// AuthUserKey is the key for user in context
+	AuthUserKey ContextValueKey = "X-Bcs-User"
+	// LaneKey is the key for lane
+	LaneKey ContextValueKey = "X-Lane"
+	// InnerClientHeaderKey is the key for client in header
+	InnerClientHeaderKey = "X-Bcs-Client"
+	// AuthorizationHeaderKey is the key for authorization in header
+	AuthorizationHeaderKey = "Authorization"
+	// CustomUsernameHeaderKey is the key for custom username in heade
+	CustomUsernameHeaderKey = "X-Bcs-Username"
+
+	// RequestIDKey x-request-id header key
+	RequestIDKey = "X-Request-Id"
+	// UsernameKey x-project-username header key
+	UsernameKey = "X-Project-Username"
+
+	// TraceparentKey traceparent header key
+	TraceparentKey = "Traceparent"
+	// TracestateKey tracestate header key
+	TracestateKey = "Tracestate"
+
+	// LaneIDPrefix 染色的header前缀
+	LaneIDPrefix = "X-Lane-"
+)
+
+var (
+	passthroughHeaderKeys = []string{RequestIDKey, TraceparentKey, TracestateKey}
+)
+
+// CustomHeaderMatcher for http header
+func CustomHeaderMatcher(key string) (string, bool) {
+	key = textproto.CanonicalMIMEHeaderKey(key)
+	if strings.HasPrefix(key, LaneIDPrefix) {
+		return key, true
+	}
+	switch key {
+	case RequestIDKey:
+		return RequestIDKey, true
+	case UsernameKey:
+		return UsernameKey, true
+	case InnerClientHeaderKey:
+		return InnerClientHeaderKey, true
+	case CustomUsernameHeaderKey:
+		return CustomUsernameHeaderKey, true
+	case TraceparentKey:
+		return TraceparentKey, true
+	case TracestateKey:
+		return TracestateKey, true
+	default:
+		return runtime.DefaultHeaderMatcher(key)
+	}
+}
+
+// LaneHeaderInterceptor 透传泳道header
+// x-lane-xxx,x-request-id,traceparent
+// 从incoming context获取request-id,以及包含泳道前缀的header
+func LaneHeaderInterceptor() grpc.UnaryClientInterceptor {
+	// 创建一个集合用于快速查找需要透传的键
+	passthroughHeaderKeysSet := make(map[string]struct{}, len(passthroughHeaderKeys))
+	for _, key := range passthroughHeaderKeys {
+		passthroughHeaderKeysSet[textproto.CanonicalMIMEHeaderKey(key)] = struct{}{}
+	}
+	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
+		invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+
+		passthroughHeader := map[string]string{}
+		if md, ok := metadata.FromIncomingContext(ctx); ok {
+			for k, v := range md {
+				if len(v) > 0 {
+					tmpKey := textproto.CanonicalMIMEHeaderKey(k)
+					// 检查是否包含laneID前缀
+					if strings.HasPrefix(tmpKey, LaneIDPrefix) {
+						passthroughHeader[k] = v[0]
+						continue // 已找到需要透传的，无需继续检查此键
+					}
+					// 检查是否在需要透传的键集合中
+					if _, exists := passthroughHeaderKeysSet[tmpKey]; exists {
+						passthroughHeader[k] = v[0]
+					}
+				}
+			}
+		}
+		klog.Infof("passthrough header: %v", passthroughHeader)
+		outgoingCtx := metadata.NewOutgoingContext(ctx, metadata.New(passthroughHeader))
+		// 调用下一个处理器
+		return invoker(outgoingCtx, method, req, reply, cc, opts...)
+	}
+}
+
+// GetLaneIDByCtx get lane id by ctx
+func GetLaneIDByCtx(ctx context.Context) map[string]string {
+	// http 格式的以key value方式存放，eg: key: X-Lane value: X-Lane-xxx:xxx
+	v, ok := ctx.Value(LaneKey).(string)
+	if ok || v != "" {
+		result := strings.Split(v, ":")
+		if len(result) != 2 {
+			return nil
+		}
+		return map[string]string{result[0]: result[1]}
+	}
+	if !ok || v == "" {
+		return grpcLaneIDValue(ctx)
+	}
+	return nil
+}
+
+// grpcLaneIDValue grpc lane id 处理
+func grpcLaneIDValue(ctx context.Context) map[string]string {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if ok {
+		for k, v := range md {
+			tmpKey := textproto.CanonicalMIMEHeaderKey(k)
+			if strings.HasPrefix(tmpKey, LaneIDPrefix) && len(v) > 0 {
+				return map[string]string{tmpKey: md.Get(k)[0]}
+			}
+		}
+	}
+	return nil
+}
+
+// GetLaneIDByHeader get lane id by header
+func GetLaneIDByHeader(h http.Header) (string, string) {
+	for k, v := range h {
+		if strings.HasPrefix(k, LaneIDPrefix) && len(v) > 0 {
+			return k, v[0]
+		}
+	}
+	return "", ""
+}
+
+// WithLaneIdCtx ctx lane id
+func WithLaneIdCtx(ctx context.Context, h http.Header) context.Context {
+	for k, v := range h {
+		if strings.HasPrefix(k, LaneIDPrefix) && len(v) > 0 {
+			ctx = context.WithValue(ctx, LaneKey, k+":"+v[0])
+		}
+	}
+	return ctx
+}