TRKJ/EdgeNode
2
0
Fork 0
mirror of https://github.com/TeaOSLab/EdgeNode.git synced 2025-11-06 01:50:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

改进流量限制

Browse Source
This commit is contained in:
GoEdgeLab
2021-11-10 14:39:02 +08:00
parent d69e4e8fc4
commit 940517e660
4 changed files with 35 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
internal/nodes/http_request.go
View File

@@ -271,12 +271,12 @@ func (this *HTTPRequest) doEnd() {
// TODO 增加是否开启开关 // TODO 增加是否开启开关
if this.Server != nil { if this.Server != nil {
if this.isCached { if this.isCached {
stats.SharedTrafficStatManager.Add(this.Server.Id, this.Host, this.writer.sentBodyBytes, this.writer.sentBodyBytes, 1, 1, 0, 0) stats.SharedTrafficStatManager.Add(this.Server.Id, this.Host, this.writer.sentBodyBytes, this.writer.sentBodyBytes, 1, 1, 0, 0, this.Server.ShouldCheckTrafficLimit())
} else { } else {
if this.isAttack { if this.isAttack {
stats.SharedTrafficStatManager.Add(this.Server.Id, this.Host, this.writer.sentBodyBytes, 0, 1, 0, 1, this.writer.sentBodyBytes) stats.SharedTrafficStatManager.Add(this.Server.Id, this.Host, this.writer.sentBodyBytes, 0, 1, 0, 1, this.writer.sentBodyBytes, this.Server.ShouldCheckTrafficLimit())
} else { } else {
stats.SharedTrafficStatManager.Add(this.Server.Id, this.Host, this.writer.sentBodyBytes, 0, 1, 0, 0, 0) stats.SharedTrafficStatManager.Add(this.Server.Id, this.Host, this.writer.sentBodyBytes, 0, 1, 0, 0, 0, this.Server.ShouldCheckTrafficLimit())
} }
} }
} }

8
internal/nodes/listener_tcp.go
View File

@@ -65,14 +65,14 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
var serverName = tlsConn.ConnectionState().ServerName var serverName = tlsConn.ConnectionState().ServerName
if len(serverName) > 0 { if len(serverName) > 0 {
// 统计 // 统计
stats.SharedTrafficStatManager.Add(firstServer.Id, serverName, 0, 0, 1, 0, 0, 0) stats.SharedTrafficStatManager.Add(firstServer.Id, serverName, 0, 0, 1, 0, 0, 0, firstServer.ShouldCheckTrafficLimit())
recordStat = true recordStat = true
} }
} }
// 统计 // 统计
if !recordStat { if !recordStat {
stats.SharedTrafficStatManager.Add(firstServer.Id, "", 0, 0, 1, 0, 0, 0) stats.SharedTrafficStatManager.Add(firstServer.Id, "", 0, 0, 1, 0, 0, 0, firstServer.ShouldCheckTrafficLimit())
} }
originConn, err := this.connectOrigin(firstServer.ReverseProxy, conn.RemoteAddr().String()) originConn, err := this.connectOrigin(firstServer.ReverseProxy, conn.RemoteAddr().String())
@@ -125,7 +125,9 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
} }
// 记录流量 // 记录流量
stats.SharedTrafficStatManager.Add(firstServer.Id, "", int64(n), 0, 0, 0, 0, 0) if firstServer != nil {
stats.SharedTrafficStatManager.Add(firstServer.Id, "", int64(n), 0, 0, 0, 0, 0, firstServer.ShouldCheckTrafficLimit())
}
} }
if err != nil { if err != nil {
closer() closer()

12
internal/nodes/listener_udp.go
View File

@@ -64,7 +64,7 @@ func (this *UDPListener) Serve() error {
remotelogs.Error("UDP_LISTENER", "unable to find a origin server") remotelogs.Error("UDP_LISTENER", "unable to find a origin server")
continue continue
} }
conn = NewUDPConn(firstServer.Id, addr, this.Listener, originConn.(*net.UDPConn)) conn = NewUDPConn(firstServer, addr, this.Listener, originConn.(*net.UDPConn))
this.connLocker.Lock() this.connLocker.Lock()
this.connMap[addr.String()] = conn this.connMap[addr.String()] = conn
this.connLocker.Unlock() this.connLocker.Unlock()
@@ -174,7 +174,7 @@ type UDPConn struct {
isClosed bool isClosed bool
} }
func NewUDPConn(serverId int64, addr net.Addr, proxyConn *net.UDPConn, serverConn *net.UDPConn) *UDPConn { func NewUDPConn(server *serverconfigs.ServerConfig, addr net.Addr, proxyConn *net.UDPConn, serverConn *net.UDPConn) *UDPConn {
conn := &UDPConn{ conn := &UDPConn{
addr: addr, addr: addr,
proxyConn: proxyConn, proxyConn: proxyConn,
@@ -184,7 +184,9 @@ func NewUDPConn(serverId int64, addr net.Addr, proxyConn *net.UDPConn, serverCon
} }
// 统计 // 统计
stats.SharedTrafficStatManager.Add(serverId, "", 0, 0, 1, 0, 0, 0) if server != nil {
stats.SharedTrafficStatManager.Add(server.Id, "", 0, 0, 1, 0, 0, 0, server.ShouldCheckTrafficLimit())
}
go func() { go func() {
buffer := bytePool32k.Get() buffer := bytePool32k.Get()
@@ -203,7 +205,9 @@ func NewUDPConn(serverId int64, addr net.Addr, proxyConn *net.UDPConn, serverCon
} }
// 记录流量 // 记录流量
stats.SharedTrafficStatManager.Add(serverId, "", int64(n), 0, 0, 0, 0, 0) if server != nil {
stats.SharedTrafficStatManager.Add(server.Id, "", int64(n), 0, 0, 0, 0, 0, server.ShouldCheckTrafficLimit())
}
} }
if err != nil { if err != nil {
conn.isOk = false conn.isOk = false

5
internal/stats/traffic_stat_manager.go
View File

@@ -26,6 +26,7 @@ type TrafficItem struct {
CountCachedRequests int64 CountCachedRequests int64
CountAttackRequests int64 CountAttackRequests int64
AttackBytes int64 AttackBytes int64
CheckingTrafficLimit bool
} }
// TrafficStatManager 区域流量统计 // TrafficStatManager 区域流量统计
@@ -86,7 +87,7 @@ func (this *TrafficStatManager) Start(configFunc func() *nodeconfigs.NodeConfig)
} }
// Add 添加流量 // Add 添加流量
func (this *TrafficStatManager) Add(serverId int64, domain string, bytes int64, cachedBytes int64, countRequests int64, countCachedRequests int64, countAttacks int64, attackBytes int64) { func (this *TrafficStatManager) Add(serverId int64, domain string, bytes int64, cachedBytes int64, countRequests int64, countCachedRequests int64, countAttacks int64, attackBytes int64, checkingTrafficLimit bool) {
if bytes == 0 && countRequests == 0 { if bytes == 0 && countRequests == 0 {
return return
} }
@@ -110,6 +111,7 @@ func (this *TrafficStatManager) Add(serverId int64, domain string, bytes int64,
item.CountCachedRequests += countCachedRequests item.CountCachedRequests += countCachedRequests
item.CountAttackRequests += countAttacks item.CountAttackRequests += countAttacks
item.AttackBytes += attackBytes item.AttackBytes += attackBytes
item.CheckingTrafficLimit = checkingTrafficLimit
// 单个域名流量 // 单个域名流量
var domainKey = strconv.FormatInt(timestamp, 10) + "@" + strconv.FormatInt(serverId, 10) + "@" + domain var domainKey = strconv.FormatInt(timestamp, 10) + "@" + strconv.FormatInt(serverId, 10) + "@" + domain
@@ -168,6 +170,7 @@ func (this *TrafficStatManager) Upload() error {
CountCachedRequests: item.CountCachedRequests, CountCachedRequests: item.CountCachedRequests,
CountAttackRequests: item.CountAttackRequests, CountAttackRequests: item.CountAttackRequests,
AttackBytes: item.AttackBytes, AttackBytes: item.AttackBytes,
CheckTrafficLimiting: item.CheckingTrafficLimit,
CreatedAt: timestamp, CreatedAt: timestamp,
}) })
} }