Golang metadata.FromContext函数代码示例

func (s *testMetaServer) Status(ctx context.Context, _ *pbtypes.Void) (*ServerStatus, error) {
	md, _ := metadata.FromContext(ctx)
	if want, got := md["want-access-token"], md["authorization"]; got != want {
		return nil, grpc.Errorf(codes.Unknown, "got access-token %q, want %q", got, want)
	}
	return &ServerStatus{}, nil
}

// ListSubscriptions lists matching subscriptions.
func (c *SubscriberClient) ListSubscriptions(ctx context.Context, req *pubsubpb.ListSubscriptionsRequest) *SubscriptionIterator {
	md, _ := metadata.FromContext(ctx)
	ctx = metadata.NewContext(ctx, metadata.Join(md, c.metadata))
	it := &SubscriptionIterator{}
	it.InternalFetch = func(pageSize int, pageToken string) ([]*pubsubpb.Subscription, string, error) {
		var resp *pubsubpb.ListSubscriptionsResponse
		req.PageToken = pageToken
		if pageSize > math.MaxInt32 {
			req.PageSize = math.MaxInt32
		} else {
			req.PageSize = int32(pageSize)
		}
		err := gax.Invoke(ctx, func(ctx context.Context) error {
			var err error
			resp, err = c.subscriberClient.ListSubscriptions(ctx, req)
			return err
		}, c.CallOptions.ListSubscriptions...)
		if err != nil {
			return nil, "", err
		}
		return resp.Subscriptions, resp.NextPageToken, nil
	}
	fetch := func(pageSize int, pageToken string) (string, error) {
		items, nextPageToken, err := it.InternalFetch(pageSize, pageToken)
		if err != nil {
			return "", err
		}
		it.items = append(it.items, items...)
		return nextPageToken, nil
	}
	it.pageInfo, it.nextFunc = iterator.NewPageInfo(fetch, it.bufLen, it.takeBuf)
	return it
}

func NewRaftProxyHealthServer(local HealthServer, connSelector raftpicker.Interface, cluster raftpicker.RaftCluster, ctxMod func(context.Context) (context.Context, error)) HealthServer {
	redirectChecker := func(ctx context.Context) (context.Context, error) {
		s, ok := transport.StreamFromContext(ctx)
		if !ok {
			return ctx, grpc.Errorf(codes.InvalidArgument, "remote addr is not found in context")
		}
		addr := s.ServerTransport().RemoteAddr().String()
		md, ok := metadata.FromContext(ctx)
		if ok && len(md["redirect"]) != 0 {
			return ctx, grpc.Errorf(codes.ResourceExhausted, "more than one redirect to leader from: %s", md["redirect"])
		}
		if !ok {
			md = metadata.New(map[string]string{})
		}
		md["redirect"] = append(md["redirect"], addr)
		return metadata.NewContext(ctx, md), nil
	}
	mods := []func(context.Context) (context.Context, error){redirectChecker}
	mods = append(mods, ctxMod)

	return &raftProxyHealthServer{
		local:        local,
		cluster:      cluster,
		connSelector: connSelector,
		ctxMods:      mods,
	}
}

func newStreamInterceptor(s *etcdserver.EtcdServer) grpc.StreamServerInterceptor {
	smap := monitorLeader(s)

	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		if !api.IsCapabilityEnabled(api.V3rpcCapability) {
			return rpctypes.ErrGRPCNotCapable
		}

		md, ok := metadata.FromContext(ss.Context())
		if ok {
			if ks := md[rpctypes.MetadataRequireLeaderKey]; len(ks) > 0 && ks[0] == rpctypes.MetadataHasLeader {
				if s.Leader() == types.ID(raft.None) {
					return rpctypes.ErrGRPCNoLeader
				}

				cctx, cancel := context.WithCancel(ss.Context())
				ss = serverStreamWithCtx{ctx: cctx, cancel: &cancel, ServerStream: ss}

				smap.mu.Lock()
				smap.streams[ss] = struct{}{}
				smap.mu.Unlock()

				defer func() {
					smap.mu.Lock()
					delete(smap.streams, ss)
					smap.mu.Unlock()
					cancel()
				}()

			}
		}
		return metricsStreamInterceptor(srv, ss, info, handler)
	}
}

