TRKJ/EdgeNode
2
0
Fork 0
mirror of https://github.com/TeaOSLab/EdgeNode.git synced 2025-11-04 07:40:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

重新实现套餐

Browse Source
This commit is contained in:
刘祥超
2023-09-06 16:34:11 +08:00
parent a6059ab070
commit e93275ac5c
10 changed files with 71 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
internal/nodes/client_conn.go
View File

@@ -198,9 +198,9 @@ func (this *ClientConn) Write(b []byte) (n int, err error) {
var cost = time.Since(before).Seconds()
if cost > 1 {
stats.SharedBandwidthStatManager.AddBandwidth(this.userId, this.serverId, int64(float64(n)/cost), int64(n))
stats.SharedBandwidthStatManager.AddBandwidth(this.userId, this.userPlanId, this.serverId, int64(float64(n)/cost), int64(n))
} else {
stats.SharedBandwidthStatManager.AddBandwidth(this.userId, this.serverId, int64(n), int64(n))
stats.SharedBandwidthStatManager.AddBandwidth(this.userId, this.userPlanId, this.serverId, int64(n), int64(n))
}
}
}

21
internal/nodes/client_conn_base.go
View File

@@ -16,6 +16,7 @@ type BaseClientConn struct {
isBound bool
userId int64
userPlanId int64
serverId int64
remoteAddr string
hasLimit bool
@@ -106,11 +107,31 @@ func (this *BaseClientConn) SetUserId(userId int64) {
}
}
func (this *BaseClientConn) SetUserPlanId(userPlanId int64) {
this.userPlanId = userPlanId
// 设置包装前连接
switch conn := this.rawConn.(type) {
case *tls.Conn:
nativeConn, ok := conn.NetConn().(ClientConnInterface)
if ok {
nativeConn.SetUserPlanId(userPlanId)
}
case *ClientConn:
conn.SetUserPlanId(userPlanId)
}
}
// UserId 获取当前连接所属服务的用户ID
func (this *BaseClientConn) UserId() int64 {
return this.userId
}
// UserPlanId 用户套餐ID
func (this *BaseClientConn) UserPlanId() int64 {
return this.userPlanId
}
// RawIP 原本IP
func (this *BaseClientConn) RawIP() string {
if len(this.rawIP) > 0 {

5
internal/nodes/client_conn_interface.go
View File

@@ -18,9 +18,12 @@ type ClientConnInterface interface {
// SetServerId 设置服务ID
SetServerId(serverId int64) (goNext bool)
// SetUserId 设置所属服务的用户ID
// SetUserId 设置所属网站的用户ID
SetUserId(userId int64)
// SetUserPlanId 设置
SetUserPlanId(userPlanId int64)
// UserId 获取当前连接所属服务的用户ID
UserId() int64

2
internal/nodes/http_request.go
View File

@@ -198,7 +198,7 @@ func (this *HTTPRequest) Do() {
}
// 流量限制
if this.ReqServer.TrafficLimit != nil && this.ReqServer.TrafficLimit.IsOn && !this.ReqServer.TrafficLimit.IsEmpty() && this.ReqServer.TrafficLimitStatus != nil && this.ReqServer.TrafficLimitStatus.IsValid() {
if this.ReqServer.TrafficLimitStatus != nil && this.ReqServer.TrafficLimitStatus.IsValid() {
this.doTrafficLimit()
this.doEnd()
return

10
internal/nodes/http_request_traffic_limit.go
View File

@@ -8,15 +8,17 @@ import (
// 流量限制
func (this *HTTPRequest) doTrafficLimit() {
var config = this.ReqServer.TrafficLimit
this.tags = append(this.tags, "bandwidth")
this.tags = append(this.tags, "trafficLimit")
var statusCode = 509
this.writer.statusCode = statusCode
this.ProcessResponseHeaders(this.writer.Header(), statusCode)
this.writer.Header().Set("Content-Type", "text/html; charset=utf-8")
this.writer.WriteHeader(statusCode)
if len(config.NoticePageBody) != 0 {
var config = this.ReqServer.TrafficLimit
if config != nil && len(config.NoticePageBody) != 0 {
_, _ = this.writer.WriteString(this.Format(config.NoticePageBody))
} else {
_, _ = this.writer.WriteString(this.Format(serverconfigs.DefaultTrafficLimitNoticePageBody))

6
internal/nodes/listener_http.go
View File

@@ -177,6 +177,12 @@ func (this *HTTPListener) ServeHTTP(rawWriter http.ResponseWriter, rawReq *http.
return
}
clientConn.SetUserId(server.UserId)
var userPlanId int64
if server.UserPlan != nil && server.UserPlan.Id > 0 {
userPlanId = server.UserPlan.Id
}
clientConn.SetUserPlanId(userPlanId)
}
}
}

12
internal/nodes/listener_tcp.go
View File

@@ -80,6 +80,12 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
return nil
}
clientConn.SetUserId(server.UserId)
var userPlanId int64
if server.UserPlan != nil && server.UserPlan.Id > 0 {
userPlanId = server.UserPlan.Id
}
clientConn.SetUserPlanId(userPlanId)
} else {
tlsConn, ok := conn.(*tls.Conn)
if ok {
@@ -92,6 +98,12 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
return nil
}
clientConn.SetUserId(server.UserId)
var userPlanId int64
if server.UserPlan != nil && server.UserPlan.Id > 0 {
userPlanId = server.UserPlan.Id
}
clientConn.SetUserPlanId(userPlanId)
}
}
}

6
internal/nodes/listener_udp.go
View File

@@ -404,7 +404,11 @@ func NewUDPConn(server *serverconfigs.ServerConfig, addr net.Addr, proxyListener
stats.SharedTrafficStatManager.Add(server.UserId, server.Id, "", int64(n), 0, 0, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId())
// 带宽
stats.SharedBandwidthStatManager.AddBandwidth(server.UserId, server.Id, int64(n), int64(n))
var userPlanId int64
if server.UserPlan != nil && server.UserPlan.Id > 0 {
userPlanId = server.UserPlan.Id
}
stats.SharedBandwidthStatManager.AddBandwidth(server.UserId, userPlanId, server.Id, int64(n), int64(n))
}
}
if err != nil {

5
internal/stats/bandwidth_stat_manager.go
View File

@@ -62,6 +62,7 @@ type BandwidthStat struct {
CountRequests int64 `json:"countRequests"`
CountCachedRequests int64 `json:"countCachedRequests"`
CountAttackRequests int64 `json:"countAttackRequests"`
UserPlanId int64 `json:"userPlanId"`
}
// BandwidthStatManager 服务带宽统计
@@ -153,6 +154,7 @@ func (this *BandwidthStatManager) Loop() error {
CountRequests: stat.CountRequests,
CountCachedRequests: stat.CountCachedRequests,
CountAttackRequests: stat.CountAttackRequests,
UserPlanId: stat.UserPlanId,
NodeRegionId: regionId,
})
delete(this.m, key)
@@ -178,7 +180,7 @@ func (this *BandwidthStatManager) Loop() error {
}
// AddBandwidth 添加带宽数据
func (this *BandwidthStatManager) AddBandwidth(userId int64, serverId int64, peekBytes int64, totalBytes int64) {
func (this *BandwidthStatManager) AddBandwidth(userId int64, userPlanId int64, serverId int64, peekBytes int64, totalBytes int64) {
if serverId <= 0 || (peekBytes == 0 && totalBytes == 0) {
return
}
@@ -217,6 +219,7 @@ func (this *BandwidthStatManager) AddBandwidth(userId int64, serverId int64, pee
Day: day,
TimeAt: timeAt,
UserId: userId,
UserPlanId: userPlanId,
ServerId: serverId,
CurrentBytes: peekBytes,
MaxBytes: peekBytes,

20
internal/stats/bandwidth_stat_manager_test.go
View File

@@ -12,22 +12,22 @@ import (
func TestBandwidthStatManager_Add(t *testing.T) {
var manager = stats.NewBandwidthStatManager()
manager.AddBandwidth(1, 1, 10, 10)
manager.AddBandwidth(1, 1, 10, 10)
manager.AddBandwidth(1, 1, 10, 10)
manager.AddBandwidth(1, 0, 1, 10, 10)
manager.AddBandwidth(1, 0, 1, 10, 10)
manager.AddBandwidth(1, 0, 1, 10, 10)
time.Sleep(1 * time.Second)
manager.AddBandwidth(1, 1, 85, 85)
manager.AddBandwidth(1, 0, 1, 85, 85)
time.Sleep(1 * time.Second)
manager.AddBandwidth(1, 1, 25, 25)
manager.AddBandwidth(1, 1, 75, 75)
manager.AddBandwidth(1, 0, 1, 25, 25)
manager.AddBandwidth(1, 0, 1, 75, 75)
manager.Inspect()
}
func TestBandwidthStatManager_Loop(t *testing.T) {
var manager = stats.NewBandwidthStatManager()
manager.AddBandwidth(1, 1, 10, 10)
manager.AddBandwidth(1, 1, 10, 10)
manager.AddBandwidth(1, 1, 10, 10)
manager.AddBandwidth(1, 0, 1, 10, 10)
manager.AddBandwidth(1, 0, 1, 10, 10)
manager.AddBandwidth(1, 0, 1, 10, 10)
err := manager.Loop()
if err != nil {
t.Fatal(err)
@@ -40,7 +40,7 @@ func BenchmarkBandwidthStatManager_Add(b *testing.B) {
var i int
for pb.Next() {
i++
manager.AddBandwidth(1, int64(i%100), 10, 10)
manager.AddBandwidth(1, 0, int64(i%100), 10, 10)
}
})
}