Repository: go-kratos/kratos
Pull request #1687 (2021-12-12 13:33:42)

fix:internal/host: fix network interface is enabled


Patch:
diff --git a/contrib/registry/consul/client.go b/contrib/registry/consul/client.go
index 62ae663a..b711a5bf 100644
--- a/contrib/registry/consul/client.go
+++ b/contrib/registry/consul/client.go
@@ -18,38 +18,29 @@ type Client struct {
 	cli    *api.Client
 	ctx    context.Context
 	cancel context.CancelFunc
+
+	// resolve service entry endpoints
+	resolver ServiceResolver
 }
 
 // NewClient creates consul client
 func NewClient(cli *api.Client) *Client {
-	c := &Client{cli: cli}
+	c := &Client{cli: cli, resolver: defaultResolver}
 	c.ctx, c.cancel = context.WithCancel(context.Background())
 	return c
 }
 
-// Service get services from consul
-func (d *Client) Service(ctx context.Context, service string, index uint64, passingOnly bool) ([]*registry.ServiceInstance, uint64, error) {
-	opts := &api.QueryOptions{
-		WaitIndex: index,
-		WaitTime:  time.Second * 55,
-	}
-	opts = opts.WithContext(ctx)
-	entries, meta, err := d.cli.Health().Service(service, "", passingOnly, opts)
-	if err != nil {
-		return nil, 0, err
-	}
-
-	services := make([]*registry.ServiceInstance, 0)
-
+func defaultResolver(_ context.Context, entries []*api.ServiceEntry) []*registry.ServiceInstance {
+	services := make([]*registry.ServiceInstance, 0, len(entries))
 	for _, entry := range entries {
 		var version string
 		for _, tag := range entry.Service.Tags {
-			strs := strings.SplitN(tag, "=", 2)
-			if len(strs) == 2 && strs[0] == "version" {
-				version = strs[1]
+			ss := strings.SplitN(tag, "=", 2)
+			if len(ss) == 2 && ss[0] == "version" {
+				version = ss[1]
 			}
 		}
-		var endpoints []string
+		var endpoints []string //nolint:prealloc
 		for scheme, addr := range entry.Service.TaggedAddresses {
 			if scheme == "lan_ipv4" || scheme == "wan_ipv4" || scheme == "lan_ipv6" || scheme == "wan_ipv6" {
 				continue
@@ -64,11 +55,29 @@ func (d *Client) Service(ctx context.Context, service string, index uint64, pass
 			Endpoints: endpoints,
 		})
 	}
-	return services, meta.LastIndex, nil
+
+	return services
+}
+
+// ServiceResolver is used to resolve service endpoints
+type ServiceResolver func(ctx context.Context, entries []*api.ServiceEntry) []*registry.ServiceInstance
+
+// Service get services from consul
+func (c *Client) Service(ctx context.Context, service string, index uint64, passingOnly bool) ([]*registry.ServiceInstance, uint64, error) {
+	opts := &api.QueryOptions{
+		WaitIndex: index,
+		WaitTime:  time.Second * 55,
+	}
+	opts = opts.WithContext(ctx)
+	entries, meta, err := c.cli.Health().Service(service, "", passingOnly, opts)
+	if err != nil {
+		return nil, 0, err
+	}
+	return c.resolver(ctx, entries), meta.LastIndex, nil
 }
 
-// Register register service instacen to consul
-func (d *Client) Register(ctx context.Context, svc *registry.ServiceInstance, enableHealthCheck bool) error {
+// Register register service instance to consul
+func (c *Client) Register(_ context.Context, svc *registry.ServiceInstance, enableHealthCheck bool) error {
 	addresses := make(map[string]api.ServiceAddress)
 	checkAddresses := make([]string, 0, len(svc.Endpoints))
 	for _, endpoint := range svc.Endpoints {
@@ -103,7 +112,7 @@ func (d *Client) Register(ctx context.Context, svc *registry.ServiceInstance, en
 			})
 		}
 	}
-	err := d.cli.Agent().ServiceRegister(asr)
+	err := c.cli.Agent().ServiceRegister(asr)
 	if err != nil {
 		return err
 	}
@@ -113,8 +122,8 @@ func (d *Client) Register(ctx context.Context, svc *registry.ServiceInstance, en
 		for {
 			select {
 			case <-ticker.C:
-				_ = d.cli.Agent().UpdateTTL("service:"+svc.ID, "pass", "pass")
-			case <-d.ctx.Done():
+				_ = c.cli.Agent().UpdateTTL("service:"+svc.ID, "pass", "pass")
+			case <-c.ctx.Done():
 				return
 			}
 		}
@@ -123,7 +132,7 @@ func (d *Client) Register(ctx context.Context, svc *registry.ServiceInstance, en
 }
 
 // Deregister deregister service by service ID
-func (d *Client) Deregister(ctx context.Context, serviceID string) error {
-	d.cancel()
-	return d.cli.Agent().ServiceDeregister(serviceID)
+func (c *Client) Deregister(_ context.Context, serviceID string) error {
+	c.cancel()
+	return c.cli.Agent().ServiceDeregister(serviceID)
 }
diff --git a/contrib/registry/consul/registry.go b/contrib/registry/consul/registry.go
index d2348b90..41cdd8c7 100644
--- a/contrib/registry/consul/registry.go
+++ b/contrib/registry/consul/registry.go
@@ -26,6 +26,15 @@ func WithHealthCheck(enable bool) Option {
 	}
 }
 
+// WithServiceResolver with endpoint function option.
+func WithServiceResolver(fn ServiceResolver) Option {
+	return func(o *Registry) {
+		if o.cli != nil {
+			o.cli.resolver = fn
+		}
+	}
+}
+
 // Config is consul registry config
 type Config struct {
 	*api.Config
diff --git a/internal/host/host.go b/internal/host/host.go
index ac0cf212..078660be 100644
--- a/internal/host/host.go
+++ b/internal/host/host.go
@@ -53,7 +53,17 @@ func Extract(hostPort string, lis net.Listener) (string, error) {
 	if err != nil {
 		return "", err
 	}
+	lowest := int(^uint(0) >> 1)
+	var result net.IP
 	for _, iface := range ifaces {
+		if (iface.Flags & net.FlagUp) == 0 {
+			continue
+		}
+		if iface.Index < lowest || result == nil {
+			lowest = iface.Index
+		} else if result != nil {
+			continue
+		}
 		addrs, err := iface.Addrs()
 		if err != nil {
 			continue
@@ -69,9 +79,12 @@ func Extract(hostPort string, lis net.Listener) (string, error) {
 				continue
 			}
 			if isValidIP(ip.String()) {
-				return net.JoinHostPort(ip.String(), port), nil
+				result = ip
 			}
 		}
 	}
+	if result != nil {
+		return net.JoinHostPort(result.String(), port), nil
+	}
 	return "", nil
 }
diff --git a/middleware/selector/selector.go b/middleware/selector/selector.go
index 53eaf4b6..06b4d2f9 100644
--- a/middleware/selector/selector.go
+++ b/middleware/selector/selector.go
@@ -11,7 +11,7 @@ import (
 
 type (
 	transporter func(ctx context.Context) (transport.Transporter, bool)
-	MatchFunc   func(operation string) bool
+	MatchFunc   func(ctx context.Context, operation string) bool
 )
 
 var (
@@ -83,7 +83,13 @@ func (b *Builder) Build() middleware.Middleware {
 }
 
 // matchs is match operation compliance Builder
-func (b *Builder) matchs(operation string) bool {
+func (b *Builder) matchs(ctx context.Context, transporter transporter) bool {
+	info, ok := transporter(ctx)
+	if !ok {
+		return false
+	}
+
+	operation := info.Operation()
 	for _, prefix := range b.prefix {
 		if prefixMatch(prefix, operation) {
 			return true
@@ -101,21 +107,19 @@ func (b *Builder) matchs(operation string) bool {
 	}
 
 	if b.match != nil {
-		return b.match(operation)
+		if b.match(ctx, operation) {
+			return true
+		}
 	}
+
 	return false
 }
 
 // selector middleware
-func selector(transporter transporter, match MatchFunc, ms ...middleware.Middleware) middleware.Middleware {
+func selector(transporter transporter, match func(context.Context, transporter) bool, ms ...middleware.Middleware) middleware.Middleware {
 	return func(handler middleware.Handler) middleware.Handler {
 		return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
-			info, ok := transporter(ctx)
-			if !ok {
-				return handler(ctx, req)
-			}
-
-			if !match(info.Operation()) {
+			if !match(ctx, transporter) {
 				return handler(ctx, req)
 			}
 			return middleware.Chain(ms...)(handler)(ctx, req)
diff --git a/transport/http/client.go b/transport/http/client.go
index c40c43bb..17cd34e3 100644
--- a/transport/http/client.go
+++ b/transport/http/client.go
@@ -231,7 +231,7 @@ func (client *Client) Invoke(ctx context.Context, method, path string, args inte
 
 func (client *Client) invoke(ctx context.Context, req *http.Request, args interface{}, reply interface{}, c callInfo, opts ...CallOption) error {
 	h := func(ctx context.Context, in interface{}) (interface{}, error) {
-		res, err := client.do(ctx, req, c)
+		res, err := client.do(req.WithContext(ctx))
 		if res != nil {
 			cs := csAttempt{res: res}
 			for _, o := range opts {
@@ -263,19 +263,18 @@ func (client *Client) Do(req *http.Request, opts ...CallOption) (*http.Response,
 			return nil, err
 		}
 	}
-	ctx := req.Context()
 
-	return client.do(ctx, req, c)
+	return client.do(req)
 }
 
-func (client *Client) do(ctx context.Context, req *http.Request, c callInfo) (*http.Response, error) {
+func (client *Client) do(req *http.Request) (*http.Response, error) {
 	var done func(context.Context, selector.DoneInfo)
 	if client.r != nil {
 		var (
 			err  error
 			node selector.Node
 		)
-		if node, done, err = client.opts.selector.Select(ctx); err != nil {
+		if node, done, err = client.opts.selector.Select(req.Context()); err != nil {
 			return nil, errors.ServiceUnavailable("NODE_NOT_FOUND", err.Error())
 		}
 		if client.insecure {
@@ -288,14 +287,14 @@ func (client *Client) do(ctx context.Context, req *http.Request, c callInfo) (*h
 	}
 	resp, err := client.cc.Do(req)
 	if err == nil {
-		err = client.opts.errorDecoder(ctx, resp)
+		err = client.opts.errorDecoder(req.Context(), resp)
 	}
 
 	if err != nil {
 		return nil, err
 	}
 	if done != nil {
-		done(ctx, selector.DoneInfo{Err: err})
+		done(req.Context(), selector.DoneInfo{Err: err})
 	}
 	return resp, nil
 }


Test patch:
diff --git a/internal/host/host_test.go b/internal/host/host_test.go
index 102be651..e383ec37 100644
--- a/internal/host/host_test.go
+++ b/internal/host/host_test.go
@@ -128,3 +128,13 @@ func TestExtractHostPort(t *testing.T) {
 	}
 	t.Logf("host port: %s,  %d", host, port)
 }
+
+func TestIpIsUp(t *testing.T) {
+	interfaces, err := net.Interfaces()
+	if err != nil {
+		t.Fail()
+	}
+	for i := range interfaces {
+		println(interfaces[i].Name, interfaces[i].Flags&net.FlagUp)
+	}
+}
diff --git a/middleware/selector/selector_test.go b/middleware/selector/selector_test.go
index aec1a45a..b66c3064 100644
--- a/middleware/selector/selector_test.go
+++ b/middleware/selector/selector_test.go
@@ -17,6 +17,7 @@ type Transport struct {
 	kind      transport.Kind
 	endpoint  string
 	operation string
+	headers   *mockHeader
 }
 
 func (tr *Transport) Kind() transport.Kind {
@@ -32,13 +33,33 @@ func (tr *Transport) Operation() string {
 }
 
 func (tr *Transport) RequestHeader() transport.Header {
-	return nil
+	return tr.headers
 }
 
 func (tr *Transport) ReplyHeader() transport.Header {
 	return nil
 }
 
+type mockHeader struct {
+	m map[string]string
+}
+
+func (m *mockHeader) Get(key string) string {
+	return m.m[key]
+}
+
+func (m *mockHeader) Set(key, value string) {
+	m.m[key] = value
+}
+
+func (m *mockHeader) Keys() []string {
+	keys := make([]string, 0, len(m.m))
+	for k := range m.m {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
 func TestMatch(t *testing.T) {
 	tests := []struct {
 		name string
@@ -139,7 +160,7 @@ func TestFunc(t *testing.T) {
 				t.Log(req)
 				return "reply", nil
 			}
-			next = Server(testMiddleware).Match(func(operation string) bool {
+			next = Server(testMiddleware).Match(func(ctx context.Context, operation string) bool {
 				if strings.HasPrefix(operation, "/go-kratos.dev") || strings.HasSuffix(operation, "world") {
 					return true
 				}
@@ -152,6 +173,66 @@ func TestFunc(t *testing.T) {
 	}
 }
 
+func TestHeaderFunc(t *testing.T) {
+	tests := []struct {
+		name string
+		ctx  context.Context
+	}{
+		{
+			name: "/hello.Update/world",
+			ctx: transport.NewServerContext(context.Background(), &Transport{
+				operation: "/hello.Update/world",
+				headers:   &mockHeader{map[string]string{"X-Test": "test"}},
+			}),
+		},
+		{
+			name: "/hi.Create/world",
+			ctx: transport.NewServerContext(context.Background(), &Transport{
+				operation: "/hi.Create/world",
+				headers:   &mockHeader{map[string]string{"X-Test": "test2", "go-kratos": "kratos"}},
+			}),
+		},
+		{
+			name: "/test.Name/1234",
+			ctx: transport.NewServerContext(context.Background(), &Transport{
+				operation: "/test.Name/1234",
+				headers:   &mockHeader{map[string]string{"X-Test": "test3"}},
+			}),
+		},
+		{
+			name: "/go-kratos.dev/kratos",
+			ctx: transport.NewServerContext(context.Background(), &Transport{
+				operation: "/go-kratos.dev/kratos",
+				headers:   &mockHeader{map[string]string{"X-Test": "test"}},
+			}),
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			next := func(ctx context.Context, req interface{}) (interface{}, error) {
+				t.Log(req)
+				return "reply", nil
+			}
+			next = Server(testMiddleware).Match(func(ctx context.Context, operation string) bool {
+				tr, ok := transport.FromServerContext(ctx)
+				if !ok {
+					return false
+				}
+				if tr.RequestHeader().Get("X-Test") == "test" {
+					return true
+				}
+				if tr.RequestHeader().Get("go-kratos") == "kratos" {
+					return true
+				}
+				return false
+			}).Build()(next)
+			reply, err := next(test.ctx, test.name)
+			assert.Equal(t, reply, "reply")
+			assert.Nil(t, err)
+		})
+	}
+}
+
 func testMiddleware(handler middleware.Handler) middleware.Handler {
 	return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
 		fmt.Println("before")