func getGRPCTraceID(ctx context.Context) (string, error) {

	md, ok := metadata.FromContext(ctx)
	if !ok {
		id, err := uuid.NewV4()
		if err != nil {
			return "", err
		}
		return id.String(), nil
	}

	tokens := md["trace"]
	if len(tokens) == 0 {
		id, err := uuid.NewV4()
		if err != nil {
			return "", err
		}
		return id.String(), nil
	}

	if tokens[0] != "" {
		return tokens[0], nil
	}

	id, err := uuid.NewV4()
	if err != nil {
		return "", err
	}
	return id.String(), nil
}

func MetadataFromContext(ctx context.Context) (metadata.MD, error) {
	md, ok := metadata.FromContext(ctx)
	if !ok {
		return md, ErrContext
	}
	return md, nil
}

func TestAnnotateContext_ForwardsGrpcMetadata(t *testing.T) {
	ctx := context.Background()
	request, err := http.NewRequest("GET", "http://www.example.com", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(%q, %q, nil) failed with %v; want success", "GET", "http://www.example.com", err)
	}
	request.Header.Add("Some-Irrelevant-Header", "some value")
	request.Header.Add("Grpc-Metadata-FooBar", "Value1")
	request.Header.Add("Grpc-Metadata-Foo-BAZ", "Value2")
	request.Header.Add("Grpc-Metadata-foo-bAz", "Value3")
	request.Header.Add("Authorization", "Token 1234567890")
	annotated, err := runtime.AnnotateContext(ctx, request)
	if err != nil {
		t.Errorf("runtime.AnnotateContext(ctx, %#v) failed with %v; want success", request, err)
		return
	}
	md, ok := metadata.FromContext(annotated)
	if got, want := len(md), emptyForwardMetaCount+3; !ok || got != want {
		t.Errorf("Expected %d metadata items in context; got %d", got, want)
	}
	if got, want := md["foobar"], []string{"Value1"}; !reflect.DeepEqual(got, want) {
		t.Errorf(`md["foobar"] = %q; want %q`, got, want)
	}
	if got, want := md["foo-baz"], []string{"Value2", "Value3"}; !reflect.DeepEqual(got, want) {
		t.Errorf(`md["foo-baz"] = %q want %q`, got, want)
	}
	if got, want := md["authorization"], []string{"Token 1234567890"}; !reflect.DeepEqual(got, want) {
		t.Errorf(`md["authorization"] = %q want %q`, got, want)
	}
}

// LogRPCWithFields will feed any request context into a logrus Entry.
func LogRPCWithFields(log *logrus.Logger, ctx context.Context) *logrus.Entry {
	md, ok := metadata.FromContext(ctx)
	if !ok {
		return logrus.NewEntry(log)
	}
	return log.WithFields(MetadataToFields(md))
}

func NewRaftProxyResourceAllocatorServer(local ResourceAllocatorServer, connSelector raftselector.ConnProvider, localCtxMod, remoteCtxMod func(context.Context) (context.Context, error)) ResourceAllocatorServer {
	redirectChecker := func(ctx context.Context) (context.Context, error) {
		s, ok := transport.StreamFromContext(ctx)
		if !ok {
			return ctx, grpc.Errorf(codes.InvalidArgument, "remote addr is not found in context")
		}
		addr := s.ServerTransport().RemoteAddr().String()
		md, ok := metadata.FromContext(ctx)
		if ok && len(md["redirect"]) != 0 {
			return ctx, grpc.Errorf(codes.ResourceExhausted, "more than one redirect to leader from: %s", md["redirect"])
		}
		if !ok {
			md = metadata.New(map[string]string{})
		}
		md["redirect"] = append(md["redirect"], addr)
		return metadata.NewContext(ctx, md), nil
	}
	remoteMods := []func(context.Context) (context.Context, error){redirectChecker}
	remoteMods = append(remoteMods, remoteCtxMod)

	var localMods []func(context.Context) (context.Context, error)
	if localCtxMod != nil {
		localMods = []func(context.Context) (context.Context, error){localCtxMod}
	}

	return &raftProxyResourceAllocatorServer{
		local:         local,
		connSelector:  connSelector,
		localCtxMods:  localMods,
		remoteCtxMods: remoteMods,
	}
}

