提交

header.go

@@ -0,0 +1,93 @@
+package grpcx
+
+import (
+	"context"
+	"encoding/json"
+	"xiaoxin-plus/pkg/errorx"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+type GrpcReqHeader struct {
+	Key   string
+	Value string
+}
+
+// 客户端写入请求header
+// 注意事项
+func WriteHeader(ctx context.Context, headers ...GrpcReqHeader) context.Context {
+
+	// 先读取，然后替换
+	oldHeader, _ := readHeaderAll(ctx)
+	header := map[string]string{}
+	for _, val := range oldHeader {
+		header[val.Key] = val.Value
+	}
+	for _, val := range headers {
+		header["grpcx_"+val.Key] = val.Value
+	}
+
+	pp := []string{}
+	for key, val := range header {
+		pp = append(pp, key, val)
+	}
+
+	headerData := metadata.Pairs(pp...)
+	return metadata.NewOutgoingContext(ctx, headerData)
+}
+
+func readHeaderAll(ctx context.Context) (header []GrpcReqHeader, err error) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		// 本地的方式
+		md, ok = metadata.FromOutgoingContext(ctx)
+		if !ok {
+			err = errorx.New("metadata.FromIncomingContext error")
+			return
+		}
+	}
+
+	for k, v := range md {
+		header = append(header, GrpcReqHeader{Key: k, Value: v[0]})
+	}
+	return
+}
+
+// 服务端读取请求header
+func ReadHeader(ctx context.Context, keys ...string) (header []GrpcReqHeader, err error) {
+	// rpc的方式
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		// 本地的方式
+		md, ok = metadata.FromOutgoingContext(ctx)
+		if !ok {
+			err = errorx.New("metadata.FromIncomingContext error")
+			return
+		}
+	}
+	for _, key := range keys {
+		if v, ok := md["grpcx_"+key]; ok {
+			header = append(header, GrpcReqHeader{Key: key, Value: v[0]})
+		}
+	}
+	// if v, ok := md["grpc_client_header"]; ok {
+	// 	_ = json.Unmarshal([]byte(v[0]), &header)
+	// }
+	return
+}
+
+type GrpcRespHeader struct{}
+
+func ServerWriteRespHeader(ctx context.Context, header GrpcRespHeader) {
+	h, _ := json.Marshal(header)
+	headerData := metadata.Pairs("grpc_server_header", string(h))
+	grpc.SendHeader(ctx, headerData)
+}
+
+// 需要添加到请求中
+// func ClientReadRespHeader(ctx context.Context, h *GrpcRespHeader) grpc.CallOption {
+// 	var header, trailer metadata.MD
+// 	r, err := l.svcCtx.RpcOrgansLogin.In(l.ctx, in, grpc.Header(&header), grpc.Trailer(&trailer))
+// 	fmt.Println(header, trailer)
+// }