Authenticate on service start

registry/service/service.go

@@ -22,8 +22,8 @@ type serviceRegistry struct {
 	name string
 	// address
 	address []string
-	// client to call registry
-	client pb.RegistryService
+	// registry is the proto client
+	registry pb.RegistryService
 }
 
 func (s *serviceRegistry) callOpts() []client.CallOption {
@@ -46,6 +46,17 @@ func (s *serviceRegistry) Init(opts ...registry.Option) error {
 	for _, o := range opts {
 		o(&s.opts)
 	}
+
+	// extract the client from the context, fallback to grpc
+	var cli client.Client
+	if c, ok := s.opts.Context.Value(clientKey{}).(client.Client); ok {
+		cli = c
+	} else {
+		cli = grpc.NewClient()
+	}
+
+	s.registry = pb.NewRegistryService(DefaultService, cli)
+
 	return nil
 }
 
@@ -67,7 +78,7 @@ func (s *serviceRegistry) Register(srv *registry.Service, opts ...registry.Regis
 	pbSrv.Options.Ttl = int64(options.TTL.Seconds())
 
 	// register the service
-	_, err := s.client.Register(options.Context, pbSrv, s.callOpts()...)
+	_, err := s.registry.Register(options.Context, pbSrv, s.callOpts()...)
 	if err != nil {
 		return err
 	}
@@ -85,7 +96,7 @@ func (s *serviceRegistry) Deregister(srv *registry.Service, opts ...registry.Der
 	}
 
 	// deregister the service
-	_, err := s.client.Deregister(options.Context, ToProto(srv), s.callOpts()...)
+	_, err := s.registry.Deregister(options.Context, ToProto(srv), s.callOpts()...)
 	if err != nil {
 		return err
 	}
@@ -101,7 +112,7 @@ func (s *serviceRegistry) GetService(name string, opts ...registry.GetOption) ([
 		options.Context = context.TODO()
 	}
 
-	rsp, err := s.client.GetService(options.Context, &pb.GetRequest{
+	rsp, err := s.registry.GetService(options.Context, &pb.GetRequest{
 		Service: name,
 	}, s.callOpts()...)
 
@@ -125,7 +136,7 @@ func (s *serviceRegistry) ListServices(opts ...registry.ListOption) ([]*registry
 		options.Context = context.TODO()
 	}
 
-	rsp, err := s.client.ListServices(options.Context, &pb.ListRequest{}, s.callOpts()...)
+	rsp, err := s.registry.ListServices(options.Context, &pb.ListRequest{}, s.callOpts()...)
 	if err != nil {
 		return nil, err
 	}
@@ -147,7 +158,7 @@ func (s *serviceRegistry) Watch(opts ...registry.WatchOption) (registry.Watcher,
 		options.Context = context.TODO()
 	}
 
-	stream, err := s.client.Watch(options.Context, &pb.WatchRequest{
+	stream, err := s.registry.Watch(options.Context, &pb.WatchRequest{
 		Service: options.Service,
 	}, s.callOpts()...)
 
@@ -171,27 +182,29 @@ func NewRegistry(opts ...registry.Option) registry.Registry {
 
 	// the registry address
 	addrs := options.Addrs
-
 	if len(addrs) == 0 {
 		addrs = []string{"127.0.0.1:8000"}
 	}
 
-	// use mdns as a fall back in case its used
-	mReg := registry.NewRegistry()
+	if options.Context == nil {
+		options.Context = context.TODO()
+	}
 
-	// create new client with mdns
-	cli := grpc.NewClient(
-		client.Registry(mReg),
-	)
+	// extract the client from the context, fallback to grpc
+	var cli client.Client
+	if c, ok := options.Context.Value(clientKey{}).(client.Client); ok {
+		cli = c
+	} else {
+		cli = grpc.NewClient()
+	}
 
-	// service name
-	// TODO: accept option
+	// service name. TODO: accept option
 	name := DefaultService
 
 	return &serviceRegistry{
-		opts:    options,
-		name:    name,
-		address: addrs,
-		client:  pb.NewRegistryService(name, cli),
+		opts:     options,
+		name:     name,
+		address:  addrs,
+		registry: pb.NewRegistryService(name, cli),
 	}
 }