github.com/weaveworks/common@v0.0.0-20230728070032-dd9e68f319d5/user/grpc.go (about)

     1  package user
     2  
     3  import (
     4  	"golang.org/x/net/context"
     5  	"google.golang.org/grpc/metadata"
     6  )
     7  
     8  // ExtractFromGRPCRequest extracts the user ID from the request metadata and returns
     9  // the user ID and a context with the user ID injected.
    10  func ExtractFromGRPCRequest(ctx context.Context) (string, context.Context, error) {
    11  	md, ok := metadata.FromIncomingContext(ctx)
    12  	if !ok {
    13  		return "", ctx, ErrNoOrgID
    14  	}
    15  
    16  	orgIDs, ok := md[lowerOrgIDHeaderName]
    17  	if !ok || len(orgIDs) != 1 {
    18  		return "", ctx, ErrNoOrgID
    19  	}
    20  
    21  	return orgIDs[0], InjectOrgID(ctx, orgIDs[0]), nil
    22  }
    23  
    24  // InjectIntoGRPCRequest injects the orgID from the context into the request metadata.
    25  func InjectIntoGRPCRequest(ctx context.Context) (context.Context, error) {
    26  	orgID, err := ExtractOrgID(ctx)
    27  	if err != nil {
    28  		return ctx, err
    29  	}
    30  
    31  	md, ok := metadata.FromOutgoingContext(ctx)
    32  	if !ok {
    33  		md = metadata.New(map[string]string{})
    34  	}
    35  	newCtx := ctx
    36  	if orgIDs, ok := md[lowerOrgIDHeaderName]; ok {
    37  		if len(orgIDs) == 1 {
    38  			if orgIDs[0] != orgID {
    39  				return ctx, ErrDifferentOrgIDPresent
    40  			}
    41  		} else {
    42  			return ctx, ErrTooManyOrgIDs
    43  		}
    44  	} else {
    45  		md = md.Copy()
    46  		md[lowerOrgIDHeaderName] = []string{orgID}
    47  		newCtx = metadata.NewOutgoingContext(ctx, md)
    48  	}
    49  
    50  	return newCtx, nil
    51  }