From 8a537d1cd3a7b825bfab2273c2a188cab57ebe84 Mon Sep 17 00:00:00 2001
From: Kong Jun <kongjun18@outlook.com>
Date: Sun, 6 Oct 2024 15:54:51 +0800
Subject: [PATCH] feat: update vmess inbound users dynamically

---
 box.go           | 11 ++++++++---
 inbound/vmess.go | 43 ++++++++++++++++++++++++++++++++++---------
 2 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/box.go b/box.go
index 9e334a97..2c34fa37 100644
--- a/box.go
+++ b/box.go
@@ -437,15 +437,20 @@ func (s *Box) Router() adapter.Router {
 func (s *Box) UpdateUser(ctx context.Context, req *pb.UpdateUserRequest) (*pb.UpdateUserReply, error) {
 	for _, inbound := range s.inbounds {
 		if inbound.Tag() == req.Tag {
-			switch inbound.(type) {
+			switch inbound := inbound.(type) {
 			case *pkginbound.ShadowsocksMulti:
-				in := inbound.(*pkginbound.ShadowsocksMulti)
-				err := in.UpdateUserPassword(req.Users, req.Passwords)
+				err := inbound.UpdateUserPassword(req.Users, req.Passwords)
 				if err != nil {
 					return nil, grpc.Errorf(codes.Internal, err.Error())
 				}
 				return &pb.UpdateUserReply{}, nil
 			}
+			case *pkginbound.VMess:
+				err := inbound.Updateusers(req.Users, req.Passwords)
+				if err != nil {
+					return nil, grpc.Errorf(codes.Internal, err.Error())
+				}
+				return &pb.UpdateUserReply{}, nil
 			break
 		}
 	}
diff --git a/inbound/vmess.go b/inbound/vmess.go
index 15451275..87d2ae34 100644
--- a/inbound/vmess.go
+++ b/inbound/vmess.go
@@ -4,6 +4,17 @@ import (
 	"context"
 	"net"
 	"os"
+	"sync"
+
+	vmess "github.com/sagernet/sing-vmess"
+	"github.com/sagernet/sing-vmess/packetaddr"
+	"github.com/sagernet/sing/common"
+	"github.com/sagernet/sing/common/auth"
+	E "github.com/sagernet/sing/common/exceptions"
+	F "github.com/sagernet/sing/common/format"
+	M "github.com/sagernet/sing/common/metadata"
+	N "github.com/sagernet/sing/common/network"
+	"github.com/sagernet/sing/common/ntp"
 
 	"github.com/sagernet/sing-box/adapter"
 	"github.com/sagernet/sing-box/common/mux"
@@ -13,15 +24,6 @@ import (
 	"github.com/sagernet/sing-box/log"
 	"github.com/sagernet/sing-box/option"
 	"github.com/sagernet/sing-box/transport/v2ray"
-	"github.com/sagernet/sing-vmess"
-	"github.com/sagernet/sing-vmess/packetaddr"
-	"github.com/sagernet/sing/common"
-	"github.com/sagernet/sing/common/auth"
-	E "github.com/sagernet/sing/common/exceptions"
-	F "github.com/sagernet/sing/common/format"
-	M "github.com/sagernet/sing/common/metadata"
-	N "github.com/sagernet/sing/common/network"
-	"github.com/sagernet/sing/common/ntp"
 )
 
 var (
@@ -36,6 +38,7 @@ type VMess struct {
 	users     []option.VMessUser
 	tlsConfig tls.ServerConfig
 	transport adapter.V2RayServerTransport
+	mu        sync.RWMutex
 }
 
 func NewVMess(ctx context.Context, router adapter.Router, logger log.ContextLogger, tag string, options option.VMessInboundOptions) (*VMess, error) {
@@ -155,6 +158,8 @@ func (h *VMess) NewConnection(ctx context.Context, conn net.Conn, metadata adapt
 			return err
 		}
 	}
+	h.mu.RLock()
+	defer h.mu.RUnlock()
 	return h.service.NewConnection(adapter.WithContext(log.ContextWithNewID(ctx), &metadata), conn, adapter.UpstreamMetadata(metadata))
 }
 
@@ -198,6 +203,26 @@ func (h *VMess) newPacketConnection(ctx context.Context, conn N.PacketConn, meta
 	return h.router.RoutePacketConnection(ctx, conn, metadata)
 }
 
+func (v *VMess) UpdateUsers(nameList []string, userIdList []string, alterIdList []int) error {
+	users := make([]option.VMessUser, 0, len(nameList))
+	for i, name := range nameList {
+		users = append(users, option.VMessUser{Name: name, UUID: userIdList[i], AlterId: alterIdList[i]})
+	}
+
+	v.mu.Lock()
+	defer v.mu.Unlock()
+
+	err := v.service.UpdateUsers(common.MapIndexed(users, func(index int, it option.VMessUser) int {
+		return index
+	}), userIdList, alterIdList)
+	if err != nil {
+		return err
+	}
+	v.users = users
+
+	return nil
+}
+
 var _ adapter.V2RayServerTransportHandler = (*vmessTransportHandler)(nil)
 
 type vmessTransportHandler VMess