Golang metadata.FromContext函數代碼示例

// getSegment creates a new segment based on whether we're part of an existing flow
func getSegment(name string, ctx context.Context) *awsxray.Segment {
	md, _ := metadata.FromContext(ctx)
	parentId := getParentId(md)
	traceId := getTraceId(md)

	// try get existing segment for parent Id
	if s, ok := awsxray.FromContext(ctx); ok {
		// only set existing segment as parent if its not a subsegment itself
		if len(parentId) == 0 && len(s.Type) == 0 {
			parentId = s.Id
		}
		if len(traceId) == 0 {
			traceId = s.TraceId
		}
	}

	// create segment
	s := &awsxray.Segment{
		Id:        fmt.Sprintf("%x", getRandom(8)),
		Name:      name,
		TraceId:   traceId,
		StartTime: float64(time.Now().Truncate(time.Millisecond).UnixNano()) / 1e9,
	}

	// we have a parent so subsegment
	if len(parentId) > 0 {
		s.ParentId = parentId
		s.Type = "subsegment"
	}

	return s
}

func (g *grpcClient) Publish(ctx context.Context, p client.Publication, opts ...client.PublishOption) error {
	md, ok := metadata.FromContext(ctx)
	if !ok {
		md = make(map[string]string)
	}
	md["Content-Type"] = p.ContentType()

	cf, err := g.newCodec(p.ContentType())
	if err != nil {
		return errors.InternalServerError("go.micro.client", err.Error())
	}

	b := &buffer{bytes.NewBuffer(nil)}
	if err := cf(b).Write(&codec.Message{Type: codec.Publication}, p.Message()); err != nil {
		return errors.InternalServerError("go.micro.client", err.Error())
	}

	g.once.Do(func() {
		g.opts.Broker.Connect()
	})

	return g.opts.Broker.Publish(p.Topic(), &broker.Message{
		Header: md,
		Body:   b.Bytes(),
	})
}

// Nearby returns all hotels within a given distance.
func (s *Geo) Nearby(ctx context.Context, req *geo.Request, rsp *geo.Result) error {
	md, _ := metadata.FromContext(ctx)
	traceID := md["traceID"]

	if tr, ok := trace.FromContext(ctx); ok {
		tr.LazyPrintf("traceID %s", traceID)
	}

	// create center point for query
	center := &geoindex.GeoPoint{
		Pid:  "",
		Plat: float64(req.Lat),
		Plon: float64(req.Lon),
	}

	// find points around center point
	points := s.index.KNearest(center, maxSearchResults, geoindex.Km(maxSearchRadius), func(p geoindex.Point) bool {
		return true
	})

	for _, p := range points {
		rsp.HotelIds = append(rsp.HotelIds, p.Id())
	}
	return nil
}

func TestRequestToContext(t *testing.T) {
	testData := []struct {
		request *http.Request
		expect  metadata.Metadata
	}{
		{
			&http.Request{
				Header: http.Header{
					"foo1": []string{"bar"},
					"foo2": []string{"bar", "baz"},
				},
			},
			metadata.Metadata{
				"foo1": "bar",
				"foo2": "bar,baz",
			},
		},
	}

	for _, d := range testData {
		ctx := RequestToContext(d.request)
		md, ok := metadata.FromContext(ctx)
		if !ok {
			t.Fatalf("Expected metadata for request %+v", d.request)
		}
		for k, v := range d.expect {
			if val := md[k]; val != v {
				t.Fatal("Expected %s for key %s for expected md %+v, got md %+v", v, k, d.expect, md)
			}
		}
	}
}

func (p *platform) Introspect(ctx context.Context) (*Token, error) {
	t, ok := p.FromContext(ctx)
	if !ok {
		md, kk := metadata.FromContext(ctx)
		if !kk {
			return nil, ErrInvalidToken
		}
		t, ok = p.FromHeader(md)
		if !ok {
			return nil, ErrInvalidToken
		}
	}

	rsp, err := p.c.Introspect(context.TODO(), &oauth2.IntrospectRequest{
		AccessToken: t.AccessToken,
	})
	if err != nil {
		return nil, err
	}

	// if its not active just err?
	if !rsp.Active {
		return nil, ErrInvalidToken
	}

	return &Token{
		AccessToken:  rsp.Token.AccessToken,
		RefreshToken: rsp.Token.RefreshToken,
		TokenType:    rsp.Token.TokenType,
		ExpiresAt:    time.Unix(rsp.Token.ExpiresAt, 0),
		Scopes:       rsp.Token.Scopes,
		Metadata:     rsp.Token.Metadata,
	}, nil
}

