metadata.go

Documentation: github.com/containerd/ttrpc

     1  /*
     2     Copyright The containerd Authors.
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package ttrpc
    18  
    19  import (
    20  	"context"
    21  	"strings"
    22  )
    23  
    24  // MD is the user type for ttrpc metadata
    25  type MD map[string][]string
    26  
    27  // Get returns the metadata for a given key when they exist.
    28  // If there is no metadata, a nil slice and false are returned.
    29  func (m MD) Get(key string) ([]string, bool) {
    30  	key = strings.ToLower(key)
    31  	list, ok := m[key]
    32  	if !ok || len(list) == 0 {
    33  		return nil, false
    34  	}
    35  
    36  	return list, true
    37  }
    38  
    39  // Set sets the provided values for a given key.
    40  // The values will overwrite any existing values.
    41  // If no values provided, a key will be deleted.
    42  func (m MD) Set(key string, values ...string) {
    43  	key = strings.ToLower(key)
    44  	if len(values) == 0 {
    45  		delete(m, key)
    46  		return
    47  	}
    48  	m[key] = values
    49  }
    50  
    51  // Append appends additional values to the given key.
    52  func (m MD) Append(key string, values ...string) {
    53  	key = strings.ToLower(key)
    54  	if len(values) == 0 {
    55  		return
    56  	}
    57  	current, ok := m[key]
    58  	if ok {
    59  		m.Set(key, append(current, values...)...)
    60  	} else {
    61  		m.Set(key, values...)
    62  	}
    63  }
    64  
    65  func (m MD) setRequest(r *Request) {
    66  	for k, values := range m {
    67  		for _, v := range values {
    68  			r.Metadata = append(r.Metadata, &KeyValue{
    69  				Key:   k,
    70  				Value: v,
    71  			})
    72  		}
    73  	}
    74  }
    75  
    76  func (m MD) fromRequest(r *Request) {
    77  	for _, kv := range r.Metadata {
    78  		m[kv.Key] = append(m[kv.Key], kv.Value)
    79  	}
    80  }
    81  
    82  type metadataKey struct{}
    83  
    84  // GetMetadata retrieves metadata from context.Context (previously attached with WithMetadata)
    85  func GetMetadata(ctx context.Context) (MD, bool) {
    86  	metadata, ok := ctx.Value(metadataKey{}).(MD)
    87  	return metadata, ok
    88  }
    89  
    90  // GetMetadataValue gets a specific metadata value by name from context.Context
    91  func GetMetadataValue(ctx context.Context, name string) (string, bool) {
    92  	metadata, ok := GetMetadata(ctx)
    93  	if !ok {
    94  		return "", false
    95  	}
    96  
    97  	if list, ok := metadata.Get(name); ok {
    98  		return list[0], true
    99  	}
   100  
   101  	return "", false
   102  }
   103  
   104  // WithMetadata attaches metadata map to a context.Context
   105  func WithMetadata(ctx context.Context, md MD) context.Context {
   106  	return context.WithValue(ctx, metadataKey{}, md)
   107  }
   108
View as plain text