<commit_before>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2020

// Package instalambda provides Instana tracing instrumentation for
// AWS Lambda functions
package instalambda

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-lambda-go/events"
	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-lambda-go/lambdacontext"
	instana "github.com/instana/go-sensor"
	"github.com/opentracing/opentracing-go"
	otlog "github.com/opentracing/opentracing-go/log"
)

const (
	awsLambdaFlushMaxRetries  = 5
	awsLambdaFlushRetryPeriod = 50 * time.Millisecond
)

type wrappedHandler struct {
	lambda.Handler

	sensor      *instana.Sensor
	onColdStart sync.Once
}

// NewHandler creates a new instrumented handler that can be used with `lambda.StartHandler()` from a handler function
func NewHandler(handlerFunc interface{}, sensor *instana.Sensor) *wrappedHandler {
	return WrapHandler(lambda.NewHandler(handlerFunc), sensor)
}

// WrapHandler instruments a lambda.Handler to trace the invokations with Instana
func WrapHandler(h lambda.Handler, sensor *instana.Sensor) *wrappedHandler {
	return &wrappedHandler{
		Handler: h,
		sensor:  sensor,
	}
}

// Invoke is a handler function for a wrapped handler
func (h *wrappedHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
	lc, ok := lambdacontext.FromContext(ctx)
	if !ok {
		return h.Handler.Invoke(ctx, payload)
	}

	opts := append([]opentracing.StartSpanOption{opentracing.Tags{
		"lambda.arn":     lc.InvokedFunctionArn + ":" + lambdacontext.FunctionVersion,
		"lambda.name":    lambdacontext.FunctionName,
		"lambda.version": lambdacontext.FunctionVersion,
	}}, h.triggerEventSpanOptions(payload, lc.ClientContext)...)
	sp := h.sensor.Tracer().StartSpan("aws.lambda.entry", opts...)

	h.onColdStart.Do(func() {
		sp.SetTag("lambda.coldStart", true)
	})

	done := make(chan struct{})
	timeoutChannel := make(<-chan time.Time)

	originalDeadline, deadlineDefined := ctx.Deadline()

	if deadlineDefined {
		deadline := originalDeadline.Add(-100 * time.Millisecond)
		timeoutChannel = time.After(time.Until(deadline))
	}

	var resp []byte
	var err error

	go func() {
		resp, err = h.Handler.Invoke(instana.ContextWithSpan(ctx, sp), payload)
		if err != nil {
			sp.LogFields(otlog.Error(err))
		}

		done <- struct{}{}
	}()

	select {
	case <-done:
		h.sensor.Logger().Debug("no timeout")
	case <-timeoutChannel:
		h.sensor.Logger().Debug("timeout")

		remainingTime := originalDeadline.Sub(time.Now())
		remainingTimeInMilliseconds := remainingTime.Nanoseconds() / 1000000
		sp.SetTag("lambda.msleft", remainingTimeInMilliseconds)
		sp.SetTag("lambda.error", fmt.Sprintf(`The Lambda function was still running when only %d ms were left, it might have ended in a timeout.`, remainingTimeInMilliseconds))

		sp.LogFields(otlog.Error(errors.New("Timeout")))
	}

	h.finishSpanAndFlush(sp)

	return resp, err
}

func (h *wrappedHandler) finishSpanAndFlush(sp opentracing.Span) {
	sp.Finish()

	// ensure that all collected data has been sent before the invocation is finished
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		var i int
		for {
			if err := tr.Flush(context.Background()); err != nil {
				if err == instana.ErrAgentNotReady && i < awsLambdaFlushMaxRetries {
					i++
					time.Sleep(awsLambdaFlushRetryPeriod)
					continue
				}

				h.sensor.Logger().Error("failed to send traces:", err)
			}

			break
		}
	}
}

func (h *wrappedHandler) triggerEventSpanOptions(payload []byte, lcc lambdacontext.ClientContext) []opentracing.StartSpanOption {
	switch detectTriggerEventType(payload) {
	case apiGatewayEventType:
		var v events.APIGatewayProxyRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal API Gateway event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractAPIGatewayTriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case apiGatewayV2EventType:
		var v events.APIGatewayV2HTTPRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal API Gateway v2.0 event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractAPIGatewayV2TriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case albEventType:
		var v events.ALBTargetGroupRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal ALB event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractALBTriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case cloudWatchEventType:
		var v events.CloudWatchEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal CloudWatch event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractCloudWatchTriggerTags(v)}
	case cloudWatchLogsEventType:
		var v events.CloudwatchLogsEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal CloudWatch Logs event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractCloudWatchLogsTriggerTags(v)}
	case s3EventType:
		var v events.S3Event
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal S3 event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractS3TriggerTags(v)}
	case sqsEventType:
		var v events.SQSEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal SQS event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractSQSTriggerTags(v)}
	case invokeRequestType:
		tags := opentracing.Tags{
			"lambda.trigger": "aws:lambda.invoke",
		}

		opts := []opentracing.StartSpanOption{tags}
		if parentCtx, ok := h.extractParentContext(lcc.Custom); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}
		return opts

	default:
		h.sensor.Logger().Info("unsupported AWS Lambda trigger event type, the entry span will include generic tags only")
		return []opentracing.StartSpanOption{opentracing.Tags{}}
	}
}

func (h *wrappedHandler) extractParentContext(headers map[string]string) (opentracing.SpanContext, bool) {
	hdrs := http.Header{}
	for k, v := range headers {
		hdrs.Set(k, v)
	}

	switch parentCtx, err := h.sensor.Tracer().Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(hdrs)); err {
	case nil:
		return parentCtx, true
	case opentracing.ErrSpanContextNotFound:
		h.sensor.Logger().Debug("lambda invoke event did not provide trace context")
	case opentracing.ErrUnsupportedFormat:
		h.sensor.Logger().Info("lambda invoke event provided trace context in unsupported format")
	default:
		h.sensor.Logger().Warn("failed to extract span context from the lambda invoke event:", err)
	}

	return nil, false
}

func (h *wrappedHandler) extractAPIGatewayTriggerTags(evt events.APIGatewayProxyRequest) opentracing.Tags {
	tags := opentracing.Tags{
		"lambda.trigger": "aws:api.gateway",
		"http.method":    evt.HTTPMethod,
		"http.url":       evt.Path,
		"http.path_tpl":  evt.Resource,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, evt.MultiValueQueryStringParameters).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, evt.MultiValueHeaders); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractAPIGatewayV2TriggerTags(evt events.APIGatewayV2HTTPRequest) opentracing.Tags {
	routeKeyPath := evt.RouteKey
	// Strip any trailing HTTP request method
	if i := strings.Index(routeKeyPath, " "); i >= 0 {
		routeKeyPath = evt.RouteKey[i+1:]
	}

	tags := opentracing.Tags{
		"lambda.trigger": "aws:api.gateway",
		"http.method":    evt.RequestContext.HTTP.Method,
		"http.url":       evt.RequestContext.HTTP.Path,
		"http.path_tpl":  routeKeyPath,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, nil).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, nil); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractALBTriggerTags(evt events.ALBTargetGroupRequest) opentracing.Tags {
	tags := opentracing.Tags{
		"lambda.trigger": "aws:application.load.balancer",
		"http.method":    evt.HTTPMethod,
		"http.url":       evt.Path,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, evt.MultiValueQueryStringParameters).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, evt.MultiValueHeaders); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractCloudWatchTriggerTags(evt events.CloudWatchEvent) opentracing.Tags {
	return opentracing.Tags{
		"lambda.trigger":              "aws:cloudwatch.events",
		"cloudwatch.events.id":        evt.ID,
		"cloudwatch.events.resources": evt.Resources,
	}
}

func (h *wrappedHandler) extractCloudWatchLogsTriggerTags(evt events.CloudwatchLogsEvent) opentracing.Tags {
	logs, err := evt.AWSLogs.Parse()
	if err != nil {
		return opentracing.Tags{
			"lambda.trigger":                "aws:cloudwatch.logs",
			"cloudwatch.logs.decodingError": err,
		}
	}

	var events []string
	for _, event := range logs.LogEvents {
		events = append(events, event.Message)
	}

	return opentracing.Tags{
		"lambda.trigger":         "aws:cloudwatch.logs",
		"cloudwatch.logs.group":  logs.LogGroup,
		"cloudwatch.logs.stream": logs.LogStream,
		"cloudwatch.logs.events": events,
	}
}

func (h *wrappedHandler) extractS3TriggerTags(evt events.S3Event) opentracing.Tags {
	var events []instana.AWSS3EventTags
	for _, rec := range evt.Records {
		events = append(events, instana.AWSS3EventTags{
			Name:   rec.EventName,
			Bucket: rec.S3.Bucket.Name,
			Object: rec.S3.Object.Key,
		})
	}

	return opentracing.Tags{
		"lambda.trigger": "aws:s3",
		"s3.events":      events,
	}
}

func (h *wrappedHandler) extractSQSTriggerTags(evt events.SQSEvent) opentracing.Tags {
	var msgs []instana.AWSSQSMessageTags
	for _, rec := range evt.Records {
		msgs = append(msgs, instana.AWSSQSMessageTags{
			Queue: rec.EventSourceARN,
		})
	}

	return opentracing.Tags{
		"lambda.trigger": "aws:sqs",
		"sqs.messages":   msgs,
	}
}

func (h *wrappedHandler) sanitizeHTTPParams(
	queryStringParams map[string]string,
	multiValueQueryStringParams map[string][]string,
) url.Values {
	secretMatcher := instana.DefaultSecretsMatcher()
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		secretMatcher = tr.Options().Secrets
	}

	params := url.Values{}

	for k, v := range queryStringParams {
		if secretMatcher.Match(k) {
			v = "<redacted>"
		}
		params.Set(k, v)
	}

	for k, vv := range multiValueQueryStringParams {
		isSecret := secretMatcher.Match(k)
		for _, v := range vv {
			if isSecret {
				v = "<redacted>"
			}
			params.Add(k, v)
		}
	}

	return params
}

func (h *wrappedHandler) collectHTTPHeaders(headers map[string]string, multiValueHeaders map[string][]string) map[string]string {
	var collectableHeaders []string
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		collectableHeaders = tr.Options().CollectableHTTPHeaders
	}

	if len(collectableHeaders) == 0 {
		return nil
	}

	// Normalize header names first by bringing them to the canonical MIME format to avoid missing headers because of mismatching case
	hdrs := http.Header{}
	for k, v := range headers {
		hdrs.Set(k, v)
	}

	for k, vv := range multiValueHeaders {
		for _, v := range vv {
			hdrs.Add(k, v)
		}
	}

	collected := make(map[string]string)
	for _, k := range collectableHeaders {
		if v := hdrs.Get(k); v != "" {
			collected[k] = v
		}
	}

	return collected
}
<commit_msg>Extract AWS Lambda timeout threshold into a constant<commit_after>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2020

// Package instalambda provides Instana tracing instrumentation for
// AWS Lambda functions
package instalambda

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-lambda-go/events"
	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-lambda-go/lambdacontext"
	instana "github.com/instana/go-sensor"
	"github.com/opentracing/opentracing-go"
	otlog "github.com/opentracing/opentracing-go/log"
)

const (
	awsLambdaFlushMaxRetries  = 5
	awsLambdaFlushRetryPeriod = 50 * time.Millisecond
	awsLambdaTimeoutThreshold = 100 * time.Millisecond
)

type wrappedHandler struct {
	lambda.Handler

	sensor      *instana.Sensor
	onColdStart sync.Once
}

