go.undefinedlabs.com/scopeagent@v0.4.2/tracer/propagation_env_var.go (about)

     1  package tracer
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"strconv"
     7  	"strings"
     8  
     9  	"github.com/google/uuid"
    10  	"github.com/opentracing/opentracing-go"
    11  )
    12  
    13  const (
    14  	environmentKeyPrefix      = "CTX_"
    15  	EnvironmentVariableFormat = 10
    16  )
    17  
    18  type envVarPropagator struct {
    19  	tracer *tracerImpl
    20  }
    21  
    22  func (p *envVarPropagator) Inject(
    23  	spanContext opentracing.SpanContext,
    24  	opaqueCarrier interface{},
    25  ) error {
    26  	sc, ok := spanContext.(SpanContext)
    27  	if !ok {
    28  		return opentracing.ErrInvalidSpanContext
    29  	}
    30  	opaqueValue := opaqueCarrier.(*[]string)
    31  	carrier := envVarCarrier(*opaqueValue)
    32  	if carrier == nil {
    33  		return opentracing.ErrInvalidCarrier
    34  	}
    35  
    36  	carrier.Set(fieldNameTraceID, UUIDToString(sc.TraceID))
    37  	carrier.Set(fieldNameSpanID, strconv.FormatUint(sc.SpanID, 16))
    38  	carrier.Set(fieldNameSampled, strconv.FormatBool(sc.Sampled))
    39  	for k, v := range sc.Baggage {
    40  		carrier.Set(prefixBaggage+k, v)
    41  	}
    42  	appendScopeEnvVars(&carrier)
    43  	*opaqueValue = carrier
    44  	return nil
    45  }
    46  
    47  func (p *envVarPropagator) Extract(
    48  	opaqueCarrier interface{},
    49  ) (opentracing.SpanContext, error) {
    50  	opaqueValue := opaqueCarrier.(*[]string)
    51  	carrier := envVarCarrier(*opaqueValue)
    52  	if carrier == nil {
    53  		return nil, opentracing.ErrInvalidCarrier
    54  	}
    55  	requiredFieldCount := 0
    56  	var traceID uuid.UUID
    57  	var spanID uint64
    58  	var sampled bool
    59  	var err error
    60  	decodedBaggage := make(map[string]string)
    61  	err = carrier.ForeachKey(func(k, v string) error {
    62  		switch strings.ToLower(k) {
    63  		case fieldNameTraceID:
    64  			traceID, err = StringToUUID(v)
    65  			if err != nil {
    66  				return opentracing.ErrSpanContextCorrupted
    67  			}
    68  		case fieldNameSpanID:
    69  			spanID, err = strconv.ParseUint(v, 16, 64)
    70  			if err != nil {
    71  				return opentracing.ErrSpanContextCorrupted
    72  			}
    73  		case fieldNameSampled:
    74  			sampled, err = strconv.ParseBool(v)
    75  			if err != nil {
    76  				return opentracing.ErrSpanContextCorrupted
    77  			}
    78  		default:
    79  			lowercaseK := strings.ToLower(k)
    80  			if strings.HasPrefix(lowercaseK, prefixBaggage) {
    81  				decodedBaggage[strings.TrimPrefix(lowercaseK, prefixBaggage)] = v
    82  			}
    83  			// Balance off the requiredFieldCount++ just below...
    84  			requiredFieldCount--
    85  		}
    86  		requiredFieldCount++
    87  		return nil
    88  	})
    89  	if err != nil {
    90  		return nil, err
    91  	}
    92  	if requiredFieldCount < tracerStateFieldCount {
    93  		if requiredFieldCount == 0 {
    94  			return nil, opentracing.ErrSpanContextNotFound
    95  		}
    96  		return nil, opentracing.ErrSpanContextCorrupted
    97  	}
    98  
    99  	return SpanContext{
   100  		TraceID: traceID,
   101  		SpanID:  spanID,
   102  		Sampled: sampled,
   103  		Baggage: decodedBaggage,
   104  	}, nil
   105  }
   106  
   107  // Environment variable names used by the utilities in the Shell and Utilities volume of IEEE Std 1003.1-2001
   108  // consist solely of uppercase letters, digits, and the '_' (underscore)
   109  var escapeMap = map[string]string{
   110  	".": "__dt__",
   111  	"-": "__dh__",
   112  }
   113  
   114  // Environment Variables Carrier
   115  type envVarCarrier []string
   116  
   117  // Set implements Set() of opentracing.TextMapWriter
   118  func (carrier *envVarCarrier) Set(key, val string) {
   119  	var newCarrier []string
   120  	keyUpper := strings.ToUpper(key)
   121  	ctxKey := escape(environmentKeyPrefix + keyUpper)
   122  	if carrier != nil {
   123  		for _, item := range *carrier {
   124  			if strings.Index(item, ctxKey) < 0 {
   125  				newCarrier = append(newCarrier, item)
   126  			}
   127  		}
   128  	}
   129  	newCarrier = append(newCarrier, fmt.Sprintf("%s=%s", ctxKey, val))
   130  	*carrier = newCarrier
   131  }
   132  
   133  // ForeachKey conforms to the TextMapReader interface.
   134  func (carrier *envVarCarrier) ForeachKey(handler func(key, val string) error) error {
   135  	if carrier != nil {
   136  		for _, item := range *carrier {
   137  			if strings.Index(item, environmentKeyPrefix) >= 0 {
   138  				kv := strings.Split(item, "=")
   139  				err := handler(unescape(kv[0][4:]), kv[1])
   140  				if err != nil {
   141  					return err
   142  				}
   143  			}
   144  		}
   145  	}
   146  	return nil
   147  }
   148  
   149  // We need to sanitize the env vars due:
   150  // Environment variable names used by the utilities in the Shell and Utilities volume of IEEE Std 1003.1-2001
   151  // consist solely of uppercase letters, digits, and the '_' (underscore)
   152  func escape(value string) string {
   153  	for key, val := range escapeMap {
   154  		value = strings.Replace(value, key, val, -1)
   155  	}
   156  	return value
   157  }
   158  func unescape(value string) string {
   159  	for key, val := range escapeMap {
   160  		value = strings.Replace(value, val, key, -1)
   161  	}
   162  	return value
   163  }
   164  
   165  // Append all SCOPE_* environment variables to a child command
   166  func appendScopeEnvVars(carrier *envVarCarrier) {
   167  	envVars := os.Environ()
   168  	for _, item := range envVars {
   169  		if strings.Index(item, "SCOPE_") == 0 {
   170  			kv := strings.Split(item, "=")
   171  			carrier.Set(kv[0], kv[1])
   172  		}
   173  	}
   174  }