grpcx/header.go

package grpcx

import (
	"context"
	"encoding/json"
	"errors"

	"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 = errors.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 = errors.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)
// }