// NewHandler creates a new instrumented handler that can be used with `lambda.StartHandler()` from a handler function
func NewHandler(handlerFunc interface{}, sensor *instana.Sensor) *wrappedHandler {
	return WrapHandler(lambda.NewHandler(handlerFunc), sensor)
}

// WrapHandler instruments a lambda.Handler to trace the invokations with Instana
func WrapHandler(h lambda.Handler, sensor *instana.Sensor) *wrappedHandler {
	return &wrappedHandler{
		Handler: h,
		sensor:  sensor,
	}
}

// Invoke is a handler function for a wrapped handler
func (h *wrappedHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
	lc, ok := lambdacontext.FromContext(ctx)
	if !ok {
		return h.Handler.Invoke(ctx, payload)
	}

	opts := append([]opentracing.StartSpanOption{opentracing.Tags{
		"lambda.arn":     lc.InvokedFunctionArn + ":" + lambdacontext.FunctionVersion,
		"lambda.name":    lambdacontext.FunctionName,
		"lambda.version": lambdacontext.FunctionVersion,
	}}, h.triggerEventSpanOptions(payload, lc.ClientContext)...)
	sp := h.sensor.Tracer().StartSpan("aws.lambda.entry", opts...)

	h.onColdStart.Do(func() {
		sp.SetTag("lambda.coldStart", true)
	})

	done := make(chan struct{})
	timeoutChannel := make(<-chan time.Time)

	originalDeadline, deadlineDefined := ctx.Deadline()

	if deadlineDefined {
		deadline := originalDeadline.Add(-awsLambdaTimeoutThreshold)
		timeoutChannel = time.After(time.Until(deadline))
	}

	var resp []byte
	var err error

	go func() {
		resp, err = h.Handler.Invoke(instana.ContextWithSpan(ctx, sp), payload)
		if err != nil {
			sp.LogFields(otlog.Error(err))
		}

		done <- struct{}{}
	}()

	select {
	case <-done:
		h.sensor.Logger().Debug("no timeout")
	case <-timeoutChannel:
		h.sensor.Logger().Debug("timeout")

		remainingTime := originalDeadline.Sub(time.Now())
		remainingTimeInMilliseconds := remainingTime.Nanoseconds() / 1000000
		sp.SetTag("lambda.msleft", remainingTimeInMilliseconds)
		sp.SetTag("lambda.error", fmt.Sprintf(`The Lambda function was still running when only %d ms were left, it might have ended in a timeout.`, remainingTimeInMilliseconds))

		sp.LogFields(otlog.Error(errors.New("Timeout")))
	}

	h.finishSpanAndFlush(sp)

	return resp, err
}

func (h *wrappedHandler) finishSpanAndFlush(sp opentracing.Span) {
	sp.Finish()

	// ensure that all collected data has been sent before the invocation is finished
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		var i int
		for {
			if err := tr.Flush(context.Background()); err != nil {
				if err == instana.ErrAgentNotReady && i < awsLambdaFlushMaxRetries {
					i++
					time.Sleep(awsLambdaFlushRetryPeriod)
					continue
				}

				h.sensor.Logger().Error("failed to send traces:", err)
			}

			break
		}
	}
}

func (h *wrappedHandler) triggerEventSpanOptions(payload []byte, lcc lambdacontext.ClientContext) []opentracing.StartSpanOption {
	switch detectTriggerEventType(payload) {
	case apiGatewayEventType:
		var v events.APIGatewayProxyRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal API Gateway event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractAPIGatewayTriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case apiGatewayV2EventType:
		var v events.APIGatewayV2HTTPRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal API Gateway v2.0 event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractAPIGatewayV2TriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case albEventType:
		var v events.ALBTargetGroupRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal ALB event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractALBTriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case cloudWatchEventType:
		var v events.CloudWatchEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal CloudWatch event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractCloudWatchTriggerTags(v)}
	case cloudWatchLogsEventType:
		var v events.CloudwatchLogsEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal CloudWatch Logs event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractCloudWatchLogsTriggerTags(v)}
	case s3EventType:
		var v events.S3Event
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal S3 event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractS3TriggerTags(v)}
	case sqsEventType:
		var v events.SQSEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal SQS event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractSQSTriggerTags(v)}
	case invokeRequestType:
		tags := opentracing.Tags{
			"lambda.trigger": "aws:lambda.invoke",
		}

		opts := []opentracing.StartSpanOption{tags}
		if parentCtx, ok := h.extractParentContext(lcc.Custom); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}
		return opts

	default:
		h.sensor.Logger().Info("unsupported AWS Lambda trigger event type, the entry span will include generic tags only")
		return []opentracing.StartSpanOption{opentracing.Tags{}}
	}
}

func (h *wrappedHandler) extractParentContext(headers map[string]string) (opentracing.SpanContext, bool) {
	hdrs := http.Header{}
	for k, v := range headers {
		hdrs.Set(k, v)
	}

	switch parentCtx, err := h.sensor.Tracer().Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(hdrs)); err {
	case nil:
		return parentCtx, true
	case opentracing.ErrSpanContextNotFound:
		h.sensor.Logger().Debug("lambda invoke event did not provide trace context")
	case opentracing.ErrUnsupportedFormat:
		h.sensor.Logger().Info("lambda invoke event provided trace context in unsupported format")
	default:
		h.sensor.Logger().Warn("failed to extract span context from the lambda invoke event:", err)
	}

	return nil, false
}

func (h *wrappedHandler) extractAPIGatewayTriggerTags(evt events.APIGatewayProxyRequest) opentracing.Tags {
	tags := opentracing.Tags{
		"lambda.trigger": "aws:api.gateway",
		"http.method":    evt.HTTPMethod,
		"http.url":       evt.Path,
		"http.path_tpl":  evt.Resource,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, evt.MultiValueQueryStringParameters).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, evt.MultiValueHeaders); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractAPIGatewayV2TriggerTags(evt events.APIGatewayV2HTTPRequest) opentracing.Tags {
	routeKeyPath := evt.RouteKey
	// Strip any trailing HTTP request method
	if i := strings.Index(routeKeyPath, " "); i >= 0 {
		routeKeyPath = evt.RouteKey[i+1:]
	}

	tags := opentracing.Tags{
		"lambda.trigger": "aws:api.gateway",
		"http.method":    evt.RequestContext.HTTP.Method,
		"http.url":       evt.RequestContext.HTTP.Path,
		"http.path_tpl":  routeKeyPath,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, nil).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, nil); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractALBTriggerTags(evt events.ALBTargetGroupRequest) opentracing.Tags {
	tags := opentracing.Tags{
		"lambda.trigger": "aws:application.load.balancer",
		"http.method":    evt.HTTPMethod,
		"http.url":       evt.Path,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, evt.MultiValueQueryStringParameters).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, evt.MultiValueHeaders); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractCloudWatchTriggerTags(evt events.CloudWatchEvent) opentracing.Tags {
	return opentracing.Tags{
		"lambda.trigger":              "aws:cloudwatch.events",
		"cloudwatch.events.id":        evt.ID,
		"cloudwatch.events.resources": evt.Resources,
	}
}

func (h *wrappedHandler) extractCloudWatchLogsTriggerTags(evt events.CloudwatchLogsEvent) opentracing.Tags {
	logs, err := evt.AWSLogs.Parse()
	if err != nil {
		return opentracing.Tags{
			"lambda.trigger":                "aws:cloudwatch.logs",
			"cloudwatch.logs.decodingError": err,
		}
	}

	var events []string
	for _, event := range logs.LogEvents {
		events = append(events, event.Message)
	}

	return opentracing.Tags{
		"lambda.trigger":         "aws:cloudwatch.logs",
		"cloudwatch.logs.group":  logs.LogGroup,
		"cloudwatch.logs.stream": logs.LogStream,
		"cloudwatch.logs.events": events,
	}
}

func (h *wrappedHandler) extractS3TriggerTags(evt events.S3Event) opentracing.Tags {
	var events []instana.AWSS3EventTags
	for _, rec := range evt.Records {
		events = append(events, instana.AWSS3EventTags{
			Name:   rec.EventName,
			Bucket: rec.S3.Bucket.Name,
			Object: rec.S3.Object.Key,
		})
	}

	return opentracing.Tags{
		"lambda.trigger": "aws:s3",
		"s3.events":      events,
	}
}

func (h *wrappedHandler) extractSQSTriggerTags(evt events.SQSEvent) opentracing.Tags {
	var msgs []instana.AWSSQSMessageTags
	for _, rec := range evt.Records {
		msgs = append(msgs, instana.AWSSQSMessageTags{
			Queue: rec.EventSourceARN,
		})
	}

	return opentracing.Tags{
		"lambda.trigger": "aws:sqs",
		"sqs.messages":   msgs,
	}
}

func (h *wrappedHandler) sanitizeHTTPParams(
	queryStringParams map[string]string,
	multiValueQueryStringParams map[string][]string,
) url.Values {
	secretMatcher := instana.DefaultSecretsMatcher()
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		secretMatcher = tr.Options().Secrets
	}

	params := url.Values{}

	for k, v := range queryStringParams {
		if secretMatcher.Match(k) {
			v = "<redacted>"
		}
		params.Set(k, v)
	}

	for k, vv := range multiValueQueryStringParams {
		isSecret := secretMatcher.Match(k)
		for _, v := range vv {
			if isSecret {
				v = "<redacted>"
			}
			params.Add(k, v)
		}
	}

	return params
}

func (h *wrappedHandler) collectHTTPHeaders(headers map[string]string, multiValueHeaders map[string][]string) map[string]string {
	var collectableHeaders []string
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		collectableHeaders = tr.Options().CollectableHTTPHeaders
	}

	if len(collectableHeaders) == 0 {
		return nil
	}

	// Normalize header names first by bringing them to the canonical MIME format to avoid missing headers because of mismatching case
	hdrs := http.Header{}
	for k, v := range headers {
		hdrs.Set(k, v)
	}

	for k, vv := range multiValueHeaders {
		for _, v := range vv {
			hdrs.Add(k, v)
		}
	}

	collected := make(map[string]string)
	for _, k := range collectableHeaders {
		if v := hdrs.Get(k); v != "" {
			collected[k] = v
		}
	}

	return collected
}
<|endoftext|>


<commit_before>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2020

// Package instalambda provides Instana tracing instrumentation for
// AWS Lambda functions
package instalambda

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-lambda-go/events"
	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-lambda-go/lambdacontext"
	instana "github.com/instana/go-sensor"
	"github.com/opentracing/opentracing-go"
	otlog "github.com/opentracing/opentracing-go/log"
)

const (
	awsLambdaFlushMaxRetries  = 5
	awsLambdaFlushRetryPeriod = 50 * time.Millisecond
)

type wrappedHandler struct {
	lambda.Handler

	sensor      *instana.Sensor
	onColdStart sync.Once
}

// NewHandler creates a new instrumented handler that can be used with `lambda.StartHandler()` from a handler function
func NewHandler(handlerFunc interface{}, sensor *instana.Sensor) *wrappedHandler {
	return WrapHandler(lambda.NewHandler(handlerFunc), sensor)
}

// WrapHandler instruments a lambda.Handler to trace the invokations with Instana
func WrapHandler(h lambda.Handler, sensor *instana.Sensor) *wrappedHandler {
	return &wrappedHandler{
		Handler: h,
		sensor:  sensor,
	}
}