func handlerWrapper(fn server.HandlerFunc, t Trace, s *registry.Service) server.HandlerFunc {
	return func(ctx context.Context, req server.Request, rsp interface{}) error {
		// embed trace instance
		newCtx := NewContext(ctx, t)

		var span *Span
		var err error

		// Expectation is that we're the initiator of tracing
		// So get trace info from metadata
		md, ok := metadata.FromContext(ctx)
		if !ok {
			// this is a new span
			span = t.NewSpan(nil)
			span.Debug = true
		} else {
			// can we gt the span from the header?
			span, ok = t.FromHeader(md)
			if !ok {
				// no, ok create one
				span = t.NewSpan(nil)
			}
			span.Timestamp = time.Time{}
			span.Debug = true
		}

		// mark client request
		span.Annotations = append(span.Annotations, &Annotation{
			Timestamp: time.Now(),
			Type:      AnnServerRequest,
			Service:   s,
		})

		// embed the span in the context
		newCtx = t.NewContext(newCtx, span)

		// defer the completion of the span
		defer func() {
			var debug map[string]string
			if err != nil {
				debug = map[string]string{"error": err.Error()}
			}
			// mark server response
			span.Annotations = append(span.Annotations, &Annotation{
				Timestamp: time.Now(),
				Type:      AnnServerResponse,
				Service:   s,
				Debug:     debug,
			})
			// flush the span to the collector on return
			t.Collect(span)
		}()
		err = fn(newCtx, req, rsp)
		return err
	}
}

func (r *rpcClient) stream(ctx context.Context, address string, req Request, opts CallOptions) (Streamer, error) {
	msg := &transport.Message{
		Header: make(map[string]string),
	}

	md, ok := metadata.FromContext(ctx)
	if ok {
		for k, v := range md {
			msg.Header[k] = v
		}
	}

	// set timeout in nanoseconds
	msg.Header["Timeout"] = fmt.Sprintf("%d", opts.RequestTimeout)
	// set the content type for the request
	msg.Header["Content-Type"] = req.ContentType()

	cf, err := r.newCodec(req.ContentType())
	if err != nil {
		return nil, errors.InternalServerError("go.micro.client", err.Error())
	}

	c, err := r.opts.Transport.Dial(address, transport.WithStream(), transport.WithTimeout(opts.DialTimeout))
	if err != nil {
		return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error sending request: %v", err))
	}

	stream := &rpcStream{
		context: ctx,
		request: req,
		closed:  make(chan bool),
		codec:   newRpcPlusCodec(msg, c, cf),
	}

	ch := make(chan error, 1)

	go func() {
		ch <- stream.Send(req.Request())
	}()

	var grr error

	select {
	case err := <-ch:
		grr = err
	case <-ctx.Done():
		grr = errors.New("go.micro.client", fmt.Sprintf("%v", ctx.Err()), 408)
	}

	if grr != nil {
		stream.Close()
		return nil, grr
	}

	return stream, nil
}

func (l *labelWrapper) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...client.CallOption) error {
	md, _ := metadata.FromContext(ctx)

	labels := make(map[string]string)

	for k, v := range md {
		if strings.HasPrefix(k, LabelPrefix) {
			key := strings.TrimPrefix(k, LabelPrefix)
			key = strings.ToLower(key)
			labels[key] = strings.ToLower(v)
		}
	}

	filter := func(services []*registry.Service) []*registry.Service {
		var filtered []*registry.Service

	SERVICE:
		for _, service := range services {
			// check labels for the service
			for lk, lv := range labels {
				// if we match then append the service and return
				if v := service.Metadata[lk]; v == lv {
					filtered = append(filtered, service)
					continue SERVICE
				}
			}

			var nodes []*registry.Node

		NODE:
			for _, node := range service.Nodes {
				// check labels against the node
				for lk, lv := range labels {
					// if it matches then append node
					if v := node.Metadata[lk]; v == lv {
						nodes = append(nodes, node)
						continue NODE
					}
				}
			}

			if len(nodes) > 0 {
				service.Nodes = nodes
				filtered = append(filtered, service)
			}
		}

		return filtered
	}

	callOptions := append(opts, client.WithSelectOption(
		selector.Filter(filter),
	))

	return l.Client.Call(ctx, req, rsp, callOptions...)
}

// GetProfiles returns hotel profiles for requested IDs
func (s *Profile) GetProfiles(ctx context.Context, req *profile.Request, rsp *profile.Result) error {
	md, _ := metadata.FromContext(ctx)
	traceID := md["traceID"]
	if tr, ok := trace.FromContext(ctx); ok {
		tr.LazyPrintf("traceID %s", traceID)
	}

	for _, i := range req.HotelIds {
		rsp.Hotels = append(rsp.Hotels, s.hotels[i])
	}
	return nil
}

