github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/nphttp2/meta_api.go

/*
 * Copyright 2022 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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 nphttp2

import (
	"context"
	"net"

	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/codes"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/metadata"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/status"
	"github.com/cloudwego/kitex/pkg/streaming"
)

// SetHeader sets the header metadata to be sent from the server to the client.
// The context provided must be the context passed to the server's handler.
//
// Streaming RPCs should prefer the SetHeader method of the ServerStream.
//
// When called multiple times, all the provided metadata will be merged.  All
// the metadata will be sent out when one of the following happens:
//
//   - grpc.SendHeader is called, or for streaming handlers, stream.SendHeader.
//   - The first response message is sent.  For unary handlers, this occurs when
//     the handler returns; for streaming handlers, this can happen when stream's
//     SendMsg method is called.
//   - An RPC status is sent out (error or success).  This occurs when the handler
//     returns.
//
// SetHeader will fail if called after any of the events above.
//
// The error returned is compatible with the status package.  However, the
// status code will often not match the RPC status as seen by the client
// application, and therefore, should not be relied upon for this purpose.
func SetHeader(ctx context.Context, md metadata.MD) error {
	if md.Len() == 0 {
		return nil
	}
	stream := serverTransportStreamFromContext(ctx)
	if stream == nil {
		return status.Errorf(codes.Internal, "grpc: failed to fetch the stream from the context %v", ctx)
	}
	return stream.SetHeader(md)
}

// SendHeader sends header metadata. It may be called at most once, and may not
// be called after any event that causes headers to be sent (see SetHeader for
// a complete list).  The provided md and headers set by SetHeader() will be
// sent.
//
// The error returned is compatible with the status package.  However, the
// status code will often not match the RPC status as seen by the client
// application, and therefore, should not be relied upon for this purpose.
func SendHeader(ctx context.Context, md metadata.MD) error {
	stream := serverTransportStreamFromContext(ctx)
	if stream == nil {
		return status.Errorf(codes.Internal, "grpc: failed to fetch the stream from the context %v", ctx)
	}
	if err := stream.SendHeader(md); err != nil {
		return err
	}
	return nil
}

// SetTrailer sets the trailer metadata that will be sent when an RPC returns.
// When called more than once, all the provided metadata will be merged.
//
// The error returned is compatible with the status package.  However, the
// status code will often not match the RPC status as seen by the client
// application, and therefore, should not be relied upon for this purpose.
func SetTrailer(ctx context.Context, md metadata.MD) error {
	if md.Len() == 0 {
		return nil
	}
	stream := serverTransportStreamFromContext(ctx)
	if stream == nil {
		return status.Errorf(codes.Internal, "grpc: failed to fetch the stream from the context %v", ctx)
	}
	stream.SetTrailer(md)
	return nil
}

func serverTransportStreamFromContext(ctx context.Context) streaming.Stream {
	return streaming.GetStream(ctx)
}

type (
	headerKey  struct{}
	trailerKey struct{}
)

// GRPCHeader is used for unary call client to get header from response.
func GRPCHeader(ctx context.Context, md *metadata.MD) context.Context {
	return context.WithValue(ctx, headerKey{}, md)
}

// GRPCTrailer is used for unary call client to get taiiler from response.
func GRPCTrailer(ctx context.Context, md *metadata.MD) context.Context {
	return context.WithValue(ctx, trailerKey{}, md)
}

// GetHeaderMetadataFromCtx is used to get the metadata of stream Header from ctx.
func GetHeaderMetadataFromCtx(ctx context.Context) *metadata.MD {
	header := ctx.Value(headerKey{})
	if header != nil {
		return header.(*metadata.MD)
	}
	return nil
}

// set header and trailer to the ctx by default.
func receiveHeaderAndTrailer(ctx context.Context, conn net.Conn) context.Context {
	if md, err := conn.(hasHeader).Header(); err == nil {
		if h := ctx.Value(headerKey{}); h != nil {
			// If using GRPCHeader(), set the value directly
			hd := h.(*metadata.MD)
			*hd = md
		} else {
			ctx = context.WithValue(ctx, headerKey{}, &md)
		}
	}
	if md := conn.(hasTrailer).Trailer(); md != nil {
		if t := ctx.Value(trailerKey{}); t != nil {
			// If using GRPCTrailer(), set the value directly
			tr := t.(*metadata.MD)
			*tr = md
		} else {
			ctx = context.WithValue(ctx, trailerKey{}, &md)
		}
	}
	return ctx
}