platform: Refactor CommandClient & Connections

experimental/libbox/command_client.go

@@ -27,6 +27,7 @@ type CommandClient struct {
 	ctx         context.Context
 	cancel      context.CancelFunc
 	clientMutex sync.RWMutex
+	standalone  bool
 }
 
 type CommandClientOptions struct {
@@ -48,7 +49,7 @@ type CommandClientHandler interface {
 	WriteGroups(message OutboundGroupIterator)
 	InitializeClashMode(modeList StringIterator, currentMode string)
 	UpdateClashMode(newMode string)
-	WriteConnections(message *Connections)
+	WriteConnectionEvents(events *ConnectionEvents)
 }
 
 type LogEntry struct {
@@ -73,7 +74,7 @@ func SetXPCDialer(dialer XPCDialer) {
 }
 
 func NewStandaloneCommandClient() *CommandClient {
-	return new(CommandClient)
+	return &CommandClient{standalone: true}
 }
 
 func NewCommandClient(handler CommandClientHandler, options *CommandClientOptions) *CommandClient {
@@ -97,147 +98,135 @@ func streamClientAuthInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc
 	return streamer(ctx, desc, cc, method, opts...)
 }
 
-func (c *CommandClient) grpcDial() (*grpc.ClientConn, error) {
-	var target string
-	if sCommandServerListenPort == 0 {
-		target = "unix://" + filepath.Join(sBasePath, "command.sock")
-	} else {
-		target = net.JoinHostPort("127.0.0.1", strconv.Itoa(int(sCommandServerListenPort)))
-	}
-	var (
-		conn *grpc.ClientConn
-		err  error
-	)
-	clientOptions := []grpc.DialOption{
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
-		grpc.WithUnaryInterceptor(unaryClientAuthInterceptor),
-		grpc.WithStreamInterceptor(streamClientAuthInterceptor),
-	}
-	for i := 0; i < 10; i++ {
-		conn, err = grpc.NewClient(target, clientOptions...)
-		if err == nil {
-			return conn, nil
+const (
+	commandClientDialAttempts  = 10
+	commandClientDialBaseDelay = 100 * time.Millisecond
+	commandClientDialStepDelay = 50 * time.Millisecond
+)
+
+func commandClientDialDelay(attempt int) time.Duration {
+	return commandClientDialBaseDelay + time.Duration(attempt)*commandClientDialStepDelay
+}
+
+func dialTarget() (string, func(context.Context, string) (net.Conn, error)) {
+	if sXPCDialer != nil {
+		return "passthrough:///xpc", func(ctx context.Context, _ string) (net.Conn, error) {
+			fileDescriptor, err := sXPCDialer.DialXPC()
+			if err != nil {
+				return nil, err
+			}
+			return networkConnectionFromFileDescriptor(fileDescriptor)
 		}
-		time.Sleep(time.Duration(100+i*50) * time.Millisecond)
 	}
-	return nil, err
+	if sCommandServerListenPort == 0 {
+		return "unix://" + filepath.Join(sBasePath, "command.sock"), nil
+	}
+	return net.JoinHostPort("127.0.0.1", strconv.Itoa(int(sCommandServerListenPort))), nil
+}
+
+func networkConnectionFromFileDescriptor(fileDescriptor int32) (net.Conn, error) {
+	file := os.NewFile(uintptr(fileDescriptor), "xpc-command-socket")
+	if file == nil {
+		return nil, E.New("invalid file descriptor")
+	}
+	networkConnection, err := net.FileConn(file)
+	if err != nil {
+		file.Close()
+		return nil, E.Cause(err, "create connection from fd")
+	}
+	file.Close()
+	return networkConnection, nil
+}
+
+func (c *CommandClient) dialWithRetry(target string, contextDialer func(context.Context, string) (net.Conn, error), retryDial bool) (*grpc.ClientConn, daemon.StartedServiceClient, error) {
+	var connection *grpc.ClientConn
+	var client daemon.StartedServiceClient
+	var lastError error
+
+	for attempt := 0; attempt < commandClientDialAttempts; attempt++ {
+		if connection == nil {
+			options := []grpc.DialOption{
+				grpc.WithTransportCredentials(insecure.NewCredentials()),
+				grpc.WithUnaryInterceptor(unaryClientAuthInterceptor),
+				grpc.WithStreamInterceptor(streamClientAuthInterceptor),
+			}
+			if contextDialer != nil {
+				options = append(options, grpc.WithContextDialer(contextDialer))
+			}
+			var err error
+			connection, err = grpc.NewClient(target, options...)
+			if err != nil {
+				lastError = err
+				if !retryDial {
+					return nil, nil, err
+				}
+				time.Sleep(commandClientDialDelay(attempt))
+				continue
+			}
+			client = daemon.NewStartedServiceClient(connection)
+		}
+		waitDuration := commandClientDialDelay(attempt)
+		ctx, cancel := context.WithTimeout(context.Background(), waitDuration)
+		_, err := client.GetStartedAt(ctx, &emptypb.Empty{}, grpc.WaitForReady(true))
+		cancel()
+		if err == nil {
+			return connection, client, nil
+		}
+		lastError = err
+	}
+
+	if connection != nil {
+		connection.Close()
+	}
+	return nil, nil, lastError
 }
 
 func (c *CommandClient) Connect() error {
 	c.clientMutex.Lock()
 	common.Close(common.PtrOrNil(c.grpcConn))
 
-	if sXPCDialer != nil {
-		fd, err := sXPCDialer.DialXPC()
-		if err != nil {
-			c.clientMutex.Unlock()
-			return err
-		}
-		file := os.NewFile(uintptr(fd), "xpc-command-socket")
-		if file == nil {
-			c.clientMutex.Unlock()
-			return E.New("invalid file descriptor")
-		}
-		netConn, err := net.FileConn(file)
-		if err != nil {
-			file.Close()
-			c.clientMutex.Unlock()
-			return E.Cause(err, "create connection from fd")
-		}
-		file.Close()
-
-		clientOptions := []grpc.DialOption{
-			grpc.WithTransportCredentials(insecure.NewCredentials()),
-			grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
-				return netConn, nil
-			}),
-			grpc.WithUnaryInterceptor(unaryClientAuthInterceptor),
-			grpc.WithStreamInterceptor(streamClientAuthInterceptor),
-		}
-
-		grpcConn, err := grpc.NewClient("passthrough:///xpc", clientOptions...)
-		if err != nil {
-			netConn.Close()
-			c.clientMutex.Unlock()
-			return err
-		}
-
-		c.grpcConn = grpcConn
-		c.grpcClient = daemon.NewStartedServiceClient(grpcConn)
-		c.ctx, c.cancel = context.WithCancel(context.Background())
-		c.clientMutex.Unlock()
-	} else {
-		conn, err := c.grpcDial()
-		if err != nil {
-			c.clientMutex.Unlock()
-			return err
-		}
-		c.grpcConn = conn
-		c.grpcClient = daemon.NewStartedServiceClient(conn)
-		c.ctx, c.cancel = context.WithCancel(context.Background())
+	target, contextDialer := dialTarget()
+	connection, client, err := c.dialWithRetry(target, contextDialer, true)
+	if err != nil {
 		c.clientMutex.Unlock()
+		return err
 	}
+	c.grpcConn = connection
+	c.grpcClient = client
+	c.ctx, c.cancel = context.WithCancel(context.Background())
+	c.clientMutex.Unlock()
 
 	c.handler.Connected()
-	for _, command := range c.options.commands {
-		switch command {
-		case CommandLog:
-			go c.handleLogStream()
-		case CommandStatus:
-			go c.handleStatusStream()
-		case CommandGroup:
-			go c.handleGroupStream()
-		case CommandClashMode:
-			go c.handleClashModeStream()
-		case CommandConnections:
-			go c.handleConnectionsStream()
-		default:
-			return E.New("unknown command: ", command)
-		}
-	}
-	return nil
+	return c.dispatchCommands()
 }
 
 func (c *CommandClient) ConnectWithFD(fd int32) error {
 	c.clientMutex.Lock()
 	common.Close(common.PtrOrNil(c.grpcConn))
 
-	file := os.NewFile(uintptr(fd), "xpc-command-socket")
-	if file == nil {
-		c.clientMutex.Unlock()
-		return E.New("invalid file descriptor")
-	}
-
-	netConn, err := net.FileConn(file)
+	networkConnection, err := networkConnectionFromFileDescriptor(fd)
 	if err != nil {
-		file.Close()
-		c.clientMutex.Unlock()
-		return E.Cause(err, "create connection from fd")
-	}
-	file.Close()
-
-	clientOptions := []grpc.DialOption{
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
-		grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
-			return netConn, nil
-		}),
-		grpc.WithUnaryInterceptor(unaryClientAuthInterceptor),
-		grpc.WithStreamInterceptor(streamClientAuthInterceptor),
-	}
-
-	grpcConn, err := grpc.NewClient("passthrough:///xpc", clientOptions...)
-	if err != nil {
-		netConn.Close()
 		c.clientMutex.Unlock()
 		return err
 	}
-
-	c.grpcConn = grpcConn
-	c.grpcClient = daemon.NewStartedServiceClient(grpcConn)
+	connection, client, err := c.dialWithRetry("passthrough:///xpc", func(ctx context.Context, _ string) (net.Conn, error) {
+		return networkConnection, nil
+	}, false)
+	if err != nil {
+		networkConnection.Close()
+		c.clientMutex.Unlock()
+		return err
+	}
+	c.grpcConn = connection
+	c.grpcClient = client
 	c.ctx, c.cancel = context.WithCancel(context.Background())
 	c.clientMutex.Unlock()
 
 	c.handler.Connected()
+	return c.dispatchCommands()
+}
+
+func (c *CommandClient) dispatchCommands() error {
 	for _, command := range c.options.commands {
 		switch command {
 		case CommandLog:
@@ -281,57 +270,41 @@ func (c *CommandClient) getClientForCall() (daemon.StartedServiceClient, error)
 		return c.grpcClient, nil
 	}
 
-	if sXPCDialer != nil {
-		fd, err := sXPCDialer.DialXPC()
-		if err != nil {
-			return nil, err
-		}
-		file := os.NewFile(uintptr(fd), "xpc-command-socket")
-		if file == nil {
-			return nil, E.New("invalid file descriptor")
-		}
-		netConn, err := net.FileConn(file)
-		if err != nil {
-			file.Close()
-			return nil, E.Cause(err, "create connection from fd")
-		}
-		file.Close()
-
-		clientOptions := []grpc.DialOption{
-			grpc.WithTransportCredentials(insecure.NewCredentials()),
-			grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
-				return netConn, nil
-			}),
-			grpc.WithUnaryInterceptor(unaryClientAuthInterceptor),
-			grpc.WithStreamInterceptor(streamClientAuthInterceptor),
-		}
-
-		grpcConn, err := grpc.NewClient("passthrough:///xpc", clientOptions...)
-		if err != nil {
-			netConn.Close()
-			return nil, err
-		}
-
-		c.grpcConn = grpcConn
-		c.grpcClient = daemon.NewStartedServiceClient(grpcConn)
-		if c.ctx == nil {
-			c.ctx, c.cancel = context.WithCancel(context.Background())
-		}
-		return c.grpcClient, nil
-	}
-
-	conn, err := c.grpcDial()
+	target, contextDialer := dialTarget()
+	connection, client, err := c.dialWithRetry(target, contextDialer, true)
 	if err != nil {
 		return nil, err
 	}
-	c.grpcConn = conn
-	c.grpcClient = daemon.NewStartedServiceClient(conn)
+	c.grpcConn = connection
+	c.grpcClient = client
 	if c.ctx == nil {
 		c.ctx, c.cancel = context.WithCancel(context.Background())
 	}
 	return c.grpcClient, nil
 }
 
+func (c *CommandClient) closeConnection() {
+	c.clientMutex.Lock()
+	defer c.clientMutex.Unlock()
+	if c.grpcConn != nil {
+		c.grpcConn.Close()
+		c.grpcConn = nil
+		c.grpcClient = nil
+	}
+}
+
+func callWithResult[T any](c *CommandClient, call func(client daemon.StartedServiceClient) (T, error)) (T, error) {
+	client, err := c.getClientForCall()
+	if err != nil {
+		var zero T
+		return zero, err
+	}
+	if c.standalone {
+		defer c.closeConnection()
+	}
+	return call(client)
+}
+
 func (c *CommandClient) getStreamContext() (daemon.StartedServiceClient, context.Context) {
 	c.clientMutex.RLock()
 	defer c.clientMutex.RUnlock()
@@ -468,175 +441,134 @@ func (c *CommandClient) handleConnectionsStream() {
 		return
 	}
 
-	var connections Connections
 	for {
-		conns, err := stream.Recv()
+		events, err := stream.Recv()
 		if err != nil {
 			c.handler.Disconnected(err.Error())
 			return
 		}
-		connections.input = ConnectionsFromGRPC(conns)
-		c.handler.WriteConnections(&connections)
+		libboxEvents := ConnectionEventsFromGRPC(events)
+		c.handler.WriteConnectionEvents(libboxEvents)
 	}
 }
 
 func (c *CommandClient) SelectOutbound(groupTag string, outboundTag string) error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.SelectOutbound(context.Background(), &daemon.SelectOutboundRequest{
-		GroupTag:    groupTag,
-		OutboundTag: outboundTag,
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.SelectOutbound(context.Background(), &daemon.SelectOutboundRequest{
+			GroupTag:    groupTag,
+			OutboundTag: outboundTag,
+		})
 	})
 	return err
 }
 
 func (c *CommandClient) URLTest(groupTag string) error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.URLTest(context.Background(), &daemon.URLTestRequest{
-		OutboundTag: groupTag,
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.URLTest(context.Background(), &daemon.URLTestRequest{
+			OutboundTag: groupTag,
+		})
 	})
 	return err
 }
 
 func (c *CommandClient) SetClashMode(newMode string) error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.SetClashMode(context.Background(), &daemon.ClashMode{
-		Mode: newMode,
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.SetClashMode(context.Background(), &daemon.ClashMode{
+			Mode: newMode,
+		})
 	})
 	return err
 }
 
 func (c *CommandClient) CloseConnection(connId string) error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.CloseConnection(context.Background(), &daemon.CloseConnectionRequest{
-		Id: connId,
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.CloseConnection(context.Background(), &daemon.CloseConnectionRequest{
+			Id: connId,
+		})
 	})
 	return err
 }
 
 func (c *CommandClient) CloseConnections() error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.CloseAllConnections(context.Background(), &emptypb.Empty{})
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.CloseAllConnections(context.Background(), &emptypb.Empty{})
+	})
 	return err
 }
 
 func (c *CommandClient) ServiceReload() error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.ReloadService(context.Background(), &emptypb.Empty{})
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.ReloadService(context.Background(), &emptypb.Empty{})
+	})
 	return err
 }
 
 func (c *CommandClient) ServiceClose() error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.StopService(context.Background(), &emptypb.Empty{})
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.StopService(context.Background(), &emptypb.Empty{})
+	})
 	return err
 }
 
 func (c *CommandClient) ClearLogs() error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.ClearLogs(context.Background(), &emptypb.Empty{})
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.ClearLogs(context.Background(), &emptypb.Empty{})
+	})
 	return err
 }
 
 func (c *CommandClient) GetSystemProxyStatus() (*SystemProxyStatus, error) {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return nil, err
-	}
-
-	status, err := client.GetSystemProxyStatus(context.Background(), &emptypb.Empty{})
-	if err != nil {
-		return nil, err
-	}
-	return SystemProxyStatusFromGRPC(status), nil
+	return callWithResult(c, func(client daemon.StartedServiceClient) (*SystemProxyStatus, error) {
+		status, err := client.GetSystemProxyStatus(context.Background(), &emptypb.Empty{})
+		if err != nil {
+			return nil, err
+		}
+		return SystemProxyStatusFromGRPC(status), nil
+	})
 }
 
 func (c *CommandClient) SetSystemProxyEnabled(isEnabled bool) error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.SetSystemProxyEnabled(context.Background(), &daemon.SetSystemProxyEnabledRequest{
-		Enabled: isEnabled,
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.SetSystemProxyEnabled(context.Background(), &daemon.SetSystemProxyEnabledRequest{
+			Enabled: isEnabled,
+		})
 	})
 	return err
 }
 
 func (c *CommandClient) GetDeprecatedNotes() (DeprecatedNoteIterator, error) {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return nil, err
-	}
-
-	warnings, err := client.GetDeprecatedWarnings(context.Background(), &emptypb.Empty{})
-	if err != nil {
-		return nil, err
-	}
-
-	var notes []*DeprecatedNote
-	for _, warning := range warnings.Warnings {
-		notes = append(notes, &DeprecatedNote{
-			Description:   warning.Message,
-			MigrationLink: warning.MigrationLink,
-		})
-	}
-	return newIterator(notes), nil
+	return callWithResult(c, func(client daemon.StartedServiceClient) (DeprecatedNoteIterator, error) {
+		warnings, err := client.GetDeprecatedWarnings(context.Background(), &emptypb.Empty{})
+		if err != nil {
+			return nil, err
+		}
+		var notes []*DeprecatedNote
+		for _, warning := range warnings.Warnings {
+			notes = append(notes, &DeprecatedNote{
+				Description:   warning.Message,
+				MigrationLink: warning.MigrationLink,
+			})
+		}
+		return newIterator(notes), nil
+	})
 }
 
 func (c *CommandClient) GetStartedAt() (int64, error) {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return 0, err
-	}
-
-	startedAt, err := client.GetStartedAt(context.Background(), &emptypb.Empty{})
-	if err != nil {
-		return 0, err
-	}
-	return startedAt.StartedAt, nil
+	return callWithResult(c, func(client daemon.StartedServiceClient) (int64, error) {
+		startedAt, err := client.GetStartedAt(context.Background(), &emptypb.Empty{})
+		if err != nil {
+			return 0, err
+		}
+		return startedAt.StartedAt, nil
+	})
 }
 
 func (c *CommandClient) SetGroupExpand(groupTag string, isExpand bool) error {
-	client, err := c.getClientForCall()
-	if err != nil {
-		return err
-	}
-
-	_, err = client.SetGroupExpand(context.Background(), &daemon.SetGroupExpandRequest{
-		GroupTag: groupTag,
-		IsExpand: isExpand,
+	_, err := callWithResult(c, func(client daemon.StartedServiceClient) (*emptypb.Empty, error) {
+		return client.SetGroupExpand(context.Background(), &daemon.SetGroupExpandRequest{
+			GroupTag: groupTag,
+			IsExpand: isExpand,
+		})
 	})
 	return err
 }