func (c *clientWrapper) setHeaders(ctx context.Context) context.Context {
	md, ok := metadata.FromContext(ctx)
	if !ok {
		md = metadata.Metadata{}
	}
	for k, v := range c.headers {
		if _, ok := md[k]; !ok {
			md[k] = v
		}
	}
	return metadata.NewContext(ctx, md)
}

// VerifyToken returns a customer from authentication token.
func (s *Auth) VerifyToken(ctx context.Context, req *auth.Request, rsp *auth.Result) error {
	md, _ := metadata.FromContext(ctx)
	traceID := md["traceID"]

	if tr, ok := trace.FromContext(ctx); ok {
		tr.LazyPrintf("traceID %s", traceID)
	}

	customer := s.customers[req.AuthToken]
	if customer == nil {
		return errors.New("Invalid Token")
	}

	rsp.Customer = customer
	return nil
}

func (g *grpcClient) stream(ctx context.Context, address string, req client.Request, opts client.CallOptions) (client.Streamer, error) {
	header := make(map[string]string)
	if md, ok := metadata.FromContext(ctx); ok {
		for k, v := range md {
			header[k] = v
		}
	}

	// set timeout in nanoseconds
	header["timeout"] = fmt.Sprintf("%d", opts.RequestTimeout)
	// set the content type for the request
	header["x-content-type"] = req.ContentType()

	md := gmetadata.New(header)
	ctx = gmetadata.NewContext(ctx, md)

	cf, err := g.newGRPCCodec(req.ContentType())
	if err != nil {
		return nil, errors.InternalServerError("go.micro.client", err.Error())
	}

	// TODO: do not use insecure
	cc, err := grpc.Dial(address, grpc.WithCodec(cf), grpc.WithTimeout(opts.DialTimeout), grpc.WithInsecure())
	if err != nil {
		return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error sending request: %v", err))
	}

	desc := &grpc.StreamDesc{
		StreamName:    req.Service() + req.Method(),
		ClientStreams: true,
		ServerStreams: true,
	}

	st, err := grpc.NewClientStream(ctx, desc, cc, req.Method())
	if err != nil {
		return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error creating stream: %v", err))
	}

	return &grpcStream{
		context: ctx,
		request: req,
		closed:  make(chan bool),
		stream:  st,
		conn:    cc,
	}, nil
}

func (t *TestHandler) Exec(ctx context.Context, req *TestRequest, rsp *TestResponse) error {
	md, ok := metadata.FromContext(ctx)
	if !ok {
		return fmt.Errorf("Expected metadata got %t", ok)
	}

	for k, v := range t.expect {
		if val := md[k]; val != v {
			return fmt.Errorf("Expected %s for key %s got %s", v, k, val)
		}
	}

	t.t.Logf("Received request %+v", req)
	t.t.Logf("Received metadata %+v", md)

	return nil
}

func (o *otWrapper) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...client.CallOption) error {
	md, _ := metadata.FromContext(ctx)
	name := fmt.Sprintf("%s.%s", req.Service(), req.Method())
	var sp opentracing.Span
	wireContext, err := o.ot.Extract(opentracing.TextMap, opentracing.TextMapCarrier(md))
	if err != nil {
		sp = o.ot.StartSpan(name)
	} else {
		sp = o.ot.StartSpan(name, opentracing.ChildOf(wireContext))
	}
	defer sp.Finish()
	if err := sp.Tracer().Inject(sp.Context(), opentracing.TextMap, opentracing.TextMapCarrier(md)); err != nil {
		return err
	}
	ctx = metadata.NewContext(ctx, md)
	return o.Client.Call(ctx, req, rsp, opts...)
}

func TestWrapper(t *testing.T) {
	testData := []struct {
		existing  metadata.Metadata
		headers   metadata.Metadata
		overwrite bool
	}{
		{
			existing: metadata.Metadata{},
			headers: metadata.Metadata{
				"foo": "bar",
			},
			overwrite: true,
		},
		{
			existing: metadata.Metadata{
				"foo": "bar",
			},
			headers: metadata.Metadata{
				"foo": "baz",
			},
			overwrite: false,
		},
	}

	for _, d := range testData {
		c := &clientWrapper{
			headers: d.headers,
		}

		ctx := metadata.NewContext(context.Background(), d.existing)
		c.setHeaders(ctx)

		md, _ := metadata.FromContext(ctx)

		for k, v := range d.headers {
			if d.overwrite && md[k] != v {
				t.Fatalf("Expected %s=%s got %s=%s", k, v, k, md[k])
			}
			if !d.overwrite && md[k] != d.existing[k] {
				t.Fatalf("Expected %s=%s got %s=%s", k, d.existing[k], k, md[k])
			}
		}
	}

}