// Invoke is a handler function for a wrapped handler
func (h *wrappedHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
	lc, ok := lambdacontext.FromContext(ctx)
	if !ok {
		return h.Handler.Invoke(ctx, payload)
	}

	opts := append([]opentracing.StartSpanOption{opentracing.Tags{
		"lambda.arn":     lc.InvokedFunctionArn + ":" + lambdacontext.FunctionVersion,
		"lambda.name":    lambdacontext.FunctionName,
		"lambda.version": lambdacontext.FunctionVersion,
	}}, h.triggerEventSpanOptions(payload, lc.ClientContext)...)
	sp := h.sensor.Tracer().StartSpan("aws.lambda.entry", opts...)

	h.onColdStart.Do(func() {
		sp.SetTag("lambda.coldStart", true)
	})

	done := make(chan struct{})
	timeoutChannel := make(<-chan time.Time)

	originalDeadline, deadlineDefined := ctx.Deadline()

	if deadlineDefined {
		deadline := originalDeadline.Add(-100 * time.Millisecond)
		timeoutChannel = time.After(time.Until(deadline))
	}

	var resp []byte
	var err error

	go func() {
		resp, err = h.Handler.Invoke(instana.ContextWithSpan(ctx, sp), payload)
		if err != nil {
			sp.LogFields(otlog.Error(err))
		}

		done <- struct{}{}
	}()

	select {
	case <-done:
		h.sensor.Logger().Debug("no timeout")
		if deadlineDefined {
			remainingTime := originalDeadline.Sub(time.Now())
			sp.SetTag("lambda.msleft", remainingTime.Milliseconds())
		}
	case <-timeoutChannel:
		h.sensor.Logger().Debug("timeout")

		remainingTime := originalDeadline.Sub(time.Now())
		sp.SetTag("lambda.msleft", remainingTime.Milliseconds())
		sp.SetTag("lambda.error", fmt.Sprintf(`The Lambda function was still running when only %d ms were left, it might have ended in a timeout.`, remainingTime.Milliseconds()))
	}

	h.finishSpanAndFlush(sp)

	return resp, err
}

func (h *wrappedHandler) finishSpanAndFlush(sp opentracing.Span) {
	sp.Finish()

	// ensure that all collected data has been sent before the invocation is finished
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		var i int
		for {
			if err := tr.Flush(context.Background()); err != nil {
				if err == instana.ErrAgentNotReady && i < awsLambdaFlushMaxRetries {
					i++
					time.Sleep(awsLambdaFlushRetryPeriod)
					continue
				}

				h.sensor.Logger().Error("failed to send traces:", err)
			}

			break
		}
	}
}

func (h *wrappedHandler) triggerEventSpanOptions(payload []byte, lcc lambdacontext.ClientContext) []opentracing.StartSpanOption {
	switch detectTriggerEventType(payload) {
	case apiGatewayEventType:
		var v events.APIGatewayProxyRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal API Gateway event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractAPIGatewayTriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case apiGatewayV2EventType:
		var v events.APIGatewayV2HTTPRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal API Gateway v2.0 event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractAPIGatewayV2TriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case albEventType:
		var v events.ALBTargetGroupRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal ALB event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractALBTriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case cloudWatchEventType:
		var v events.CloudWatchEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal CloudWatch event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractCloudWatchTriggerTags(v)}
	case cloudWatchLogsEventType:
		var v events.CloudwatchLogsEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal CloudWatch Logs event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractCloudWatchLogsTriggerTags(v)}
	case s3EventType:
		var v events.S3Event
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal S3 event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractS3TriggerTags(v)}
	case sqsEventType:
		var v events.SQSEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal SQS event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractSQSTriggerTags(v)}
	case invokeRequestType:
		tags := opentracing.Tags{
			"lambda.trigger": "aws:lambda.invoke",
		}

		opts := []opentracing.StartSpanOption{tags}
		if parentCtx, ok := h.extractParentContext(lcc.Custom); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}
		return opts

	default:
		h.sensor.Logger().Info("unsupported AWS Lambda trigger event type, the entry span will include generic tags only")
		return []opentracing.StartSpanOption{opentracing.Tags{}}
	}
}

func (h *wrappedHandler) extractParentContext(headers map[string]string) (opentracing.SpanContext, bool) {
	hdrs := http.Header{}
	for k, v := range headers {
		hdrs.Set(k, v)
	}

	switch parentCtx, err := h.sensor.Tracer().Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(hdrs)); err {
	case nil:
		return parentCtx, true
	case opentracing.ErrSpanContextNotFound:
		h.sensor.Logger().Debug("lambda invoke event did not provide trace context")
	case opentracing.ErrUnsupportedFormat:
		h.sensor.Logger().Info("lambda invoke event provided trace context in unsupported format")
	default:
		h.sensor.Logger().Warn("failed to extract span context from the lambda invoke event:", err)
	}

	return nil, false
}

func (h *wrappedHandler) extractAPIGatewayTriggerTags(evt events.APIGatewayProxyRequest) opentracing.Tags {
	tags := opentracing.Tags{
		"lambda.trigger": "aws:api.gateway",
		"http.method":    evt.HTTPMethod,
		"http.url":       evt.Path,
		"http.path_tpl":  evt.Resource,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, evt.MultiValueQueryStringParameters).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, evt.MultiValueHeaders); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractAPIGatewayV2TriggerTags(evt events.APIGatewayV2HTTPRequest) opentracing.Tags {
	routeKeyPath := evt.RouteKey
	// Strip any trailing HTTP request method
	if i := strings.Index(routeKeyPath, " "); i >= 0 {
		routeKeyPath = evt.RouteKey[i+1:]
	}

	tags := opentracing.Tags{
		"lambda.trigger": "aws:api.gateway",
		"http.method":    evt.RequestContext.HTTP.Method,
		"http.url":       evt.RequestContext.HTTP.Path,
		"http.path_tpl":  routeKeyPath,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, nil).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, nil); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractALBTriggerTags(evt events.ALBTargetGroupRequest) opentracing.Tags {
	tags := opentracing.Tags{
		"lambda.trigger": "aws:application.load.balancer",
		"http.method":    evt.HTTPMethod,
		"http.url":       evt.Path,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, evt.MultiValueQueryStringParameters).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, evt.MultiValueHeaders); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractCloudWatchTriggerTags(evt events.CloudWatchEvent) opentracing.Tags {
	return opentracing.Tags{
		"lambda.trigger":              "aws:cloudwatch.events",
		"cloudwatch.events.id":        evt.ID,
		"cloudwatch.events.resources": evt.Resources,
	}
}

func (h *wrappedHandler) extractCloudWatchLogsTriggerTags(evt events.CloudwatchLogsEvent) opentracing.Tags {
	logs, err := evt.AWSLogs.Parse()
	if err != nil {
		return opentracing.Tags{
			"lambda.trigger":                "aws:cloudwatch.logs",
			"cloudwatch.logs.decodingError": err,
		}
	}

	var events []string
	for _, event := range logs.LogEvents {
		events = append(events, event.Message)
	}

	return opentracing.Tags{
		"lambda.trigger":         "aws:cloudwatch.logs",
		"cloudwatch.logs.group":  logs.LogGroup,
		"cloudwatch.logs.stream": logs.LogStream,
		"cloudwatch.logs.events": events,
	}
}

func (h *wrappedHandler) extractS3TriggerTags(evt events.S3Event) opentracing.Tags {
	var events []instana.AWSS3EventTags
	for _, rec := range evt.Records {
		events = append(events, instana.AWSS3EventTags{
			Name:   rec.EventName,
			Bucket: rec.S3.Bucket.Name,
			Object: rec.S3.Object.Key,
		})
	}

	return opentracing.Tags{
		"lambda.trigger": "aws:s3",
		"s3.events":      events,
	}
}

func (h *wrappedHandler) extractSQSTriggerTags(evt events.SQSEvent) opentracing.Tags {
	var msgs []instana.AWSSQSMessageTags
	for _, rec := range evt.Records {
		msgs = append(msgs, instana.AWSSQSMessageTags{
			Queue: rec.EventSourceARN,
		})
	}

	return opentracing.Tags{
		"lambda.trigger": "aws:sqs",
		"sqs.messages":   msgs,
	}
}

func (h *wrappedHandler) sanitizeHTTPParams(
	queryStringParams map[string]string,
	multiValueQueryStringParams map[string][]string,
) url.Values {
	secretMatcher := instana.DefaultSecretsMatcher()
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		secretMatcher = tr.Options().Secrets
	}

	params := url.Values{}

	for k, v := range queryStringParams {
		if secretMatcher.Match(k) {
			v = "<redacted>"
		}
		params.Set(k, v)
	}

	for k, vv := range multiValueQueryStringParams {
		isSecret := secretMatcher.Match(k)
		for _, v := range vv {
			if isSecret {
				v = "<redacted>"
			}
			params.Add(k, v)
		}
	}

	return params
}

func (h *wrappedHandler) collectHTTPHeaders(headers map[string]string, multiValueHeaders map[string][]string) map[string]string {
	var collectableHeaders []string
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		collectableHeaders = tr.Options().CollectableHTTPHeaders
	}

	if len(collectableHeaders) == 0 {
		return nil
	}

	// Normalize header names first by bringing them to the canonical MIME format to avoid missing headers because of mismatching case
	hdrs := http.Header{}
	for k, v := range headers {
		hdrs.Set(k, v)
	}

	for k, vv := range multiValueHeaders {
		for _, v := range vv {
			hdrs.Add(k, v)
		}
	}

	collected := make(map[string]string)
	for _, k := range collectableHeaders {
		if v := hdrs.Get(k); v != "" {
			collected[k] = v
		}
	}

	return collected
}
<commit_msg>simplify Lambda timeout handling<commit_after>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2020

// Package instalambda provides Instana tracing instrumentation for
// AWS Lambda functions
package instalambda

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-lambda-go/events"
	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-lambda-go/lambdacontext"
	instana "github.com/instana/go-sensor"
	"github.com/opentracing/opentracing-go"
	otlog "github.com/opentracing/opentracing-go/log"
)

const (
	awsLambdaFlushMaxRetries  = 5
	awsLambdaFlushRetryPeriod = 50 * time.Millisecond
)

type wrappedHandler struct {
	lambda.Handler

	sensor      *instana.Sensor
	onColdStart sync.Once
}

// NewHandler creates a new instrumented handler that can be used with `lambda.StartHandler()` from a handler function
func NewHandler(handlerFunc interface{}, sensor *instana.Sensor) *wrappedHandler {
	return WrapHandler(lambda.NewHandler(handlerFunc), sensor)
}

// WrapHandler instruments a lambda.Handler to trace the invokations with Instana
func WrapHandler(h lambda.Handler, sensor *instana.Sensor) *wrappedHandler {
	return &wrappedHandler{
		Handler: h,
		sensor:  sensor,
	}
}

// Invoke is a handler function for a wrapped handler
func (h *wrappedHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
	lc, ok := lambdacontext.FromContext(ctx)
	if !ok {
		return h.Handler.Invoke(ctx, payload)
	}

	opts := append([]opentracing.StartSpanOption{opentracing.Tags{
		"lambda.arn":     lc.InvokedFunctionArn + ":" + lambdacontext.FunctionVersion,
		"lambda.name":    lambdacontext.FunctionName,
		"lambda.version": lambdacontext.FunctionVersion,
	}}, h.triggerEventSpanOptions(payload, lc.ClientContext)...)
	sp := h.sensor.Tracer().StartSpan("aws.lambda.entry", opts...)

	h.onColdStart.Do(func() {
		sp.SetTag("lambda.coldStart", true)
	})

	done := make(chan struct{})
	timeoutChannel := make(<-chan time.Time)

	originalDeadline, deadlineDefined := ctx.Deadline()

	if deadlineDefined {
		deadline := originalDeadline.Add(-100 * time.Millisecond)
		timeoutChannel = time.After(time.Until(deadline))
	}

	var resp []byte
	var err error

	go func() {
		resp, err = h.Handler.Invoke(instana.ContextWithSpan(ctx, sp), payload)
		if err != nil {
			sp.LogFields(otlog.Error(err))
		}

		done <- struct{}{}
	}()

	select {
	case <-done:
		h.sensor.Logger().Debug("no timeout")
	case <-timeoutChannel:
		h.sensor.Logger().Debug("timeout")

		remainingTime := originalDeadline.Sub(time.Now())
		sp.SetTag("lambda.msleft", remainingTime.Milliseconds())
		sp.SetTag("lambda.error", fmt.Sprintf(`The Lambda function was still running when only %d ms were left, it might have ended in a timeout.`, remainingTime.Milliseconds()))

		sp.LogFields(otlog.Error(errors.New("Timeout")))
	}

	h.finishSpanAndFlush(sp)

	return resp, err
}