func (s *testServer) UnaryCall(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
	md, ok := metadata.FromContext(ctx)
	if ok {
		if err := grpc.SendHeader(ctx, md); err != nil {
			grpclog.Fatalf("grpc.SendHeader(%v, %v) = %v, want %v", ctx, md, err, nil)
		}
		grpc.SetTrailer(ctx, md)
	}
	if s.security != "" {
		// Check Auth info
		authInfo, ok := credentials.FromContext(ctx)
		if !ok {
			grpclog.Fatalf("Failed to get AuthInfo from ctx.")
		}
		var authType string
		switch info := authInfo.(type) {
		case credentials.TLSInfo:
			authType = info.AuthType()
		default:
			grpclog.Fatalf("Unknown AuthInfo type")
		}
		if authType != s.security {
			grpclog.Fatalf("Wrong auth type: got %q, want %q", authType, s.security)
		}
	}

	// Simulate some service delay.
	time.Sleep(time.Second)
	return &testpb.SimpleResponse{
		Payload: newPayload(in.GetResponseType(), in.GetResponseSize()),
	}, nil
}

func (c Client) GetRatePlans(ctx context.Context, hotelIDs []int32, inDate string, outDate string) RatePlanReply {
	md, _ := metadata.FromContext(ctx)

	trace.Req(md["traceID"], md["from"], "service.rate", "GetRatePlans")
	defer trace.Rep(md["traceID"], "service.rate", md["from"], time.Now())

	args := &pb.Args{
		HotelIds: hotelIDs,
		InDate:   inDate,
		OutDate:  outDate,
	}

	reply, err := c.client.GetRates(ctx, args)

	if err != nil {
		return RatePlanReply{
			RatePlans: []*pb.RatePlan{},
			Err:       err,
		}
	}

	return RatePlanReply{
		RatePlans: reply.RatePlans,
		Err:       nil,
	}
}

func (s *networkServerRPC) ValidateContext(ctx context.Context) error {
	md, ok := metadata.FromContext(ctx)
	if !ok {
		return errors.NewErrInternal("Could not get metadata from context")
	}
	var id, token string
	if ids, ok := md["id"]; ok && len(ids) == 1 {
		id = ids[0]
	}
	if id == "" {
		return errors.NewErrInvalidArgument("Metadata", "id missing")
	}
	if tokens, ok := md["token"]; ok && len(tokens) == 1 {
		token = tokens[0]
	}
	if token == "" {
		return errors.NewErrInvalidArgument("Metadata", "token missing")
	}
	var claims *jwt.StandardClaims
	claims, err := security.ValidateJWT(token, []byte(s.networkServer.(*networkServer).Identity.PublicKey))
	if err != nil {
		return err
	}
	if claims.Subject != id {
		return errors.NewErrInvalidArgument("Metadata", "token was issued for a different component id")
	}
	return nil
}

func (s *_ABitOfEverythingServer) List(_ *empty.Empty, stream examples.StreamService_ListServer) error {
	s.m.Lock()
	defer s.m.Unlock()

	err := stream.SendHeader(metadata.New(map[string]string{
		"count": fmt.Sprintf("%d", len(s.v)),
	}))
	if err != nil {
		return nil
	}

	for _, msg := range s.v {
		if err := stream.Send(msg); err != nil {
			return err
		}
	}

	// return error when metadata includes error header
	if header, ok := metadata.FromContext(stream.Context()); ok {
		if v, ok := header["error"]; ok {
			stream.SetTrailer(metadata.New(map[string]string{
				"foo": "foo2",
				"bar": "bar2",
			}))
			return grpc.Errorf(codes.InvalidArgument, "error metadata: %v", v)
		}
	}
	return nil
}

func isForwardedRequest(ctx context.Context) bool {
	md, _ := metadata.FromContext(ctx)
	if len(md[certForwardedKey]) != 1 {
		return false
	}
	return md[certForwardedKey][0] == "true"
}

// Reads the OIDC JWT passed in the context and verifies it using the given OIDC client.
// Returns the verified identity on success, error otherwise.
func VerifiedIdentityFromContext(client *gooidc.Client, ctx context.Context) (*gooidc.Identity, error) {
	md, ok := metadata.FromContext(ctx)
	if !ok {
		return nil, errors.New("missing RPC credentials")
	}
	rawJWT, ok := md["jwt"]
	if !ok {
		return nil, errors.New("missing OIDC credentials")
	}
	if len(rawJWT) != 1 {
		return nil, errors.New("incorrect JWT data sent")
	}
	jwt, err := jose.ParseJWT(rawJWT[0])
	if err != nil {
		return nil, err
	}
	if err := client.VerifyJWT(jwt); err != nil {
		return nil, err
	}
	claims, err := jwt.Claims()
	if err != nil {
		return nil, err
	}
	return gooidc.IdentityFromClaims(claims)
}