mirror of
				https://github.com/TeaOSLab/EdgeNode.git
				synced 2025-11-04 07:40:56 +08:00 
			
		
		
		
	优化用户服务整体启用和禁用
This commit is contained in:
		@@ -160,7 +160,7 @@ func (this *BaseListener) findNamedServer(name string) (serverConfig *serverconf
 | 
			
		||||
 | 
			
		||||
// 严格查找域名
 | 
			
		||||
func (this *BaseListener) findNamedServerMatched(name string) (serverConfig *serverconfigs.ServerConfig, serverName string) {
 | 
			
		||||
	group := this.Group
 | 
			
		||||
	var group = this.Group
 | 
			
		||||
	if group == nil {
 | 
			
		||||
		return nil, ""
 | 
			
		||||
	}
 | 
			
		||||
 
 | 
			
		||||
@@ -175,6 +175,15 @@ func (this *HTTPListener) ServeHTTP(rawWriter http.ResponseWriter, rawReq *http.
 | 
			
		||||
		}
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	// 检查用户
 | 
			
		||||
	if server != nil && server.UserId > 0 {
 | 
			
		||||
		if !SharedUserManager.CheckUserServersIsEnabled(server.UserId) {
 | 
			
		||||
			rawWriter.WriteHeader(http.StatusNotFound)
 | 
			
		||||
			_, _ = rawWriter.Write([]byte("The site owner is unavailable."))
 | 
			
		||||
			return
 | 
			
		||||
		}
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	// 包装新请求对象
 | 
			
		||||
	var req = &HTTPRequest{
 | 
			
		||||
		RawReq:     rawReq,
 | 
			
		||||
 
 | 
			
		||||
@@ -92,7 +92,7 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	// 是否已达到流量限制
 | 
			
		||||
	if this.reachedTrafficLimit() {
 | 
			
		||||
	if this.reachedTrafficLimit() || (server.UserId > 0 && !SharedUserManager.CheckUserServersIsEnabled(server.UserId)) {
 | 
			
		||||
		// 关闭连接
 | 
			
		||||
		tcpConn, ok := conn.(LingerConn)
 | 
			
		||||
		if ok {
 | 
			
		||||
 
 | 
			
		||||
@@ -170,6 +170,11 @@ func (this *UDPListener) servePacketListener(listener UDPPacketListener) error {
 | 
			
		||||
			return nil
 | 
			
		||||
		}
 | 
			
		||||
 | 
			
		||||
		// 检查用户状态
 | 
			
		||||
		if firstServer.UserId > 0 && !SharedUserManager.CheckUserServersIsEnabled(firstServer.UserId) {
 | 
			
		||||
			return nil
 | 
			
		||||
		}
 | 
			
		||||
 | 
			
		||||
		n, cm, clientAddr, err := listener.ReadFrom(buffer)
 | 
			
		||||
		if err != nil {
 | 
			
		||||
			if this.isClosed {
 | 
			
		||||
 
 | 
			
		||||
@@ -430,6 +430,22 @@ func (this *Node) execTask(rpcClient *rpc.RPCClient, nodeCtx context.Context, ta
 | 
			
		||||
				}
 | 
			
		||||
			}
 | 
			
		||||
		}
 | 
			
		||||
	case "userServersStateChanged":
 | 
			
		||||
		if task.UserId > 0 {
 | 
			
		||||
			resp, err := rpcClient.UserRPC.CheckUserServersState(nodeCtx, &pb.CheckUserServersStateRequest{UserId: task.UserId})
 | 
			
		||||
			if err != nil {
 | 
			
		||||
				return err
 | 
			
		||||
			}
 | 
			
		||||
 | 
			
		||||
			SharedUserManager.UpdateUserServersIsEnabled(task.UserId, resp.IsEnabled)
 | 
			
		||||
 | 
			
		||||
			if resp.IsEnabled {
 | 
			
		||||
				err = this.syncUserServersConfig(task.UserId)
 | 
			
		||||
				if err != nil {
 | 
			
		||||
					return err
 | 
			
		||||
				}
 | 
			
		||||
			}
 | 
			
		||||
		}
 | 
			
		||||
	default:
 | 
			
		||||
		remotelogs.Error("NODE", "task '"+types.String(task.Id)+"', type '"+task.Type+"' has not been handled")
 | 
			
		||||
	}
 | 
			
		||||
@@ -615,6 +631,36 @@ func (this *Node) syncServerConfig(serverId int64) error {
 | 
			
		||||
	return nil
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// 同步某个用户下的所有服务配置
 | 
			
		||||
func (this *Node) syncUserServersConfig(userId int64) error {
 | 
			
		||||
	rpcClient, err := rpc.SharedRPC()
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return err
 | 
			
		||||
	}
 | 
			
		||||
	serverConfigsResp, err := rpcClient.ServerRPC.ComposeAllUserServersConfig(rpcClient.Context(), &pb.ComposeAllUserServersConfigRequest{
 | 
			
		||||
		UserId: userId,
 | 
			
		||||
	})
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return err
 | 
			
		||||
	}
 | 
			
		||||
	if len(serverConfigsResp.ServersConfigJSON) == 0 {
 | 
			
		||||
		return nil
 | 
			
		||||
	}
 | 
			
		||||
	var serverConfigs = []*serverconfigs.ServerConfig{}
 | 
			
		||||
	err = json.Unmarshal(serverConfigsResp.ServersConfigJSON, &serverConfigs)
 | 
			
		||||
	if err != nil {
 | 
			
		||||
		return err
 | 
			
		||||
	}
 | 
			
		||||
	this.locker.Lock()
 | 
			
		||||
	defer this.locker.Unlock()
 | 
			
		||||
 | 
			
		||||
	for _, config := range serverConfigs {
 | 
			
		||||
		this.updatingServerMap[config.Id] = config
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
	return nil
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// 启动同步计时器
 | 
			
		||||
func (this *Node) startSyncTimer() {
 | 
			
		||||
	// TODO 这个时间间隔可以自行设置
 | 
			
		||||
 
 | 
			
		||||
							
								
								
									
										49
									
								
								internal/nodes/user_manager.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										49
									
								
								internal/nodes/user_manager.go
									
									
									
									
									
										Normal file
									
								
							@@ -0,0 +1,49 @@
 | 
			
		||||
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
 | 
			
		||||
 | 
			
		||||
package nodes
 | 
			
		||||
 | 
			
		||||
import (
 | 
			
		||||
	"sync"
 | 
			
		||||
)
 | 
			
		||||
 | 
			
		||||
var SharedUserManager = NewUserManager()
 | 
			
		||||
 | 
			
		||||
type User struct {
 | 
			
		||||
	ServersEnabled bool
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
type UserManager struct {
 | 
			
		||||
	userMap map[int64]*User // id => *User
 | 
			
		||||
 | 
			
		||||
	locker sync.RWMutex
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
func NewUserManager() *UserManager {
 | 
			
		||||
	return &UserManager{
 | 
			
		||||
		userMap: map[int64]*User{},
 | 
			
		||||
	}
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
func (this *UserManager) UpdateUserServersIsEnabled(userId int64, isEnabled bool) {
 | 
			
		||||
	this.locker.Lock()
 | 
			
		||||
	u, ok := this.userMap[userId]
 | 
			
		||||
	if ok {
 | 
			
		||||
		u.ServersEnabled = isEnabled
 | 
			
		||||
	} else {
 | 
			
		||||
		u = &User{ServersEnabled: isEnabled}
 | 
			
		||||
		this.userMap[userId] = u
 | 
			
		||||
	}
 | 
			
		||||
	this.locker.Unlock()
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
func (this *UserManager) CheckUserServersIsEnabled(userId int64) (isEnabled bool) {
 | 
			
		||||
	this.locker.RLock()
 | 
			
		||||
	u, ok := this.userMap[userId]
 | 
			
		||||
	if ok {
 | 
			
		||||
		isEnabled = u.ServersEnabled
 | 
			
		||||
	} else {
 | 
			
		||||
		isEnabled = true
 | 
			
		||||
	}
 | 
			
		||||
	this.locker.RUnlock()
 | 
			
		||||
	return
 | 
			
		||||
}
 | 
			
		||||
@@ -49,6 +49,7 @@ type RPCClient struct {
 | 
			
		||||
	FirewallRPC            pb.FirewallServiceClient
 | 
			
		||||
	SSLCertRPC             pb.SSLCertServiceClient
 | 
			
		||||
	ScriptRPC              pb.ScriptServiceClient
 | 
			
		||||
	UserRPC                pb.UserServiceClient
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) {
 | 
			
		||||
@@ -81,6 +82,7 @@ func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) {
 | 
			
		||||
	client.FirewallRPC = pb.NewFirewallServiceClient(client)
 | 
			
		||||
	client.SSLCertRPC = pb.NewSSLCertServiceClient(client)
 | 
			
		||||
	client.ScriptRPC = pb.NewScriptServiceClient(client)
 | 
			
		||||
	client.UserRPC = pb.NewUserServiceClient(client)
 | 
			
		||||
 | 
			
		||||
	err := client.init()
 | 
			
		||||
	if err != nil {
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user