func (h *wrappedHandler) finishSpanAndFlush(sp opentracing.Span) {
	sp.Finish()

	// ensure that all collected data has been sent before the invocation is finished
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		var i int
		for {
			if err := tr.Flush(context.Background()); err != nil {
				if err == instana.ErrAgentNotReady && i < awsLambdaFlushMaxRetries {
					i++
					time.Sleep(awsLambdaFlushRetryPeriod)
					continue
				}

				h.sensor.Logger().Error("failed to send traces:", err)
			}

			break
		}
	}
}

func (h *wrappedHandler) triggerEventSpanOptions(payload []byte, lcc lambdacontext.ClientContext) []opentracing.StartSpanOption {
	switch detectTriggerEventType(payload) {
	case apiGatewayEventType:
		var v events.APIGatewayProxyRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal API Gateway event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractAPIGatewayTriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case apiGatewayV2EventType:
		var v events.APIGatewayV2HTTPRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal API Gateway v2.0 event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractAPIGatewayV2TriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case albEventType:
		var v events.ALBTargetGroupRequest
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal ALB event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		opts := []opentracing.StartSpanOption{h.extractALBTriggerTags(v)}
		if parentCtx, ok := h.extractParentContext(v.Headers); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}

		return opts
	case cloudWatchEventType:
		var v events.CloudWatchEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal CloudWatch event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractCloudWatchTriggerTags(v)}
	case cloudWatchLogsEventType:
		var v events.CloudwatchLogsEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal CloudWatch Logs event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractCloudWatchLogsTriggerTags(v)}
	case s3EventType:
		var v events.S3Event
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal S3 event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractS3TriggerTags(v)}
	case sqsEventType:
		var v events.SQSEvent
		if err := json.Unmarshal(payload, &v); err != nil {
			h.sensor.Logger().Warn("failed to unmarshal SQS event payload: ", err)
			return []opentracing.StartSpanOption{opentracing.Tags{}}
		}

		return []opentracing.StartSpanOption{h.extractSQSTriggerTags(v)}
	case invokeRequestType:
		tags := opentracing.Tags{
			"lambda.trigger": "aws:lambda.invoke",
		}

		opts := []opentracing.StartSpanOption{tags}
		if parentCtx, ok := h.extractParentContext(lcc.Custom); ok {
			opts = append(opts, opentracing.ChildOf(parentCtx))
		}
		return opts

	default:
		h.sensor.Logger().Info("unsupported AWS Lambda trigger event type, the entry span will include generic tags only")
		return []opentracing.StartSpanOption{opentracing.Tags{}}
	}
}

func (h *wrappedHandler) extractParentContext(headers map[string]string) (opentracing.SpanContext, bool) {
	hdrs := http.Header{}
	for k, v := range headers {
		hdrs.Set(k, v)
	}

	switch parentCtx, err := h.sensor.Tracer().Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(hdrs)); err {
	case nil:
		return parentCtx, true
	case opentracing.ErrSpanContextNotFound:
		h.sensor.Logger().Debug("lambda invoke event did not provide trace context")
	case opentracing.ErrUnsupportedFormat:
		h.sensor.Logger().Info("lambda invoke event provided trace context in unsupported format")
	default:
		h.sensor.Logger().Warn("failed to extract span context from the lambda invoke event:", err)
	}

	return nil, false
}

func (h *wrappedHandler) extractAPIGatewayTriggerTags(evt events.APIGatewayProxyRequest) opentracing.Tags {
	tags := opentracing.Tags{
		"lambda.trigger": "aws:api.gateway",
		"http.method":    evt.HTTPMethod,
		"http.url":       evt.Path,
		"http.path_tpl":  evt.Resource,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, evt.MultiValueQueryStringParameters).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, evt.MultiValueHeaders); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractAPIGatewayV2TriggerTags(evt events.APIGatewayV2HTTPRequest) opentracing.Tags {
	routeKeyPath := evt.RouteKey
	// Strip any trailing HTTP request method
	if i := strings.Index(routeKeyPath, " "); i >= 0 {
		routeKeyPath = evt.RouteKey[i+1:]
	}

	tags := opentracing.Tags{
		"lambda.trigger": "aws:api.gateway",
		"http.method":    evt.RequestContext.HTTP.Method,
		"http.url":       evt.RequestContext.HTTP.Path,
		"http.path_tpl":  routeKeyPath,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, nil).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, nil); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractALBTriggerTags(evt events.ALBTargetGroupRequest) opentracing.Tags {
	tags := opentracing.Tags{
		"lambda.trigger": "aws:application.load.balancer",
		"http.method":    evt.HTTPMethod,
		"http.url":       evt.Path,
		"http.params":    h.sanitizeHTTPParams(evt.QueryStringParameters, evt.MultiValueQueryStringParameters).Encode(),
	}

	if headers := h.collectHTTPHeaders(evt.Headers, evt.MultiValueHeaders); len(headers) > 0 {
		tags["http.header"] = headers
	}

	return tags
}

func (h *wrappedHandler) extractCloudWatchTriggerTags(evt events.CloudWatchEvent) opentracing.Tags {
	return opentracing.Tags{
		"lambda.trigger":              "aws:cloudwatch.events",
		"cloudwatch.events.id":        evt.ID,
		"cloudwatch.events.resources": evt.Resources,
	}
}

func (h *wrappedHandler) extractCloudWatchLogsTriggerTags(evt events.CloudwatchLogsEvent) opentracing.Tags {
	logs, err := evt.AWSLogs.Parse()
	if err != nil {
		return opentracing.Tags{
			"lambda.trigger":                "aws:cloudwatch.logs",
			"cloudwatch.logs.decodingError": err,
		}
	}

	var events []string
	for _, event := range logs.LogEvents {
		events = append(events, event.Message)
	}

	return opentracing.Tags{
		"lambda.trigger":         "aws:cloudwatch.logs",
		"cloudwatch.logs.group":  logs.LogGroup,
		"cloudwatch.logs.stream": logs.LogStream,
		"cloudwatch.logs.events": events,
	}
}

func (h *wrappedHandler) extractS3TriggerTags(evt events.S3Event) opentracing.Tags {
	var events []instana.AWSS3EventTags
	for _, rec := range evt.Records {
		events = append(events, instana.AWSS3EventTags{
			Name:   rec.EventName,
			Bucket: rec.S3.Bucket.Name,
			Object: rec.S3.Object.Key,
		})
	}

	return opentracing.Tags{
		"lambda.trigger": "aws:s3",
		"s3.events":      events,
	}
}

func (h *wrappedHandler) extractSQSTriggerTags(evt events.SQSEvent) opentracing.Tags {
	var msgs []instana.AWSSQSMessageTags
	for _, rec := range evt.Records {
		msgs = append(msgs, instana.AWSSQSMessageTags{
			Queue: rec.EventSourceARN,
		})
	}

	return opentracing.Tags{
		"lambda.trigger": "aws:sqs",
		"sqs.messages":   msgs,
	}
}

func (h *wrappedHandler) sanitizeHTTPParams(
	queryStringParams map[string]string,
	multiValueQueryStringParams map[string][]string,
) url.Values {
	secretMatcher := instana.DefaultSecretsMatcher()
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		secretMatcher = tr.Options().Secrets
	}

	params := url.Values{}

	for k, v := range queryStringParams {
		if secretMatcher.Match(k) {
			v = "<redacted>"
		}
		params.Set(k, v)
	}

	for k, vv := range multiValueQueryStringParams {
		isSecret := secretMatcher.Match(k)
		for _, v := range vv {
			if isSecret {
				v = "<redacted>"
			}
			params.Add(k, v)
		}
	}

	return params
}

func (h *wrappedHandler) collectHTTPHeaders(headers map[string]string, multiValueHeaders map[string][]string) map[string]string {
	var collectableHeaders []string
	if tr, ok := h.sensor.Tracer().(instana.Tracer); ok {
		collectableHeaders = tr.Options().CollectableHTTPHeaders
	}

	if len(collectableHeaders) == 0 {
		return nil
	}

	// Normalize header names first by bringing them to the canonical MIME format to avoid missing headers because of mismatching case
	hdrs := http.Header{}
	for k, v := range headers {
		hdrs.Set(k, v)
	}

	for k, vv := range multiValueHeaders {
		for _, v := range vv {
			hdrs.Add(k, v)
		}
	}

	collected := make(map[string]string)
	for _, k := range collectableHeaders {
		if v := hdrs.Get(k); v != "" {
			collected[k] = v
		}
	}

	return collected
}
<|endoftext|>


<commit_before>package main

// createuvm does what it says on the tin. Simple test utility for looking
// at startup timing.

import (
	"bufio"
	"fmt"
	"os"

	"github.com/Microsoft/hcsshim/internal/uvm"
)

func main() {

	fmt.Println("Creating...")
	lcowUVM, err := uvm.Create(&uvm.UVMOptions{OperatingSystem: "linux", ID: "uvm"})
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create utility VM: %s", err)
		os.Exit(-1)
	}

	fmt.Print("Created. Press 'Enter' to start...")
	bufio.NewReader(os.Stdin).ReadBytes('\n')

	fmt.Println("Starting...")
	if err := lcowUVM.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to start utility VM: %s", err)
		os.Exit(-1)
	}

	fmt.Print("Started. Use `hcsdiag console -uvm uvm`. Press 'Enter' to terminate...")
	bufio.NewReader(os.Stdin).ReadBytes('\n')

	lcowUVM.Terminate()
	os.Exit(0)
}
<commit_msg>Remove unused test cmd<commit_after><|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 gcsx

