github.com/annwntech/go-micro/v2@v2.9.5/metadata/context.go (about)

     1  // Package metadata is a way of defining message headers
     2  package metadata
     3  
     4  import (
     5  	"context"
     6  	"strings"
     7  )
     8  
     9  type mdIncomingKey struct{}
    10  type mdOutgoingKey struct{}
    11  type metadataKey struct{}
    12  
    13  type rawMetadata struct {
    14  	md Metadata
    15  }
    16  
    17  // FromIncomingContext returns metadata from incoming ctx
    18  // returned metadata shoud not be modified or race condition happens
    19  func FromIncomingContext(ctx context.Context) (Metadata, bool) {
    20  	if ctx == nil {
    21  		return nil, false
    22  	}
    23  	md, ok := ctx.Value(mdIncomingKey{}).(*rawMetadata)
    24  	if !ok || md.md == nil {
    25  		return nil, false
    26  	}
    27  	return md.md, ok
    28  }
    29  
    30  // FromOutgoingContext returns metadata from outgoing ctx
    31  // returned metadata shoud not be modified or race condition happens
    32  func FromOutgoingContext(ctx context.Context) (Metadata, bool) {
    33  	if ctx == nil {
    34  		return nil, false
    35  	}
    36  	md, ok := ctx.Value(mdOutgoingKey{}).(*rawMetadata)
    37  	if !ok || md.md == nil {
    38  		return nil, false
    39  	}
    40  	return md.md, ok
    41  }
    42  
    43  // FromContext returns metadata from the given context
    44  // returned metadata shoud not be modified or race condition happens
    45  //
    46  // Deprecated: use FromIncomingContext or FromOutgoingContext
    47  func FromContext(ctx context.Context) (Metadata, bool) {
    48  	if ctx == nil {
    49  		return nil, false
    50  	}
    51  	md, ok := ctx.Value(metadataKey{}).(Metadata)
    52  	if !ok || md == nil {
    53  		return nil, false
    54  	}
    55  	// capitalise all values
    56  	newMD := make(Metadata, len(md))
    57  	for k, v := range md {
    58  		newMD[strings.Title(k)] = v
    59  	}
    60  	return newMD, ok
    61  }
    62  
    63  // NewContext creates a new context with the given metadata
    64  //
    65  // Deprecated: use NewIncomingContext or NewOutgoingContext
    66  func NewContext(ctx context.Context, md Metadata) context.Context {
    67  	if ctx == nil {
    68  		ctx = context.Background()
    69  	}
    70  	ctx = context.WithValue(ctx, metadataKey{}, md)
    71  	ctx = context.WithValue(ctx, mdIncomingKey{}, &rawMetadata{})
    72  	ctx = context.WithValue(ctx, mdOutgoingKey{}, &rawMetadata{})
    73  	return ctx
    74  }
    75  
    76  // SetOutgoingContext modify outgoing context with given metadata
    77  func SetOutgoingContext(ctx context.Context, md Metadata) bool {
    78  	if ctx == nil {
    79  		return false
    80  	}
    81  	if omd, ok := ctx.Value(mdOutgoingKey{}).(*rawMetadata); ok {
    82  		omd.md = md
    83  		return true
    84  	}
    85  	return false
    86  }
    87  
    88  // SetIncomingContext modify incoming context with given metadata
    89  func SetIncomingContext(ctx context.Context, md Metadata) bool {
    90  	if ctx == nil {
    91  		return false
    92  	}
    93  	if omd, ok := ctx.Value(mdIncomingKey{}).(*rawMetadata); ok {
    94  		omd.md = md
    95  		return true
    96  	}
    97  	return false
    98  }
    99  
   100  // NewIncomingContext creates a new context with incoming metadata attached
   101  func NewIncomingContext(ctx context.Context, md Metadata) context.Context {
   102  	if ctx == nil {
   103  		ctx = context.Background()
   104  	}
   105  	ctx = context.WithValue(ctx, mdIncomingKey{}, &rawMetadata{md})
   106  	ctx = context.WithValue(ctx, mdOutgoingKey{}, &rawMetadata{})
   107  	return ctx
   108  }
   109  
   110  // NewOutgoingContext creates a new context with outcoming metadata attached
   111  func NewOutgoingContext(ctx context.Context, md Metadata) context.Context {
   112  	if ctx == nil {
   113  		ctx = context.Background()
   114  	}
   115  	ctx = context.WithValue(ctx, mdOutgoingKey{}, &rawMetadata{md})
   116  	ctx = context.WithValue(ctx, mdIncomingKey{}, &rawMetadata{})
   117  	return ctx
   118  }