Procházet zdrojové kódy

test case for removing handler on the fly

Darien Raymond před 7 roky
rodič
revize
deaee9fa65

+ 27 - 1
app/commander/commander.go

@@ -9,6 +9,7 @@ import (
 
 	"google.golang.org/grpc"
 	"v2ray.com/core"
+	"v2ray.com/core/common"
 )
 
 type Commander struct {
@@ -19,6 +20,21 @@ type Commander struct {
 	callbacks []core.ServiceRegistryCallback
 }
 
+func NewCommander(ctx context.Context, config *Config) (*Commander, error) {
+	v := core.FromContext(ctx)
+	if v == nil {
+		return nil, newError("V is not in context.")
+	}
+	c := &Commander{
+		config: *config,
+		ohm:    v.OutboundHandlerManager(),
+	}
+	if err := v.RegisterFeature((*core.Commander)(nil), c); err != nil {
+		return nil, err
+	}
+	return c, nil
+}
+
 func (c *Commander) RegisterService(callback core.ServiceRegistryCallback) {
 	c.Lock()
 	defer c.Unlock()
@@ -42,7 +58,11 @@ func (c *Commander) Start() error {
 		buffer: make(chan net.Conn, 4),
 	}
 
-	c.server.Serve(listener)
+	go func() {
+		if err := c.server.Serve(listener); err != nil {
+			newError("failed to start grpc server").Base(err).AtError().WriteToLog()
+		}
+	}()
 
 	c.ohm.RemoveHandler(context.Background(), c.config.Tag)
 	c.ohm.AddHandler(context.Background(), &CommanderOutbound{
@@ -61,3 +81,9 @@ func (c *Commander) Close() {
 		c.server = nil
 	}
 }
+
+func init() {
+	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, cfg interface{}) (interface{}, error) {
+		return NewCommander(ctx, cfg.(*Config))
+	}))
+}

+ 5 - 2
app/proxyman/outbound/outbound.go

@@ -14,8 +14,9 @@ import (
 // Manager is to manage all outbound handlers.
 type Manager struct {
 	sync.RWMutex
-	defaultHandler core.OutboundHandler
-	taggedHandler  map[string]core.OutboundHandler
+	defaultHandler   core.OutboundHandler
+	taggedHandler    map[string]core.OutboundHandler
+	untaggedHandlers []core.OutboundHandler
 }
 
 // New creates a new Manager.
@@ -68,6 +69,8 @@ func (m *Manager) AddHandler(ctx context.Context, handler core.OutboundHandler)
 	tag := handler.Tag()
 	if len(tag) > 0 {
 		m.taggedHandler[tag] = handler
+	} else {
+		m.untaggedHandlers = append(m.untaggedHandlers, handler)
 	}
 
 	return nil

+ 126 - 0
testing/scenarios/command_test.go

@@ -0,0 +1,126 @@
+package scenarios
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"google.golang.org/grpc"
+	"v2ray.com/core"
+	"v2ray.com/core/app/commander"
+	"v2ray.com/core/app/proxyman"
+	"v2ray.com/core/app/proxyman/command"
+	"v2ray.com/core/app/router"
+	"v2ray.com/core/common/net"
+	"v2ray.com/core/common/serial"
+	"v2ray.com/core/proxy/dokodemo"
+	"v2ray.com/core/proxy/freedom"
+	"v2ray.com/core/testing/servers/tcp"
+	. "v2ray.com/ext/assert"
+)
+
+func TestCommanderRemoveHandler(t *testing.T) {
+	assert := With(t)
+
+	tcpServer := tcp.Server{
+		MsgProcessor: xor,
+	}
+	dest, err := tcpServer.Start()
+	assert(err, IsNil)
+	defer tcpServer.Close()
+
+	clientPort := pickPort()
+	cmdPort := pickPort()
+	clientConfig := &core.Config{
+		App: []*serial.TypedMessage{
+			serial.ToTypedMessage(&commander.Config{Tag: "api"}),
+			serial.ToTypedMessage(&router.Config{
+				Rule: []*router.RoutingRule{
+					{
+						InboundTag: []string{"api"},
+						Tag:        "api",
+					},
+				},
+			}),
+			serial.ToTypedMessage(&command.Config{}),
+		},
+		Inbound: []*core.InboundHandlerConfig{
+			{
+				Tag: "d",
+				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
+					PortRange: net.SinglePortRange(clientPort),
+					Listen:    net.NewIPOrDomain(net.LocalHostIP),
+				}),
+				ProxySettings: serial.ToTypedMessage(&dokodemo.Config{
+					Address: net.NewIPOrDomain(dest.Address),
+					Port:    uint32(dest.Port),
+					NetworkList: &net.NetworkList{
+						Network: []net.Network{net.Network_TCP},
+					},
+				}),
+			},
+			{
+				Tag: "api",
+				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
+					PortRange: net.SinglePortRange(cmdPort),
+					Listen:    net.NewIPOrDomain(net.LocalHostIP),
+				}),
+				ProxySettings: serial.ToTypedMessage(&dokodemo.Config{
+					Address: net.NewIPOrDomain(dest.Address),
+					Port:    uint32(dest.Port),
+					NetworkList: &net.NetworkList{
+						Network: []net.Network{net.Network_TCP},
+					},
+				}),
+			},
+		},
+		Outbound: []*core.OutboundHandlerConfig{
+			{
+				Tag:           "default-outbound",
+				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
+			},
+		},
+	}
+
+	servers, err := InitializeServerConfigs(clientConfig)
+	assert(err, IsNil)
+
+	{
+		conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
+			IP:   []byte{127, 0, 0, 1},
+			Port: int(clientPort),
+		})
+		assert(err, IsNil)
+
+		payload := "commander request."
+		nBytes, err := conn.Write([]byte(payload))
+		assert(err, IsNil)
+		assert(nBytes, Equals, len(payload))
+
+		response := make([]byte, 1024)
+		nBytes, err = conn.Read(response)
+		assert(err, IsNil)
+		assert(response[:nBytes], Equals, xor([]byte(payload)))
+		assert(conn.Close(), IsNil)
+	}
+
+	cmdConn, err := grpc.Dial(fmt.Sprintf("127.0.0.1:%d", cmdPort), grpc.WithInsecure())
+	assert(err, IsNil)
+
+	hsClient := command.NewHandlerServiceClient(cmdConn)
+	resp, err := hsClient.RemoveInbound(context.Background(), &command.RemoveInboundRequest{
+		Tag: "d",
+	})
+	assert(err, IsNil)
+	assert(resp, IsNotNil)
+
+	{
+		_, err := net.DialTCP("tcp", nil, &net.TCPAddr{
+			IP:   []byte{127, 0, 0, 1},
+			Port: int(clientPort),
+		})
+		assert(err, IsNotNil)
+	}
+
+	CloseAllServers(servers)
+}