import (
	"testing"

	"github.com/jacobsa/gcloud/gcs/mock_gcs"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

func TestRandomReader(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Invariant-checking random reader
////////////////////////////////////////////////////////////////////////

type checkingRandomReader struct {
	ctx     context.Context
	wrapped RandomReader
}

func (rr *checkingRandomReader) ReadAt(p []byte, offset int64) (int, error) {
	rr.wrapped.CheckInvariants()
	defer rr.wrapped.CheckInvariants()
	return rr.wrapped.ReadAt(rr.ctx, p, offset)
}

func (rr *checkingRandomReader) Destroy() {
	rr.wrapped.CheckInvariants()
	rr.wrapped.Destroy()
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type RandomReaderTest struct {
	bucket mock_gcs.MockBucket
	rr     checkingRandomReader
}

func init() { RegisterTestSuite(&RandomReaderTest{}) }

var _ SetUpInterface = &RandomReaderTest{}
var _ TearDownInterface = &RandomReaderTest{}

func (t *RandomReaderTest) SetUp(ti *TestInfo) {
	panic("TODO")
}

func (t *RandomReaderTest) TearDown() {
	t.rr.Destroy()
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *RandomReaderTest) DoesFoo() {
	AssertTrue(false, "TODO")
}
<commit_msg>Added test names.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 gcsx

import (
	"testing"

	"github.com/jacobsa/gcloud/gcs/mock_gcs"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

func TestRandomReader(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Invariant-checking random reader
////////////////////////////////////////////////////////////////////////

type checkingRandomReader struct {
	ctx     context.Context
	wrapped RandomReader
}

func (rr *checkingRandomReader) ReadAt(p []byte, offset int64) (int, error) {
	rr.wrapped.CheckInvariants()
	defer rr.wrapped.CheckInvariants()
	return rr.wrapped.ReadAt(rr.ctx, p, offset)
}

func (rr *checkingRandomReader) Destroy() {
	rr.wrapped.CheckInvariants()
	rr.wrapped.Destroy()
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type RandomReaderTest struct {
	bucket mock_gcs.MockBucket
	rr     checkingRandomReader
}

func init() { RegisterTestSuite(&RandomReaderTest{}) }

var _ SetUpInterface = &RandomReaderTest{}
var _ TearDownInterface = &RandomReaderTest{}

func (t *RandomReaderTest) SetUp(ti *TestInfo) {
	panic("TODO")
}

func (t *RandomReaderTest) TearDown() {
	t.rr.Destroy()
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *RandomReaderTest) EmptyRead() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) NoExistingReader() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) ExistingReader_WrongOffset() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) NewReaderReturnsError() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) ReaderFails() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) ReaderOvershootsRange() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) ReaderExhausted_ReadFinished() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) ReaderExhausted_ReadNotFinished() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) PropagatesCancellation() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) DoesntPropagateCancellationAfterReturning() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) UpgradesReadsToMinimumSize() {
	AssertTrue(false, "TODO")
}

func (t *RandomReaderTest) UpgradesSequentialReads() {
	AssertTrue(false, "TODO")
}
<|endoftext|>


<commit_before>package sctp

import (
	"encoding/binary"
	"fmt"
)

/*
chunkPayloadData represents an SCTP Chunk of type DATA

 0                   1                   2                   3
 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|   Type = 0    | Reserved|U|B|E|    Length                     |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                              TSN                              |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|      Stream Identifier S      |   Stream Sequence Number n    |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                  Payload Protocol Identifier                  |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                                                               |
|                 User Data (seq n of Stream S)                 |
|                                                               |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

An unfragmented user message shall have both the B and E bits set to
'1'.  Setting both B and E bits to '0' indicates a middle fragment of
a multi-fragment user message, as summarized in the following table:
   B E                  Description
============================================================
|  1 0 | First piece of a fragmented user message          |
+----------------------------------------------------------+
|  0 0 | Middle piece of a fragmented user message         |
+----------------------------------------------------------+
|  0 1 | Last piece of a fragmented user message           |
+----------------------------------------------------------+
|  1 1 | Unfragmented message                              |
============================================================
|             Table 1: Fragment Description Flags          |
============================================================
*/
type chunkPayloadData struct {
	chunkHeader

	unordered        bool
	beginingFragment bool
	endingFragment   bool
	immediateSack    bool

	tsn                  uint32
	streamIdentifier     uint16
	streamSequenceNumber uint16
	payloadType          PayloadProtocolIdentifier
	userData             []byte
}

const (
	payloadDataEndingFragmentBitmask   = 1
	payloadDataBeginingFragmentBitmask = 2
	payloadDataUnorderedBitmask        = 4
	payloadDataImmediateSACK           = 8

	payloadDataHeaderSize = 12
)

// PayloadProtocolIdentifier is an enum for DataChannel payload types
type PayloadProtocolIdentifier uint32

// PayloadProtocolIdentifier enums
const (
	PayloadTypeWebRTCDCEP        PayloadProtocolIdentifier = 50
	PayloadTypeWebRTCString      PayloadProtocolIdentifier = 51
	PayloadTypeWebRTCBinary      PayloadProtocolIdentifier = 52
	PayloadTypeWebRTCStringEmpty PayloadProtocolIdentifier = 56
	PayloadTypeWebRTCBinaryEmpty PayloadProtocolIdentifier = 57
)

func (p PayloadProtocolIdentifier) String() string {
	switch p {
	case PayloadTypeWebRTCDCEP:
		return "WebRTC DCEP"
	case PayloadTypeWebRTCString:
		return "WebRTC String"
	case PayloadTypeWebRTCBinary:
		return "WebRTC Binary"
	case PayloadTypeWebRTCStringEmpty:
		return "WebRTC String (Empty)"
	case PayloadTypeWebRTCBinaryEmpty:
		return "WebRTC Binary (Empty)"
	default:
		return fmt.Sprintf("Unknown Payload Protocol Identifier: %d", p)
	}
}

func (p *chunkPayloadData) unmarshal(raw []byte) error {
	if err := p.chunkHeader.unmarshal(raw); err != nil {
		return err
	}

	p.immediateSack = p.flags&payloadDataImmediateSACK != 0
	p.unordered = p.flags&payloadDataUnorderedBitmask != 0
	p.beginingFragment = p.flags&payloadDataBeginingFragmentBitmask != 0
	p.endingFragment = p.flags&payloadDataEndingFragmentBitmask != 0

	p.tsn = binary.BigEndian.Uint32(p.raw[0:])
	p.streamIdentifier = binary.BigEndian.Uint16(p.raw[4:])
	p.streamSequenceNumber = binary.BigEndian.Uint16(p.raw[6:])
	p.payloadType = PayloadProtocolIdentifier(binary.BigEndian.Uint32(p.raw[8:]))
	p.userData = p.raw[payloadDataHeaderSize:]

	return nil
}

func (p *chunkPayloadData) marshal() ([]byte, error) {

	payRaw := make([]byte, payloadDataHeaderSize+len(p.userData))

	binary.BigEndian.PutUint32(payRaw[0:], p.tsn)
	binary.BigEndian.PutUint16(payRaw[4:], p.streamIdentifier)
	binary.BigEndian.PutUint16(payRaw[6:], p.streamSequenceNumber)
	binary.BigEndian.PutUint32(payRaw[8:], uint32(p.payloadType))
	copy(payRaw[payloadDataHeaderSize:], p.userData)

	flags := uint8(0)
	if p.endingFragment {
		flags = 1
	}
	if p.beginingFragment {
		flags |= 1 << 1
	}
	if p.unordered {
		flags |= 1 << 2
	}
	if p.immediateSack {
		flags |= 1 << 3
	}

	p.chunkHeader.flags = flags
	p.chunkHeader.typ = PAYLOADDATA
	p.chunkHeader.raw = payRaw
	return p.chunkHeader.marshal()
}

func (p *chunkPayloadData) check() (abort bool, err error) {
	return false, nil
}
<commit_msg>Fix PayloadTypeWebRTCBinary value<commit_after>package sctp

import (
	"encoding/binary"
	"fmt"
)

/*
chunkPayloadData represents an SCTP Chunk of type DATA

 0                   1                   2                   3
 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|   Type = 0    | Reserved|U|B|E|    Length                     |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                              TSN                              |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|      Stream Identifier S      |   Stream Sequence Number n    |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                  Payload Protocol Identifier                  |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                                                               |
|                 User Data (seq n of Stream S)                 |
|                                                               |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

An unfragmented user message shall have both the B and E bits set to
'1'.  Setting both B and E bits to '0' indicates a middle fragment of
a multi-fragment user message, as summarized in the following table:
   B E                  Description
============================================================
|  1 0 | First piece of a fragmented user message          |
+----------------------------------------------------------+
|  0 0 | Middle piece of a fragmented user message         |
+----------------------------------------------------------+
|  0 1 | Last piece of a fragmented user message           |
+----------------------------------------------------------+
|  1 1 | Unfragmented message                              |
============================================================
|             Table 1: Fragment Description Flags          |
============================================================
*/
type chunkPayloadData struct {
	chunkHeader

	unordered        bool
	beginingFragment bool
	endingFragment   bool
	immediateSack    bool

	tsn                  uint32
	streamIdentifier     uint16
	streamSequenceNumber uint16
	payloadType          PayloadProtocolIdentifier
	userData             []byte
}

const (
	payloadDataEndingFragmentBitmask   = 1
	payloadDataBeginingFragmentBitmask = 2
	payloadDataUnorderedBitmask        = 4
	payloadDataImmediateSACK           = 8

	payloadDataHeaderSize = 12
)

// PayloadProtocolIdentifier is an enum for DataChannel payload types
type PayloadProtocolIdentifier uint32

// PayloadProtocolIdentifier enums
const (
	PayloadTypeWebRTCDCEP        PayloadProtocolIdentifier = 50
	PayloadTypeWebRTCString      PayloadProtocolIdentifier = 51
	PayloadTypeWebRTCBinary      PayloadProtocolIdentifier = 53
	PayloadTypeWebRTCStringEmpty PayloadProtocolIdentifier = 56
	PayloadTypeWebRTCBinaryEmpty PayloadProtocolIdentifier = 57
)

func (p PayloadProtocolIdentifier) String() string {
	switch p {
	case PayloadTypeWebRTCDCEP:
		return "WebRTC DCEP"
	case PayloadTypeWebRTCString:
		return "WebRTC String"
	case PayloadTypeWebRTCBinary:
		return "WebRTC Binary"
	case PayloadTypeWebRTCStringEmpty:
		return "WebRTC String (Empty)"
	case PayloadTypeWebRTCBinaryEmpty:
		return "WebRTC Binary (Empty)"
	default:
		return fmt.Sprintf("Unknown Payload Protocol Identifier: %d", p)
	}
}

func (p *chunkPayloadData) unmarshal(raw []byte) error {
	if err := p.chunkHeader.unmarshal(raw); err != nil {
		return err
	}

	p.immediateSack = p.flags&payloadDataImmediateSACK != 0
	p.unordered = p.flags&payloadDataUnorderedBitmask != 0
	p.beginingFragment = p.flags&payloadDataBeginingFragmentBitmask != 0
	p.endingFragment = p.flags&payloadDataEndingFragmentBitmask != 0

	p.tsn = binary.BigEndian.Uint32(p.raw[0:])
	p.streamIdentifier = binary.BigEndian.Uint16(p.raw[4:])
	p.streamSequenceNumber = binary.BigEndian.Uint16(p.raw[6:])
	p.payloadType = PayloadProtocolIdentifier(binary.BigEndian.Uint32(p.raw[8:]))
	p.userData = p.raw[payloadDataHeaderSize:]

	return nil
}

func (p *chunkPayloadData) marshal() ([]byte, error) {

	payRaw := make([]byte, payloadDataHeaderSize+len(p.userData))

	binary.BigEndian.PutUint32(payRaw[0:], p.tsn)
	binary.BigEndian.PutUint16(payRaw[4:], p.streamIdentifier)
	binary.BigEndian.PutUint16(payRaw[6:], p.streamSequenceNumber)
	binary.BigEndian.PutUint32(payRaw[8:], uint32(p.payloadType))
	copy(payRaw[payloadDataHeaderSize:], p.userData)

	flags := uint8(0)
	if p.endingFragment {
		flags = 1
	}
	if p.beginingFragment {
		flags |= 1 << 1
	}
	if p.unordered {
		flags |= 1 << 2
	}
	if p.immediateSack {
		flags |= 1 << 3
	}

	p.chunkHeader.flags = flags
	p.chunkHeader.typ = PAYLOADDATA
	p.chunkHeader.raw = payRaw
	return p.chunkHeader.marshal()
}

func (p *chunkPayloadData) check() (abort bool, err error) {
	return false, nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"mime"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/progrium/go-basher"
)

var Version string

func assert(err error) {
	if err != nil {
		println("!!", err.Error())
		os.Exit(2)
	}
}

func fatal(msg string) {
	println("!!", msg)
	os.Exit(2)
}

func UploadUrl(args []string) {
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var release map[string]interface{}
	assert(json.Unmarshal(bytes, &release))
	url, ok := release["upload_url"].(string)
	if !ok {
		os.Exit(2)
	}
	url = strings.Replace(url, "{", "", 1)
	url = strings.Replace(url, "}", "", 1)
	fmt.Println(url)
}

func ReleaseIdFromTagname(args []string) {
	tagname := args[0]
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var releases []map[string]interface{}
	assert(json.Unmarshal(bytes, &releases))
	for _, release := range releases {
		if release["tag_name"].(string) == tagname {
			fmt.Println(strconv.Itoa(int(release["id"].(float64))))
			return
		}
	}
	os.Exit(2)
}

func MimeType(args []string) {
	filename := args[0]
	ext := filepath.Ext(filename)
	mime.AddExtensionType(".gz", "application/gzip")
	mime.AddExtensionType(".tgz", "application/gzip")
	mime.AddExtensionType(".tar", "application/tar")
	mime.AddExtensionType(".zip", "application/zip")
	mimetype := mime.TypeByExtension(ext)
	if mimetype != "" {
		fmt.Println(mimetype)
	} else {
		fmt.Println("application/octet-stream")
	}
}

func Checksum(args []string) {
	if len(args) < 1 {
		fatal("No algorithm specified")
	}
	var h hash.Hash
	switch args[0] {
	case "md5":
		h = md5.New()
	case "sha1":
		h = sha1.New()
	case "sha256":
		h = sha256.New()
	default:
		fatal("Algorithm '" + args[0] + "' is unsupported")
	}
	io.Copy(h, os.Stdin)
	fmt.Printf("%x\n", h.Sum(nil))
}

func main() {
	os.Setenv("VERSION", Version)
	basher.Application(map[string]func([]string){
		"upload-url":              UploadUrl,
		"release-id-from-tagname": ReleaseIdFromTagname,
		"mimetype":                MimeType,
		"checksum":                Checksum,
	}, []string{
		"bash/gh-release.bash",
	}, Asset, true)
}
<commit_msg>Fixes #12: Upload url parsing now fixed<commit_after>package main

import (
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"mime"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/progrium/go-basher"
)

var Version string

func assert(err error) {
	if err != nil {
		println("!!", err.Error())
		os.Exit(2)
	}
}

func fatal(msg string) {
	println("!!", msg)
	os.Exit(2)
}

func UploadUrl(args []string) {
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var release map[string]interface{}
	assert(json.Unmarshal(bytes, &release))
	url, ok := release["upload_url"].(string)
	if !ok {
		os.Exit(2)
	}
	i := strings.Index(url, "{")
	if i > -1 {
		url = url[:i]
	}
	i = strings.Index(url, "?")
	if i > -1 {
		url = url[:i]
	}
	fmt.Println(url + "?name=")
}

func ReleaseIdFromTagname(args []string) {
	tagname := args[0]
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var releases []map[string]interface{}
	assert(json.Unmarshal(bytes, &releases))
	for _, release := range releases {
		if release["tag_name"].(string) == tagname {
			fmt.Println(strconv.Itoa(int(release["id"].(float64))))
			return
		}
	}
	os.Exit(2)
}

func MimeType(args []string) {
	filename := args[0]
	ext := filepath.Ext(filename)
	mime.AddExtensionType(".gz", "application/gzip")
	mime.AddExtensionType(".tgz", "application/gzip")
	mime.AddExtensionType(".tar", "application/tar")
	mime.AddExtensionType(".zip", "application/zip")
	mimetype := mime.TypeByExtension(ext)
	if mimetype != "" {
		fmt.Println(mimetype)
	} else {
		fmt.Println("application/octet-stream")
	}
}

func Checksum(args []string) {
	if len(args) < 1 {
		fatal("No algorithm specified")
	}
	var h hash.Hash
	switch args[0] {
	case "md5":
		h = md5.New()
	case "sha1":
		h = sha1.New()
	case "sha256":
		h = sha256.New()
	default:
		fatal("Algorithm '" + args[0] + "' is unsupported")
	}
	io.Copy(h, os.Stdin)
	fmt.Printf("%x\n", h.Sum(nil))
}

func main() {
	os.Setenv("VERSION", Version)
	basher.Application(map[string]func([]string){
		"upload-url":              UploadUrl,
		"release-id-from-tagname": ReleaseIdFromTagname,
		"mimetype":                MimeType,
		"checksum":                Checksum,
	}, []string{
		"bash/gh-release.bash",
	}, Asset, true)
}
<|endoftext|>


<commit_before>package main

import (
    "github.com/hoisie/web.go"
)

func hello(val string) string { return "hello " + val }

func main() {
    web.Get("/(.*)", hello)
    web.Run("0.0.0.0:9999")
}
<commit_msg>Update example<commit_after>package main

import (
    "github.com/rday/web"
)

type Message struct {
    Greeting string
    Response string
}

func hello(val string) (Message, error) {
    msg := Message{val, "Hello " + val}
    return msg, nil
}

func plain(val string) ([]byte, error) {
    return []byte("Plain " + val), nil
}

func main() {
    web.Get("/plain/(.*)", plain)
    web.Get("/(.*)", hello)
    web.Run("0.0.0.0:9999")
}

<|endoftext|>


<commit_before>package main

import (
    //"fmt"
    "image"
    "image/color"
    //"image/gif"
    "os"
    //"bufio"
    "io"
    "compress/lzw"
    "bytes"
    "math"
)

func writeHeader(w io.Writer, m image.Image) {
    b := m.Bounds()

    header := make([]byte, 0x320)

    header[0] = 'G'
    header[1] = 'I'
    header[2] = 'F'
    header[3] = '8'
    header[4] = '9'
    header[5] = 'a'

    header[7] = byte(b.Max.X / 255)
    header[6] = byte(b.Max.X % 255)
    header[9] = byte(b.Max.Y / 255)
    header[8] = byte(b.Max.Y % 255)

    header[0x0B] = byte(0x00) // Background color.
    header[0x0C] = byte(0x00) // Default pixel aspect ratio.

    // Grayscale color table.
    for i := 0; i < 255; i++ {
        header[0x0F + i * 3] = byte(i)
        header[0x0E + i * 3] = byte(i)
        header[0x0D + i * 3] = byte(i)
    }

    header[0x30D] = byte(0x21) // GCE data header.
    header[0x30E] = byte(0xF9) // GCE data header (cont).
    header[0x30F] = byte(0x04) // Next 4 bytes are GCE data.
    header[0x310] = byte(0x01) // There is a transparent pixel.
    header[0x311] = byte(0x00) // Animation delay, LSB.
    header[0x312] = byte(0x00) // Animation delay, MSB.
    header[0x313] = byte(0x10) // And it is color #16 (0x10).
    header[0x314] = byte(0x00) // End of GCE data.

    header[0x315] = byte(0x2C) // Start of Image Descriptor.

    header[0x316] = byte(b.Min.X / 255)
    header[0x315] = byte(b.Min.X % 255)
    header[0x318] = byte(b.Min.Y / 255)
    header[0x317] = byte(b.Min.Y % 255)

    header[0x31B] = byte(b.Max.X / 255)
    header[0x31A] = byte(b.Max.X % 255)
    header[0x31D] = byte(b.Max.Y / 255)
    header[0x31C] = byte(b.Max.Y % 255)

    header[0x31E] = byte(0x00) // No local color table.

    header[0x31F] = byte(0x08) // Start of LZW with minimum code size 8.

    w.Write(header)
}

func compressImage(m image.Image) *bytes.Buffer {
    b := m.Bounds()

    compressedImageBuffer := bytes.NewBuffer(make([]byte, 0, 255))
    lzww := lzw.NewWriter(compressedImageBuffer, lzw.LSB, int(8))

    for y := b.Min.Y; y < b.Max.Y; y++ {
        for x := b.Min.X; x < b.Max.X; x++ {
            c := color.GrayModel.Convert(m.At(x, y)).(color.Gray)
            lzww.Write([]byte{c.Y})
            //lzww.Write([]byte{byte(x ^ y)})
            //lzww.Write([]byte{byte(0x00)})
        }
    }
    lzww.Close()

    return compressedImageBuffer
}

func writeBlocks(w io.Writer, compressedImage *bytes.Buffer) {
    const maxBlockSize = 255
    bytesSoFar := 0
    bytesRemaining := compressedImage.Len()
    for bytesRemaining > 0 {
        if bytesSoFar == 0 {
            blockSize := math.Min(maxBlockSize, float64(bytesRemaining))
            w.Write([]byte{byte(blockSize)})
        }

        b, _ :=  compressedImage.ReadByte()
        w.Write([]byte{b})

        bytesSoFar = (bytesSoFar + 1) % maxBlockSize
        bytesRemaining--
    }
}

func Encode(w io.Writer, m image.Image) error {
    writeHeader(w, m)
    writeBlocks(w, compressImage(m))
    w.Write([]byte{0, ';'})

    return nil
}

func main() {
    m := image.NewRGBA(image.Rect(0, 0, 100, 100))
    m.Set(1, 1, color.RGBA{0x00, 0xFF, 0x00, 0xFF})
    file, _ := os.Create("new_image.gif")
    Encode(file, m)
}
<commit_msg>Add GCT reference to header<commit_after>package main

import (
    //"fmt"
    "image"
    "image/color"
    //"image/gif"
    "os"
    //"bufio"
    "io"
    "compress/lzw"
    "bytes"
    "math"
)

func writeHeader(w io.Writer, m image.Image) {
    b := m.Bounds()

    header := make([]byte, 0x320)

    header[0] = 'G'
    header[1] = 'I'
    header[2] = 'F'
    header[3] = '8'
    header[4] = '9'
    header[5] = 'a'

    header[7] = byte(b.Max.X / 255)
    header[6] = byte(b.Max.X % 255)
    header[9] = byte(b.Max.Y / 255)
    header[8] = byte(b.Max.Y % 255)

    header[0x0A] = byte(0xF7) // GCT follows for 256 colors with resolution
                              // 3 x 8 bits/primary

    header[0x0B] = byte(0x00) // Background color.
    header[0x0C] = byte(0x00) // Default pixel aspect ratio.

    // Grayscale color table.
    for i := 0; i < 255; i++ {
        header[0x0F + i * 3] = byte(i)
        header[0x0E + i * 3] = byte(i)
        header[0x0D + i * 3] = byte(i)
    }

    header[0x30D] = byte(0x21) // GCE data header.
    header[0x30E] = byte(0xF9) // GCE data header (cont).
    header[0x30F] = byte(0x04) // Next 4 bytes are GCE data.
    header[0x310] = byte(0x01) // There is a transparent pixel.
    header[0x311] = byte(0x00) // Animation delay, LSB.
    header[0x312] = byte(0x00) // Animation delay, MSB.
    header[0x313] = byte(0x10) // And it is color #16 (0x10).
    header[0x314] = byte(0x00) // End of GCE data.

    header[0x315] = byte(0x2C) // Start of Image Descriptor.

    header[0x316] = byte(b.Min.X / 255)
    header[0x315] = byte(b.Min.X % 255)
    header[0x318] = byte(b.Min.Y / 255)
    header[0x317] = byte(b.Min.Y % 255)

    header[0x31B] = byte(b.Max.X / 255)
    header[0x31A] = byte(b.Max.X % 255)
    header[0x31D] = byte(b.Max.Y / 255)
    header[0x31C] = byte(b.Max.Y % 255)

    header[0x31E] = byte(0x00) // No local color table.

    header[0x31F] = byte(0x08) // Start of LZW with minimum code size 8.

    w.Write(header)
}

func compressImage(m image.Image) *bytes.Buffer {
    b := m.Bounds()

    compressedImageBuffer := bytes.NewBuffer(make([]byte, 0, 255))
    lzww := lzw.NewWriter(compressedImageBuffer, lzw.LSB, int(8))

    for y := b.Min.Y; y < b.Max.Y; y++ {
        for x := b.Min.X; x < b.Max.X; x++ {
            c := color.GrayModel.Convert(m.At(x, y)).(color.Gray)
            lzww.Write([]byte{c.Y})
            //lzww.Write([]byte{byte(x ^ y)})
            //lzww.Write([]byte{byte(0x00)})
        }
    }
    lzww.Close()

    return compressedImageBuffer
}

func writeBlocks(w io.Writer, compressedImage *bytes.Buffer) {
    const maxBlockSize = 255
    bytesSoFar := 0
    bytesRemaining := compressedImage.Len()
    for bytesRemaining > 0 {
        if bytesSoFar == 0 {
            blockSize := math.Min(maxBlockSize, float64(bytesRemaining))
            w.Write([]byte{byte(blockSize)})
        }

        b, _ :=  compressedImage.ReadByte()
        w.Write([]byte{b})

        bytesSoFar = (bytesSoFar + 1) % maxBlockSize
        bytesRemaining--
    }
}

func Encode(w io.Writer, m image.Image) error {
    writeHeader(w, m)
    writeBlocks(w, compressImage(m))
    w.Write([]byte{0, ';'})

    return nil
}

func main() {
    m := image.NewRGBA(image.Rect(0, 0, 100, 100))
    m.Set(1, 1, color.RGBA{0x00, 0xFF, 0x00, 0xFF})
    file, _ := os.Create("new_image.gif")
    Encode(file, m)
}
<|endoftext|>


<commit_before>package tests_test

import (
	"io"
	"os/exec"
	"strings"
	"time"

	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/client-go/kubecli"

	"kubevirt.io/kubevirt/tests"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	// Define relevant k8s versions
	relevantk8sVer = "1.16.2"

	// Define a timeout for read opeartions in order to prevent test hanging
	readTimeout     = 1 * time.Minute
	processWaitTime = 2 * time.Minute

	bufferSize = 1024
)

var _ = Describe("[rfe_id:3423][crit:high][vendor:cnv-qe@redhat.com][level:component]VmWatch", func() {
	tests.FlagParse()

	virtCli, err := kubecli.GetKubevirtClient()
	tests.PanicOnError(err)

	type vmStatus struct {
		name,
		age,
		running,
		volume string
	}

	type vmiStatus struct {
		name,
		age,
		phase,
		ip,
		node string
	}

	newVMStatus := func(fields []string) *vmStatus {
		flen := len(fields)
		stat := &vmStatus{}

		switch {
		case flen > 3:
			stat.volume = fields[3]
			fallthrough
		case flen > 2:
			stat.running = fields[2]
			fallthrough
		case flen > 1:
			stat.age = fields[1]
			fallthrough
		case flen > 0:
			stat.name = fields[0]
		}

		return stat
	}

	newVMIStatus := func(fields []string) *vmiStatus {
		flen := len(fields)
		stat := &vmiStatus{}

		switch {
		case flen > 4:
			stat.node = fields[4]
			fallthrough
		case flen > 3:
			stat.ip = fields[3]
			fallthrough
		case flen > 2:
			stat.phase = fields[2]
			fallthrough
		case flen > 1:
			stat.age = fields[1]
			fallthrough
		case flen > 0:
			stat.name = fields[0]
		}

		return stat
	}

	// Fail the test if stderr has something to read
	failOnError := func(rc io.ReadCloser) {
		defer GinkgoRecover()

		buf := make([]byte, bufferSize)

		n, err := rc.Read(buf)

		if err != nil && n > 0 {
			rc.Close()
			Fail(string(buf[:n]))
		}
	}

	// Reads from stdin until a newline character is found
	readLine := func(rc io.ReadCloser, timeout time.Duration) string {
		lineChan := make(chan string)

		go func() {
			defer GinkgoRecover()

			var line strings.Builder
			buf := make([]byte, 1)
			defer close(lineChan)

			for {
				n, err := rc.Read(buf)

				if err != nil && err != io.EOF && !strings.Contains(err.Error(), "file already closed") {
					Fail(err.Error())
					return
				}

				if n > 0 {
					if buf[0] != '\n' {
						line.WriteByte(buf[0])
					} else {
						break
					}
				}
			}

			lineChan <- line.String()
		}()

		select {
		case line := <-lineChan:
			return line
		case <-time.After(timeout):
			err := rc.Close()
			Expect(err).ToNot(HaveOccurred(), "stdout should have been closed properly")

			Fail("Timeout reached on read operation")

			return ""
		}
	}

	// Reads VM status from the given pipe (stdin in this case) and
	// returns a new status.
	// if old_status is non-nil, the function will read status lines until
	// new_status.running != old_status.running in order to skip duplicated status lines
	readVMStatus := func(rc io.ReadCloser, old_status *vmStatus, timeout time.Duration) *vmStatus {
		new_stat := newVMStatus(strings.Fields(readLine(rc, timeout)))

		for old_status != nil && new_stat.running == old_status.running {
			new_stat = newVMStatus(strings.Fields(readLine(rc, timeout)))
		}

		return new_stat
	}

	// Reads VMI status from the given pipe (stdin in this case) and
	// returns a new status.
	// if old_status is non-nil, the function will read status lines until
	// new_status.phase != old_status.phase in order to skip duplicated lines
	readVMIStatus := func(rc io.ReadCloser, old_status *vmiStatus, timeout time.Duration) *vmiStatus {
		newStat := newVMIStatus(strings.Fields(readLine(rc, timeout)))

		for old_status != nil && newStat.phase == old_status.phase {
			newStat = newVMIStatus(strings.Fields(readLine(rc, timeout)))
		}

		return newStat
	}

	// Create a command with output/error redirection.
	// Returns (cmd, stdout, stderr)
	createCommandWithNSAndRedirect := func(namespace, cmdName string, args ...string) (*exec.Cmd, io.ReadCloser, io.ReadCloser) {
		cmdName, cmd, err := tests.CreateCommandWithNS(namespace, cmdName, args...)

		Expect(cmdName).ToNot(Equal(""))
		Expect(cmd).ToNot(BeNil())
		Expect(err).ToNot(HaveOccurred(), "Command should have been created with proper kubectl/oc arguments")

		// Output redirection
		stdOut, err := cmd.StdoutPipe()
		Expect(err).ToNot(HaveOccurred(), "stdout should have been redirected")
		Expect(stdOut).ToNot(BeNil())

		stdErr, err := cmd.StderrPipe()
		Expect(err).ToNot(HaveOccurred(), "stderr should have been redirected")
		Expect(stdErr).ToNot(BeNil())

		return cmd, stdOut, stdErr
	}

	BeforeEach(func() {
		tests.SkipIfVersionBelow("Printing format for `kubectl get -w` on custom resources is only relevant for 1.16.2+", relevantk8sVer)
	})

	It("[test_id:3468]Should update vm status with the proper columns using 'kubectl get vm -w'", func() {
		By("Creating a new VM spec")
		vm := tests.NewRandomVMWithEphemeralDisk(tests.ContainerDiskFor(tests.ContainerDiskCirros))
		Expect(vm).ToNot(BeNil())

		By("Setting up the kubectl command")
		cmd, stdout, stderr :=
			createCommandWithNSAndRedirect(vm.ObjectMeta.Namespace, tests.GetK8sCmdClient(), "get", "vm", "-w")
		Expect(cmd).ToNot(BeNil())

		err = cmd.Start()
		Expect(err).ToNot(HaveOccurred(), "Command should have started successfully")

		defer cmd.Process.Kill()

		time.Sleep(processWaitTime)

		go failOnError(stderr)

		By("Applying the VM to the cluster")
		vm, err := virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Create(vm)
		Expect(err).ToNot(HaveOccurred(), "VM should have been added to the cluster")

		// Read column titles
		vmStatus := readVMStatus(stdout, nil, readTimeout)
		Expect(vmStatus.name).To(Equal("NAME"), "Output should have the NAME column")
		Expect(vmStatus.age).To(Equal("AGE"), "Output should have the AGE column")
		Expect(vmStatus.running).To(Equal("RUNNING"), "Output should have the RUNNING column")
		Expect(vmStatus.volume).To(Equal("VOLUME"), "Output should have the VOLUME column")

		// Read first status of the vm
		vmStatus = readVMStatus(stdout, vmStatus, readTimeout)
		Expect(vmStatus.name).To(Equal(vm.Name))
		By("Expecting vm.running == false")
		Expect(vmStatus.running).To(Equal("false"))

		By("Starting the VM")
		vm = tests.StartVirtualMachine(vm)

		vmStatus = readVMStatus(stdout, vmStatus, readTimeout)
		By("Expecting vm.running == true")
		Expect(vmStatus.running).To(Equal("true"))

		By("Restarting the VM")
		err = virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Restart(vm.ObjectMeta.Name)
		Expect(err).ToNot(HaveOccurred(), "VM should have been restarted")

		vmStatus = readVMStatus(stdout, nil, readTimeout)
		By("Expecting vm.running == true")
		Expect(vmStatus.running).To(Equal("true"))

		By("Stopping the VM")
		vm = tests.StopVirtualMachine(vm)

		vmStatus = readVMStatus(stdout, vmStatus, readTimeout)
		By("Expecting vm.running == false")
		Expect(vmStatus.running).To(Equal("false"))

		By("Deleting the VM")
		err = virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Delete(vm.ObjectMeta.Name, &v1.DeleteOptions{})
		Expect(err).ToNot(HaveOccurred(), "VM should have been deleted from the cluster")
	})

	It("[test_id:3466]Should update vmi status with the proper columns using 'kubectl get vmi -w'", func() {
		By("Creating a random VMI spec")
		vm := tests.NewRandomVMWithEphemeralDisk(tests.ContainerDiskFor(tests.ContainerDiskCirros))

		Expect(vm).ToNot(BeNil())

		By("Setting up the kubectl command")
		cmd, stdout, stderr :=
			createCommandWithNSAndRedirect(vm.ObjectMeta.Namespace, tests.GetK8sCmdClient(), "get", "vmi", "-w")
		Expect(cmd).ToNot(BeNil())

		err = cmd.Start()
		Expect(err).ToNot(HaveOccurred(), "Command should have stared successfully")

		defer cmd.Process.Kill()

		time.Sleep(processWaitTime)

		go failOnError(stderr)

		By("Applying vmi to the cluster")
		vm, err = virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Create(vm)
		Expect(err).ToNot(HaveOccurred(), "VMI should have been added to the cluster")

		// Start a VMI
		vm = tests.StartVirtualMachine(vm)

		// Read the column titles
		vmiStatus := readVMIStatus(stdout, nil, readTimeout)
		Expect(vmiStatus.name).To(Equal("NAME"), "Output should have the NAME column")
		Expect(vmiStatus.age).To(Equal("AGE"), "Output should have the AGE column")
		Expect(vmiStatus.phase).To(Equal("PHASE"), "Output should have the PHASE column")
		Expect(vmiStatus.ip).To(Equal("IP"), "Output should have the IP column")
		Expect(vmiStatus.node).To(Equal("NODENAME"), "Output should have the NODENAME column")

		By("Expecting vmi.phase == ''")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal(""))

		By("Expecting vmi.phase == Pending")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Pending"))

		By("Expecting vmi.phase == Scheduling")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Scheduling"))

		By("Expecting vmi.phase == Scheduled")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Scheduled"))

		By("Expecting vmi.phase == Running")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Running"))

		// Restart the VMI
		err = virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Restart(vm.ObjectMeta.Name)
		Expect(err).ToNot(HaveOccurred(), "VMI should have been restarted")

		By("Expecting vmi.phase == Failed")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Failed"))

		By("Expecting vmi.phase == ''")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal(""))

		By("Expecting vmi.phase == Pending")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Pending"))

		By("Expecting vmi.phase == Scheduling")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Scheduling"))

		By("Expecting vmi.phase == Scheduled")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Scheduled"))

		By("Expecting vmi.phase == Running")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Running"))
	})
})
<commit_msg>Renamed variables<commit_after>package tests_test

import (
	"io"
	"os/exec"
	"strings"
	"time"

	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/client-go/kubecli"

	"kubevirt.io/kubevirt/tests"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	// Define relevant k8s versions
	relevantk8sVer = "1.16.2"

	// Define a timeout for read opeartions in order to prevent test hanging
	readTimeout     = 1 * time.Minute
	processWaitTime = 2 * time.Minute

	bufferSize = 1024
)

var _ = FDescribe("[rfe_id:3423][crit:high][vendor:cnv-qe@redhat.com][level:component]VmWatch", func() {
	tests.FlagParse()

	virtCli, err := kubecli.GetKubevirtClient()
	tests.PanicOnError(err)

	type vmStatus struct {
		name,
		age,
		running,
		volume string
	}

	type vmiStatus struct {
		name,
		age,
		phase,
		ip,
		node string
	}

	newVMStatus := func(fields []string) *vmStatus {
		flen := len(fields)
		stat := &vmStatus{}

		switch {
		case flen > 3:
			stat.volume = fields[3]
			fallthrough
		case flen > 2:
			stat.running = fields[2]
			fallthrough
		case flen > 1:
			stat.age = fields[1]
			fallthrough
		case flen > 0:
			stat.name = fields[0]
		}

		return stat
	}

	newVMIStatus := func(fields []string) *vmiStatus {
		flen := len(fields)
		stat := &vmiStatus{}

		switch {
		case flen > 4:
			stat.node = fields[4]
			fallthrough
		case flen > 3:
			stat.ip = fields[3]
			fallthrough
		case flen > 2:
			stat.phase = fields[2]
			fallthrough
		case flen > 1:
			stat.age = fields[1]
			fallthrough
		case flen > 0:
			stat.name = fields[0]
		}

		return stat
	}

	// Fail the test if stderr has something to read
	failOnError := func(rc io.ReadCloser) {
		defer GinkgoRecover()

		buf := make([]byte, bufferSize)

		n, err := rc.Read(buf)

		if err != nil && n > 0 {
			rc.Close()
			Fail(string(buf[:n]))
		}
	}

	// Reads from stdin until a newline character is found
	readLine := func(rc io.ReadCloser, timeout time.Duration) string {
		lineChan := make(chan string)

		go func() {
			defer GinkgoRecover()

			var line strings.Builder
			buf := make([]byte, 1)
			defer close(lineChan)

			for {
				n, err := rc.Read(buf)

				if err != nil && err != io.EOF && !strings.Contains(err.Error(), "file already closed") {
					Fail(err.Error())
					return
				}

				if n > 0 {
					if buf[0] != '\n' {
						line.WriteByte(buf[0])
					} else {
						break
					}
				}
			}

			lineChan <- line.String()
		}()

		select {
		case line := <-lineChan:
			return line
		case <-time.After(timeout):
			err := rc.Close()
			Expect(err).ToNot(HaveOccurred(), "stdout should have been closed properly")

			Fail("Timeout reached on read operation")

			return ""
		}
	}

	// Reads VM status from the given pipe (stdin in this case) and
	// returns a new status.
	// if old_status is non-nil, the function will read status lines until
	// new_status.running != old_status.running in order to skip duplicated status lines
	readVMStatus := func(rc io.ReadCloser, oldStatus *vmStatus, timeout time.Duration) *vmStatus {
		newStat := newVMStatus(strings.Fields(readLine(rc, timeout)))

		for oldStatus != nil && newStat.running == oldStatus.running {
			newStat = newVMStatus(strings.Fields(readLine(rc, timeout)))
		}

		return newStat
	}

	// Reads VMI status from the given pipe (stdin in this case) and
	// returns a new status.
	// if old_status is non-nil, the function will read status lines until
	// new_status.phase != old_status.phase in order to skip duplicated lines
	readVMIStatus := func(rc io.ReadCloser, oldStatus *vmiStatus, timeout time.Duration) *vmiStatus {
		newStat := newVMIStatus(strings.Fields(readLine(rc, timeout)))

		for oldStatus != nil && newStat.phase == oldStatus.phase {
			newStat = newVMIStatus(strings.Fields(readLine(rc, timeout)))
		}

		return newStat
	}

	// Create a command with output/error redirection.
	// Returns (cmd, stdout, stderr)
	createCommandWithNSAndRedirect := func(namespace, cmdName string, args ...string) (*exec.Cmd, io.ReadCloser, io.ReadCloser) {
		cmdName, cmd, err := tests.CreateCommandWithNS(namespace, cmdName, args...)

		Expect(cmdName).ToNot(Equal(""))
		Expect(cmd).ToNot(BeNil())
		Expect(err).ToNot(HaveOccurred(), "Command should have been created with proper kubectl/oc arguments")

		// Output redirection
		stdOut, err := cmd.StdoutPipe()
		Expect(err).ToNot(HaveOccurred(), "stdout should have been redirected")
		Expect(stdOut).ToNot(BeNil())

		stdErr, err := cmd.StderrPipe()
		Expect(err).ToNot(HaveOccurred(), "stderr should have been redirected")
		Expect(stdErr).ToNot(BeNil())

		return cmd, stdOut, stdErr
	}

	BeforeEach(func() {
		tests.SkipIfVersionBelow("Printing format for `kubectl get -w` on custom resources is only relevant for 1.16.2+", relevantk8sVer)
		tests.BeforeTestCleanup()
	})

	It("[test_id:3468]Should update vm status with the proper columns using 'kubectl get vm -w'", func() {
		By("Creating a new VM spec")
		vm := tests.NewRandomVMWithEphemeralDisk(tests.ContainerDiskFor(tests.ContainerDiskCirros))
		Expect(vm).ToNot(BeNil())

		By("Setting up the kubectl command")
		cmd, stdout, stderr :=
			createCommandWithNSAndRedirect(vm.ObjectMeta.Namespace, tests.GetK8sCmdClient(), "get", "vm", "-w")
		Expect(cmd).ToNot(BeNil())

		err = cmd.Start()
		Expect(err).ToNot(HaveOccurred(), "Command should have started successfully")

		defer cmd.Process.Kill()

		time.Sleep(processWaitTime)

		go failOnError(stderr)

		By("Applying the VM to the cluster")
		vm, err := virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Create(vm)
		Expect(err).ToNot(HaveOccurred(), "VM should have been added to the cluster")

		// Read column titles
		vmStatus := readVMStatus(stdout, nil, readTimeout)
		Expect(vmStatus.name).To(Equal("NAME"), "Output should have the NAME column")
		Expect(vmStatus.age).To(Equal("AGE"), "Output should have the AGE column")
		Expect(vmStatus.running).To(Equal("RUNNING"), "Output should have the RUNNING column")
		Expect(vmStatus.volume).To(Equal("VOLUME"), "Output should have the VOLUME column")

		// Read first status of the vm
		vmStatus = readVMStatus(stdout, vmStatus, readTimeout)
		Expect(vmStatus.name).To(Equal(vm.Name))
		By("Expecting vm.running == false")
		Expect(vmStatus.running).To(Equal("false"))

		By("Starting the VM")
		vm = tests.StartVirtualMachine(vm)

		vmStatus = readVMStatus(stdout, vmStatus, readTimeout)
		By("Expecting vm.running == true")
		Expect(vmStatus.running).To(Equal("true"))

		By("Restarting the VM")
		err = virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Restart(vm.ObjectMeta.Name)
		Expect(err).ToNot(HaveOccurred(), "VM should have been restarted")

		vmStatus = readVMStatus(stdout, nil, readTimeout)
		By("Expecting vm.running == true")
		Expect(vmStatus.running).To(Equal("true"))

		By("Stopping the VM")
		vm = tests.StopVirtualMachine(vm)

		vmStatus = readVMStatus(stdout, vmStatus, readTimeout)
		By("Expecting vm.running == false")
		Expect(vmStatus.running).To(Equal("false"))

		By("Deleting the VM")
		err = virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Delete(vm.ObjectMeta.Name, &v1.DeleteOptions{})
		Expect(err).ToNot(HaveOccurred(), "VM should have been deleted from the cluster")
	})

	It("[test_id:3466]Should update vmi status with the proper columns using 'kubectl get vmi -w'", func() {
		By("Creating a random VMI spec")
		vm := tests.NewRandomVMWithEphemeralDisk(tests.ContainerDiskFor(tests.ContainerDiskCirros))

		Expect(vm).ToNot(BeNil())

		By("Setting up the kubectl command")
		cmd, stdout, stderr :=
			createCommandWithNSAndRedirect(vm.ObjectMeta.Namespace, tests.GetK8sCmdClient(), "get", "vmi", "-w")
		Expect(cmd).ToNot(BeNil())

		err = cmd.Start()
		Expect(err).ToNot(HaveOccurred(), "Command should have stared successfully")

		defer cmd.Process.Kill()

		time.Sleep(processWaitTime)

		go failOnError(stderr)

		By("Applying vmi to the cluster")
		vm, err = virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Create(vm)
		Expect(err).ToNot(HaveOccurred(), "VMI should have been added to the cluster")

		// Start a VMI
		vm = tests.StartVirtualMachine(vm)

		// Read the column titles
		vmiStatus := readVMIStatus(stdout, nil, readTimeout)
		Expect(vmiStatus.name).To(Equal("NAME"), "Output should have the NAME column")
		Expect(vmiStatus.age).To(Equal("AGE"), "Output should have the AGE column")
		Expect(vmiStatus.phase).To(Equal("PHASE"), "Output should have the PHASE column")
		Expect(vmiStatus.ip).To(Equal("IP"), "Output should have the IP column")
		Expect(vmiStatus.node).To(Equal("NODENAME"), "Output should have the NODENAME column")

		By("Expecting vmi.phase == ''")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal(""))

		By("Expecting vmi.phase == Pending")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Pending"))

		By("Expecting vmi.phase == Scheduling")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Scheduling"))

		By("Expecting vmi.phase == Scheduled")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Scheduled"))

		By("Expecting vmi.phase == Running")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Running"))

		// Restart the VMI
		err = virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Restart(vm.ObjectMeta.Name)
		Expect(err).ToNot(HaveOccurred(), "VMI should have been restarted")

		By("Expecting vmi.phase == Failed")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Failed"))

		By("Expecting vmi.phase == ''")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal(""))

		By("Expecting vmi.phase == Pending")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Pending"))

		By("Expecting vmi.phase == Scheduling")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Scheduling"))

		By("Expecting vmi.phase == Scheduled")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Scheduled"))

		By("Expecting vmi.phase == Running")
		vmiStatus = readVMIStatus(stdout, vmiStatus, readTimeout)
		Expect(vmiStatus.phase).To(Equal("Running"))

		By("Deleting the VM")
		err = virtCli.VirtualMachine(vm.ObjectMeta.Namespace).Delete(vm.ObjectMeta.Name, &v1.DeleteOptions{})
		Expect(err).ToNot(HaveOccurred(), "VM should have been deleted from the cluster")
	})
})
<|endoftext|>