Compare commits

...

70 Commits

Author SHA1 Message Date
刘祥超
6ff030dbd8 编译时加入configs/cluster.template.yaml文件 2022-11-27 14:52:48 +08:00
刘祥超
0ddeef6986 支持使用域名中含有通配符清除缓存数据 2022-11-26 11:05:46 +08:00
刘祥超
976bd3600b 优化OpenFileCache功能 2022-11-25 14:52:04 +08:00
刘祥超
a64047a934 优化配置重载程序 2022-11-25 10:50:57 +08:00
刘祥超
e82f207935 统计API调用时低于一半的采样率返回总统计 2022-11-23 20:23:46 +08:00
刘祥超
61b5316a1f 优化代码 2022-11-23 20:13:34 +08:00
刘祥超
82329aa8b0 修复一处编译错误 2022-11-22 18:40:03 +08:00
刘祥超
7dabd9c19c 在监控系统运行时上报API连接状况 2022-11-22 11:23:39 +08:00
刘祥超
9437acd18c 优化代码 2022-11-21 21:08:47 +08:00
刘祥超
9da7a34edf 节点可以单独设置所使用的API节点地址 2022-11-21 19:55:28 +08:00
刘祥超
b6a5491dcc 优化Partial Content兼容性 2022-11-20 18:07:46 +08:00
刘祥超
bcee658567 优化Partial Content配置编码速度 2022-11-19 23:11:05 +08:00
刘祥超
afc8f7b703 优化Partial Content缓存 2022-11-19 21:20:53 +08:00
刘祥超
7a4b89d2fb 缓存Header中忽略Set-Cookie 2022-11-19 17:35:23 +08:00
刘祥超
c6299a2fb0 减少文件缓存写入次数 2022-11-19 17:23:45 +08:00
刘祥超
8b5d74af9b 进一步提升文件缓存写入速度 2022-11-19 15:55:05 +08:00
刘祥超
a194360a56 Update go.mod 2022-11-18 17:39:49 +08:00
刘祥超
b12f7f69ba 优化代码 2022-11-17 20:28:55 +08:00
刘祥超
06ec4d3fba 优化代码 2022-11-17 10:38:20 +08:00
刘祥超
c209ab912f 优化代码 2022-11-17 10:35:43 +08:00
刘祥超
32720d772d 优化代码 2022-11-17 10:32:26 +08:00
刘祥超
a89c02fd10 请求变量增加${cname},WAF checkpoint增加cname和isCNAME 2022-11-16 15:01:10 +08:00
刘祥超
37ef86b92f 接收HTTP请求时去除域名后面的点符号 2022-11-16 11:25:11 +08:00
刘祥超
4c19c37f49 写入缓存时减少对缓存目录的检查频率 2022-11-15 22:25:49 +08:00
刘祥超
1bb818b5b0 边缘节点支持设置多个缓存目录 2022-11-15 20:42:25 +08:00
刘祥超
825e46458f 优化代码 2022-11-15 10:06:57 +08:00
刘祥超
a42737bd28 缩短节点运行日志队列长度 2022-11-14 16:42:50 +08:00
刘祥超
5f76be2cfd 优化代码 2022-11-13 10:32:12 +08:00
刘祥超
dbddf8a91a 优化代码 2022-11-08 21:37:20 +08:00
刘祥超
6c457f41f6 优化代码 2022-11-08 20:58:17 +08:00
刘祥超
e4b2a650f0 优化代码 2022-11-08 20:19:51 +08:00
刘祥超
913ba95801 优化缓存相关代码 2022-11-08 11:03:37 +08:00
刘祥超
a9f8e39703 修复节点设置的“缓存磁盘容量”不起作用的问题 2022-11-07 21:32:20 +08:00
刘祥超
534f013f59 使用版本号来读取节点任务,提升任务同步稳定性 2022-11-06 12:07:26 +08:00
刘祥超
258380f75c 修复无法回报任务执行失败的问题 2022-11-05 14:56:57 +08:00
刘祥超
8c0e51ec46 域名跳转增加是否忽略端口选项 2022-11-05 14:30:29 +08:00
刘祥超
4c37c7ab84 时钟同步增加是否检查chrony选项 2022-11-03 14:59:26 +08:00
刘祥超
f005da1d5f 防盗链提示增加缓存时间,以提升性能 2022-11-02 15:24:30 +08:00
刘祥超
e99acc4694 版本修改为0.5.8 2022-11-02 15:11:55 +08:00
刘祥超
408357dfcf 优化DDoS防护相关错误提示信息 2022-11-01 17:37:40 +08:00
刘祥超
0109a27c06 上传访问日志发生网络错误时不提交 2022-11-01 14:55:06 +08:00
刘祥超
e6e2dccc42 版本号修改为0.5.7 2022-10-31 19:14:03 +08:00
刘祥超
09dcf0d712 集群全局服务配置中增加多个访问日志相关选项 2022-10-26 17:51:16 +08:00
刘祥超
60aebd9306 URL跳转中增加域名跳转、端口跳转 2022-10-26 16:14:37 +08:00
刘祥超
04191d04d3 节点设置中增加“通过IP名单”选项 2022-10-26 10:42:16 +08:00
刘祥超
b80a5c525f 节点缓存目录所在磁盘空间不足时(<5G),暂停缓存写入,同时启动LFU清理 2022-10-25 15:14:28 +08:00
刘祥超
265c1e5312 WAF参数定义增加优先级,可以让“轻”任务优先执行 2022-10-24 17:57:07 +08:00
刘祥超
2723f705b6 修复在iptables中加入ipv6的错误 2022-10-24 16:37:54 +08:00
刘祥超
b4cddd6341 集群服务设置--访问日志中可以设置是否只记录通用Header 2022-10-24 14:39:18 +08:00
刘祥超
5636a81d48 防盗链功能增加禁止的来源域名 2022-10-24 10:21:23 +08:00
刘祥超
d8059960de 文件缓存索引表取消UNIQUE索引,尽可能避免 sqlite malformed 错误 2022-10-23 20:45:41 +08:00
刘祥超
17af4064af 带宽和流量提交失败时,将在一定时间内重试 2022-10-23 19:41:21 +08:00
刘祥超
15f37d2c93 优化用户服务整体启用和禁用 2022-10-23 16:21:11 +08:00
刘祥超
6dc3aa8cb7 单请求写入时间从1个小时增加到2个小时 2022-10-23 09:52:50 +08:00
刘祥超
900cccf2f1 修复源站Websocket源站读取失败导致的异常错误 2022-10-18 19:43:53 +08:00
刘祥超
1fec88dfc6 优化代码 2022-10-14 15:00:05 +08:00
刘祥超
7da9363336 上传带宽信息时附带区域ID信息 2022-10-11 18:57:35 +08:00
刘祥超
d82e633bba 时钟同步程序每天只提示一次警告信息 2022-10-11 11:31:00 +08:00
刘祥超
b363bbaafd 版本修改为0.5.6 2022-10-01 08:50:12 +08:00
刘祥超
92a20e3c9a 修复Websocket无法正常交互的问题 2022-09-30 16:34:21 +08:00
刘祥超
5742dfb263 修复Websocket响应可能被缓存的问题 2022-09-30 14:55:42 +08:00
刘祥超
0ae63511d5 版本调整为v0.5.5 2022-09-28 18:57:27 +08:00
刘祥超
aa60092c20 修复开启WAF后,自动记录请求Body的Bug 2022-09-28 16:46:05 +08:00
刘祥超
54fc265d24 systemd服务增加BEGIN INIT INFO 2022-09-28 08:17:25 +08:00
刘祥超
a5ac900784 版本修改为0.5.4 2022-09-27 08:05:20 +08:00
刘祥超
4053f1da32 程序延时100ms退出 2022-09-26 16:27:51 +08:00
刘祥超
0374ccd8a8 版本号改为0.5.3.2 2022-09-26 16:27:28 +08:00
刘祥超
1d46c446cf 程序退出时关闭sqlite数据库 2022-09-26 16:14:24 +08:00
刘祥超
54b66805f9 将版本修改为0.5.4 2022-09-26 15:17:06 +08:00
刘祥超
f7afcbde92 版本修改为0.5.3.1 2022-09-26 13:02:46 +08:00
87 changed files with 2761 additions and 631 deletions

View File

@@ -50,6 +50,7 @@ function build() {
fi fi
cp "$ROOT"/configs/api.template.yaml "$DIST"/configs cp "$ROOT"/configs/api.template.yaml "$DIST"/configs
cp "$ROOT"/configs/cluster.template.yaml "$DIST"/configs
cp -R "$ROOT"/www "$DIST"/ cp -R "$ROOT"/www "$DIST"/
cp -R "$ROOT"/pages "$DIST"/ cp -R "$ROOT"/pages "$DIST"/

View File

@@ -3,6 +3,7 @@
package caches package caches
const ( const (
SuffixAll = "@GOEDGE_" // 通用后缀
SuffixWebP = "@GOEDGE_WEBP" // WebP后缀 SuffixWebP = "@GOEDGE_WEBP" // WebP后缀
SuffixCompression = "@GOEDGE_" // 压缩后缀 SuffixCompression + Encoding SuffixCompression = "@GOEDGE_" // 压缩后缀 SuffixCompression + Encoding
SuffixMethod = "@GOEDGE_" // 请求方法后缀 SuffixMethod + RequestMethod SuffixMethod = "@GOEDGE_" // 请求方法后缀 SuffixMethod + RequestMethod

View File

@@ -0,0 +1,11 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package caches
import "github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/shared"
type FileDir struct {
Path string
Capacity *shared.SizeCapacity
IsFull bool
}

View File

@@ -2,6 +2,7 @@ package caches
import ( import (
"github.com/TeaOSLab/EdgeNode/internal/utils" "github.com/TeaOSLab/EdgeNode/internal/utils"
"strings"
"time" "time"
) )
@@ -59,3 +60,17 @@ func (this *Item) IncreaseHit(week int32) {
this.Week = week this.Week = week
} }
} }
func (this *Item) RequestURI() string {
var schemeIndex = strings.Index(this.Key, "://")
if schemeIndex <= 0 {
return ""
}
var firstSlashIndex = strings.Index(this.Key[schemeIndex+3:], "/")
if firstSlashIndex <= 0 {
return ""
}
return this.Key[schemeIndex+3+firstSlashIndex:]
}

View File

@@ -81,3 +81,14 @@ func TestItems_Memory2(t *testing.T) {
t.Log(w, len(i)) t.Log(w, len(i))
} }
} }
func TestItem_RequestURI(t *testing.T) {
for _, u := range []string{
"https://goedge.cn/hello/world",
"https://goedge.cn:8080/hello/world",
"https://goedge.cn/hello/world?v=1&t=123",
} {
var item = &Item{Key: u}
t.Log(u, "=>", item.RequestURI())
}
}

View File

@@ -160,6 +160,7 @@ func (this *FileList) CleanPrefix(prefix string) error {
} }
defer func() { defer func() {
// TODO 需要优化
this.memoryCache.Clean() this.memoryCache.Clean()
}() }()
@@ -172,6 +173,46 @@ func (this *FileList) CleanPrefix(prefix string) error {
return nil return nil
} }
// CleanMatchKey 清理通配符匹配的缓存数据,类似于 https://*.example.com/hello
func (this *FileList) CleanMatchKey(key string) error {
if len(key) == 0 {
return nil
}
defer func() {
// TODO 需要优化
this.memoryCache.Clean()
}()
for _, db := range this.dbList {
err := db.CleanMatchKey(key)
if err != nil {
return err
}
}
return nil
}
// CleanMatchPrefix 清理通配符匹配的缓存数据,类似于 https://*.example.com/prefix/
func (this *FileList) CleanMatchPrefix(prefix string) error {
if len(prefix) == 0 {
return nil
}
defer func() {
// TODO 需要优化
this.memoryCache.Clean()
}()
for _, db := range this.dbList {
err := db.CleanMatchPrefix(prefix)
if err != nil {
return err
}
}
return nil
}
func (this *FileList) Remove(hash string) error { func (this *FileList) Remove(hash string) error {
_, err := this.remove(hash) _, err := this.remove(hash)
return err return err

View File

@@ -13,7 +13,10 @@ import (
"github.com/iwind/TeaGo/logs" "github.com/iwind/TeaGo/logs"
"github.com/iwind/TeaGo/types" "github.com/iwind/TeaGo/types"
timeutil "github.com/iwind/TeaGo/utils/time" timeutil "github.com/iwind/TeaGo/utils/time"
"net"
"net/url"
"os" "os"
"path/filepath"
"runtime" "runtime"
"strings" "strings"
"time" "time"
@@ -108,7 +111,7 @@ func (this *FileListDB) Open(dbPath string) error {
this.writeBatch = dbs.NewBatch(writeDB, 4) this.writeBatch = dbs.NewBatch(writeDB, 4)
this.writeBatch.OnFail(func(err error) { this.writeBatch.OnFail(func(err error) {
remotelogs.Warn("LIST_FILE_DB", "run batch failed: "+err.Error()) remotelogs.Warn("LIST_FILE_DB", "run batch failed: "+err.Error()+" ("+filepath.Base(this.dbPath)+")")
}) })
goman.New(func() { goman.New(func() {
@@ -388,6 +391,85 @@ func (this *FileListDB) CleanPrefix(prefix string) error {
} }
} }
func (this *FileListDB) CleanMatchKey(key string) error {
if !this.isReady {
return nil
}
// 忽略 @GOEDGE_
if strings.Contains(key, SuffixAll) {
return nil
}
u, err := url.Parse(key)
if err != nil {
return nil
}
var host = u.Host
hostPart, _, err := net.SplitHostPort(host)
if err == nil && len(hostPart) > 0 {
host = hostPart
}
if len(host) == 0 {
return nil
}
// 转义
var queryKey = strings.ReplaceAll(key, "%", "\\%")
queryKey = strings.ReplaceAll(queryKey, "_", "\\_")
queryKey = strings.Replace(queryKey, "*", "%", 1)
// TODO 检查大批量数据下的操作性能
var staleLife = 600 // TODO 需要可以设置
var unixTime = utils.UnixTime() // 只删除当前的,不删除新的
_, err = this.writeDB.Exec(`UPDATE "`+this.itemsTableName+`" SET "expiredAt"=0, "staleAt"=? WHERE "host" GLOB ? AND "host" NOT GLOB ? AND "key" LIKE ? ESCAPE '\'`, unixTime+int64(staleLife), host, "*."+host, queryKey)
if err != nil {
return err
}
_, err = this.writeDB.Exec(`UPDATE "`+this.itemsTableName+`" SET "expiredAt"=0, "staleAt"=? WHERE "host" GLOB ? AND "host" NOT GLOB ? AND "key" LIKE ? ESCAPE '\'`, unixTime+int64(staleLife), host, "*."+host, queryKey+SuffixAll+"%")
if err != nil {
return err
}
return nil
}
func (this *FileListDB) CleanMatchPrefix(prefix string) error {
if !this.isReady {
return nil
}
u, err := url.Parse(prefix)
if err != nil {
return nil
}
var host = u.Host
hostPart, _, err := net.SplitHostPort(host)
if err == nil && len(hostPart) > 0 {
host = hostPart
}
if len(host) == 0 {
return nil
}
// 转义
var queryPrefix = strings.ReplaceAll(prefix, "%", "\\%")
queryPrefix = strings.ReplaceAll(queryPrefix, "_", "\\_")
queryPrefix = strings.Replace(queryPrefix, "*", "%", 1)
queryPrefix += "%"
// TODO 检查大批量数据下的操作性能
var staleLife = 600 // TODO 需要可以设置
var unixTime = utils.UnixTime() // 只删除当前的,不删除新的
_, err = this.writeDB.Exec(`UPDATE "`+this.itemsTableName+`" SET "expiredAt"=0, "staleAt"=? WHERE "host" GLOB ? AND "host" NOT GLOB ? AND "key" LIKE ? ESCAPE '\'`, unixTime+int64(staleLife), host, "*."+host, queryPrefix)
return err
}
func (this *FileListDB) CleanAll() error { func (this *FileListDB) CleanAll() error {
if !this.isReady { if !this.isReady {
return nil return nil
@@ -473,6 +555,7 @@ func (this *FileListDB) initTables(times int) error {
{ {
// expiredAt - 过期时间,用来判断有无过期 // expiredAt - 过期时间,用来判断有无过期
// staleAt - 过时缓存最大时间,用来清理缓存 // staleAt - 过时缓存最大时间,用来清理缓存
// 不对 hash 增加 unique 参数,是尽可能避免产生 malformed 错误
_, err := this.writeDB.Exec(`CREATE TABLE IF NOT EXISTS "` + this.itemsTableName + `" ( _, err := this.writeDB.Exec(`CREATE TABLE IF NOT EXISTS "` + this.itemsTableName + `" (
"id" integer NOT NULL PRIMARY KEY AUTOINCREMENT, "id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
"hash" varchar(32), "hash" varchar(32),
@@ -498,7 +581,7 @@ ON "` + this.itemsTableName + `" (
"staleAt" ASC "staleAt" ASC
); );
CREATE UNIQUE INDEX IF NOT EXISTS "hash" CREATE INDEX IF NOT EXISTS "hash"
ON "` + this.itemsTableName + `" ( ON "` + this.itemsTableName + `" (
"hash" ASC "hash" ASC
); );

View File

@@ -47,3 +47,41 @@ func TestFileListDB_IncreaseHitAsync(t *testing.T) {
// wait transaction // wait transaction
time.Sleep(1 * time.Second) time.Sleep(1 * time.Second)
} }
func TestFileListDB_CleanMatchKey(t *testing.T) {
var db = caches.NewFileListDB()
err := db.Open(Tea.Root + "/data/cache-db-large.db")
if err != nil {
t.Fatal(err)
}
err = db.Init()
err = db.CleanMatchKey("https://*.goedge.cn/large-text")
if err != nil {
t.Fatal(err)
}
err = db.CleanMatchKey("https://*.goedge.cn:1234/large-text?%2B____")
if err != nil {
t.Fatal(err)
}
}
func TestFileListDB_CleanMatchPrefix(t *testing.T) {
var db = caches.NewFileListDB()
err := db.Open(Tea.Root + "/data/cache-db-large.db")
if err != nil {
t.Fatal(err)
}
err = db.Init()
err = db.CleanMatchPrefix("https://*.goedge.cn/large-text")
if err != nil {
t.Fatal(err)
}
err = db.CleanMatchPrefix("https://*.goedge.cn:1234/large-text?%2B____")
if err != nil {
t.Fatal(err)
}
}

View File

@@ -18,6 +18,12 @@ type ListInterface interface {
// CleanPrefix 清除某个前缀的缓存 // CleanPrefix 清除某个前缀的缓存
CleanPrefix(prefix string) error CleanPrefix(prefix string) error
// CleanMatchKey 清除通配符匹配的Key
CleanMatchKey(key string) error
// CleanMatchPrefix 清除通配符匹配的前缀
CleanMatchPrefix(prefix string) error
// Remove 删除内容 // Remove 删除内容
Remove(hash string) error Remove(hash string) error

View File

@@ -1,8 +1,11 @@
package caches package caches
import ( import (
"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
"github.com/TeaOSLab/EdgeNode/internal/zero" "github.com/TeaOSLab/EdgeNode/internal/zero"
"github.com/iwind/TeaGo/logs" "github.com/iwind/TeaGo/logs"
"net"
"net/url"
"strconv" "strconv"
"strings" "strings"
"sync" "sync"
@@ -146,6 +149,82 @@ func (this *MemoryList) CleanPrefix(prefix string) error {
return nil return nil
} }
// CleanMatchKey 清理通配符匹配的缓存数据,类似于 https://*.example.com/hello
func (this *MemoryList) CleanMatchKey(key string) error {
if strings.Contains(key, SuffixAll) {
return nil
}
u, err := url.Parse(key)
if err != nil {
return nil
}
var host = u.Host
hostPart, _, err := net.SplitHostPort(host)
if err == nil && len(hostPart) > 0 {
host = hostPart
}
if len(host) == 0 {
return nil
}
var requestURI = u.RequestURI()
this.locker.RLock()
defer this.locker.RUnlock()
// TODO 需要优化性能支持千万级数据低于1s的处理速度
for _, itemMap := range this.itemMaps {
for _, item := range itemMap {
if configutils.MatchDomain(host, item.Host) {
var itemRequestURI = item.RequestURI()
if itemRequestURI == requestURI || strings.HasPrefix(itemRequestURI, requestURI+SuffixAll) {
item.ExpiredAt = 0
}
}
}
}
return nil
}
// CleanMatchPrefix 清理通配符匹配的缓存数据,类似于 https://*.example.com/prefix/
func (this *MemoryList) CleanMatchPrefix(prefix string) error {
u, err := url.Parse(prefix)
if err != nil {
return nil
}
var host = u.Host
hostPart, _, err := net.SplitHostPort(host)
if err == nil && len(hostPart) > 0 {
host = hostPart
}
if len(host) == 0 {
return nil
}
var requestURI = u.RequestURI()
var isRootPath = requestURI == "/"
this.locker.RLock()
defer this.locker.RUnlock()
// TODO 需要优化性能支持千万级数据低于1s的处理速度
for _, itemMap := range this.itemMaps {
for _, item := range itemMap {
if configutils.MatchDomain(host, item.Host) {
var itemRequestURI = item.RequestURI()
if isRootPath || strings.HasPrefix(itemRequestURI, requestURI) {
item.ExpiredAt = 0
}
}
}
}
return nil
}
func (this *MemoryList) Remove(hash string) error { func (this *MemoryList) Remove(hash string) error {
this.locker.Lock() this.locker.Lock()

View File

@@ -24,7 +24,8 @@ func init() {
type Manager struct { type Manager struct {
// 全局配置 // 全局配置
MaxDiskCapacity *shared.SizeCapacity MaxDiskCapacity *shared.SizeCapacity
DiskDir string MainDiskDir string
SubDiskDirs []*serverconfigs.CacheDir
MaxMemoryCapacity *shared.SizeCapacity MaxMemoryCapacity *shared.SizeCapacity
policyMap map[int64]*serverconfigs.HTTPCachePolicy // policyId => []*Policy policyMap map[int64]*serverconfigs.HTTPCachePolicy // policyId => []*Policy
@@ -47,12 +48,10 @@ func (this *Manager) UpdatePolicies(newPolicies []*serverconfigs.HTTPCachePolicy
this.locker.Lock() this.locker.Lock()
defer this.locker.Unlock() defer this.locker.Unlock()
newPolicyIds := []int64{} var newPolicyIds = []int64{}
for _, policy := range newPolicies { for _, policy := range newPolicies {
// 使用节点单独的缓存目录 // 使用节点单独的缓存目录
if len(this.DiskDir) > 0 { policy.UpdateDiskDir(this.MainDiskDir, this.SubDiskDirs)
policy.UpdateDiskDir(this.DiskDir)
}
newPolicyIds = append(newPolicyIds, policy.Id) newPolicyIds = append(newPolicyIds, policy.Id)
} }

View File

@@ -19,7 +19,7 @@ type OpenFileCache struct {
poolList *linkedlist.List poolList *linkedlist.List
watcher *fsnotify.Watcher watcher *fsnotify.Watcher
locker sync.Mutex locker sync.RWMutex
maxSize int maxSize int
count int count int
@@ -54,13 +54,18 @@ func NewOpenFileCache(maxSize int) (*OpenFileCache, error) {
} }
func (this *OpenFileCache) Get(filename string) *OpenFile { func (this *OpenFileCache) Get(filename string) *OpenFile {
this.locker.Lock() this.locker.RLock()
defer this.locker.Unlock()
pool, ok := this.poolMap[filename] pool, ok := this.poolMap[filename]
this.locker.RUnlock()
if ok { if ok {
file, consumed := pool.Get() file, consumed := pool.Get()
if consumed { if consumed {
this.locker.Lock()
this.count-- this.count--
// pool如果为空也不需要从列表中删除避免put时需要重新创建
this.locker.Unlock()
} }
return file return file
} }
@@ -146,6 +151,7 @@ func (this *OpenFileCache) CloseAll() {
this.poolMap = map[string]*OpenFilePool{} this.poolMap = map[string]*OpenFilePool{}
this.poolList.Reset() this.poolList.Reset()
_ = this.watcher.Close() _ = this.watcher.Close()
this.count = 0
this.locker.Unlock() this.locker.Unlock()
} }

View File

@@ -0,0 +1,43 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package caches_test
import (
"github.com/TeaOSLab/EdgeNode/internal/caches"
"testing"
"time"
)
func TestNewOpenFileCache_Close(t *testing.T) {
cache, err := caches.NewOpenFileCache(1024)
if err != nil {
t.Fatal(err)
}
cache.Debug()
cache.Put("a.txt", caches.NewOpenFile(nil, nil, nil, 0))
cache.Put("b.txt", caches.NewOpenFile(nil, nil, nil, 0))
cache.Put("b.txt", caches.NewOpenFile(nil, nil, nil, 0))
cache.Put("b.txt", caches.NewOpenFile(nil, nil, nil, 0))
cache.Put("c.txt", caches.NewOpenFile(nil, nil, nil, 0))
cache.Get("b.txt")
cache.Get("d.txt")
cache.Close("a.txt")
time.Sleep(100 * time.Second)
}
func TestNewOpenFileCache_CloseAll(t *testing.T) {
cache, err := caches.NewOpenFileCache(1024)
if err != nil {
t.Fatal(err)
}
cache.Debug()
cache.Put("a.txt", caches.NewOpenFile(nil, nil, nil, 0))
cache.Put("b.txt", caches.NewOpenFile(nil, nil, nil, 0))
cache.Put("c.txt", caches.NewOpenFile(nil, nil, nil, 0))
cache.Get("b.txt")
cache.Get("d.txt")
cache.CloseAll()
time.Sleep(6 * time.Second)
}

View File

@@ -4,6 +4,8 @@ package caches_test
import ( import (
"github.com/TeaOSLab/EdgeNode/internal/caches" "github.com/TeaOSLab/EdgeNode/internal/caches"
"github.com/iwind/TeaGo/rands"
"sync"
"testing" "testing"
) )
@@ -15,3 +17,30 @@ func TestOpenFilePool_Get(t *testing.T) {
t.Log(pool.Get()) t.Log(pool.Get())
t.Log(pool.Get()) t.Log(pool.Get())
} }
func TestOpenFilePool_Close(t *testing.T) {
var pool = caches.NewOpenFilePool("a")
pool.Put(caches.NewOpenFile(nil, nil, nil, 0))
pool.Put(caches.NewOpenFile(nil, nil, nil, 0))
pool.Close()
}
func TestOpenFilePool_Concurrent(t *testing.T) {
var pool = caches.NewOpenFilePool("a")
var concurrent = 1000
var wg = &sync.WaitGroup{}
wg.Add(concurrent)
for i := 0; i < concurrent; i++ {
go func() {
defer wg.Done()
if rands.Int(0, 1) == 1 {
pool.Put(caches.NewOpenFile(nil, nil, nil, 0))
}
if rands.Int(0, 1) == 0 {
pool.Get()
}
}()
}
wg.Wait()
}

View File

@@ -3,38 +3,88 @@
package caches package caches
import ( import (
"bytes"
"encoding/json" "encoding/json"
"errors" "github.com/iwind/TeaGo/types"
"os" "os"
"strconv"
) )
// PartialRanges 内容分区范围定义 // PartialRanges 内容分区范围定义
type PartialRanges struct { type PartialRanges struct {
Ranges [][2]int64 `json:"ranges"` Version int `json:"version"` // 版本号
Ranges [][2]int64 `json:"ranges"` // 范围
BodySize int64 `json:"bodySize"` // 总长度
} }
// NewPartialRanges 获取新对象 // NewPartialRanges 获取新对象
func NewPartialRanges() *PartialRanges { func NewPartialRanges(expiresAt int64) *PartialRanges {
return &PartialRanges{Ranges: [][2]int64{}} return &PartialRanges{
Ranges: [][2]int64{},
Version: 1,
}
}
// NewPartialRangesFromData 从数据中解析范围
func NewPartialRangesFromData(data []byte) (*PartialRanges, error) {
var rs = NewPartialRanges(0)
for {
var index = bytes.IndexRune(data, '\n')
if index < 0 {
break
}
var line = data[:index]
var colonIndex = bytes.IndexRune(line, ':')
if colonIndex > 0 {
switch string(line[:colonIndex]) {
case "v": // 版本号
rs.Version = types.Int(line[colonIndex+1:])
case "b": // 总长度
rs.BodySize = types.Int64(line[colonIndex+1:])
case "r": // 范围信息
var commaIndex = bytes.IndexRune(line, ',')
if commaIndex > 0 {
rs.Ranges = append(rs.Ranges, [2]int64{types.Int64(line[colonIndex+1 : commaIndex]), types.Int64(line[commaIndex+1:])})
}
}
}
data = data[index+1:]
if len(data) == 0 {
break
}
}
return rs, nil
} }
// NewPartialRangesFromJSON 从JSON中解析范围 // NewPartialRangesFromJSON 从JSON中解析范围
func NewPartialRangesFromJSON(data []byte) (*PartialRanges, error) { func NewPartialRangesFromJSON(data []byte) (*PartialRanges, error) {
var rs = NewPartialRanges() var rs = NewPartialRanges(0)
err := json.Unmarshal(data, &rs) err := json.Unmarshal(data, &rs)
if err != nil { if err != nil {
return nil, err return nil, err
} }
rs.Version = 0
return rs, nil return rs, nil
} }
// NewPartialRangesFromFile 从文件中加载范围信息
func NewPartialRangesFromFile(path string) (*PartialRanges, error) { func NewPartialRangesFromFile(path string) (*PartialRanges, error) {
data, err := os.ReadFile(path) data, err := os.ReadFile(path)
if err != nil { if err != nil {
return nil, err return nil, err
} }
return NewPartialRangesFromJSON(data) if len(data) == 0 {
return NewPartialRanges(0), nil
}
// 兼容老的JSON格式
if data[0] == '{' {
return NewPartialRangesFromJSON(data)
}
// 新的格式
return NewPartialRangesFromData(data)
} }
// Add 添加新范围 // Add 添加新范围
@@ -105,29 +155,27 @@ func (this *PartialRanges) Nearest(begin int64, end int64) (r [2]int64, ok bool)
return return
} }
// AsJSON 转换为JSON // 转换为字符串
func (this *PartialRanges) AsJSON() ([]byte, error) { func (this *PartialRanges) String() string {
return json.Marshal(this) var s = "v:" + strconv.Itoa(this.Version) + "\n" + // version
"b:" + this.formatInt64(this.BodySize) + "\n" // bodySize
for _, r := range this.Ranges {
s += "r:" + this.formatInt64(r[0]) + "," + this.formatInt64(r[1]) + "\n" // range
}
return s
}
// Bytes 将内容转换为字节
func (this *PartialRanges) Bytes() []byte {
return []byte(this.String())
} }
// WriteToFile 写入到文件中 // WriteToFile 写入到文件中
func (this *PartialRanges) WriteToFile(path string) error { func (this *PartialRanges) WriteToFile(path string) error {
data, err := this.AsJSON() return os.WriteFile(path, this.Bytes(), 0666)
if err != nil {
return errors.New("convert to json failed: " + err.Error())
}
return os.WriteFile(path, data, 0666)
}
// ReadFromFile 从文件中读取
func (this *PartialRanges) ReadFromFile(path string) (*PartialRanges, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, err
}
return NewPartialRangesFromJSON(data)
} }
// Max 获取最大位置
func (this *PartialRanges) Max() int64 { func (this *PartialRanges) Max() int64 {
if len(this.Ranges) > 0 { if len(this.Ranges) > 0 {
return this.Ranges[len(this.Ranges)-1][1] return this.Ranges[len(this.Ranges)-1][1]
@@ -135,6 +183,11 @@ func (this *PartialRanges) Max() int64 {
return 0 return 0
} }
// Reset 重置范围信息
func (this *PartialRanges) Reset() {
this.Ranges = [][2]int64{}
}
func (this *PartialRanges) merge(index int) { func (this *PartialRanges) merge(index int) {
// forward // forward
var lastIndex = index var lastIndex = index
@@ -187,3 +240,7 @@ func (this *PartialRanges) max(n1 int64, n2 int64) int64 {
} }
return n2 return n2
} }
func (this *PartialRanges) formatInt64(i int64) string {
return strconv.FormatInt(i, 10)
}

View File

@@ -3,14 +3,16 @@
package caches_test package caches_test
import ( import (
"encoding/json"
"github.com/TeaOSLab/EdgeNode/internal/caches" "github.com/TeaOSLab/EdgeNode/internal/caches"
"github.com/iwind/TeaGo/assert" "github.com/iwind/TeaGo/assert"
"github.com/iwind/TeaGo/logs" "github.com/iwind/TeaGo/logs"
"testing" "testing"
"time"
) )
func TestNewPartialRanges(t *testing.T) { func TestNewPartialRanges(t *testing.T) {
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
r.Add(1, 100) r.Add(1, 100)
r.Add(50, 300) r.Add(50, 300)
@@ -28,7 +30,7 @@ func TestNewPartialRanges(t *testing.T) {
func TestNewPartialRanges1(t *testing.T) { func TestNewPartialRanges1(t *testing.T) {
var a = assert.NewAssertion(t) var a = assert.NewAssertion(t)
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
r.Add(1, 100) r.Add(1, 100)
r.Add(1, 101) r.Add(1, 101)
r.Add(1, 102) r.Add(1, 102)
@@ -47,7 +49,7 @@ func TestNewPartialRanges1(t *testing.T) {
func TestNewPartialRanges2(t *testing.T) { func TestNewPartialRanges2(t *testing.T) {
// low -> high // low -> high
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
r.Add(1, 100) r.Add(1, 100)
r.Add(1, 101) r.Add(1, 101)
r.Add(1, 102) r.Add(1, 102)
@@ -63,7 +65,7 @@ func TestNewPartialRanges2(t *testing.T) {
func TestNewPartialRanges3(t *testing.T) { func TestNewPartialRanges3(t *testing.T) {
// high -> low // high -> low
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
r.Add(301, 302) r.Add(301, 302)
r.Add(303, 304) r.Add(303, 304)
r.Add(200, 300) r.Add(200, 300)
@@ -75,7 +77,7 @@ func TestNewPartialRanges3(t *testing.T) {
func TestNewPartialRanges4(t *testing.T) { func TestNewPartialRanges4(t *testing.T) {
// nearby // nearby
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
r.Add(301, 302) r.Add(301, 302)
r.Add(303, 304) r.Add(303, 304)
r.Add(305, 306) r.Add(305, 306)
@@ -90,7 +92,7 @@ func TestNewPartialRanges4(t *testing.T) {
} }
func TestNewPartialRanges5(t *testing.T) { func TestNewPartialRanges5(t *testing.T) {
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
for j := 0; j < 1000; j++ { for j := 0; j < 1000; j++ {
r.Add(int64(j), int64(j+100)) r.Add(int64(j), int64(j+100))
} }
@@ -100,7 +102,7 @@ func TestNewPartialRanges5(t *testing.T) {
func TestNewPartialRanges_Nearest(t *testing.T) { func TestNewPartialRanges_Nearest(t *testing.T) {
{ {
// nearby // nearby
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
r.Add(301, 400) r.Add(301, 400)
r.Add(401, 500) r.Add(401, 500)
r.Add(501, 600) r.Add(501, 600)
@@ -112,7 +114,7 @@ func TestNewPartialRanges_Nearest(t *testing.T) {
{ {
// nearby // nearby
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
r.Add(301, 400) r.Add(301, 400)
r.Add(450, 500) r.Add(450, 500)
r.Add(550, 600) r.Add(550, 600)
@@ -131,45 +133,100 @@ func TestNewPartialRanges_Large_Range(t *testing.T) {
var largeSize int64 = 10000000000000 var largeSize int64 = 10000000000000
t.Log(largeSize/1024/1024/1024, "G") t.Log(largeSize/1024/1024/1024, "G")
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
r.Add(1, largeSize) r.Add(1, largeSize)
jsonData, err := r.AsJSON() var s = r.String()
if err != nil { t.Log(s)
t.Fatal(err)
}
t.Log(string(jsonData))
r2, err := caches.NewPartialRangesFromJSON(jsonData) r2, err := caches.NewPartialRangesFromData([]byte(s))
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
a.IsTrue(largeSize == r2.Ranges[0][1]) a.IsTrue(largeSize == r2.Ranges[0][1])
logs.PrintAsJSON(r, t)
} }
func TestNewPartialRanges_AsJSON(t *testing.T) { func TestPartialRanges_Encode_JSON(t *testing.T) {
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
for j := 0; j < 1000; j++ { for i := 0; i < 10; i++ {
r.Add(int64(j), int64(j+100)) r.Ranges = append(r.Ranges, [2]int64{int64(i * 100), int64(i*100 + 100)})
} }
data, err := r.AsJSON() var before = time.Now()
data, err := json.Marshal(r)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
t.Log(string(data)) t.Log(time.Since(before).Seconds()*1000, "ms")
t.Log(len(data))
}
r2, err := caches.NewPartialRangesFromJSON(data) func TestPartialRanges_Encode_String(t *testing.T) {
var r = caches.NewPartialRanges(0)
r.BodySize = 1024
for i := 0; i < 10; i++ {
r.Ranges = append(r.Ranges, [2]int64{int64(i * 100), int64(i*100 + 100)})
}
var before = time.Now()
var data = r.String()
t.Log(time.Since(before).Seconds()*1000, "ms")
t.Log(len(data))
r2, err := caches.NewPartialRangesFromData([]byte(data))
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
t.Log(r2.Ranges) logs.PrintAsJSON(r2, t)
}
func TestPartialRanges_Version(t *testing.T) {
{
ranges, err := caches.NewPartialRangesFromData([]byte(`e:1668928495
r:0,1048576
r:1140260864,1140295164`))
if err != nil {
t.Fatal(err)
}
t.Log("version:", ranges.Version)
}
{
ranges, err := caches.NewPartialRangesFromData([]byte(`e:1668928495
r:0,1048576
r:1140260864,1140295164
v:0
`))
if err != nil {
t.Fatal(err)
}
t.Log("version:", ranges.Version)
}
{
ranges, err := caches.NewPartialRangesFromJSON([]byte(`{}`))
if err != nil {
t.Fatal(err)
}
t.Log("version:", ranges.Version)
}
} }
func BenchmarkNewPartialRanges(b *testing.B) { func BenchmarkNewPartialRanges(b *testing.B) {
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
var r = caches.NewPartialRanges() var r = caches.NewPartialRanges(0)
for j := 0; j < 1000; j++ { for j := 0; j < 1000; j++ {
r.Add(int64(j), int64(j+100)) r.Add(int64(j), int64(j+100))
} }
} }
} }
func BenchmarkPartialRanges_String(b *testing.B) {
var r = caches.NewPartialRanges(0)
r.BodySize = 1024
for i := 0; i < 10; i++ {
r.Ranges = append(r.Ranges, [2]int64{int64(i * 100), int64(i*100 + 100)})
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
_ = r.String()
}
}

View File

@@ -42,7 +42,7 @@ func (this *FileReader) InitAutoDiscard(autoDiscard bool) error {
this.header = this.openFile.header this.header = this.openFile.header
} }
isOk := false var isOk = false
if autoDiscard { if autoDiscard {
defer func() { defer func() {
@@ -67,17 +67,17 @@ func (this *FileReader) InitAutoDiscard(autoDiscard bool) error {
this.expiresAt = int64(binary.BigEndian.Uint32(buf[:SizeExpiresAt])) this.expiresAt = int64(binary.BigEndian.Uint32(buf[:SizeExpiresAt]))
status := types.Int(string(buf[OffsetStatus : OffsetStatus+SizeStatus])) var status = types.Int(string(buf[OffsetStatus : OffsetStatus+SizeStatus]))
if status < 100 || status > 999 { if status < 100 || status > 999 {
return errors.New("invalid status") return errors.New("invalid status")
} }
this.status = status this.status = status
// URL // URL
urlLength := binary.BigEndian.Uint32(buf[OffsetURLLength : OffsetURLLength+SizeURLLength]) var urlLength = binary.BigEndian.Uint32(buf[OffsetURLLength : OffsetURLLength+SizeURLLength])
// header // header
headerSize := int(binary.BigEndian.Uint32(buf[OffsetHeaderLength : OffsetHeaderLength+SizeHeaderLength])) var headerSize = int(binary.BigEndian.Uint32(buf[OffsetHeaderLength : OffsetHeaderLength+SizeHeaderLength]))
if headerSize == 0 { if headerSize == 0 {
return nil return nil
} }
@@ -86,7 +86,7 @@ func (this *FileReader) InitAutoDiscard(autoDiscard bool) error {
// body // body
this.bodyOffset = this.headerOffset + int64(headerSize) this.bodyOffset = this.headerOffset + int64(headerSize)
bodySize := int(binary.BigEndian.Uint64(buf[OffsetBodyLength : OffsetBodyLength+SizeBodyLength])) var bodySize = int(binary.BigEndian.Uint64(buf[OffsetBodyLength : OffsetBodyLength+SizeBodyLength]))
if bodySize == 0 { if bodySize == 0 {
isOk = true isOk = true
return nil return nil
@@ -158,7 +158,7 @@ func (this *FileReader) ReadHeader(buf []byte, callback ReaderFunc) error {
return nil return nil
} }
isOk := false var isOk = false
defer func() { defer func() {
if !isOk { if !isOk {
@@ -171,7 +171,7 @@ func (this *FileReader) ReadHeader(buf []byte, callback ReaderFunc) error {
return err return err
} }
headerSize := this.headerSize var headerSize = this.headerSize
for { for {
n, err := this.fp.Read(buf) n, err := this.fp.Read(buf)
@@ -215,7 +215,7 @@ func (this *FileReader) ReadHeader(buf []byte, callback ReaderFunc) error {
} }
func (this *FileReader) ReadBody(buf []byte, callback ReaderFunc) error { func (this *FileReader) ReadBody(buf []byte, callback ReaderFunc) error {
isOk := false var isOk = false
defer func() { defer func() {
if !isOk { if !isOk {
@@ -261,11 +261,12 @@ func (this *FileReader) Read(buf []byte) (n int, err error) {
if err != nil && err != io.EOF { if err != nil && err != io.EOF {
_ = this.discard() _ = this.discard()
} }
return return
} }
func (this *FileReader) ReadBodyRange(buf []byte, start int64, end int64, callback ReaderFunc) error { func (this *FileReader) ReadBodyRange(buf []byte, start int64, end int64, callback ReaderFunc) error {
isOk := false var isOk = false
defer func() { defer func() {
if !isOk { if !isOk {
@@ -273,7 +274,7 @@ func (this *FileReader) ReadBodyRange(buf []byte, start int64, end int64, callba
} }
}() }()
offset := start var offset = start
if start < 0 { if start < 0 {
offset = this.bodyOffset + this.bodySize + end offset = this.bodyOffset + this.bodySize + end
end = this.bodyOffset + this.bodySize - 1 end = this.bodyOffset + this.bodySize - 1
@@ -296,7 +297,7 @@ func (this *FileReader) ReadBodyRange(buf []byte, start int64, end int64, callba
for { for {
n, err := this.fp.Read(buf) n, err := this.fp.Read(buf)
if n > 0 { if n > 0 {
n2 := int(end-offset) + 1 var n2 = int(end-offset) + 1
if n2 <= n { if n2 <= n {
_, e := callback(n2) _, e := callback(n2)
if e != nil { if e != nil {
@@ -344,12 +345,12 @@ func (this *FileReader) FP() *os.File {
} }
func (this *FileReader) Close() error { func (this *FileReader) Close() error {
if this.openFileCache != nil { if this.isClosed {
if this.isClosed { return nil
return nil }
} this.isClosed = true
this.isClosed = true
if this.openFileCache != nil {
if this.openFile != nil { if this.openFile != nil {
this.openFileCache.Put(this.fp.Name(), this.openFile) this.openFileCache.Put(this.fp.Name(), this.openFile)
} else { } else {
@@ -359,6 +360,7 @@ func (this *FileReader) Close() error {
} }
return nil return nil
} }
return this.fp.Close() return this.fp.Close()
} }

View File

@@ -19,7 +19,7 @@ func TestFileReader(t *testing.T) {
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
_, path := storage.keyPath("my-key") _, path, _ := storage.keyPath("my-key")
fp, err := os.Open(path) fp, err := os.Open(path)
if err != nil { if err != nil {
@@ -105,7 +105,7 @@ func TestFileReader_Range(t *testing.T) {
} }
_ = writer.Close()**/ _ = writer.Close()**/
_, path := storage.keyPath("my-number") _, path, _ := storage.keyPath("my-number")
fp, err := os.Open(path) fp, err := os.Open(path)
if err != nil { if err != nil {

View File

@@ -117,13 +117,10 @@ func (this *PartialFileReader) ContainsRange(r rangeutils.Range) (r2 rangeutils.
r2, ok = this.ranges.Nearest(r.Start(), r.End()) r2, ok = this.ranges.Nearest(r.Start(), r.End())
if ok && this.bodySize > 0 { if ok && this.bodySize > 0 {
// 考虑可配置 // 考虑可配置
var span int64 = 512 * 1024 const minSpan = 128 << 10
if this.bodySize > 1<<30 {
span = 1 << 20
}
// 这里限制返回的最小缓存,防止因为返回的内容过小而导致请求过多 // 这里限制返回的最小缓存,防止因为返回的内容过小而导致请求过多
if r2.Length() < r.Length() && r2.Length() < span { if r2.Length() < r.Length() && r2.Length() < minSpan {
ok = false ok = false
} }
} }
@@ -138,6 +135,10 @@ func (this *PartialFileReader) MaxLength() int64 {
return this.ranges.Max() + 1 return this.ranges.Max() + 1
} }
func (this *PartialFileReader) Ranges() *PartialRanges {
return this.ranges
}
func (this *PartialFileReader) discard() error { func (this *PartialFileReader) discard() error {
_ = os.Remove(this.rangePath) _ = os.Remove(this.rangePath)
return this.FileReader.discard() return this.FileReader.discard()

View File

@@ -21,6 +21,7 @@ import (
"github.com/iwind/TeaGo/rands" "github.com/iwind/TeaGo/rands"
"github.com/iwind/TeaGo/types" "github.com/iwind/TeaGo/types"
stringutil "github.com/iwind/TeaGo/utils/string" stringutil "github.com/iwind/TeaGo/utils/string"
"golang.org/x/sys/unix"
"golang.org/x/text/language" "golang.org/x/text/language"
"golang.org/x/text/message" "golang.org/x/text/message"
"math" "math"
@@ -48,8 +49,7 @@ const (
SizeBodyLength = 8 SizeBodyLength = 8
OffsetBodyLength = OffsetHeaderLength + SizeHeaderLength OffsetBodyLength = OffsetHeaderLength + SizeHeaderLength
SizeMeta = SizeExpiresAt + SizeStatus + SizeURLLength + SizeHeaderLength + SizeBodyLength SizeMeta = SizeExpiresAt + SizeStatus + SizeURLLength + SizeHeaderLength + SizeBodyLength
OffsetKey = SizeMeta
) )
const ( const (
@@ -58,6 +58,7 @@ const (
HotItemLifeSeconds int64 = 3600 // 热点数据生命周期 HotItemLifeSeconds int64 = 3600 // 热点数据生命周期
FileToMemoryMaxSize = 32 * sizes.M // 可以从文件写入到内存的最大文件尺寸 FileToMemoryMaxSize = 32 * sizes.M // 可以从文件写入到内存的最大文件尺寸
FileTmpSuffix = ".tmp" FileTmpSuffix = ".tmp"
MinDiskSpace = 5 << 30 // 当前磁盘最小剩余空间
) )
var sharedWritingFileKeyMap = map[string]zero.Zero{} // key => bool var sharedWritingFileKeyMap = map[string]zero.Zero{} // key => bool
@@ -90,6 +91,11 @@ type FileStorage struct {
ignoreKeys *setutils.FixedSet ignoreKeys *setutils.FixedSet
openFileCache *OpenFileCache openFileCache *OpenFileCache
mainDir string
mainDiskIsFull bool
subDirs []*FileDir
} }
func NewFileStorage(policy *serverconfigs.HTTPCachePolicy) *FileStorage { func NewFileStorage(policy *serverconfigs.HTTPCachePolicy) *FileStorage {
@@ -153,6 +159,16 @@ func (this *FileStorage) UpdatePolicy(newPolicy *serverconfigs.HTTPCachePolicy)
return return
} }
var subDirs = []*FileDir{}
for _, subDir := range newOptions.SubDirs {
subDirs = append(subDirs, &FileDir{
Path: subDir.Path,
Capacity: subDir.Capacity,
IsFull: false,
})
}
this.checkDiskSpace()
err = newOptions.Init() err = newOptions.Init()
if err != nil { if err != nil {
remotelogs.Error("CACHE", "update policy '"+types.String(this.policy.Id)+"' failed: init options failed: "+err.Error()) remotelogs.Error("CACHE", "update policy '"+types.String(this.policy.Id)+"' failed: init options failed: "+err.Error())
@@ -215,6 +231,19 @@ func (this *FileStorage) Init() error {
this.options.Dir = filepath.Clean(this.options.Dir) this.options.Dir = filepath.Clean(this.options.Dir)
var dir = this.options.Dir var dir = this.options.Dir
var subDirs = []*FileDir{}
for _, subDir := range this.options.SubDirs {
subDirs = append(subDirs, &FileDir{
Path: subDir.Path,
Capacity: subDir.Capacity,
IsFull: false,
})
}
this.subDirs = subDirs
if len(subDirs) > 0 {
this.checkDiskSpace()
}
if len(dir) == 0 { if len(dir) == 0 {
return errors.New("[CACHE]cache storage dir can not be empty") return errors.New("[CACHE]cache storage dir can not be empty")
} }
@@ -287,6 +316,9 @@ func (this *FileStorage) Init() error {
// open file cache // open file cache
this.initOpenFileCache() this.initOpenFileCache()
// 检查磁盘空间
this.checkDiskSpace()
return nil return nil
} }
@@ -314,7 +346,7 @@ func (this *FileStorage) openReader(key string, allowMemory bool, useStale bool,
} }
} }
hash, path := this.keyPath(key) hash, path, _ := this.keyPath(key)
// 检查文件记录是否已过期 // 检查文件记录是否已过期
if !useStale { if !useStale {
@@ -382,16 +414,16 @@ func (this *FileStorage) openReader(key string, allowMemory bool, useStale bool,
} }
// OpenWriter 打开缓存文件等待写入 // OpenWriter 打开缓存文件等待写入
func (this *FileStorage) OpenWriter(key string, expiresAt int64, status int, size int64, maxSize int64, isPartial bool) (Writer, error) { func (this *FileStorage) OpenWriter(key string, expiresAt int64, status int, headerSize int, bodySize int64, maxSize int64, isPartial bool) (Writer, error) {
return this.openWriter(key, expiresAt, status, size, maxSize, isPartial, false) return this.openWriter(key, expiresAt, status, headerSize, bodySize, maxSize, isPartial, false)
} }
// OpenFlushWriter 打开从其他媒介直接刷入的写入器 // OpenFlushWriter 打开从其他媒介直接刷入的写入器
func (this *FileStorage) OpenFlushWriter(key string, expiresAt int64, status int) (Writer, error) { func (this *FileStorage) OpenFlushWriter(key string, expiresAt int64, status int, headerSize int, bodySize int64) (Writer, error) {
return this.openWriter(key, expiresAt, status, -1, -1, false, true) return this.openWriter(key, expiresAt, status, headerSize, bodySize, -1, false, true)
} }
func (this *FileStorage) openWriter(key string, expiredAt int64, status int, size int64, maxSize int64, isPartial bool, isFlushing bool) (Writer, error) { func (this *FileStorage) openWriter(key string, expiredAt int64, status int, headerSize int, bodySize int64, maxSize int64, isPartial bool, isFlushing bool) (Writer, error) {
// 是否正在退出 // 是否正在退出
if teaconst.IsQuiting { if teaconst.IsQuiting {
return nil, ErrWritingUnavailable return nil, ErrWritingUnavailable
@@ -409,8 +441,8 @@ func (this *FileStorage) openWriter(key string, expiredAt int64, status int, siz
maxMemorySize = maxSize maxMemorySize = maxSize
} }
var memoryStorage = this.memoryStorage var memoryStorage = this.memoryStorage
if !isFlushing && !isPartial && memoryStorage != nil && ((size > 0 && size < maxMemorySize) || size < 0) { if !isFlushing && !isPartial && memoryStorage != nil && ((bodySize > 0 && bodySize < maxMemorySize) || bodySize < 0) {
writer, err := memoryStorage.OpenWriter(key, expiredAt, status, size, maxMemorySize, false) writer, err := memoryStorage.OpenWriter(key, expiredAt, status, headerSize, bodySize, maxMemorySize, false)
if err == nil { if err == nil {
return writer, nil return writer, nil
} }
@@ -463,17 +495,9 @@ func (this *FileStorage) openWriter(key string, expiredAt int64, status int, siz
var hash = stringutil.Md5(key) var hash = stringutil.Md5(key)
// TODO 可以只stat一次 dir, diskIsFull := this.subDir(hash)
var dir = this.options.Dir + "/p" + strconv.FormatInt(this.policy.Id, 10) + "/" + hash[:2] + "/" + hash[2:4] if diskIsFull {
_, err = os.Stat(dir) return nil, NewCapacityError("the disk is full")
if err != nil {
if !os.IsNotExist(err) {
return nil, err
}
err = os.MkdirAll(dir, 0777)
if err != nil {
return nil, err
}
} }
// 检查缓存是否已经生成 // 检查缓存是否已经生成
@@ -520,19 +544,38 @@ func (this *FileStorage) openWriter(key string, expiredAt int64, status int, siz
// 从已经存储的内容中读取信息 // 从已经存储的内容中读取信息
var isNewCreated = true var isNewCreated = true
var partialBodyOffset int64 var partialBodyOffset int64
var partialRanges *PartialRanges
if isPartial { if isPartial {
readerFp, err := os.OpenFile(tmpPath, os.O_RDONLY, 0444) // 数据库中是否存在
if err == nil { existsCacheItem, _ := this.list.Exist(hash)
var partialReader = NewPartialFileReader(readerFp) if existsCacheItem {
err = partialReader.Init() readerFp, err := os.OpenFile(tmpPath, os.O_RDONLY, 0444)
_ = partialReader.Close() if err == nil {
if err == nil && partialReader.bodyOffset > 0 { var partialReader = NewPartialFileReader(readerFp)
isNewCreated = false err = partialReader.Init()
partialBodyOffset = partialReader.bodyOffset _ = partialReader.Close()
} else { if err == nil && partialReader.bodyOffset > 0 {
_ = this.removeCacheFile(tmpPath) partialRanges = partialReader.Ranges()
if bodySize > 0 && partialRanges != nil && partialRanges.BodySize > 0 && bodySize != partialRanges.BodySize {
_ = this.removeCacheFile(tmpPath)
} else {
isNewCreated = false
partialBodyOffset = partialReader.bodyOffset
}
} else {
_ = this.removeCacheFile(tmpPath)
}
} }
} }
if isNewCreated {
err = this.list.Remove(hash)
if err != nil {
return nil, err
}
}
if partialRanges == nil {
partialRanges = NewPartialRanges(expiredAt)
}
} }
var flags = os.O_CREATE | os.O_WRONLY var flags = os.O_CREATE | os.O_WRONLY
@@ -542,7 +585,16 @@ func (this *FileStorage) openWriter(key string, expiredAt int64, status int, siz
var before = time.Now() var before = time.Now()
writer, err := os.OpenFile(tmpPath, flags, 0666) writer, err := os.OpenFile(tmpPath, flags, 0666)
if err != nil { if err != nil {
return nil, err // TODO 检查在各个系统中的稳定性
if os.IsNotExist(err) {
_ = os.MkdirAll(dir, 0777)
// open file again
writer, err = os.OpenFile(tmpPath, flags, 0666)
}
if err != nil {
return nil, err
}
} }
if !isFlushing { if !isFlushing {
if time.Since(before) >= maxOpenFilesSlowCost { if time.Since(before) >= maxOpenFilesSlowCost {
@@ -574,9 +626,12 @@ func (this *FileStorage) openWriter(key string, expiredAt int64, status int, siz
return nil, ErrFileIsWriting return nil, ErrFileIsWriting
} }
var metaBodySize int64 = -1
var metaHeaderSize = -1
if isNewCreated { if isNewCreated {
// 写入过期时间 // 写入meta
var metaBytes = make([]byte, SizeMeta+len(key)) // 从v0.5.8开始不再在meta中写入Key
var metaBytes = make([]byte, SizeMeta)
binary.BigEndian.PutUint32(metaBytes[OffsetExpiresAt:], uint32(expiredAt)) binary.BigEndian.PutUint32(metaBytes[OffsetExpiresAt:], uint32(expiredAt))
// 写入状态码 // 写入状态码
@@ -585,17 +640,17 @@ func (this *FileStorage) openWriter(key string, expiredAt int64, status int, siz
} }
copy(metaBytes[OffsetStatus:], strconv.Itoa(status)) copy(metaBytes[OffsetStatus:], strconv.Itoa(status))
// 写入URL长度
binary.BigEndian.PutUint32(metaBytes[OffsetURLLength:], uint32(len(key)))
// 写入Header Length // 写入Header Length
binary.BigEndian.PutUint32(metaBytes[OffsetHeaderLength:], uint32(0)) if headerSize > 0 {
binary.BigEndian.PutUint32(metaBytes[OffsetHeaderLength:], uint32(headerSize))
metaHeaderSize = headerSize
}
// 写入Body Length // 写入Body Length
binary.BigEndian.PutUint64(metaBytes[OffsetBodyLength:], uint64(0)) if bodySize > 0 {
binary.BigEndian.PutUint64(metaBytes[OffsetBodyLength:], uint64(bodySize))
// 写入URL metaBodySize = bodySize
copy(metaBytes[OffsetKey:], key) }
_, err = writer.Write(metaBytes) _, err = writer.Write(metaBytes)
if err != nil { if err != nil {
@@ -605,12 +660,7 @@ func (this *FileStorage) openWriter(key string, expiredAt int64, status int, siz
isOk = true isOk = true
if isPartial { if isPartial {
ranges, err := NewPartialRangesFromFile(cachePathName + "@ranges.cache") return NewPartialFileWriter(writer, key, expiredAt, metaHeaderSize, metaBodySize, isNewCreated, isPartial, partialBodyOffset, partialRanges, func() {
if err != nil {
ranges = NewPartialRanges()
}
return NewPartialFileWriter(writer, key, expiredAt, isNewCreated, isPartial, partialBodyOffset, ranges, func() {
sharedWritingFileKeyLocker.Lock() sharedWritingFileKeyLocker.Lock()
delete(sharedWritingFileKeyMap, key) delete(sharedWritingFileKeyMap, key)
if len(sharedWritingFileKeyMap) == 0 { if len(sharedWritingFileKeyMap) == 0 {
@@ -619,7 +669,7 @@ func (this *FileStorage) openWriter(key string, expiredAt int64, status int, siz
sharedWritingFileKeyLocker.Unlock() sharedWritingFileKeyLocker.Unlock()
}), nil }), nil
} else { } else {
return NewFileWriter(this, writer, key, expiredAt, -1, func() { return NewFileWriter(this, writer, key, expiredAt, metaHeaderSize, metaBodySize, -1, func() {
sharedWritingFileKeyLocker.Lock() sharedWritingFileKeyLocker.Lock()
delete(sharedWritingFileKeyMap, key) delete(sharedWritingFileKeyMap, key)
if len(sharedWritingFileKeyMap) == 0 { if len(sharedWritingFileKeyMap) == 0 {
@@ -646,7 +696,7 @@ func (this *FileStorage) AddToList(item *Item) {
} }
item.MetaSize = SizeMeta + 128 item.MetaSize = SizeMeta + 128
hash := stringutil.Md5(item.Key) var hash = stringutil.Md5(item.Key)
err := this.list.Add(hash, item) err := this.list.Add(hash, item)
if err != nil && !strings.Contains(err.Error(), "UNIQUE constraint failed") { if err != nil && !strings.Contains(err.Error(), "UNIQUE constraint failed") {
remotelogs.Error("CACHE", "add to list failed: "+err.Error()) remotelogs.Error("CACHE", "add to list failed: "+err.Error())
@@ -668,7 +718,7 @@ func (this *FileStorage) Delete(key string) error {
_ = memoryStorage.Delete(key) _ = memoryStorage.Delete(key)
}) })
hash, path := this.keyPath(key) hash, path, _ := this.keyPath(key)
err := this.list.Remove(hash) err := this.list.Remove(hash)
if err != nil { if err != nil {
return err return err
@@ -708,56 +758,67 @@ func (this *FileStorage) CleanAll() error {
// 删除缓存和目录 // 删除缓存和目录
// 不能直接删除子目录,比较危险 // 不能直接删除子目录,比较危险
dir := this.dir()
fp, err := os.Open(dir)
if err != nil {
return err
}
defer func() {
_ = fp.Close()
}()
stat, err := fp.Stat() var rootDirs = []string{this.options.Dir}
if err != nil { var subDirs = this.subDirs // copy slice
return err if len(subDirs) > 0 {
for _, subDir := range subDirs {
rootDirs = append(rootDirs, subDir.Path)
}
} }
if !stat.IsDir() { for _, rootDir := range rootDirs {
return nil var dir = rootDir + "/p" + types.String(this.policy.Id)
} fp, err := os.Open(dir)
// 改成待删除
subDirs, err := fp.Readdir(-1)
if err != nil {
return err
}
for _, info := range subDirs {
subDir := info.Name()
// 检查目录名
ok, err := regexp.MatchString(`^[0-9a-f]{2}$`, subDir)
if err != nil { if err != nil {
return err return err
} }
if !ok { defer func() {
continue _ = fp.Close()
} }()
// 修改目录名 stat, err := fp.Stat()
tmpDir := dir + "/" + subDir + "-deleted"
err = os.Rename(dir+"/"+subDir, tmpDir)
if err != nil { if err != nil {
return err return err
} }
}
// 重新遍历待删除 if !stat.IsDir() {
goman.New(func() { return nil
err = this.cleanDeletedDirs(dir)
if err != nil {
remotelogs.Warn("CACHE", "delete '*-deleted' dirs failed: "+err.Error())
} }
})
// 改成待删除
subDirs, err := fp.Readdir(-1)
if err != nil {
return err
}
for _, info := range subDirs {
subDir := info.Name()
// 检查目录名
ok, err := regexp.MatchString(`^[0-9a-f]{2}$`, subDir)
if err != nil {
return err
}
if !ok {
continue
}
// 修改目录名
tmpDir := dir + "/" + subDir + "-deleted"
err = os.Rename(dir+"/"+subDir, tmpDir)
if err != nil {
return err
}
}
// 重新遍历待删除
goman.New(func() {
err = this.cleanDeletedDirs(dir)
if err != nil {
remotelogs.Warn("CACHE", "delete '*-deleted' dirs failed: "+err.Error())
}
})
}
return nil return nil
} }
@@ -780,6 +841,19 @@ func (this *FileStorage) Purge(keys []string, urlType string) error {
// 目录 // 目录
if urlType == "dir" { if urlType == "dir" {
for _, key := range keys { for _, key := range keys {
// 检查是否有通配符 http(s)://*.example.com
var schemeIndex = strings.Index(key, "://")
if schemeIndex > 0 {
var keyRight = key[schemeIndex+3:]
if strings.HasPrefix(keyRight, "*.") {
err := this.list.CleanMatchPrefix(key)
if err != nil {
return err
}
continue
}
}
err := this.list.CleanPrefix(key) err := this.list.CleanPrefix(key)
if err != nil { if err != nil {
return err return err
@@ -790,7 +864,21 @@ func (this *FileStorage) Purge(keys []string, urlType string) error {
// URL // URL
for _, key := range keys { for _, key := range keys {
hash, path := this.keyPath(key) // 检查是否有通配符 http(s)://*.example.com
var schemeIndex = strings.Index(key, "://")
if schemeIndex > 0 {
var keyRight = key[schemeIndex+3:]
if strings.HasPrefix(keyRight, "*.") {
err := this.list.CleanMatchKey(key)
if err != nil {
return err
}
continue
}
}
// 普通的Key
hash, path, _ := this.keyPath(key)
err := this.removeCacheFile(path) err := this.removeCacheFile(path)
if err != nil && !os.IsNotExist(err) { if err != nil && !os.IsNotExist(err) {
return err return err
@@ -861,25 +949,22 @@ func (this *FileStorage) CanSendfile() bool {
return this.options.EnableSendfile return this.options.EnableSendfile
} }
// 绝对路径
func (this *FileStorage) dir() string {
return this.options.Dir + "/p" + strconv.FormatInt(this.policy.Id, 10) + "/"
}
// 获取Key对应的文件路径 // 获取Key对应的文件路径
func (this *FileStorage) keyPath(key string) (hash string, path string) { func (this *FileStorage) keyPath(key string) (hash string, path string, diskIsFull bool) {
hash = stringutil.Md5(key) hash = stringutil.Md5(key)
dir := this.options.Dir + "/p" + strconv.FormatInt(this.policy.Id, 10) + "/" + hash[:2] + "/" + hash[2:4] var dir string
dir, diskIsFull = this.subDir(hash)
path = dir + "/" + hash + ".cache" path = dir + "/" + hash + ".cache"
return return
} }
// 获取Hash对应的文件路径 // 获取Hash对应的文件路径
func (this *FileStorage) hashPath(hash string) (path string) { func (this *FileStorage) hashPath(hash string) (path string, diskIsFull bool) {
if len(hash) != 32 { if len(hash) != 32 {
return "" return "", false
} }
dir := this.options.Dir + "/p" + strconv.FormatInt(this.policy.Id, 10) + "/" + hash[:2] + "/" + hash[2:4] var dir string
dir, diskIsFull = this.subDir(hash)
path = dir + "/" + hash + ".cache" path = dir + "/" + hash + ".cache"
return return
} }
@@ -937,19 +1022,39 @@ func (this *FileStorage) initList() error {
} }
// 清理任务 // 清理任务
// TODO purge每个分区
func (this *FileStorage) purgeLoop() { func (this *FileStorage) purgeLoop() {
// 检查磁盘剩余空间
this.checkDiskSpace()
// 计算是否应该开启LFU清理 // 计算是否应该开启LFU清理
var capacityBytes = this.policy.CapacityBytes() var capacityBytes = this.diskCapacityBytes()
var startLFU = false var startLFU = false
var usedPercent = float32(this.TotalDiskSize()*100) / float32(capacityBytes)
var lfuFreePercent = this.policy.PersistenceLFUFreePercent var lfuFreePercent = this.policy.PersistenceLFUFreePercent
if lfuFreePercent <= 0 { if lfuFreePercent <= 0 {
lfuFreePercent = 5 lfuFreePercent = 5
} }
if capacityBytes > 0 {
if lfuFreePercent < 100 { var hasFullDisk = this.mainDiskIsFull
if usedPercent >= 100-lfuFreePercent { if !hasFullDisk {
startLFU = true var subDirs = this.subDirs // copy slice
for _, subDir := range subDirs {
if subDir.IsFull {
hasFullDisk = true
break
}
}
}
if hasFullDisk {
startLFU = true
} else {
var usedPercent = float32(this.TotalDiskSize()*100) / float32(capacityBytes)
if capacityBytes > 0 {
if lfuFreePercent < 100 {
if usedPercent >= 100-lfuFreePercent {
startLFU = true
}
} }
} }
} }
@@ -974,7 +1079,7 @@ func (this *FileStorage) purgeLoop() {
} }
for i := 0; i < times; i++ { for i := 0; i < times; i++ {
countFound, err := this.list.Purge(purgeCount, func(hash string) error { countFound, err := this.list.Purge(purgeCount, func(hash string) error {
path := this.hashPath(hash) path, _ := this.hashPath(hash)
err := this.removeCacheFile(path) err := this.removeCacheFile(path)
if err != nil && !os.IsNotExist(err) { if err != nil && !os.IsNotExist(err) {
remotelogs.Error("CACHE", "purge '"+path+"' error: "+err.Error()) remotelogs.Error("CACHE", "purge '"+path+"' error: "+err.Error())
@@ -1008,7 +1113,7 @@ func (this *FileStorage) purgeLoop() {
remotelogs.Println("CACHE", "LFU purge policy '"+this.policy.Name+"' id: "+types.String(this.policy.Id)+", count: "+types.String(count)) remotelogs.Println("CACHE", "LFU purge policy '"+this.policy.Name+"' id: "+types.String(this.policy.Id)+", count: "+types.String(count))
err := this.list.PurgeLFU(count, func(hash string) error { err := this.list.PurgeLFU(count, func(hash string) error {
path := this.hashPath(hash) path, _ := this.hashPath(hash)
err := this.removeCacheFile(path) err := this.removeCacheFile(path)
if err != nil && !os.IsNotExist(err) { if err != nil && !os.IsNotExist(err) {
remotelogs.Error("CACHE", "purge '"+path+"' error: "+err.Error()) remotelogs.Error("CACHE", "purge '"+path+"' error: "+err.Error())
@@ -1089,7 +1194,7 @@ func (this *FileStorage) hotLoop() {
expiresAt = bestExpiresAt expiresAt = bestExpiresAt
} }
writer, err := memoryStorage.openWriter(item.Key, expiresAt, reader.Status(), reader.BodySize(), -1, false) writer, err := memoryStorage.openWriter(item.Key, expiresAt, reader.Status(), types.Int(reader.HeaderSize()), reader.BodySize(), -1, false)
if err != nil { if err != nil {
if !CanIgnoreErr(err) { if !CanIgnoreErr(err) {
remotelogs.Error("CACHE", "transfer hot item failed: "+err.Error()) remotelogs.Error("CACHE", "transfer hot item failed: "+err.Error())
@@ -1128,6 +1233,7 @@ func (this *FileStorage) hotLoop() {
memoryStorage.AddToList(&Item{ memoryStorage.AddToList(&Item{
Type: writer.ItemType(), Type: writer.ItemType(),
Key: item.Key, Key: item.Key,
Host: ParseHost(item.Key),
ExpiredAt: expiresAt, ExpiredAt: expiresAt,
HeaderSize: writer.HeaderSize(), HeaderSize: writer.HeaderSize(),
BodySize: writer.BodySize(), BodySize: writer.BodySize(),
@@ -1327,3 +1433,63 @@ func (this *FileStorage) runMemoryStorageSafety(f func(memoryStorage *MemoryStor
f(memoryStorage) f(memoryStorage)
} }
} }
// 检查磁盘剩余空间
func (this *FileStorage) checkDiskSpace() {
if this.options != nil && len(this.options.Dir) > 0 {
var stat unix.Statfs_t
err := unix.Statfs(this.options.Dir, &stat)
if err == nil {
var availableBytes = stat.Bavail * uint64(stat.Bsize)
this.mainDiskIsFull = availableBytes < MinDiskSpace
}
}
var subDirs = this.subDirs // copy slice
for _, subDir := range subDirs {
var stat unix.Statfs_t
err := unix.Statfs(subDir.Path, &stat)
if err == nil {
var availableBytes = stat.Bavail * uint64(stat.Bsize)
subDir.IsFull = availableBytes < MinDiskSpace
}
}
}
// 获取目录
func (this *FileStorage) subDir(hash string) (dirPath string, dirIsFull bool) {
var suffix = "/p" + types.String(this.policy.Id) + "/" + hash[:2] + "/" + hash[2:4]
if len(hash) < 4 {
return this.options.Dir + suffix, this.mainDiskIsFull
}
var subDirs = this.subDirs // copy slice
var countSubDirs = len(subDirs)
if countSubDirs == 0 {
return this.options.Dir + suffix, this.mainDiskIsFull
}
countSubDirs++ // add main dir
// 最多只支持16个目录
if countSubDirs > 16 {
countSubDirs = 16
}
var dirIndex = this.charCode(hash[0]) % uint8(countSubDirs)
if dirIndex == 0 {
return this.options.Dir + suffix, this.mainDiskIsFull
}
var subDir = subDirs[dirIndex-1]
return subDir.Path + suffix, subDir.IsFull
}
func (this *FileStorage) charCode(r byte) uint8 {
if r >= '0' && r <= '9' {
return r - '0'
}
if r >= 'a' && r <= 'z' {
return r - 'a' + 10
}
return 0
}

View File

@@ -62,7 +62,7 @@ func TestFileStorage_OpenWriter(t *testing.T) {
header := []byte("Header") header := []byte("Header")
body := []byte("This is Body") body := []byte("This is Body")
writer, err := storage.OpenWriter("my-key", time.Now().Unix()+86400, 200, -1, -1, false) writer, err := storage.OpenWriter("my-key", time.Now().Unix()+86400, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -100,7 +100,7 @@ func TestFileStorage_OpenWriter_Partial(t *testing.T) {
t.Fatal(err) t.Fatal(err)
} }
writer, err := storage.OpenWriter("my-key", time.Now().Unix()+86400, 200, -1, -1, true) writer, err := storage.OpenWriter("my-key", time.Now().Unix()+86400, 200, -1, -1, -1, true)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -139,7 +139,7 @@ func TestFileStorage_OpenWriter_HTTP(t *testing.T) {
t.Log(time.Since(now).Seconds()*1000, "ms") t.Log(time.Since(now).Seconds()*1000, "ms")
}() }()
writer, err := storage.OpenWriter("my-http-response", time.Now().Unix()+86400, 200, -1, -1, false) writer, err := storage.OpenWriter("my-http-response", time.Now().Unix()+86400, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -212,7 +212,7 @@ func TestFileStorage_Concurrent_Open_DifferentFile(t *testing.T) {
go func(i int) { go func(i int) {
defer wg.Done() defer wg.Done()
writer, err := storage.OpenWriter("abc"+strconv.Itoa(i), time.Now().Unix()+3600, 200, -1, -1, false) writer, err := storage.OpenWriter("abc"+strconv.Itoa(i), time.Now().Unix()+3600, 200, -1, -1, -1, false)
if err != nil { if err != nil {
if err != ErrFileIsWriting { if err != ErrFileIsWriting {
t.Error(err) t.Error(err)
@@ -267,7 +267,7 @@ func TestFileStorage_Concurrent_Open_SameFile(t *testing.T) {
go func(i int) { go func(i int) {
defer wg.Done() defer wg.Done()
writer, err := storage.OpenWriter("abc"+strconv.Itoa(0), time.Now().Unix()+3600, 200, -1, -1, false) writer, err := storage.OpenWriter("abc"+strconv.Itoa(0), time.Now().Unix()+3600, 200, -1, -1, -1, false)
if err != nil { if err != nil {
if err != ErrFileIsWriting { if err != ErrFileIsWriting {
t.Error(err) t.Error(err)
@@ -522,7 +522,7 @@ func TestFileStorage_DecodeFile(t *testing.T) {
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
_, path := storage.keyPath("my-key") _, path, _ := storage.keyPath("my-key")
t.Log(path) t.Log(path)
} }
@@ -569,6 +569,6 @@ func BenchmarkFileStorage_KeyPath(b *testing.B) {
} }
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
_, _ = storage.keyPath(strconv.Itoa(i)) _, _, _ = storage.keyPath(strconv.Itoa(i))
} }
} }

View File

@@ -14,10 +14,10 @@ type StorageInterface interface {
// OpenWriter 打开缓存写入器等待写入 // OpenWriter 打开缓存写入器等待写入
// size 和 maxSize 可能为-1 // size 和 maxSize 可能为-1
OpenWriter(key string, expiresAt int64, status int, size int64, maxSize int64, isPartial bool) (Writer, error) OpenWriter(key string, expiresAt int64, status int, headerSize int, bodySize int64, maxSize int64, isPartial bool) (Writer, error)
// OpenFlushWriter 打开从其他媒介直接刷入的写入器 // OpenFlushWriter 打开从其他媒介直接刷入的写入器
OpenFlushWriter(key string, expiresAt int64, status int) (Writer, error) OpenFlushWriter(key string, expiresAt int64, status int, headerSize int, bodySize int64) (Writer, error)
// Delete 删除某个键值对应的缓存 // Delete 删除某个键值对应的缓存
Delete(key string) error Delete(key string) error

View File

@@ -1,7 +1,6 @@
package caches package caches
import ( import (
"fmt"
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs" "github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
"github.com/TeaOSLab/EdgeNode/internal/goman" "github.com/TeaOSLab/EdgeNode/internal/goman"
"github.com/TeaOSLab/EdgeNode/internal/remotelogs" "github.com/TeaOSLab/EdgeNode/internal/remotelogs"
@@ -17,6 +16,7 @@ import (
"math" "math"
"runtime" "runtime"
"strconv" "strconv"
"strings"
"sync" "sync"
"sync/atomic" "sync/atomic"
"time" "time"
@@ -149,7 +149,7 @@ func (this *MemoryStorage) OpenReader(key string, useStale bool, isPartial bool)
} }
// OpenWriter 打开缓存写入器等待写入 // OpenWriter 打开缓存写入器等待写入
func (this *MemoryStorage) OpenWriter(key string, expiredAt int64, status int, size int64, maxSize int64, isPartial bool) (Writer, error) { func (this *MemoryStorage) OpenWriter(key string, expiredAt int64, status int, headerSize int, bodySize int64, maxSize int64, isPartial bool) (Writer, error) {
if this.ignoreKeys.Has(key) { if this.ignoreKeys.Has(key) {
return nil, ErrEntityTooLarge return nil, ErrEntityTooLarge
} }
@@ -158,15 +158,15 @@ func (this *MemoryStorage) OpenWriter(key string, expiredAt int64, status int, s
if isPartial { if isPartial {
return nil, ErrFileIsWriting return nil, ErrFileIsWriting
} }
return this.openWriter(key, expiredAt, status, size, maxSize, true) return this.openWriter(key, expiredAt, status, headerSize, bodySize, maxSize, true)
} }
// OpenFlushWriter 打开从其他媒介直接刷入的写入器 // OpenFlushWriter 打开从其他媒介直接刷入的写入器
func (this *MemoryStorage) OpenFlushWriter(key string, expiresAt int64, status int) (Writer, error) { func (this *MemoryStorage) OpenFlushWriter(key string, expiresAt int64, status int, headerSize int, bodySize int64) (Writer, error) {
return this.openWriter(key, expiresAt, status, -1, -1, true) return this.openWriter(key, expiresAt, status, headerSize, bodySize, -1, true)
} }
func (this *MemoryStorage) openWriter(key string, expiresAt int64, status int, size int64, maxSize int64, isDirty bool) (Writer, error) { func (this *MemoryStorage) openWriter(key string, expiresAt int64, status int, headerSize int, bodySize int64, maxSize int64, isDirty bool) (Writer, error) {
// 待写入队列是否已满 // 待写入队列是否已满
if isDirty && if isDirty &&
this.parentStorage != nil && this.parentStorage != nil &&
@@ -207,10 +207,10 @@ func (this *MemoryStorage) openWriter(key string, expiresAt int64, status int, s
return nil, NewCapacityError("write memory cache failed: too many keys in cache storage") return nil, NewCapacityError("write memory cache failed: too many keys in cache storage")
} }
capacityBytes := this.memoryCapacityBytes() capacityBytes := this.memoryCapacityBytes()
if size < 0 { if bodySize < 0 {
size = 0 bodySize = 0
} }
if capacityBytes > 0 && capacityBytes <= this.totalSize+size { if capacityBytes > 0 && capacityBytes <= this.totalSize+bodySize {
return nil, NewCapacityError("write memory cache failed: over memory size: " + strconv.FormatInt(capacityBytes, 10) + ", current size: " + strconv.FormatInt(this.totalSize, 10) + " bytes") return nil, NewCapacityError("write memory cache failed: over memory size: " + strconv.FormatInt(capacityBytes, 10) + ", current size: " + strconv.FormatInt(this.totalSize, 10) + " bytes")
} }
@@ -230,10 +230,10 @@ func (this *MemoryStorage) openWriter(key string, expiresAt int64, status int, s
// Delete 删除某个键值对应的缓存 // Delete 删除某个键值对应的缓存
func (this *MemoryStorage) Delete(key string) error { func (this *MemoryStorage) Delete(key string) error {
hash := this.hash(key) var hash = this.hash(key)
this.locker.Lock() this.locker.Lock()
delete(this.valuesMap, hash) delete(this.valuesMap, hash)
_ = this.list.Remove(fmt.Sprintf("%d", hash)) _ = this.list.Remove(types.String(hash))
this.locker.Unlock() this.locker.Unlock()
return nil return nil
} }
@@ -263,6 +263,19 @@ func (this *MemoryStorage) Purge(keys []string, urlType string) error {
// 目录 // 目录
if urlType == "dir" { if urlType == "dir" {
for _, key := range keys { for _, key := range keys {
// 检查是否有通配符 http(s)://*.example.com
var schemeIndex = strings.Index(key, "://")
if schemeIndex > 0 {
var keyRight = key[schemeIndex+3:]
if strings.HasPrefix(keyRight, "*.") {
err := this.list.CleanMatchPrefix(key)
if err != nil {
return err
}
continue
}
}
err := this.list.CleanPrefix(key) err := this.list.CleanPrefix(key)
if err != nil { if err != nil {
return err return err
@@ -273,6 +286,19 @@ func (this *MemoryStorage) Purge(keys []string, urlType string) error {
// URL // URL
for _, key := range keys { for _, key := range keys {
// 检查是否有通配符 http(s)://*.example.com
var schemeIndex = strings.Index(key, "://")
if schemeIndex > 0 {
var keyRight = key[schemeIndex+3:]
if strings.HasPrefix(keyRight, "*.") {
err := this.list.CleanMatchKey(key)
if err != nil {
return err
}
continue
}
}
err := this.Delete(key) err := this.Delete(key)
if err != nil { if err != nil {
return err return err
@@ -336,7 +362,12 @@ func (this *MemoryStorage) CanUpdatePolicy(newPolicy *serverconfigs.HTTPCachePol
// AddToList 将缓存添加到列表 // AddToList 将缓存添加到列表
func (this *MemoryStorage) AddToList(item *Item) { func (this *MemoryStorage) AddToList(item *Item) {
item.MetaSize = int64(len(item.Key)) + 128 /** 128是我们评估的数据结构的长度 **/ item.MetaSize = int64(len(item.Key)) + 128 /** 128是我们评估的数据结构的长度 **/
hash := fmt.Sprintf("%d", this.hash(item.Key)) var hash = types.String(this.hash(item.Key))
if len(item.Host) == 0 {
item.Host = ParseHost(item.Key)
}
_ = this.list.Add(hash, item) _ = this.list.Add(hash, item)
} }
@@ -433,7 +464,7 @@ func (this *MemoryStorage) startFlush() {
var statCount = 0 var statCount = 0
var writeDelayMS float64 = 0 var writeDelayMS float64 = 0
for hash := range this.dirtyChan { for key := range this.dirtyChan {
statCount++ statCount++
if statCount == 100 { if statCount == 100 {
@@ -455,7 +486,7 @@ func (this *MemoryStorage) startFlush() {
} }
} }
this.flushItem(hash) this.flushItem(key)
if writeDelayMS > 0 { if writeDelayMS > 0 {
time.Sleep(time.Duration(writeDelayMS) * time.Millisecond) time.Sleep(time.Duration(writeDelayMS) * time.Millisecond)
@@ -481,7 +512,7 @@ func (this *MemoryStorage) flushItem(key string) {
return return
} }
writer, err := this.parentStorage.OpenFlushWriter(key, item.ExpiresAt, item.Status) writer, err := this.parentStorage.OpenFlushWriter(key, item.ExpiresAt, item.Status, len(item.HeaderValue), int64(len(item.BodyValue)))
if err != nil { if err != nil {
if !CanIgnoreErr(err) { if !CanIgnoreErr(err) {
remotelogs.Error("CACHE", "flush items failed: open writer failed: "+err.Error()) remotelogs.Error("CACHE", "flush items failed: open writer failed: "+err.Error())
@@ -513,6 +544,7 @@ func (this *MemoryStorage) flushItem(key string) {
this.parentStorage.AddToList(&Item{ this.parentStorage.AddToList(&Item{
Type: writer.ItemType(), Type: writer.ItemType(),
Key: key, Key: key,
Host: ParseHost(key),
ExpiredAt: item.ExpiresAt, ExpiredAt: item.ExpiresAt,
HeaderSize: writer.HeaderSize(), HeaderSize: writer.HeaderSize(),
BodySize: writer.BodySize(), BodySize: writer.BodySize(),
@@ -542,7 +574,7 @@ func (this *MemoryStorage) memoryCapacityBytes() int64 {
func (this *MemoryStorage) deleteWithoutLocker(key string) error { func (this *MemoryStorage) deleteWithoutLocker(key string) error {
hash := this.hash(key) hash := this.hash(key)
delete(this.valuesMap, hash) delete(this.valuesMap, hash)
_ = this.list.Remove(fmt.Sprintf("%d", hash)) _ = this.list.Remove(types.String(hash))
return nil return nil
} }

View File

@@ -16,7 +16,7 @@ import (
func TestMemoryStorage_OpenWriter(t *testing.T) { func TestMemoryStorage_OpenWriter(t *testing.T) {
storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil) storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil)
writer, err := storage.OpenWriter("abc", time.Now().Unix()+60, 200, -1, -1, false) writer, err := storage.OpenWriter("abc", time.Now().Unix()+60, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -63,7 +63,7 @@ func TestMemoryStorage_OpenWriter(t *testing.T) {
} }
} }
writer, err = storage.OpenWriter("abc", time.Now().Unix()+60, 200, -1, -1, false) writer, err = storage.OpenWriter("abc", time.Now().Unix()+60, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -104,7 +104,7 @@ func TestMemoryStorage_OpenReaderLock(t *testing.T) {
func TestMemoryStorage_Delete(t *testing.T) { func TestMemoryStorage_Delete(t *testing.T) {
storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil) storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil)
{ {
writer, err := storage.OpenWriter("abc", time.Now().Unix()+60, 200, -1, -1, false) writer, err := storage.OpenWriter("abc", time.Now().Unix()+60, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -112,7 +112,7 @@ func TestMemoryStorage_Delete(t *testing.T) {
t.Log(len(storage.valuesMap)) t.Log(len(storage.valuesMap))
} }
{ {
writer, err := storage.OpenWriter("abc1", time.Now().Unix()+60, 200, -1, -1, false) writer, err := storage.OpenWriter("abc1", time.Now().Unix()+60, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -127,7 +127,7 @@ func TestMemoryStorage_Stat(t *testing.T) {
storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil) storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil)
expiredAt := time.Now().Unix() + 60 expiredAt := time.Now().Unix() + 60
{ {
writer, err := storage.OpenWriter("abc", expiredAt, 200, -1, -1, false) writer, err := storage.OpenWriter("abc", expiredAt, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -140,7 +140,7 @@ func TestMemoryStorage_Stat(t *testing.T) {
}) })
} }
{ {
writer, err := storage.OpenWriter("abc1", expiredAt, 200, -1, -1, false) writer, err := storage.OpenWriter("abc1", expiredAt, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -164,7 +164,7 @@ func TestMemoryStorage_CleanAll(t *testing.T) {
storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil) storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil)
expiredAt := time.Now().Unix() + 60 expiredAt := time.Now().Unix() + 60
{ {
writer, err := storage.OpenWriter("abc", expiredAt, 200, -1, -1, false) writer, err := storage.OpenWriter("abc", expiredAt, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -176,7 +176,7 @@ func TestMemoryStorage_CleanAll(t *testing.T) {
}) })
} }
{ {
writer, err := storage.OpenWriter("abc1", expiredAt, 200, -1, -1, false) writer, err := storage.OpenWriter("abc1", expiredAt, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -199,7 +199,7 @@ func TestMemoryStorage_Purge(t *testing.T) {
storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil) storage := NewMemoryStorage(&serverconfigs.HTTPCachePolicy{}, nil)
expiredAt := time.Now().Unix() + 60 expiredAt := time.Now().Unix() + 60
{ {
writer, err := storage.OpenWriter("abc", expiredAt, 200, -1, -1, false) writer, err := storage.OpenWriter("abc", expiredAt, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -211,7 +211,7 @@ func TestMemoryStorage_Purge(t *testing.T) {
}) })
} }
{ {
writer, err := storage.OpenWriter("abc1", expiredAt, 200, -1, -1, false) writer, err := storage.OpenWriter("abc1", expiredAt, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -242,7 +242,7 @@ func TestMemoryStorage_Expire(t *testing.T) {
for i := 0; i < 1000; i++ { for i := 0; i < 1000; i++ {
expiredAt := time.Now().Unix() + int64(rands.Int(0, 60)) expiredAt := time.Now().Unix() + int64(rands.Int(0, 60))
key := "abc" + strconv.Itoa(i) key := "abc" + strconv.Itoa(i)
writer, err := storage.OpenWriter(key, expiredAt, 200, -1, -1, false) writer, err := storage.OpenWriter(key, expiredAt, 200, -1, -1, -1, false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }

30
internal/caches/utils.go Normal file
View File

@@ -0,0 +1,30 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package caches
import (
"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
"net"
"strings"
)
func ParseHost(key string) string {
var schemeIndex = strings.Index(key, "://")
if schemeIndex <= 0 {
return ""
}
var firstSlashIndex = strings.Index(key[schemeIndex+3:], "/")
if firstSlashIndex <= 0 {
return ""
}
var host = key[schemeIndex+3 : schemeIndex+3+firstSlashIndex]
hostPart, _, err := net.SplitHostPort(host)
if err == nil && len(hostPart) > 0 {
host = configutils.QuoteIP(hostPart)
}
return host
}

View File

@@ -0,0 +1,51 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package caches_test
import (
"fmt"
"github.com/TeaOSLab/EdgeNode/internal/caches"
"github.com/cespare/xxhash"
"github.com/iwind/TeaGo/types"
"strconv"
"testing"
)
func TestParseHost(t *testing.T) {
for _, u := range []string{
"https://goedge.cn/hello/world",
"https://goedge.cn:8080/hello/world",
"https://goedge.cn/hello/world?v=1&t=123",
"https://[::1]:1234/hello/world?v=1&t=123",
"https://[::1]/hello/world?v=1&t=123",
"https://127.0.0.1/hello/world?v=1&t=123",
"https:/hello/world?v=1&t=123",
"123456",
} {
t.Log(u, "=>", caches.ParseHost(u))
}
}
func TestUintString(t *testing.T) {
t.Log(strconv.FormatUint(xxhash.Sum64String("https://goedge.cn/"), 10))
t.Log(strconv.FormatUint(123456789, 10))
t.Log(fmt.Sprintf("%d", 1234567890123))
}
func BenchmarkUint_String(b *testing.B) {
for i := 0; i < b.N; i++ {
_ = strconv.FormatUint(1234567890123, 10)
}
}
func BenchmarkUint_String2(b *testing.B) {
for i := 0; i < b.N; i++ {
_ = types.String(1234567890123)
}
}
func BenchmarkUint_String3(b *testing.B) {
for i := 0; i < b.N; i++ {
_ = fmt.Sprintf("%d", 1234567890123)
}
}

View File

@@ -11,25 +11,32 @@ import (
) )
type FileWriter struct { type FileWriter struct {
storage StorageInterface storage StorageInterface
rawWriter *os.File rawWriter *os.File
key string key string
headerSize int64
bodySize int64 metaHeaderSize int
expiredAt int64 headerSize int64
maxSize int64
endFunc func() metaBodySize int64 // 写入前的内容长度
once sync.Once bodySize int64
expiredAt int64
maxSize int64
endFunc func()
once sync.Once
} }
func NewFileWriter(storage StorageInterface, rawWriter *os.File, key string, expiredAt int64, maxSize int64, endFunc func()) *FileWriter { func NewFileWriter(storage StorageInterface, rawWriter *os.File, key string, expiredAt int64, metaHeaderSize int, metaBodySize int64, maxSize int64, endFunc func()) *FileWriter {
return &FileWriter{ return &FileWriter{
storage: storage, storage: storage,
key: key, key: key,
rawWriter: rawWriter, rawWriter: rawWriter,
expiredAt: expiredAt, expiredAt: expiredAt,
maxSize: maxSize, maxSize: maxSize,
endFunc: endFunc, endFunc: endFunc,
metaHeaderSize: metaHeaderSize,
metaBodySize: metaBodySize,
} }
} }
@@ -45,7 +52,10 @@ func (this *FileWriter) WriteHeader(data []byte) (n int, err error) {
// WriteHeaderLength 写入Header长度数据 // WriteHeaderLength 写入Header长度数据
func (this *FileWriter) WriteHeaderLength(headerLength int) error { func (this *FileWriter) WriteHeaderLength(headerLength int) error {
bytes4 := make([]byte, 4) if this.metaHeaderSize > 0 && this.metaHeaderSize == headerLength {
return nil
}
var bytes4 = make([]byte, 4)
binary.BigEndian.PutUint32(bytes4, uint32(headerLength)) binary.BigEndian.PutUint32(bytes4, uint32(headerLength))
_, err := this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength, io.SeekStart) _, err := this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength, io.SeekStart)
if err != nil { if err != nil {
@@ -88,7 +98,10 @@ func (this *FileWriter) WriteAt(offset int64, data []byte) error {
// WriteBodyLength 写入Body长度数据 // WriteBodyLength 写入Body长度数据
func (this *FileWriter) WriteBodyLength(bodyLength int64) error { func (this *FileWriter) WriteBodyLength(bodyLength int64) error {
bytes8 := make([]byte, 8) if this.metaBodySize >= 0 && bodyLength == this.metaBodySize {
return nil
}
var bytes8 = make([]byte, 8)
binary.BigEndian.PutUint64(bytes8, uint64(bodyLength)) binary.BigEndian.PutUint64(bytes8, uint64(bodyLength))
_, err := this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength+SizeHeaderLength, io.SeekStart) _, err := this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength+SizeHeaderLength, io.SeekStart)
if err != nil { if err != nil {
@@ -109,7 +122,7 @@ func (this *FileWriter) Close() error {
this.endFunc() this.endFunc()
}) })
path := this.rawWriter.Name() var path = this.rawWriter.Name()
err := this.WriteHeaderLength(types.Int(this.headerSize)) err := this.WriteHeaderLength(types.Int(this.headerSize))
if err != nil { if err != nil {

View File

@@ -11,13 +11,18 @@ import (
) )
type PartialFileWriter struct { type PartialFileWriter struct {
rawWriter *os.File rawWriter *os.File
key string key string
headerSize int64
bodySize int64 metaHeaderSize int
expiredAt int64 headerSize int64
endFunc func()
once sync.Once metaBodySize int64
bodySize int64
expiredAt int64
endFunc func()
once sync.Once
isNew bool isNew bool
isPartial bool isPartial bool
@@ -27,17 +32,19 @@ type PartialFileWriter struct {
rangePath string rangePath string
} }
func NewPartialFileWriter(rawWriter *os.File, key string, expiredAt int64, isNew bool, isPartial bool, bodyOffset int64, ranges *PartialRanges, endFunc func()) *PartialFileWriter { func NewPartialFileWriter(rawWriter *os.File, key string, expiredAt int64, metaHeaderSize int, metaBodySize int64, isNew bool, isPartial bool, bodyOffset int64, ranges *PartialRanges, endFunc func()) *PartialFileWriter {
return &PartialFileWriter{ return &PartialFileWriter{
key: key, key: key,
rawWriter: rawWriter, rawWriter: rawWriter,
expiredAt: expiredAt, expiredAt: expiredAt,
endFunc: endFunc, endFunc: endFunc,
isNew: isNew, isNew: isNew,
isPartial: isPartial, isPartial: isPartial,
bodyOffset: bodyOffset, bodyOffset: bodyOffset,
ranges: ranges, ranges: ranges,
rangePath: partialRangesFilePath(rawWriter.Name()), rangePath: partialRangesFilePath(rawWriter.Name()),
metaHeaderSize: metaHeaderSize,
metaBodySize: metaBodySize,
} }
} }
@@ -71,7 +78,11 @@ func (this *PartialFileWriter) AppendHeader(data []byte) error {
// WriteHeaderLength 写入Header长度数据 // WriteHeaderLength 写入Header长度数据
func (this *PartialFileWriter) WriteHeaderLength(headerLength int) error { func (this *PartialFileWriter) WriteHeaderLength(headerLength int) error {
bytes4 := make([]byte, 4) if this.metaHeaderSize > 0 && this.metaHeaderSize == headerLength {
return nil
}
var bytes4 = make([]byte, 4)
binary.BigEndian.PutUint32(bytes4, uint32(headerLength)) binary.BigEndian.PutUint32(bytes4, uint32(headerLength))
_, err := this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength, io.SeekStart) _, err := this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength, io.SeekStart)
if err != nil { if err != nil {
@@ -110,8 +121,13 @@ func (this *PartialFileWriter) WriteAt(offset int64, data []byte) error {
} }
if this.bodyOffset == 0 { if this.bodyOffset == 0 {
this.bodyOffset = SizeMeta + int64(len(this.key)) + this.headerSize var keyLength = 0
if this.ranges.Version == 0 { // 以往的版本包含有Key
keyLength = len(this.key)
}
this.bodyOffset = SizeMeta + int64(keyLength) + this.headerSize
} }
_, err := this.rawWriter.WriteAt(data, this.bodyOffset+offset) _, err := this.rawWriter.WriteAt(data, this.bodyOffset+offset)
if err != nil { if err != nil {
return err return err
@@ -129,7 +145,10 @@ func (this *PartialFileWriter) SetBodyLength(bodyLength int64) {
// WriteBodyLength 写入Body长度数据 // WriteBodyLength 写入Body长度数据
func (this *PartialFileWriter) WriteBodyLength(bodyLength int64) error { func (this *PartialFileWriter) WriteBodyLength(bodyLength int64) error {
bytes8 := make([]byte, 8) if this.metaBodySize > 0 && this.metaBodySize == bodyLength {
return nil
}
var bytes8 = make([]byte, 8)
binary.BigEndian.PutUint64(bytes8, uint64(bodyLength)) binary.BigEndian.PutUint64(bytes8, uint64(bodyLength))
_, err := this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength+SizeHeaderLength, io.SeekStart) _, err := this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength+SizeHeaderLength, io.SeekStart)
if err != nil { if err != nil {
@@ -150,8 +169,11 @@ func (this *PartialFileWriter) Close() error {
this.endFunc() this.endFunc()
}) })
this.ranges.BodySize = this.bodySize
err := this.ranges.WriteToFile(this.rangePath) err := this.ranges.WriteToFile(this.rangePath)
if err != nil { if err != nil {
_ = this.rawWriter.Close()
this.remove()
return err return err
} }

View File

@@ -26,8 +26,8 @@ func TestPartialFileWriter_Write(t *testing.T) {
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
var ranges = caches.NewPartialRanges() var ranges = caches.NewPartialRanges(0)
var writer = caches.NewPartialFileWriter(fp, "test", time.Now().Unix()+86500, true, true, 0, ranges, func() { var writer = caches.NewPartialFileWriter(fp, "test", time.Now().Unix()+86500, -1, -1, true, true, 0, ranges, func() {
t.Log("end") t.Log("end")
}) })
_, err = writer.WriteHeader([]byte("header")) _, err = writer.WriteHeader([]byte("header"))

View File

@@ -9,11 +9,15 @@ import (
// APIConfig 节点API配置 // APIConfig 节点API配置
type APIConfig struct { type APIConfig struct {
RPC struct { RPC struct {
Endpoints []string `yaml:"endpoints"` Endpoints []string `yaml:"endpoints" json:"endpoints"`
DisableUpdate bool `yaml:"disableUpdate"` DisableUpdate bool `yaml:"disableUpdate" json:"disableUpdate"`
} `yaml:"rpc"` } `yaml:"rpc" json:"rpc"`
NodeId string `yaml:"nodeId"` NodeId string `yaml:"nodeId" json:"nodeId"`
Secret string `yaml:"secret"` Secret string `yaml:"secret" json:"secret"`
}
func NewAPIConfig() *APIConfig {
return &APIConfig{}
} }
func LoadAPIConfig() (*APIConfig, error) { func LoadAPIConfig() (*APIConfig, error) {

View File

@@ -3,9 +3,9 @@ package configs
// ClusterConfig 集群配置 // ClusterConfig 集群配置
type ClusterConfig struct { type ClusterConfig struct {
RPC struct { RPC struct {
Endpoints []string `yaml:"endpoints"` Endpoints []string `yaml:"endpoints" json:"endpoints"`
DisableUpdate bool `yaml:"disableUpdate"` DisableUpdate bool `yaml:"disableUpdate" json:"disableUpdate"`
} `yaml:"rpc"` } `yaml:"rpc" json:"rpc"`
ClusterId string `yaml:"clusterId"` ClusterId string `yaml:"clusterId" json:"clusterId"`
Secret string `yaml:"secret"` Secret string `yaml:"secret" json:"secret"`
} }

View File

@@ -1,7 +1,7 @@
package teaconst package teaconst
const ( const (
Version = "0.5.3" Version = "0.5.8"
ProductName = "Edge Node" ProductName = "Edge Node"
ProcessName = "edge-node" ProcessName = "edge-node"

View File

@@ -70,7 +70,7 @@ func (this *DDoSProtectionManager) Apply(config *ddosconfigs.ProtectionConfig) e
nodeConfig, _ := nodeconfigs.SharedNodeConfig() nodeConfig, _ := nodeconfigs.SharedNodeConfig()
if nodeConfig != nil { if nodeConfig != nil {
var allowIPList = nodeConfig.AllowedIPs var allowIPList = nodeConfig.AllowedIPs
if !utils.ContainsSameStrings(allowIPList, this.lastAllowIPList) { if !utils.EqualStrings(allowIPList, this.lastAllowIPList) {
allowIPListChanged = true allowIPListChanged = true
this.lastAllowIPList = allowIPList this.lastAllowIPList = allowIPList
} }
@@ -91,6 +91,9 @@ func (this *DDoSProtectionManager) Apply(config *ddosconfigs.ProtectionConfig) e
} }
if nftablesInstance == nil { if nftablesInstance == nil {
if config == nil || !config.IsOn() {
return nil
}
return errors.New("nftables instance should not be nil") return errors.New("nftables instance should not be nil")
} }

View File

@@ -4,6 +4,7 @@ import (
"errors" "errors"
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb" "github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs" "github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
"github.com/TeaOSLab/EdgeNode/internal/utils"
executils "github.com/TeaOSLab/EdgeNode/internal/utils/exec" executils "github.com/TeaOSLab/EdgeNode/internal/utils/exec"
"os/exec" "os/exec"
"runtime" "runtime"
@@ -74,10 +75,16 @@ func (this *IPTablesAction) runAction(action string, listType IPListType, item *
} }
func (this *IPTablesAction) runActionSingleIP(action string, listType IPListType, item *pb.IPItem) error { func (this *IPTablesAction) runActionSingleIP(action string, listType IPListType, item *pb.IPItem) error {
// 暂时不支持ipv6
// TODO 将来支持ipv6
if utils.IsIPv6(item.IpFrom) {
return nil
}
if item.Type == "all" { if item.Type == "all" {
return nil return nil
} }
path := this.config.Path var path = this.config.Path
var err error var err error
if len(path) == 0 { if len(path) == 0 {
path, err = exec.LookPath("iptables") path, err = exec.LookPath("iptables")
@@ -88,6 +95,7 @@ func (this *IPTablesAction) runActionSingleIP(action string, listType IPListType
this.iptablesNotFound = true this.iptablesNotFound = true
return err return err
} }
this.config.Path = path
} }
iptablesAction := "" iptablesAction := ""
switch action { switch action {

View File

@@ -1,6 +1,7 @@
package iplibrary package iplibrary
import ( import (
"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb" "github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
teaconst "github.com/TeaOSLab/EdgeNode/internal/const" teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
"github.com/TeaOSLab/EdgeNode/internal/events" "github.com/TeaOSLab/EdgeNode/internal/events"
@@ -141,6 +142,12 @@ func (this *IPListManager) init() {
} }
func (this *IPListManager) loop() error { func (this *IPListManager) loop() error {
// 是否同步IP名单
nodeConfig, _ := nodeconfigs.SharedNodeConfig()
if nodeConfig != nil && !nodeConfig.EnableIPLists {
return nil
}
for { for {
hasNext, err := this.fetch() hasNext, err := this.fetch()
if err != nil { if err != nil {

View File

@@ -178,7 +178,7 @@ func (this *APIStream) handleWriteCache(message *pb.NodeStreamMessage) error {
} }
expiredAt := time.Now().Unix() + msg.LifeSeconds expiredAt := time.Now().Unix() + msg.LifeSeconds
writer, err := storage.OpenWriter(msg.Key, expiredAt, 200, int64(len(msg.Value)), -1, false) writer, err := storage.OpenWriter(msg.Key, expiredAt, 200, -1, int64(len(msg.Value)), -1, false)
if err != nil { if err != nil {
this.replyFail(message.RequestId, "prepare writing failed: "+err.Error()) this.replyFail(message.RequestId, "prepare writing failed: "+err.Error())
return err return err
@@ -407,7 +407,7 @@ func (this *APIStream) handleCheckLocalFirewall(message *pb.NodeStreamMessage) e
var protectionConfig = sharedNodeConfig.DDoSProtection var protectionConfig = sharedNodeConfig.DDoSProtection
err = firewalls.SharedDDoSProtectionManager.Apply(protectionConfig) err = firewalls.SharedDDoSProtectionManager.Apply(protectionConfig)
if err != nil { if err != nil {
this.replyFail(message.RequestId, dataMessage.Name+"was installed, but apply DDoS protection config failed: "+err.Error()) this.replyFail(message.RequestId, dataMessage.Name+" was installed, but apply DDoS protection config failed: "+err.Error())
} else { } else {
this.replyOk(message.RequestId, string(result.AsJSON())) this.replyOk(message.RequestId, string(result.AsJSON()))
} }

View File

@@ -43,7 +43,11 @@ func (this *HTTPAccessLogQueue) Start() {
for range ticker.C { for range ticker.C {
err := this.loop() err := this.loop()
if err != nil { if err != nil {
remotelogs.Error("ACCESS_LOG_QUEUE", err.Error()) if rpc.IsConnError(err) {
remotelogs.Debug("ACCESS_LOG_QUEUE", err.Error())
} else {
remotelogs.Error("ACCESS_LOG_QUEUE", err.Error())
}
} }
} }
} }

View File

@@ -758,6 +758,8 @@ func (this *HTTPRequest) Format(source string) string {
return strconv.FormatInt(this.requestFromTime.Unix(), 10) return strconv.FormatInt(this.requestFromTime.Unix(), 10)
case "host": case "host":
return this.ReqHost return this.ReqHost
case "cname":
return this.ReqServer.CNameDomain
case "referer": case "referer":
return this.RawReq.Referer() return this.RawReq.Referer()
case "referer.host": case "referer.host":
@@ -1603,9 +1605,25 @@ func (this *HTTPRequest) fixRequestHeader(header http.Header) {
header.Del(k) header.Del(k)
k = strings.ReplaceAll(k, "-Websocket-", "-WebSocket-") k = strings.ReplaceAll(k, "-Websocket-", "-WebSocket-")
header[k] = v header[k] = v
} else if k == "Www-Authenticate" { } else if strings.HasPrefix(k, "Sec-Ch") {
header.Del(k) header.Del(k)
header["WWW-Authenticate"] = v k = strings.ReplaceAll(k, "Sec-Ch-Ua", "Sec-CH-UA")
header[k] = v
} else {
switch k {
case "Www-Authenticate":
header.Del(k)
header["WWW-Authenticate"] = v
case "A-Im":
header.Del(k)
header["A-IM"] = v
case "Content-Md5":
header.Del(k)
header["Content-MD5"] = v
case "Sec-Gpc":
header.Del(k)
header["Content-GPC"] = v
}
} }
} }
} }
@@ -1705,10 +1723,10 @@ func (this *HTTPRequest) bytePool(contentLength int64) *utils.BytePool {
return utils.BytePool1k return utils.BytePool1k
} }
if contentLength < 32768 { // 32K if contentLength < 32768 { // 32K
return utils.BytePool4k return utils.BytePool16k
} }
if contentLength < 131072 { // 128K if contentLength < 131072 { // 128K
return utils.BytePool16k return utils.BytePool32k
} }
return utils.BytePool32k return utils.BytePool32k
} }

View File

@@ -19,6 +19,11 @@ import (
// 读取缓存 // 读取缓存
func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) { func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
// 需要动态Upgrade的不缓存
if len(this.RawReq.Header.Get("Upgrade")) > 0 {
return
}
this.cacheCanTryStale = false this.cacheCanTryStale = false
var cachePolicy = this.ReqServer.HTTPCachePolicy var cachePolicy = this.ReqServer.HTTPCachePolicy
@@ -290,34 +295,31 @@ func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
} }
} }
var pool = this.bytePool(fileSize)
var buf = pool.Get()
defer func() {
pool.Put(buf)
}()
// 读取Header // 读取Header
var headerBuf = []byte{} var headerData = []byte{}
this.writer.SetSentHeaderBytes(reader.HeaderSize()) this.writer.SetSentHeaderBytes(reader.HeaderSize())
err = reader.ReadHeader(buf, func(n int) (goNext bool, err error) { var headerPool = this.bytePool(reader.HeaderSize())
headerBuf = append(headerBuf, buf[:n]...) var headerBuf = headerPool.Get()
err = reader.ReadHeader(headerBuf, func(n int) (goNext bool, err error) {
headerData = append(headerData, headerBuf[:n]...)
for { for {
nIndex := bytes.Index(headerBuf, []byte{'\n'}) nIndex := bytes.Index(headerData, []byte{'\n'})
if nIndex >= 0 { if nIndex >= 0 {
row := headerBuf[:nIndex] row := headerData[:nIndex]
spaceIndex := bytes.Index(row, []byte{':'}) spaceIndex := bytes.Index(row, []byte{':'})
if spaceIndex <= 0 { if spaceIndex <= 0 {
return false, errors.New("invalid header '" + string(row) + "'") return false, errors.New("invalid header '" + string(row) + "'")
} }
this.writer.Header().Set(string(row[:spaceIndex]), string(row[spaceIndex+1:])) this.writer.Header().Set(string(row[:spaceIndex]), string(row[spaceIndex+1:]))
headerBuf = headerBuf[nIndex+1:] headerData = headerData[nIndex+1:]
} else { } else {
break break
} }
} }
return true, nil return true, nil
}) })
headerPool.Put(headerBuf)
if err != nil { if err != nil {
if !this.canIgnore(err) { if !this.canIgnore(err) {
remotelogs.WarnServer("HTTP_REQUEST_CACHE", this.URL()+": read from cache failed: read header failed: "+err.Error()) remotelogs.WarnServer("HTTP_REQUEST_CACHE", this.URL()+": read from cache failed: read header failed: "+err.Error())
@@ -455,13 +457,16 @@ func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
respHeader.Set("Content-Length", strconv.FormatInt(ranges[0].Length(), 10)) respHeader.Set("Content-Length", strconv.FormatInt(ranges[0].Length(), 10))
this.writer.WriteHeader(http.StatusPartialContent) this.writer.WriteHeader(http.StatusPartialContent)
err = reader.ReadBodyRange(buf, ranges[0].Start(), ranges[0].End(), func(n int) (goNext bool, err error) { var pool = this.bytePool(fileSize)
_, err = this.writer.Write(buf[:n]) var bodyBuf = pool.Get()
err = reader.ReadBodyRange(bodyBuf, ranges[0].Start(), ranges[0].End(), func(n int) (goNext bool, err error) {
_, err = this.writer.Write(bodyBuf[:n])
if err != nil { if err != nil {
return false, errWritingToClient return false, errWritingToClient
} }
return true, nil return true, nil
}) })
pool.Put(bodyBuf)
if err != nil { if err != nil {
this.varMapping["cache.status"] = "MISS" this.varMapping["cache.status"] = "MISS"
@@ -508,13 +513,16 @@ func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
} }
} }
err := reader.ReadBodyRange(buf, r.Start(), r.End(), func(n int) (goNext bool, err error) { var pool = this.bytePool(fileSize)
_, err = this.writer.Write(buf[:n]) var bodyBuf = pool.Get()
err := reader.ReadBodyRange(bodyBuf, r.Start(), r.End(), func(n int) (goNext bool, err error) {
_, err = this.writer.Write(bodyBuf[:n])
if err != nil { if err != nil {
return false, errWritingToClient return false, errWritingToClient
} }
return true, nil return true, nil
}) })
pool.Put(bodyBuf)
if err != nil { if err != nil {
if !this.canIgnore(err) { if !this.canIgnore(err) {
remotelogs.WarnServer("HTTP_REQUEST_CACHE", this.URL()+": read from cache failed: "+err.Error()) remotelogs.WarnServer("HTTP_REQUEST_CACHE", this.URL()+": read from cache failed: "+err.Error())
@@ -538,15 +546,18 @@ func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
this.writer.Prepare(resp, fileSize, reader.Status(), false) this.writer.Prepare(resp, fileSize, reader.Status(), false)
this.writer.WriteHeader(reader.Status()) this.writer.WriteHeader(reader.Status())
var pool = this.bytePool(fileSize)
var bodyBuf = pool.Get()
if storage.CanSendfile() { if storage.CanSendfile() {
if fp, canSendFile := this.writer.canSendfile(); canSendFile { if fp, canSendFile := this.writer.canSendfile(); canSendFile {
this.writer.sentBodyBytes, err = io.CopyBuffer(this.writer.rawWriter, fp, buf) this.writer.sentBodyBytes, err = io.CopyBuffer(this.writer.rawWriter, fp, bodyBuf)
} else { } else {
_, err = io.CopyBuffer(this.writer, resp.Body, buf) _, err = io.CopyBuffer(this.writer, resp.Body, bodyBuf)
} }
} else { } else {
_, err = io.CopyBuffer(this.writer, resp.Body, buf) _, err = io.CopyBuffer(this.writer, resp.Body, bodyBuf)
} }
pool.Put(bodyBuf)
if err == io.EOF { if err == io.EOF {
err = nil err = nil
} }
@@ -617,7 +628,14 @@ func (this *HTTPRequest) tryPartialReader(storage caches.StorageInterface, key s
}() }()
// 检查范围 // 检查范围
//const maxFirstSpan = 16 << 20 // TODO 可以在缓存策略中设置此值
for index, r := range ranges { for index, r := range ranges {
// 没有指定结束位置时,自动指定一个
/**if r.Start() >= 0 && r.End() == -1 {
if partialReader.MaxLength() > r.Start()+maxFirstSpan {
r[1] = r.Start() + maxFirstSpan
}
}**/
r1, ok := r.Convert(partialReader.MaxLength()) r1, ok := r.Convert(partialReader.MaxLength())
if !ok { if !ok {
return nil, nil return nil, nil

View File

@@ -1,7 +1,11 @@
package nodes package nodes
import ( import (
"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
"github.com/TeaOSLab/EdgeNode/internal/utils" "github.com/TeaOSLab/EdgeNode/internal/utils"
"github.com/iwind/TeaGo/types"
"net"
"net/http" "net/http"
"strconv" "strconv"
"strings" "strings"
@@ -13,7 +17,7 @@ func (this *HTTPRequest) doHostRedirect() (blocked bool) {
if this.web.MergeSlashes { if this.web.MergeSlashes {
urlPath = utils.CleanPath(urlPath) urlPath = utils.CleanPath(urlPath)
} }
fullURL := this.requestScheme() + "://" + this.ReqHost + urlPath var fullURL = this.requestScheme() + "://" + this.ReqHost + urlPath
for _, u := range this.web.HostRedirects { for _, u := range this.web.HostRedirects {
if !u.IsOn { if !u.IsOn {
continue continue
@@ -21,11 +25,50 @@ func (this *HTTPRequest) doHostRedirect() (blocked bool) {
if !u.MatchRequest(this.Format) { if !u.MatchRequest(this.Format) {
continue continue
} }
if u.MatchPrefix { // 匹配前缀 if len(u.Type) == 0 || u.Type == serverconfigs.HTTPHostRedirectTypeURL {
if strings.HasPrefix(fullURL, u.BeforeURL) { if u.MatchPrefix { // 匹配前缀
afterURL := u.AfterURL if strings.HasPrefix(fullURL, u.BeforeURL) {
if u.KeepRequestURI { afterURL := u.AfterURL
afterURL += this.RawReq.URL.RequestURI() if u.KeepRequestURI {
afterURL += this.RawReq.URL.RequestURI()
}
// 前后是否一致
if fullURL == afterURL {
return false
}
if u.Status <= 0 {
u.Status = http.StatusTemporaryRedirect
}
this.processResponseHeaders(this.writer.Header(), u.Status)
http.Redirect(this.RawWriter, this.RawReq, afterURL, u.Status)
return true
}
} else if u.MatchRegexp { // 正则匹配
var reg = u.BeforeURLRegexp()
if reg == nil {
continue
}
var matches = reg.FindStringSubmatch(fullURL)
if len(matches) == 0 {
continue
}
var afterURL = u.AfterURL
for i, match := range matches {
afterURL = strings.ReplaceAll(afterURL, "${"+strconv.Itoa(i)+"}", match)
}
var subNames = reg.SubexpNames()
if len(subNames) > 0 {
for _, subName := range subNames {
if len(subName) > 0 {
index := reg.SubexpIndex(subName)
if index > -1 {
afterURL = strings.ReplaceAll(afterURL, "${"+subName+"}", matches[index])
}
}
}
} }
// 前后是否一致 // 前后是否一致
@@ -33,69 +76,6 @@ func (this *HTTPRequest) doHostRedirect() (blocked bool) {
return false return false
} }
if u.Status <= 0 {
this.processResponseHeaders(this.writer.Header(), http.StatusTemporaryRedirect)
http.Redirect(this.RawWriter, this.RawReq, afterURL, http.StatusTemporaryRedirect)
} else {
this.processResponseHeaders(this.writer.Header(), u.Status)
http.Redirect(this.RawWriter, this.RawReq, afterURL, u.Status)
}
return true
}
} else if u.MatchRegexp { // 正则匹配
reg := u.BeforeURLRegexp()
if reg == nil {
continue
}
matches := reg.FindStringSubmatch(fullURL)
if len(matches) == 0 {
continue
}
afterURL := u.AfterURL
for i, match := range matches {
afterURL = strings.ReplaceAll(afterURL, "${"+strconv.Itoa(i)+"}", match)
}
subNames := reg.SubexpNames()
if len(subNames) > 0 {
for _, subName := range subNames {
if len(subName) > 0 {
index := reg.SubexpIndex(subName)
if index > -1 {
afterURL = strings.ReplaceAll(afterURL, "${"+subName+"}", matches[index])
}
}
}
}
// 前后是否一致
if fullURL == afterURL {
return false
}
if u.KeepArgs {
var qIndex = strings.Index(this.uri, "?")
if qIndex >= 0 {
afterURL += this.uri[qIndex:]
}
}
if u.Status <= 0 {
this.processResponseHeaders(this.writer.Header(), http.StatusTemporaryRedirect)
http.Redirect(this.RawWriter, this.RawReq, afterURL, http.StatusTemporaryRedirect)
} else {
this.processResponseHeaders(this.writer.Header(), u.Status)
http.Redirect(this.RawWriter, this.RawReq, afterURL, u.Status)
}
return true
} else { // 精准匹配
if fullURL == u.RealBeforeURL() {
// 前后是否一致
if fullURL == u.AfterURL {
return false
}
var afterURL = u.AfterURL
if u.KeepArgs { if u.KeepArgs {
var qIndex = strings.Index(this.uri, "?") var qIndex = strings.Index(this.uri, "?")
if qIndex >= 0 { if qIndex >= 0 {
@@ -104,12 +84,120 @@ func (this *HTTPRequest) doHostRedirect() (blocked bool) {
} }
if u.Status <= 0 { if u.Status <= 0 {
this.processResponseHeaders(this.writer.Header(), http.StatusTemporaryRedirect) u.Status = http.StatusTemporaryRedirect
http.Redirect(this.RawWriter, this.RawReq, afterURL, http.StatusTemporaryRedirect) }
} else { this.processResponseHeaders(this.writer.Header(), u.Status)
http.Redirect(this.RawWriter, this.RawReq, afterURL, u.Status)
return true
} else { // 精准匹配
if fullURL == u.RealBeforeURL() {
// 前后是否一致
if fullURL == u.AfterURL {
return false
}
var afterURL = u.AfterURL
if u.KeepArgs {
var qIndex = strings.Index(this.uri, "?")
if qIndex >= 0 {
afterURL += this.uri[qIndex:]
}
}
if u.Status <= 0 {
u.Status = http.StatusTemporaryRedirect
}
this.processResponseHeaders(this.writer.Header(), u.Status) this.processResponseHeaders(this.writer.Header(), u.Status)
http.Redirect(this.RawWriter, this.RawReq, afterURL, u.Status) http.Redirect(this.RawWriter, this.RawReq, afterURL, u.Status)
return true
} }
}
} else if u.Type == serverconfigs.HTTPHostRedirectTypeDomain {
if len(u.DomainAfter) == 0 {
continue
}
var reqHost = this.ReqHost
// 忽略跳转前端口
if u.DomainBeforeIgnorePorts {
h, _, err := net.SplitHostPort(reqHost)
if err == nil && len(h) > 0 {
reqHost = h
}
}
// 如果跳转前后域名一致,则终止
if u.DomainAfter == reqHost {
return false
}
var scheme = u.DomainAfterScheme
if len(scheme) == 0 {
scheme = this.requestScheme()
}
if u.DomainsAll || configutils.MatchDomains(u.DomainsBefore, reqHost) {
var afterURL = scheme + "://" + u.DomainAfter + urlPath
if fullURL == afterURL {
// 终止匹配
return false
}
if u.Status <= 0 {
u.Status = http.StatusTemporaryRedirect
}
this.processResponseHeaders(this.writer.Header(), u.Status)
http.Redirect(this.RawWriter, this.RawReq, afterURL, u.Status)
return true
}
} else if u.Type == serverconfigs.HTTPHostRedirectTypePort {
if u.PortAfter <= 0 {
continue
}
var scheme = u.PortAfterScheme
if len(scheme) == 0 {
scheme = this.requestScheme()
}
reqHost, reqPort, _ := net.SplitHostPort(this.ReqHost)
if len(reqHost) == 0 {
reqHost = this.ReqHost
}
if len(reqPort) == 0 {
switch this.requestScheme() {
case "http":
reqPort = "80"
case "https":
reqPort = "443"
}
}
// 如果跳转前后端口一致,则终止
if reqPort == types.String(u.PortAfter) {
return false
}
var containsPort = false
if u.PortsAll {
containsPort = true
} else {
containsPort = u.ContainsPort(types.Int(reqPort))
}
if containsPort {
var newReqHost = reqHost
if !((scheme == "http" && u.PortAfter == 80) || (scheme == "https" && u.PortAfter == 443)) {
newReqHost += ":" + types.String(u.PortAfter)
}
var afterURL = scheme + "://" + newReqHost + urlPath
if fullURL == afterURL {
// 终止匹配
return false
}
if u.Status <= 0 {
u.Status = http.StatusTemporaryRedirect
}
this.processResponseHeaders(this.writer.Header(), u.Status)
http.Redirect(this.RawWriter, this.RawReq, afterURL, u.Status)
return true return true
} }
} }

View File

@@ -51,27 +51,46 @@ func (this *HTTPRequest) log() {
addr = addr[:index] addr = addr[:index]
} }
var serverGlobalConfig = this.nodeConfig.GlobalServerConfig
// 请求Cookie // 请求Cookie
var cookies = map[string]string{} var cookies = map[string]string{}
if ref == nil || ref.ContainsField(serverconfigs.HTTPAccessLogFieldCookie) { var enableCookies = false
for _, cookie := range this.RawReq.Cookies() { if serverGlobalConfig == nil || serverGlobalConfig.HTTPAccessLog.EnableCookies {
cookies[cookie.Name] = cookie.Value enableCookies = true
if ref == nil || ref.ContainsField(serverconfigs.HTTPAccessLogFieldCookie) {
for _, cookie := range this.RawReq.Cookies() {
cookies[cookie.Name] = cookie.Value
}
} }
} }
// 请求Header // 请求Header
var pbReqHeader = map[string]*pb.Strings{} var pbReqHeader = map[string]*pb.Strings{}
if ref == nil || ref.ContainsField(serverconfigs.HTTPAccessLogFieldHeader) { if serverGlobalConfig == nil || serverGlobalConfig.HTTPAccessLog.EnableRequestHeaders {
for k, v := range this.RawReq.Header { if ref == nil || ref.ContainsField(serverconfigs.HTTPAccessLogFieldHeader) {
pbReqHeader[k] = &pb.Strings{Values: v} // 是否只记录通用Header
var commonHeadersOnly = serverGlobalConfig != nil && serverGlobalConfig.HTTPAccessLog.CommonRequestHeadersOnly
for k, v := range this.RawReq.Header {
if commonHeadersOnly && !serverconfigs.IsCommonRequestHeader(k) {
continue
}
if !enableCookies && k == "Cookie" {
continue
}
pbReqHeader[k] = &pb.Strings{Values: v}
}
} }
} }
// 响应Header // 响应Header
var pbResHeader = map[string]*pb.Strings{} var pbResHeader = map[string]*pb.Strings{}
if ref == nil || ref.ContainsField(serverconfigs.HTTPAccessLogFieldSentHeader) { if serverGlobalConfig == nil || serverGlobalConfig.HTTPAccessLog.EnableResponseHeaders {
for k, v := range this.writer.Header() { if ref == nil || ref.ContainsField(serverconfigs.HTTPAccessLogFieldSentHeader) {
pbResHeader[k] = &pb.Strings{Values: v} for k, v := range this.writer.Header() {
pbResHeader[k] = &pb.Strings{Values: v}
}
} }
} }

View File

@@ -12,6 +12,8 @@ func (this *HTTPRequest) doCheckReferers() (shouldStop bool) {
return return
} }
const cacheSeconds = "3600" // 时间不能过长,防止修改设置后长期无法生效
var refererURL = this.RawReq.Header.Get("Referer") var refererURL = this.RawReq.Header.Get("Referer")
if len(refererURL) == 0 { if len(refererURL) == 0 {
if this.web.Referers.MatchDomain(this.ReqHost, "") { if this.web.Referers.MatchDomain(this.ReqHost, "") {
@@ -19,6 +21,7 @@ func (this *HTTPRequest) doCheckReferers() (shouldStop bool) {
} }
this.tags = append(this.tags, "refererCheck") this.tags = append(this.tags, "refererCheck")
this.writer.Header().Set("Cache-Control", "max-age="+cacheSeconds)
this.writeCode(http.StatusForbidden, "The referer has been blocked.", "当前访问已被防盗链系统拦截。") this.writeCode(http.StatusForbidden, "The referer has been blocked.", "当前访问已被防盗链系统拦截。")
return true return true
@@ -31,6 +34,7 @@ func (this *HTTPRequest) doCheckReferers() (shouldStop bool) {
} }
this.tags = append(this.tags, "refererCheck") this.tags = append(this.tags, "refererCheck")
this.writer.Header().Set("Cache-Control", "max-age="+cacheSeconds)
this.writeCode(http.StatusForbidden, "The referer has been blocked.", "当前访问已被防盗链系统拦截。") this.writeCode(http.StatusForbidden, "The referer has been blocked.", "当前访问已被防盗链系统拦截。")
return true return true
@@ -38,6 +42,7 @@ func (this *HTTPRequest) doCheckReferers() (shouldStop bool) {
if !this.web.Referers.MatchDomain(this.ReqHost, u.Host) { if !this.web.Referers.MatchDomain(this.ReqHost, u.Host) {
this.tags = append(this.tags, "refererCheck") this.tags = append(this.tags, "refererCheck")
this.writer.Header().Set("Cache-Control", "max-age="+cacheSeconds)
this.writeCode(http.StatusForbidden, "The referer has been blocked.", "当前访问已被防盗链系统拦截。") this.writeCode(http.StatusForbidden, "The referer has been blocked.", "当前访问已被防盗链系统拦截。")
return true return true
} }

View File

@@ -218,7 +218,7 @@ func (this *HTTPRequest) checkWAFRequest(firewallPolicy *firewallconfigs.HTTPFir
} }
goNext, hasRequestBody, ruleGroup, ruleSet, err := w.MatchRequest(this, this.writer) goNext, hasRequestBody, ruleGroup, ruleSet, err := w.MatchRequest(this, this.writer)
if forceLog && logRequestBody && hasRequestBody { if forceLog && logRequestBody && hasRequestBody && ruleSet != nil && ruleSet.HasAttackActions() {
this.wafHasRequestBody = true this.wafHasRequestBody = true
} }
if err != nil { if err != nil {
@@ -294,7 +294,7 @@ func (this *HTTPRequest) checkWAFResponse(firewallPolicy *firewallconfigs.HTTPFi
} }
goNext, hasRequestBody, ruleGroup, ruleSet, err := w.MatchResponse(this, resp, this.writer) goNext, hasRequestBody, ruleGroup, ruleSet, err := w.MatchResponse(this, resp, this.writer)
if forceLog && logRequestBody && hasRequestBody { if forceLog && logRequestBody && hasRequestBody && ruleSet != nil && ruleSet.HasAttackActions() {
this.wafHasRequestBody = true this.wafHasRequestBody = true
} }
if err != nil { if err != nil {

View File

@@ -2,6 +2,7 @@ package nodes
import ( import (
"bufio" "bufio"
"bytes"
"errors" "errors"
"github.com/TeaOSLab/EdgeNode/internal/utils" "github.com/TeaOSLab/EdgeNode/internal/utils"
"io" "io"
@@ -9,8 +10,36 @@ import (
"net/url" "net/url"
) )
// WebsocketResponseReader Websocket响应Reader
type WebsocketResponseReader struct {
rawReader io.Reader
buf []byte
}
func NewWebsocketResponseReader(rawReader io.Reader) *WebsocketResponseReader {
return &WebsocketResponseReader{
rawReader: rawReader,
}
}
func (this *WebsocketResponseReader) Read(p []byte) (n int, err error) {
n, err = this.rawReader.Read(p)
if n > 0 {
if len(this.buf) == 0 {
this.buf = make([]byte, n)
copy(this.buf, p[:n])
} else {
this.buf = append(this.buf, p[:n]...)
}
}
return
}
// 处理Websocket请求 // 处理Websocket请求
func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shouldRetry bool) { func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shouldRetry bool) {
// 设置不缓存
this.web.Cache = nil
if this.web.WebsocketRef == nil || !this.web.WebsocketRef.IsOn || this.web.Websocket == nil || !this.web.Websocket.IsOn { if this.web.WebsocketRef == nil || !this.web.WebsocketRef.IsOn || this.web.Websocket == nil || !this.web.Websocket.IsOn {
this.writer.WriteHeader(http.StatusForbidden) this.writer.WriteHeader(http.StatusForbidden)
this.addError(errors.New("websocket have not been enabled yet")) this.addError(errors.New("websocket have not been enabled yet"))
@@ -84,14 +113,20 @@ func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shou
go func() { go func() {
// 读取第一个响应 // 读取第一个响应
resp, err := http.ReadResponse(bufio.NewReader(originConn), this.RawReq) var respReader = NewWebsocketResponseReader(originConn)
if err != nil { resp, err := http.ReadResponse(bufio.NewReader(respReader), this.RawReq)
if err != nil || resp == nil {
if resp != nil && resp.Body != nil {
_ = resp.Body.Close()
}
_ = clientConn.Close() _ = clientConn.Close()
_ = originConn.Close() _ = originConn.Close()
return return
} }
this.processResponseHeaders(resp.Header, resp.StatusCode) this.processResponseHeaders(resp.Header, resp.StatusCode)
this.writer.statusCode = resp.StatusCode
// 将响应写回客户端 // 将响应写回客户端
err = resp.Write(clientConn) err = resp.Write(clientConn)
@@ -105,6 +140,25 @@ func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shou
return return
} }
// 剩余已经从源站读取的内容
var headerBytes = respReader.buf
var headerIndex = bytes.Index(headerBytes, []byte{'\r', '\n', '\r', '\n'}) // CRLF
if headerIndex > 0 {
var leftBytes = headerBytes[headerIndex+4:]
if len(leftBytes) > 0 {
_, err = clientConn.Write(leftBytes)
if err != nil {
if resp.Body != nil {
_ = resp.Body.Close()
}
_ = clientConn.Close()
_ = originConn.Close()
return
}
}
}
if resp.Body != nil { if resp.Body != nil {
_ = resp.Body.Close() _ = resp.Body.Close()
} }

View File

@@ -303,7 +303,19 @@ func (this *HTTPWriter) PrepareCache(resp *http.Response, size int64) {
if this.isPartial { if this.isPartial {
cacheKey += caches.SuffixPartial cacheKey += caches.SuffixPartial
} }
cacheWriter, err := storage.OpenWriter(cacheKey, expiresAt, this.StatusCode(), size, cacheRef.MaxSizeBytes(), this.isPartial)
// 待写入尺寸
var totalSize = size
if totalSize < 0 && this.isPartial {
var contentRange = resp.Header.Get("Content-Range")
if len(contentRange) > 0 {
_, partialTotalSize := httpRequestParseContentRangeHeader(contentRange)
if partialTotalSize > 0 {
totalSize = partialTotalSize
}
}
}
cacheWriter, err := storage.OpenWriter(cacheKey, expiresAt, this.StatusCode(), this.calculateHeaderLength(), totalSize, cacheRef.MaxSizeBytes(), this.isPartial)
if err != nil { if err != nil {
if err == caches.ErrEntityTooLarge && addStatusHeader { if err == caches.ErrEntityTooLarge && addStatusHeader {
this.Header().Set("X-Cache", "BYPASS, entity too large") this.Header().Set("X-Cache", "BYPASS, entity too large")
@@ -324,13 +336,19 @@ func (this *HTTPWriter) PrepareCache(resp *http.Response, size int64) {
} }
// 写入Header // 写入Header
var headerBuf = utils.SharedBufferPool.Get()
for k, v := range this.Header() { for k, v := range this.Header() {
if k == "Set-Cookie" || (this.isPartial && k == "Content-Range") {
continue
}
for _, v1 := range v { for _, v1 := range v {
if this.isPartial && k == "Content-Type" && strings.Contains(v1, "multipart/byteranges") { if this.isPartial && k == "Content-Type" && strings.Contains(v1, "multipart/byteranges") {
continue continue
} }
_, err = cacheWriter.WriteHeader([]byte(k + ":" + v1 + "\n")) _, err = headerBuf.Write([]byte(k + ":" + v1 + "\n"))
if err != nil { if err != nil {
utils.SharedBufferPool.Put(headerBuf)
remotelogs.Error("HTTP_WRITER", "write cache failed: "+err.Error()) remotelogs.Error("HTTP_WRITER", "write cache failed: "+err.Error())
_ = this.cacheWriter.Discard() _ = this.cacheWriter.Discard()
this.cacheWriter = nil this.cacheWriter = nil
@@ -338,6 +356,14 @@ func (this *HTTPWriter) PrepareCache(resp *http.Response, size int64) {
} }
} }
} }
_, err = cacheWriter.WriteHeader(headerBuf.Bytes())
utils.SharedBufferPool.Put(headerBuf)
if err != nil {
remotelogs.Error("HTTP_WRITER", "write cache failed: "+err.Error())
_ = this.cacheWriter.Discard()
this.cacheWriter = nil
return
}
if this.isPartial { if this.isPartial {
// content-range // content-range
@@ -627,16 +653,21 @@ func (this *HTTPWriter) PrepareCompression(resp *http.Response, size int64) {
cacheKey += this.cacheReaderSuffix cacheKey += this.cacheReaderSuffix
} }
compressionCacheWriter, err := this.cacheStorage.OpenWriter(cacheKey+caches.SuffixCompression+compressionEncoding, expiredAt, this.StatusCode(), -1, cacheRef.MaxSizeBytes(), false) compressionCacheWriter, err := this.cacheStorage.OpenWriter(cacheKey+caches.SuffixCompression+compressionEncoding, expiredAt, this.StatusCode(), this.calculateHeaderLength(), -1, cacheRef.MaxSizeBytes(), false)
if err != nil { if err != nil {
return return
} }
// 写入Header // 写入Header
var headerBuffer = utils.SharedBufferPool.Get()
for k, v := range this.Header() { for k, v := range this.Header() {
if k == "Set-Cookie" || (this.isPartial && k == "Content-Range") {
continue
}
for _, v1 := range v { for _, v1 := range v {
_, err = compressionCacheWriter.WriteHeader([]byte(k + ":" + v1 + "\n")) _, err = headerBuffer.Write([]byte(k + ":" + v1 + "\n"))
if err != nil { if err != nil {
utils.SharedBufferPool.Put(headerBuffer)
remotelogs.Error("HTTP_WRITER", "write compression cache failed: "+err.Error()) remotelogs.Error("HTTP_WRITER", "write compression cache failed: "+err.Error())
_ = compressionCacheWriter.Discard() _ = compressionCacheWriter.Discard()
compressionCacheWriter = nil compressionCacheWriter = nil
@@ -645,6 +676,15 @@ func (this *HTTPWriter) PrepareCompression(resp *http.Response, size int64) {
} }
} }
_, err = compressionCacheWriter.WriteHeader(headerBuffer.Bytes())
utils.SharedBufferPool.Put(headerBuffer)
if err != nil {
remotelogs.Error("HTTP_WRITER", "write compression cache failed: "+err.Error())
_ = compressionCacheWriter.Discard()
compressionCacheWriter = nil
return
}
if compressionCacheWriter != nil { if compressionCacheWriter != nil {
this.compressionCacheWriter = compressionCacheWriter this.compressionCacheWriter = compressionCacheWriter
var teeWriter = writers.NewTeeWriterCloser(this.writer, compressionCacheWriter) var teeWriter = writers.NewTeeWriterCloser(this.writer, compressionCacheWriter)
@@ -942,10 +982,14 @@ func (this *HTTPWriter) finishWebP() {
expiredAt = this.cacheWriter.ExpiredAt() expiredAt = this.cacheWriter.ExpiredAt()
} }
webpCacheWriter, _ = this.cacheStorage.OpenWriter(cacheKey, expiredAt, this.StatusCode(), -1, -1, false) webpCacheWriter, _ = this.cacheStorage.OpenWriter(cacheKey, expiredAt, this.StatusCode(), -1, -1, -1, false)
if webpCacheWriter != nil { if webpCacheWriter != nil {
// 写入Header // 写入Header
for k, v := range this.Header() { for k, v := range this.Header() {
if k == "Set-Cookie" {
continue
}
// 这里是原始的数据,不需要内容编码 // 这里是原始的数据,不需要内容编码
if k == "Content-Encoding" || k == "Transfer-Encoding" { if k == "Content-Encoding" || k == "Transfer-Encoding" {
continue continue
@@ -1157,3 +1201,16 @@ func (this *HTTPWriter) finishRequest() {
_ = this.rawReader.Close() _ = this.rawReader.Close()
} }
} }
// 计算Header长度
func (this *HTTPWriter) calculateHeaderLength() (result int) {
for k, v := range this.Header() {
if k == "Set-Cookie" || (this.isPartial && k == "Content-Range") {
continue
}
for _, v1 := range v {
result += len(k) + 1 /**:**/ + len(v1) + 1 /**\n**/
}
}
return
}

View File

@@ -160,7 +160,7 @@ func (this *BaseListener) findNamedServer(name string) (serverConfig *serverconf
// 严格查找域名 // 严格查找域名
func (this *BaseListener) findNamedServerMatched(name string) (serverConfig *serverconfigs.ServerConfig, serverName string) { func (this *BaseListener) findNamedServerMatched(name string) (serverConfig *serverconfigs.ServerConfig, serverName string) {
group := this.Group var group = this.Group
if group == nil { if group == nil {
return nil, "" return nil, ""
} }

View File

@@ -45,7 +45,7 @@ func (this *HTTPListener) Serve() error {
Handler: this, Handler: this,
ReadTimeout: 1 * time.Hour, // TODO 改成可以配置 ReadTimeout: 1 * time.Hour, // TODO 改成可以配置
ReadHeaderTimeout: 3 * time.Second, // TODO 改成可以配置 ReadHeaderTimeout: 3 * time.Second, // TODO 改成可以配置
WriteTimeout: 1 * time.Hour, // TODO 改成可以配置 WriteTimeout: 2 * time.Hour, // TODO 改成可以配置
IdleTimeout: 75 * time.Second, // TODO 改成可以配置 IdleTimeout: 75 * time.Second, // TODO 改成可以配置
ConnState: func(conn net.Conn, state http.ConnState) { ConnState: func(conn net.Conn, state http.ConnState) {
switch state { switch state {
@@ -117,7 +117,7 @@ func (this *HTTPListener) Reload(group *serverconfigs.ServerAddressGroup) {
// ServerHTTP 处理HTTP请求 // ServerHTTP 处理HTTP请求
func (this *HTTPListener) ServeHTTP(rawWriter http.ResponseWriter, rawReq *http.Request) { func (this *HTTPListener) ServeHTTP(rawWriter http.ResponseWriter, rawReq *http.Request) {
// 域名 // 域名
var reqHost = rawReq.Host var reqHost = strings.TrimRight(rawReq.Host, ".")
// TLS域名 // TLS域名
if this.isIP(reqHost) { if this.isIP(reqHost) {
@@ -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{ var req = &HTTPRequest{
RawReq: rawReq, RawReq: rawReq,

View File

@@ -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) tcpConn, ok := conn.(LingerConn)
if ok { if ok {

View File

@@ -170,6 +170,11 @@ func (this *UDPListener) servePacketListener(listener UDPPacketListener) error {
return nil return nil
} }
// 检查用户状态
if firstServer.UserId > 0 && !SharedUserManager.CheckUserServersIsEnabled(firstServer.UserId) {
return nil
}
n, cm, clientAddr, err := listener.ReadFrom(buffer) n, cm, clientAddr, err := listener.ReadFrom(buffer)
if err != nil { if err != nil {
if this.isClosed { if this.isClosed {

View File

@@ -2,14 +2,15 @@ package nodes
import ( import (
"bytes" "bytes"
"context"
"encoding/json" "encoding/json"
"errors" "errors"
"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
iplib "github.com/TeaOSLab/EdgeCommon/pkg/iplibrary" iplib "github.com/TeaOSLab/EdgeCommon/pkg/iplibrary"
"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs" "github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb" "github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs" "github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/ddosconfigs" "github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/ddosconfigs"
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
"github.com/TeaOSLab/EdgeNode/internal/caches" "github.com/TeaOSLab/EdgeNode/internal/caches"
"github.com/TeaOSLab/EdgeNode/internal/configs" "github.com/TeaOSLab/EdgeNode/internal/configs"
"github.com/TeaOSLab/EdgeNode/internal/conns" "github.com/TeaOSLab/EdgeNode/internal/conns"
@@ -25,11 +26,11 @@ import (
"github.com/TeaOSLab/EdgeNode/internal/trackers" "github.com/TeaOSLab/EdgeNode/internal/trackers"
"github.com/TeaOSLab/EdgeNode/internal/utils" "github.com/TeaOSLab/EdgeNode/internal/utils"
_ "github.com/TeaOSLab/EdgeNode/internal/utils/clock" // 触发时钟更新 _ "github.com/TeaOSLab/EdgeNode/internal/utils/clock" // 触发时钟更新
"github.com/TeaOSLab/EdgeNode/internal/utils/jsonutils"
"github.com/TeaOSLab/EdgeNode/internal/waf" "github.com/TeaOSLab/EdgeNode/internal/waf"
"github.com/andybalholm/brotli" "github.com/andybalholm/brotli"
"github.com/iwind/TeaGo/Tea" "github.com/iwind/TeaGo/Tea"
"github.com/iwind/TeaGo/lists" "github.com/iwind/TeaGo/lists"
"github.com/iwind/TeaGo/logs"
"github.com/iwind/TeaGo/maps" "github.com/iwind/TeaGo/maps"
"github.com/iwind/TeaGo/types" "github.com/iwind/TeaGo/types"
"github.com/iwind/gosock/pkg/gosock" "github.com/iwind/gosock/pkg/gosock"
@@ -38,6 +39,7 @@ import (
"os" "os"
"os/exec" "os/exec"
"os/signal" "os/signal"
"path/filepath"
"runtime" "runtime"
"runtime/debug" "runtime/debug"
"sort" "sort"
@@ -59,18 +61,27 @@ type Node struct {
sock *gosock.Sock sock *gosock.Sock
locker sync.Mutex locker sync.Mutex
maxCPU int32 oldMaxCPU int32
maxThreads int oldMaxThreads int
timezone string oldTimezone string
oldHTTPCachePolicies []*serverconfigs.HTTPCachePolicy
oldHTTPFirewallPolicies []*firewallconfigs.HTTPFirewallPolicy
oldFirewallActions []*firewallconfigs.FirewallActionConfig
oldMetricItems []*serverconfigs.MetricItemConfig
updatingServerMap map[int64]*serverconfigs.ServerConfig updatingServerMap map[int64]*serverconfigs.ServerConfig
lastAPINodeVersion int64
lastAPINodeAddrs []string // 以前的API节点地址
lastTaskVersion int64
} }
func NewNode() *Node { func NewNode() *Node {
return &Node{ return &Node{
sock: gosock.NewTmpSock(teaconst.ProcessName), sock: gosock.NewTmpSock(teaconst.ProcessName),
maxThreads: -1, oldMaxThreads: -1,
maxCPU: -1, oldMaxCPU: -1,
updatingServerMap: map[int64]*serverconfigs.ServerConfig{}, updatingServerMap: map[int64]*serverconfigs.ServerConfig{},
} }
} }
@@ -190,7 +201,7 @@ func (this *Node) Start() {
} }
} }
sharedNodeConfig = nodeConfig sharedNodeConfig = nodeConfig
this.onReload(nodeConfig) this.onReload(nodeConfig, true)
// 发送事件 // 发送事件
events.Notify(events.EventLoaded) events.Notify(events.EventLoaded)
@@ -205,9 +216,7 @@ func (this *Node) Start() {
// 统计 // 统计
goman.New(func() { goman.New(func() {
stats.SharedTrafficStatManager.Start(func() *nodeconfigs.NodeConfig { stats.SharedTrafficStatManager.Start()
return sharedNodeConfig
})
}) })
goman.New(func() { goman.New(func() {
stats.SharedHTTPRequestStatManager.Start() stats.SharedHTTPRequestStatManager.Start()
@@ -313,8 +322,9 @@ func (this *Node) loop() error {
return errors.New("create rpc client failed: " + err.Error()) return errors.New("create rpc client failed: " + err.Error())
} }
var nodeCtx = rpcClient.Context() tasksResp, err := rpcClient.NodeTaskRPC.FindNodeTasks(rpcClient.Context(), &pb.FindNodeTasksRequest{
tasksResp, err := rpcClient.NodeTaskRPC.FindNodeTasks(nodeCtx, &pb.FindNodeTasksRequest{}) Version: this.lastTaskVersion,
})
if err != nil { if err != nil {
if rpc.IsConnError(err) && !Tea.IsTesting() { if rpc.IsConnError(err) && !Tea.IsTesting() {
return nil return nil
@@ -322,15 +332,18 @@ func (this *Node) loop() error {
return errors.New("read node tasks failed: " + err.Error()) return errors.New("read node tasks failed: " + err.Error())
} }
for _, task := range tasksResp.NodeTasks { for _, task := range tasksResp.NodeTasks {
err := this.execTask(rpcClient, nodeCtx, task) err := this.execTask(rpcClient, task)
this.finishTask(task.Id, err) if !this.finishTask(task.Id, task.Version, err) {
// 防止失败的任务无法重试
break
}
} }
return nil return nil
} }
// 执行任务 // 执行任务
func (this *Node) execTask(rpcClient *rpc.RPCClient, nodeCtx context.Context, task *pb.NodeTask) error { func (this *Node) execTask(rpcClient *rpc.RPCClient, task *pb.NodeTask) error {
switch task.Type { switch task.Type {
case "ipItemChanged": case "ipItemChanged":
// 防止阻塞 // 防止阻塞
@@ -360,7 +373,7 @@ func (this *Node) execTask(rpcClient *rpc.RPCClient, nodeCtx context.Context, ta
return errors.New("reload common scripts failed: " + err.Error()) return errors.New("reload common scripts failed: " + err.Error())
} }
case "nodeLevelChanged": case "nodeLevelChanged":
levelInfoResp, err := rpcClient.NodeRPC.FindNodeLevelInfo(nodeCtx, &pb.FindNodeLevelInfoRequest{}) levelInfoResp, err := rpcClient.NodeRPC.FindNodeLevelInfo(rpcClient.Context(), &pb.FindNodeLevelInfoRequest{})
if err != nil { if err != nil {
return err return err
} }
@@ -381,7 +394,7 @@ func (this *Node) execTask(rpcClient *rpc.RPCClient, nodeCtx context.Context, ta
sharedNodeConfig.ParentNodes = parentNodes sharedNodeConfig.ParentNodes = parentNodes
} }
case "ddosProtectionChanged": case "ddosProtectionChanged":
resp, err := rpcClient.NodeRPC.FindNodeDDoSProtection(nodeCtx, &pb.FindNodeDDoSProtectionRequest{}) resp, err := rpcClient.NodeRPC.FindNodeDDoSProtection(rpcClient.Context(), &pb.FindNodeDDoSProtectionRequest{})
if err != nil { if err != nil {
return err return err
} }
@@ -409,7 +422,7 @@ func (this *Node) execTask(rpcClient *rpc.RPCClient, nodeCtx context.Context, ta
return nil return nil
} }
case "globalServerConfigChanged": case "globalServerConfigChanged":
resp, err := rpcClient.NodeRPC.FindNodeGlobalServerConfig(nodeCtx, &pb.FindNodeGlobalServerConfigRequest{}) resp, err := rpcClient.NodeRPC.FindNodeGlobalServerConfig(rpcClient.Context(), &pb.FindNodeGlobalServerConfigRequest{})
if err != nil { if err != nil {
return err return err
} }
@@ -430,6 +443,22 @@ func (this *Node) execTask(rpcClient *rpc.RPCClient, nodeCtx context.Context, ta
} }
} }
} }
case "userServersStateChanged":
if task.UserId > 0 {
resp, err := rpcClient.UserRPC.CheckUserServersState(rpcClient.Context(), &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: default:
remotelogs.Error("NODE", "task '"+types.String(task.Id)+"', type '"+task.Type+"' has not been handled") remotelogs.Error("NODE", "task '"+types.String(task.Id)+"', type '"+task.Type+"' has not been handled")
} }
@@ -438,39 +467,44 @@ func (this *Node) execTask(rpcClient *rpc.RPCClient, nodeCtx context.Context, ta
} }
// 标记任务完成 // 标记任务完成
func (this *Node) finishTask(taskId int64, err error) { func (this *Node) finishTask(taskId int64, taskVersion int64, taskErr error) (success bool) {
if taskId <= 0 { if taskId <= 0 {
return return true
} }
rpcClient, err := rpc.SharedRPC() rpcClient, err := rpc.SharedRPC()
if err != nil { if err != nil {
logs.Println("[NODE]", "create rpc client failed: "+err.Error()) remotelogs.Debug("NODE", "create rpc client failed: "+err.Error())
return return false
} }
var nodeCtx = rpcClient.Context() var isOk = taskErr == nil
if isOk && taskVersion > this.lastTaskVersion {
this.lastTaskVersion = taskVersion
}
var isOk = err == nil
var errMsg = "" var errMsg = ""
if err != nil { if taskErr != nil {
errMsg = err.Error() errMsg = taskErr.Error()
} }
_, err = rpcClient.NodeTaskRPC.ReportNodeTaskDone(nodeCtx, &pb.ReportNodeTaskDoneRequest{ _, err = rpcClient.NodeTaskRPC.ReportNodeTaskDone(rpcClient.Context(), &pb.ReportNodeTaskDoneRequest{
NodeTaskId: taskId, NodeTaskId: taskId,
IsOk: isOk, IsOk: isOk,
Error: errMsg, Error: errMsg,
}) })
success = err == nil
if err != nil { if err != nil {
// 不需要上报到服务中心 // 连接错误不需要上报到服务中心
if rpc.IsConnError(err) { if rpc.IsConnError(err) {
logs.Println("[NODE]", "report task done failed: "+err.Error()) remotelogs.Debug("NODE", "report task done failed: "+err.Error())
} else { } else {
remotelogs.Error("NODE", "report task done failed: "+err.Error()) remotelogs.Error("NODE", "report task done failed: "+err.Error())
} }
} }
return success
} }
// 读取API配置 // 读取API配置
@@ -501,10 +535,8 @@ func (this *Node) syncConfig(taskVersion int64) error {
} }
// 获取同步任务 // 获取同步任务
var nodeCtx = rpcClient.Context()
// TODO 这里考虑只同步版本号有变更的 // TODO 这里考虑只同步版本号有变更的
configResp, err := rpcClient.NodeRPC.FindCurrentNodeConfig(nodeCtx, &pb.FindCurrentNodeConfigRequest{ configResp, err := rpcClient.NodeRPC.FindCurrentNodeConfig(rpcClient.Context(), &pb.FindCurrentNodeConfigRequest{
Version: -1, // 更新所有版本 Version: -1, // 更新所有版本
Compress: true, Compress: true,
NodeTaskVersion: taskVersion, NodeTaskVersion: taskVersion,
@@ -575,7 +607,7 @@ func (this *Node) syncConfig(taskVersion int64) error {
remotelogs.Println("NODE", "loading config ...") remotelogs.Println("NODE", "loading config ...")
} }
this.onReload(nodeConfig) this.onReload(nodeConfig, true)
// 发送事件 // 发送事件
events.Notify(events.EventReload) events.Notify(events.EventReload)
@@ -615,6 +647,36 @@ func (this *Node) syncServerConfig(serverId int64) error {
return nil 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() { func (this *Node) startSyncTimer() {
// TODO 这个时间间隔可以自行设置 // TODO 这个时间间隔可以自行设置
@@ -676,12 +738,12 @@ func (this *Node) checkClusterConfig() error {
return err return err
} }
logs.Println("[NODE]registering node to cluster ...") remotelogs.Debug("NODE", "registering node to cluster ...")
resp, err := rpcClient.NodeRPC.RegisterClusterNode(rpcClient.ClusterContext(config.ClusterId, config.Secret), &pb.RegisterClusterNodeRequest{Name: HOSTNAME}) resp, err := rpcClient.NodeRPC.RegisterClusterNode(rpcClient.ClusterContext(config.ClusterId, config.Secret), &pb.RegisterClusterNodeRequest{Name: HOSTNAME})
if err != nil { if err != nil {
return err return err
} }
logs.Println("[NODE]registered successfully") remotelogs.Debug("NODE", "registered successfully")
// 写入到配置文件中 // 写入到配置文件中
if len(resp.Endpoints) == 0 { if len(resp.Endpoints) == 0 {
@@ -689,8 +751,8 @@ func (this *Node) checkClusterConfig() error {
} }
var apiConfig = &configs.APIConfig{ var apiConfig = &configs.APIConfig{
RPC: struct { RPC: struct {
Endpoints []string `yaml:"endpoints"` Endpoints []string `yaml:"endpoints" json:"endpoints"`
DisableUpdate bool `yaml:"disableUpdate"` DisableUpdate bool `yaml:"disableUpdate" json:"disableUpdate"`
}{ }{
Endpoints: resp.Endpoints, Endpoints: resp.Endpoints,
DisableUpdate: false, DisableUpdate: false,
@@ -698,12 +760,12 @@ func (this *Node) checkClusterConfig() error {
NodeId: resp.UniqueId, NodeId: resp.UniqueId,
Secret: resp.Secret, Secret: resp.Secret,
} }
logs.Println("[NODE]writing 'configs/api.yaml' ...") remotelogs.Debug("NODE", "writing 'configs/api.yaml' ...")
err = apiConfig.WriteFile(Tea.ConfigFile("api.yaml")) err = apiConfig.WriteFile(Tea.ConfigFile("api.yaml"))
if err != nil { if err != nil {
return err return err
} }
logs.Println("[NODE]wrote 'configs/api.yaml' successfully") remotelogs.Debug("NODE", "wrote 'configs/api.yaml' successfully")
return nil return nil
} }
@@ -759,6 +821,7 @@ func (this *Node) listenSock() error {
// 退出主进程 // 退出主进程
events.Notify(events.EventQuit) events.Notify(events.EventQuit)
time.Sleep(100 * time.Millisecond)
utils.Exit() utils.Exit()
case "quit": case "quit":
_ = cmd.ReplyOk() _ = cmd.ReplyOk()
@@ -907,12 +970,12 @@ func (this *Node) listenSock() error {
err := this.sock.Listen() err := this.sock.Listen()
if err != nil { if err != nil {
logs.Println("NODE", err.Error()) remotelogs.Debug("NODE", err.Error())
} }
}) })
events.OnKey(events.EventQuit, this, func() { events.OnKey(events.EventQuit, this, func() {
remotelogs.Println("NODE", "quit unix sock") remotelogs.Debug("NODE", "quit unix sock")
_ = this.sock.Close() _ = this.sock.Close()
}) })
@@ -920,29 +983,94 @@ func (this *Node) listenSock() error {
} }
// 重载配置调用 // 重载配置调用
func (this *Node) onReload(config *nodeconfigs.NodeConfig) { func (this *Node) onReload(config *nodeconfigs.NodeConfig, reloadAll bool) {
nodeconfigs.ResetNodeConfig(config) nodeconfigs.ResetNodeConfig(config)
sharedNodeConfig = config sharedNodeConfig = config
// 不需要每次都全部重新加载
if !reloadAll {
return
}
// 缓存策略 // 缓存策略
var subDirs = config.CacheDiskSubDirs
for _, subDir := range subDirs {
subDir.Path = filepath.Clean(subDir.Path)
}
if len(subDirs) > 0 {
sort.Slice(subDirs, func(i, j int) bool {
return subDirs[i].Path < subDirs[j].Path
})
}
var cachePoliciesChanged = !jsonutils.Equal(caches.SharedManager.MaxDiskCapacity, config.MaxCacheDiskCapacity) ||
!jsonutils.Equal(caches.SharedManager.MaxMemoryCapacity, config.MaxCacheMemoryCapacity) ||
!jsonutils.Equal(caches.SharedManager.MainDiskDir, config.CacheDiskDir) ||
!jsonutils.Equal(caches.SharedManager.SubDiskDirs, subDirs) ||
!jsonutils.Equal(this.oldHTTPCachePolicies, config.HTTPCachePolicies)
caches.SharedManager.MaxDiskCapacity = config.MaxCacheDiskCapacity caches.SharedManager.MaxDiskCapacity = config.MaxCacheDiskCapacity
caches.SharedManager.MaxMemoryCapacity = config.MaxCacheMemoryCapacity caches.SharedManager.MaxMemoryCapacity = config.MaxCacheMemoryCapacity
caches.SharedManager.DiskDir = config.CacheDiskDir caches.SharedManager.MainDiskDir = config.CacheDiskDir
if len(config.HTTPCachePolicies) > 0 { caches.SharedManager.SubDiskDirs = subDirs
caches.SharedManager.UpdatePolicies(config.HTTPCachePolicies)
} else { if cachePoliciesChanged {
caches.SharedManager.UpdatePolicies([]*serverconfigs.HTTPCachePolicy{}) // copy
this.oldHTTPCachePolicies = []*serverconfigs.HTTPCachePolicy{}
err := jsonutils.Copy(&this.oldHTTPCachePolicies, config.HTTPCachePolicies)
if err != nil {
remotelogs.Error("NODE", "onReload: copy HTTPCachePolicies failed: "+err.Error())
}
// update
if len(config.HTTPCachePolicies) > 0 {
caches.SharedManager.UpdatePolicies(config.HTTPCachePolicies)
} else {
caches.SharedManager.UpdatePolicies([]*serverconfigs.HTTPCachePolicy{})
}
} }
// WAF策略 // WAF策略
waf.SharedWAFManager.UpdatePolicies(config.FindAllFirewallPolicies()) var allFirewallPolicies = config.FindAllFirewallPolicies()
iplibrary.SharedActionManager.UpdateActions(config.FirewallActions) if !jsonutils.Equal(allFirewallPolicies, this.oldHTTPFirewallPolicies) {
// copy
this.oldHTTPFirewallPolicies = []*firewallconfigs.HTTPFirewallPolicy{}
err := jsonutils.Copy(&this.oldHTTPFirewallPolicies, allFirewallPolicies)
if err != nil {
remotelogs.Error("NODE", "onReload: copy HTTPFirewallPolicies failed: "+err.Error())
}
// update
waf.SharedWAFManager.UpdatePolicies(allFirewallPolicies)
}
if !jsonutils.Equal(config.FirewallActions, this.oldFirewallActions) {
// copy
this.oldFirewallActions = []*firewallconfigs.FirewallActionConfig{}
err := jsonutils.Copy(&this.oldFirewallActions, config.FirewallActions)
if err != nil {
remotelogs.Error("NODE", "onReload: copy FirewallActionConfigs failed: "+err.Error())
}
// update
iplibrary.SharedActionManager.UpdateActions(config.FirewallActions)
}
// 统计指标 // 统计指标
metrics.SharedManager.Update(config.MetricItems) if !jsonutils.Equal(this.oldMetricItems, config.MetricItems) {
// copy
this.oldMetricItems = []*serverconfigs.MetricItemConfig{}
err := jsonutils.Copy(&this.oldMetricItems, config.MetricItems)
if err != nil {
remotelogs.Error("NODE", "onReload: copy MetricItemConfigs failed: "+err.Error())
}
// update
metrics.SharedManager.Update(config.MetricItems)
}
// max cpu // max cpu
if config.MaxCPU != this.maxCPU { if config.MaxCPU != this.oldMaxCPU {
if config.MaxCPU > 0 && config.MaxCPU < int32(runtime.NumCPU()) { if config.MaxCPU > 0 && config.MaxCPU < int32(runtime.NumCPU()) {
runtime.GOMAXPROCS(int(config.MaxCPU)) runtime.GOMAXPROCS(int(config.MaxCPU))
remotelogs.Println("NODE", "[CPU]set max cpu to '"+types.String(config.MaxCPU)+"'") remotelogs.Println("NODE", "[CPU]set max cpu to '"+types.String(config.MaxCPU)+"'")
@@ -952,11 +1080,11 @@ func (this *Node) onReload(config *nodeconfigs.NodeConfig) {
remotelogs.Println("NODE", "[CPU]set max cpu to '"+types.String(threads)+"'") remotelogs.Println("NODE", "[CPU]set max cpu to '"+types.String(threads)+"'")
} }
this.maxCPU = config.MaxCPU this.oldMaxCPU = config.MaxCPU
} }
// max threads // max threads
if config.MaxThreads != this.maxThreads { if config.MaxThreads != this.oldMaxThreads {
if config.MaxThreads > 0 { if config.MaxThreads > 0 {
debug.SetMaxThreads(config.MaxThreads) debug.SetMaxThreads(config.MaxThreads)
remotelogs.Println("NODE", "[THREADS]set max threads to '"+types.String(config.MaxThreads)+"'") remotelogs.Println("NODE", "[THREADS]set max threads to '"+types.String(config.MaxThreads)+"'")
@@ -964,7 +1092,7 @@ func (this *Node) onReload(config *nodeconfigs.NodeConfig) {
debug.SetMaxThreads(nodeconfigs.DefaultMaxThreads) debug.SetMaxThreads(nodeconfigs.DefaultMaxThreads)
remotelogs.Println("NODE", "[THREADS]set max threads to '"+types.String(nodeconfigs.DefaultMaxThreads)+"'") remotelogs.Println("NODE", "[THREADS]set max threads to '"+types.String(nodeconfigs.DefaultMaxThreads)+"'")
} }
this.maxThreads = config.MaxThreads this.oldMaxThreads = config.MaxThreads
} }
// timezone // timezone
@@ -973,7 +1101,7 @@ func (this *Node) onReload(config *nodeconfigs.NodeConfig) {
timeZone = "Asia/Shanghai" timeZone = "Asia/Shanghai"
} }
if this.timezone != timeZone { if this.oldTimezone != timeZone {
location, err := time.LoadLocation(timeZone) location, err := time.LoadLocation(timeZone)
if err != nil { if err != nil {
remotelogs.Error("NODE", "[TIMEZONE]change time zone failed: "+err.Error()) remotelogs.Error("NODE", "[TIMEZONE]change time zone failed: "+err.Error())
@@ -982,7 +1110,7 @@ func (this *Node) onReload(config *nodeconfigs.NodeConfig) {
remotelogs.Println("NODE", "[TIMEZONE]change time zone to '"+timeZone+"'") remotelogs.Println("NODE", "[TIMEZONE]change time zone to '"+timeZone+"'")
time.Local = location time.Local = location
this.timezone = timeZone this.oldTimezone = timeZone
} }
// product information // product information
@@ -1012,6 +1140,9 @@ func (this *Node) onReload(config *nodeconfigs.NodeConfig) {
remotelogs.Error("NODE", "[DNS_RESOLVER]set env failed: "+err.Error()) remotelogs.Error("NODE", "[DNS_RESOLVER]set env failed: "+err.Error())
} }
} }
// API Node地址这里不限制是否为空因为在为空时仍然要有对应的处理
this.changeAPINodeAddrs(config.APINodeAddrs)
} }
// reload server config // reload server config
@@ -1046,7 +1177,7 @@ func (this *Node) reloadServer() {
} }
} }
this.onReload(newNodeConfig) this.onReload(newNodeConfig, false)
err = sharedListenerManager.Start(newNodeConfig) err = sharedListenerManager.Start(newNodeConfig)
if err != nil { if err != nil {
@@ -1055,6 +1186,7 @@ func (this *Node) reloadServer() {
} }
} }
// 检查硬盘
func (this *Node) checkDisk() { func (this *Node) checkDisk() {
if runtime.GOOS != "linux" { if runtime.GOOS != "linux" {
return return
@@ -1075,3 +1207,69 @@ func (this *Node) checkDisk() {
} }
} }
} }
// 检查API节点地址
func (this *Node) changeAPINodeAddrs(apiNodeAddrs []*serverconfigs.NetworkAddressConfig) {
var addrs = []string{}
for _, addr := range apiNodeAddrs {
err := addr.Init()
if err != nil {
remotelogs.Error("NODE", "changeAPINodeAddrs: validate api node address '"+configutils.QuoteIP(addr.Host)+":"+addr.PortRange+"' failed: "+err.Error())
} else {
addrs = append(addrs, addr.FullAddresses()...)
}
}
sort.Strings(addrs)
if utils.EqualStrings(this.lastAPINodeAddrs, addrs) {
return
}
this.lastAPINodeAddrs = addrs
config, err := configs.LoadAPIConfig()
if err != nil {
remotelogs.Error("NODE", "changeAPINodeAddrs: "+err.Error())
return
}
if config == nil {
return
}
var oldEndpoints = config.RPC.Endpoints
rpcClient, err := rpc.SharedRPC()
if err != nil {
return
}
if len(addrs) > 0 {
this.lastAPINodeVersion++
var v = this.lastAPINodeVersion
// 异步检测,防止阻塞
go func(v int64) {
// 测试新的API节点地址
if rpcClient.TestEndpoints(addrs) {
config.RPC.Endpoints = addrs
} else {
config.RPC.Endpoints = oldEndpoints
this.lastAPINodeAddrs = nil // 恢复为空,以便于下次更新重试
}
// 检查测试中间有无新的变更
if v != this.lastAPINodeVersion {
return
}
err = rpcClient.UpdateConfig(config)
if err != nil {
remotelogs.Error("NODE", "changeAPINodeAddrs: update rpc config failed: "+err.Error())
}
}(v)
return
}
err = rpcClient.UpdateConfig(config)
if err != nil {
remotelogs.Error("NODE", "changeAPINodeAddrs: update rpc config failed: "+err.Error())
}
}

View File

@@ -31,12 +31,15 @@ type NodeStatusExecutor struct {
cpuLogicalCount int cpuLogicalCount int
cpuPhysicalCount int cpuPhysicalCount int
apiCallStat *rpc.CallStat
ticker *time.Ticker ticker *time.Ticker
} }
func NewNodeStatusExecutor() *NodeStatusExecutor { func NewNodeStatusExecutor() *NodeStatusExecutor {
return &NodeStatusExecutor{ return &NodeStatusExecutor{
ticker: time.NewTicker(30 * time.Second), ticker: time.NewTicker(30 * time.Second),
apiCallStat: rpc.NewCallStat(10),
} }
} }
@@ -78,6 +81,11 @@ func (this *NodeStatusExecutor) update() {
status.CacheTotalMemorySize = caches.SharedManager.TotalMemorySize() status.CacheTotalMemorySize = caches.SharedManager.TotalMemorySize()
status.TrafficInBytes = teaconst.InTrafficBytes status.TrafficInBytes = teaconst.InTrafficBytes
status.TrafficOutBytes = teaconst.OutTrafficBytes status.TrafficOutBytes = teaconst.OutTrafficBytes
apiSuccessPercent, apiAvgCostSeconds := this.apiCallStat.Sum()
status.APISuccessPercent = apiSuccessPercent
status.APIAvgCostSeconds = apiAvgCostSeconds
var localFirewall = firewalls.Firewall() var localFirewall = firewalls.Firewall()
if localFirewall != nil && !localFirewall.IsMock() { if localFirewall != nil && !localFirewall.IsMock() {
status.LocalFirewallName = localFirewall.Name() status.LocalFirewallName = localFirewall.Name()
@@ -125,9 +133,13 @@ func (this *NodeStatusExecutor) update() {
remotelogs.Error("NODE_STATUS", "failed to open rpc: "+err.Error()) remotelogs.Error("NODE_STATUS", "failed to open rpc: "+err.Error())
return return
} }
var before = time.Now()
_, err = rpcClient.NodeRPC.UpdateNodeStatus(rpcClient.Context(), &pb.UpdateNodeStatusRequest{ _, err = rpcClient.NodeRPC.UpdateNodeStatus(rpcClient.Context(), &pb.UpdateNodeStatusRequest{
StatusJSON: jsonData, StatusJSON: jsonData,
}) })
var costSeconds = time.Since(before).Seconds()
this.apiCallStat.Add(err == nil, costSeconds)
if err != nil { if err != nil {
if rpc.IsConnError(err) { if rpc.IsConnError(err) {
remotelogs.Warn("NODE_STATUS", "rpc UpdateNodeStatus() failed: "+err.Error()) remotelogs.Warn("NODE_STATUS", "rpc UpdateNodeStatus() failed: "+err.Error())
@@ -140,7 +152,7 @@ func (this *NodeStatusExecutor) update() {
// 更新CPU // 更新CPU
func (this *NodeStatusExecutor) updateCPU(status *nodeconfigs.NodeStatus) { func (this *NodeStatusExecutor) updateCPU(status *nodeconfigs.NodeStatus) {
duration := time.Duration(0) var duration = time.Duration(0)
if this.isFirstTime { if this.isFirstTime {
duration = 100 * time.Millisecond duration = 100 * time.Millisecond
} }
@@ -195,8 +207,8 @@ func (this *NodeStatusExecutor) updateDisk(status *nodeconfigs.NodeStatus) {
}) })
// 当前TeaWeb所在的fs // 当前TeaWeb所在的fs
rootFS := "" var rootFS = ""
rootTotal := uint64(0) var rootTotal = uint64(0)
if lists.ContainsString([]string{"darwin", "linux", "freebsd"}, runtime.GOOS) { if lists.ContainsString([]string{"darwin", "linux", "freebsd"}, runtime.GOOS) {
for _, p := range partitions { for _, p := range partitions {
if p.Mountpoint == "/" { if p.Mountpoint == "/" {
@@ -210,9 +222,9 @@ func (this *NodeStatusExecutor) updateDisk(status *nodeconfigs.NodeStatus) {
} }
} }
total := rootTotal var total = rootTotal
totalUsage := uint64(0) var totalUsage = uint64(0)
maxUsage := float64(0) var maxUsage = float64(0)
for _, partition := range partitions { for _, partition := range partitions {
if runtime.GOOS != "windows" && !strings.Contains(partition.Device, "/") && !strings.Contains(partition.Device, "\\") { if runtime.GOOS != "windows" && !strings.Contains(partition.Device, "/") && !strings.Contains(partition.Device, "\\") {
continue continue
@@ -252,7 +264,7 @@ func (this *NodeStatusExecutor) updateDisk(status *nodeconfigs.NodeStatus) {
// 缓存空间 // 缓存空间
func (this *NodeStatusExecutor) updateCacheSpace(status *nodeconfigs.NodeStatus) { func (this *NodeStatusExecutor) updateCacheSpace(status *nodeconfigs.NodeStatus) {
var result = []maps.Map{} var result = []maps.Map{}
cachePaths := caches.SharedManager.FindAllCachePaths() var cachePaths = caches.SharedManager.FindAllCachePaths()
for _, path := range cachePaths { for _, path := range cachePaths {
var stat unix.Statfs_t var stat unix.Statfs_t
err := unix.Statfs(path, &stat) err := unix.Statfs(path, &stat)

View File

@@ -1,5 +1,4 @@
//go:build !windows //go:build !windows
// +build !windows
package nodes package nodes

View File

@@ -1,23 +1,15 @@
package nodes package nodes
import ( import (
"context"
"crypto/tls"
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb" "github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
"github.com/TeaOSLab/EdgeNode/internal/configs" "github.com/TeaOSLab/EdgeNode/internal/configs"
"github.com/TeaOSLab/EdgeNode/internal/events" "github.com/TeaOSLab/EdgeNode/internal/events"
"github.com/TeaOSLab/EdgeNode/internal/goman" "github.com/TeaOSLab/EdgeNode/internal/goman"
"github.com/TeaOSLab/EdgeNode/internal/rpc" "github.com/TeaOSLab/EdgeNode/internal/rpc"
"github.com/TeaOSLab/EdgeNode/internal/trackers" "github.com/TeaOSLab/EdgeNode/internal/trackers"
"github.com/TeaOSLab/EdgeNode/internal/utils"
"github.com/iwind/TeaGo/Tea" "github.com/iwind/TeaGo/Tea"
"github.com/iwind/TeaGo/logs" "github.com/iwind/TeaGo/logs"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/insecure"
"net/url"
"sort"
"strings"
"sync"
"time" "time"
) )
@@ -64,6 +56,9 @@ func (this *SyncAPINodesTask) Stop() {
} }
func (this *SyncAPINodesTask) Loop() error { func (this *SyncAPINodesTask) Loop() error {
// 如果有节点定制的API节点地址
var hasCustomizedAPINodeAddrs = sharedNodeConfig != nil && len(sharedNodeConfig.APINodeAddrs) > 0
config, err := configs.LoadAPIConfig() config, err := configs.LoadAPIConfig()
if err != nil { if err != nil {
return err return err
@@ -96,21 +91,25 @@ func (this *SyncAPINodesTask) Loop() error {
} }
// 和现有的对比 // 和现有的对比
if this.isSame(newEndpoints, config.RPC.Endpoints) { if utils.EqualStrings(newEndpoints, config.RPC.Endpoints) {
return nil return nil
} }
// 测试是否有API节点可用 // 测试是否有API节点可用
var hasOk = this.testEndpoints(newEndpoints) var hasOk = rpcClient.TestEndpoints(newEndpoints)
if !hasOk { if !hasOk {
return nil return nil
} }
// 修改RPC对象配置 // 修改RPC对象配置
config.RPC.Endpoints = newEndpoints config.RPC.Endpoints = newEndpoints
err = rpcClient.UpdateConfig(config)
if err != nil { // 更新当前RPC
return err if !hasCustomizedAPINodeAddrs {
err = rpcClient.UpdateConfig(config)
if err != nil {
return err
}
} }
// 保存到文件 // 保存到文件
@@ -121,53 +120,3 @@ func (this *SyncAPINodesTask) Loop() error {
return nil return nil
} }
func (this *SyncAPINodesTask) isSame(endpoints1 []string, endpoints2 []string) bool {
sort.Strings(endpoints1)
sort.Strings(endpoints2)
return strings.Join(endpoints1, "&") == strings.Join(endpoints2, "&")
}
func (this *SyncAPINodesTask) testEndpoints(endpoints []string) bool {
if len(endpoints) == 0 {
return false
}
var wg = sync.WaitGroup{}
wg.Add(len(endpoints))
var ok = false
for _, endpoint := range endpoints {
go func(endpoint string) {
defer wg.Done()
u, err := url.Parse(endpoint)
if err != nil {
return
}
ctx, cancelFunc := context.WithTimeout(context.Background(), 5*time.Second)
defer func() {
cancelFunc()
}()
var conn *grpc.ClientConn
if u.Scheme == "http" {
conn, err = grpc.DialContext(ctx, u.Host, grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock())
} else if u.Scheme == "https" {
conn, err = grpc.DialContext(ctx, u.Host, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{
InsecureSkipVerify: true,
})), grpc.WithBlock())
}
if err != nil {
return
}
_ = conn.Close()
ok = true
}(endpoint)
}
wg.Wait()
return ok
}

View File

@@ -0,0 +1,53 @@
// 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) {
if userId <= 0 {
return true
}
this.locker.RLock()
u, ok := this.userMap[userId]
if ok {
isEnabled = u.ServersEnabled
} else {
isEnabled = true
}
this.locker.RUnlock()
return
}

View File

@@ -17,7 +17,7 @@ import (
"time" "time"
) )
var logChan = make(chan *pb.NodeLog, 1024) var logChan = make(chan *pb.NodeLog, 64) // 队列数量不需要太长,因为日志通常仅仅为调试用
func init() { func init() {
// 定期上传日志 // 定期上传日志

71
internal/rpc/call_stat.go Normal file
View File

@@ -0,0 +1,71 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package rpc
import (
"sync"
)
type callStatItem struct {
ok bool
costSeconds float64
}
type CallStat struct {
size int
items []*callStatItem
locker sync.Mutex
}
func NewCallStat(size int) *CallStat {
return &CallStat{
size: size,
}
}
func (this *CallStat) Add(ok bool, costSeconds float64) {
var size = this.size
if size <= 0 {
size = 10
}
this.locker.Lock()
this.items = append(this.items, &callStatItem{
ok: ok,
costSeconds: costSeconds,
})
if len(this.items) > size {
this.items = this.items[1:]
}
this.locker.Unlock()
}
func (this *CallStat) Sum() (successPercent float64, avgCostSeconds float64) {
this.locker.Lock()
defer this.locker.Unlock()
var size = this.size
if size <= 0 {
size = 10
}
var totalItems = len(this.items)
if totalItems <= size/2 /** 低于一半的采样率,不计入统计 **/ {
successPercent = 100
return
}
var totalOkItems = 0
var totalCostSeconds float64
for _, item := range this.items {
if item.ok {
totalOkItems++
}
totalCostSeconds += item.costSeconds
}
successPercent = float64(totalOkItems) * 100 / float64(totalItems)
avgCostSeconds = totalCostSeconds / float64(totalItems)
return
}

View File

@@ -0,0 +1,19 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package rpc_test
import (
"github.com/TeaOSLab/EdgeNode/internal/rpc"
"testing"
)
func TestNewCallStat(t *testing.T) {
var stat = rpc.NewCallStat(10)
stat.Add(true, 1)
stat.Add(true, 2)
stat.Add(true, 3)
stat.Add(false, 4)
stat.Add(true, 0)
stat.Add(true, 1)
t.Log(stat.Sum())
}

View File

@@ -49,6 +49,7 @@ type RPCClient struct {
FirewallRPC pb.FirewallServiceClient FirewallRPC pb.FirewallServiceClient
SSLCertRPC pb.SSLCertServiceClient SSLCertRPC pb.SSLCertServiceClient
ScriptRPC pb.ScriptServiceClient ScriptRPC pb.ScriptServiceClient
UserRPC pb.UserServiceClient
} }
func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) { func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) {
@@ -81,6 +82,7 @@ func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) {
client.FirewallRPC = pb.NewFirewallServiceClient(client) client.FirewallRPC = pb.NewFirewallServiceClient(client)
client.SSLCertRPC = pb.NewSSLCertServiceClient(client) client.SSLCertRPC = pb.NewSSLCertServiceClient(client)
client.ScriptRPC = pb.NewScriptServiceClient(client) client.ScriptRPC = pb.NewScriptServiceClient(client)
client.UserRPC = pb.NewUserServiceClient(client)
err := client.init() err := client.init()
if err != nil { if err != nil {
@@ -92,7 +94,6 @@ func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) {
// Context 节点上下文信息 // Context 节点上下文信息
func (this *RPCClient) Context() context.Context { func (this *RPCClient) Context() context.Context {
var ctx = context.Background()
var m = maps.Map{ var m = maps.Map{
"timestamp": time.Now().Unix(), "timestamp": time.Now().Unix(),
"type": "node", "type": "node",
@@ -109,6 +110,8 @@ func (this *RPCClient) Context() context.Context {
return context.Background() return context.Background()
} }
var token = base64.StdEncoding.EncodeToString(data) var token = base64.StdEncoding.EncodeToString(data)
var ctx = context.Background()
ctx = metadata.AppendToOutgoingContext(ctx, "nodeId", this.apiConfig.NodeId, "token", token) ctx = metadata.AppendToOutgoingContext(ctx, "nodeId", this.apiConfig.NodeId, "token", token)
return ctx return ctx
} }
@@ -157,6 +160,64 @@ func (this *RPCClient) UpdateConfig(config *configs.APIConfig) error {
return err return err
} }
// TestEndpoints 测试Endpoints是否可用
func (this *RPCClient) TestEndpoints(endpoints []string) bool {
if len(endpoints) == 0 {
return false
}
var wg = sync.WaitGroup{}
wg.Add(len(endpoints))
var ok = false
for _, endpoint := range endpoints {
go func(endpoint string) {
defer wg.Done()
u, err := url.Parse(endpoint)
if err != nil {
return
}
ctx, cancelFunc := context.WithTimeout(context.Background(), 5*time.Second)
defer func() {
cancelFunc()
}()
var conn *grpc.ClientConn
if u.Scheme == "http" {
conn, err = grpc.DialContext(ctx, u.Host, grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock())
} else if u.Scheme == "https" {
conn, err = grpc.DialContext(ctx, u.Host, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{
InsecureSkipVerify: true,
})), grpc.WithBlock())
} else {
return
}
if err != nil {
return
}
if conn == nil {
return
}
defer func() {
_ = conn.Close()
}()
var pingService = pb.NewPingServiceClient(conn)
_, err = pingService.Ping(this.Context(), &pb.PingRequest{})
if err != nil {
return
}
ok = true
}(endpoint)
}
wg.Wait()
return ok
}
// 初始化 // 初始化
func (this *RPCClient) init() error { func (this *RPCClient) init() error {
// 重新连接 // 重新连接
@@ -201,12 +262,18 @@ func (this *RPCClient) pickConn() *grpc.ClientConn {
defer this.locker.Unlock() defer this.locker.Unlock()
// 检查连接状态 // 检查连接状态
if len(this.conns) > 0 { var countConns = len(this.conns)
var availableConns = []*grpc.ClientConn{} if countConns > 0 {
if countConns == 1 {
return this.conns[0]
}
for _, stateArray := range [][2]connectivity.State{ for _, stateArray := range [][2]connectivity.State{
{connectivity.Ready, connectivity.Idle}, // 优先Ready和Idle {connectivity.Ready, connectivity.Idle}, // 优先Ready和Idle
{connectivity.Connecting, connectivity.Connecting}, {connectivity.Connecting, connectivity.Connecting},
{connectivity.TransientFailure, connectivity.TransientFailure},
} { } {
var availableConns = []*grpc.ClientConn{}
for _, conn := range this.conns { for _, conn := range this.conns {
var state = conn.GetState() var state = conn.GetState()
if state == stateArray[0] || state == stateArray[1] { if state == stateArray[0] || state == stateArray[1] {
@@ -217,26 +284,6 @@ func (this *RPCClient) pickConn() *grpc.ClientConn {
return this.randConn(availableConns) return this.randConn(availableConns)
} }
} }
if len(availableConns) > 0 {
return this.randConn(availableConns)
}
// 关闭
for _, conn := range this.conns {
_ = conn.Close()
}
}
// 重新初始化
err := this.init()
if err != nil {
// 错误提示已经在构造对象时打印过,所以这里不再重复打印
return nil
}
if len(this.conns) == 0 {
return nil
} }
return this.randConn(this.conns) return this.randConn(this.conns)
@@ -245,14 +292,15 @@ func (this *RPCClient) pickConn() *grpc.ClientConn {
func (this *RPCClient) Invoke(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error { func (this *RPCClient) Invoke(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error {
var conn = this.pickConn() var conn = this.pickConn()
if conn == nil { if conn == nil {
return errors.New("can not get available grpc connection") return errors.New("could not get available grpc connection")
} }
return conn.Invoke(ctx, method, args, reply, opts...) return conn.Invoke(ctx, method, args, reply, opts...)
} }
func (this *RPCClient) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) { func (this *RPCClient) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
var conn = this.pickConn() var conn = this.pickConn()
if conn == nil { if conn == nil {
return nil, errors.New("can not get available grpc connection") return nil, errors.New("could not get available grpc connection")
} }
return conn.NewStream(ctx, desc, method, opts...) return conn.NewStream(ctx, desc, method, opts...)
} }

62
internal/rpc/rpc_test.go Normal file
View File

@@ -0,0 +1,62 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package rpc_test
import (
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
"github.com/TeaOSLab/EdgeNode/internal/rpc"
_ "github.com/iwind/TeaGo/bootstrap"
timeutil "github.com/iwind/TeaGo/utils/time"
"sync"
"testing"
"time"
)
func TestRPCConcurrentCall(t *testing.T) {
rpcClient, err := rpc.SharedRPC()
if err != nil {
t.Fatal(err)
}
var before = time.Now()
defer func() {
t.Log("cost:", time.Since(before).Seconds()*1000, "ms")
}()
var concurrent = 3
var wg = sync.WaitGroup{}
wg.Add(concurrent)
for i := 0; i < concurrent; i++ {
go func() {
defer wg.Done()
_, err = rpcClient.NodeRPC.FindCurrentNodeConfig(rpcClient.Context(), &pb.FindCurrentNodeConfigRequest{})
if err != nil {
t.Log(err)
}
}()
}
wg.Wait()
}
func TestRPC_Retry(t *testing.T) {
rpcClient, err := rpc.SharedRPC()
if err != nil {
t.Fatal(err)
}
var ticker = time.NewTicker(1 * time.Second)
for range ticker.C {
go func() {
_, err = rpcClient.NodeRPC.FindCurrentNodeConfig(rpcClient.Context(), &pb.FindCurrentNodeConfigRequest{})
if err != nil {
t.Log(timeutil.Format("H:i:s"), err)
} else {
t.Log(timeutil.Format("H:i:s"), "success")
}
}()
}
}

View File

@@ -3,6 +3,7 @@
package stats package stats
import ( import (
"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb" "github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
"github.com/TeaOSLab/EdgeNode/internal/events" "github.com/TeaOSLab/EdgeNode/internal/events"
"github.com/TeaOSLab/EdgeNode/internal/goman" "github.com/TeaOSLab/EdgeNode/internal/goman"
@@ -42,6 +43,8 @@ type BandwidthStat struct {
type BandwidthStatManager struct { type BandwidthStatManager struct {
m map[string]*BandwidthStat // key => *BandwidthStat m map[string]*BandwidthStat // key => *BandwidthStat
pbStats []*pb.ServerBandwidthStat
lastTime string // 上一次执行的时间 lastTime string // 上一次执行的时间
ticker *time.Ticker ticker *time.Ticker
@@ -65,6 +68,12 @@ func (this *BandwidthStatManager) Start() {
} }
func (this *BandwidthStatManager) Loop() error { func (this *BandwidthStatManager) Loop() error {
var regionId int64
nodeConfig, _ := nodeconfigs.SharedNodeConfig()
if nodeConfig != nil {
regionId = nodeConfig.RegionId
}
var now = time.Now() var now = time.Now()
var day = timeutil.Format("Ymd", now) var day = timeutil.Format("Ymd", now)
var currentTime = timeutil.FormatTime("Hi", now.Unix()/300*300) var currentTime = timeutil.FormatTime("Hi", now.Unix()/300*300)
@@ -76,16 +85,29 @@ func (this *BandwidthStatManager) Loop() error {
var pbStats = []*pb.ServerBandwidthStat{} var pbStats = []*pb.ServerBandwidthStat{}
// 历史未提交记录
if len(this.pbStats) > 0 {
var expiredTime = timeutil.FormatTime("Hi", time.Now().Unix()-1200) // 只保留20分钟
for _, stat := range this.pbStats {
if stat.TimeAt > expiredTime {
pbStats = append(pbStats, stat)
}
}
this.pbStats = nil
}
this.locker.Lock() this.locker.Lock()
for key, stat := range this.m { for key, stat := range this.m {
if stat.Day < day || stat.TimeAt < currentTime { if stat.Day < day || stat.TimeAt < currentTime {
pbStats = append(pbStats, &pb.ServerBandwidthStat{ pbStats = append(pbStats, &pb.ServerBandwidthStat{
Id: 0, Id: 0,
UserId: stat.UserId, UserId: stat.UserId,
ServerId: stat.ServerId, ServerId: stat.ServerId,
Day: stat.Day, Day: stat.Day,
TimeAt: stat.TimeAt, TimeAt: stat.TimeAt,
Bytes: stat.MaxBytes / bandwidthTimestampDelim, Bytes: stat.MaxBytes / bandwidthTimestampDelim,
NodeRegionId: regionId,
}) })
delete(this.m, key) delete(this.m, key)
} }
@@ -100,6 +122,8 @@ func (this *BandwidthStatManager) Loop() error {
} }
_, err = rpcClient.ServerBandwidthStatRPC.UploadServerBandwidthStats(rpcClient.Context(), &pb.UploadServerBandwidthStatsRequest{ServerBandwidthStats: pbStats}) _, err = rpcClient.ServerBandwidthStatRPC.UploadServerBandwidthStats(rpcClient.Context(), &pb.UploadServerBandwidthStatsRequest{ServerBandwidthStats: pbStats})
if err != nil { if err != nil {
this.pbStats = pbStats
return err return err
} }
} }

View File

@@ -31,19 +31,33 @@ type TrafficItem struct {
CheckingTrafficLimit bool CheckingTrafficLimit bool
} }
func (this *TrafficItem) Add(anotherItem *TrafficItem) {
this.Bytes += anotherItem.Bytes
this.CachedBytes += anotherItem.CachedBytes
this.CountRequests += anotherItem.CountRequests
this.CountCachedRequests += anotherItem.CountCachedRequests
this.CountAttackRequests += anotherItem.CountAttackRequests
this.AttackBytes += anotherItem.AttackBytes
}
const trafficStatsMaxLife = 1200 // 最大只保存20分钟内的数据
// TrafficStatManager 区域流量统计 // TrafficStatManager 区域流量统计
type TrafficStatManager struct { type TrafficStatManager struct {
itemMap map[string]*TrafficItem // [timestamp serverId] => *TrafficItem itemMap map[string]*TrafficItem // [timestamp serverId] => *TrafficItem
domainsMap map[string]*TrafficItem // timestamp @ serverId @ domain => *TrafficItem domainsMap map[string]*TrafficItem // timestamp @ serverId @ domain => *TrafficItem
locker sync.Mutex
configFunc func() *nodeconfigs.NodeConfig pbItems []*pb.ServerDailyStat
pbDomainItems []*pb.UploadServerDailyStatsRequest_DomainStat
locker sync.Mutex
totalRequests int64 totalRequests int64
} }
// NewTrafficStatManager 获取新对象 // NewTrafficStatManager 获取新对象
func NewTrafficStatManager() *TrafficStatManager { func NewTrafficStatManager() *TrafficStatManager {
manager := &TrafficStatManager{ var manager = &TrafficStatManager{
itemMap: map[string]*TrafficItem{}, itemMap: map[string]*TrafficItem{},
domainsMap: map[string]*TrafficItem{}, domainsMap: map[string]*TrafficItem{},
} }
@@ -52,9 +66,7 @@ func NewTrafficStatManager() *TrafficStatManager {
} }
// Start 启动自动任务 // Start 启动自动任务
func (this *TrafficStatManager) Start(configFunc func() *nodeconfigs.NodeConfig) { func (this *TrafficStatManager) Start() {
this.configFunc = configFunc
// 上传请求总数 // 上传请求总数
var monitorTicker = time.NewTicker(1 * time.Minute) var monitorTicker = time.NewTicker(1 * time.Minute)
events.OnKey(events.EventQuit, this, func() { events.OnKey(events.EventQuit, this, func() {
@@ -70,7 +82,7 @@ func (this *TrafficStatManager) Start(configFunc func() *nodeconfigs.NodeConfig)
}) })
// 上传统计数据 // 上传统计数据
duration := 5 * time.Minute var duration = 5 * time.Minute
if Tea.IsTesting() { if Tea.IsTesting() {
// 测试环境缩短上传时间,方便我们调试 // 测试环境缩短上传时间,方便我们调试
duration = 30 * time.Second duration = 30 * time.Second
@@ -143,9 +155,10 @@ func (this *TrafficStatManager) Add(serverId int64, domain string, bytes int64,
// Upload 上传流量 // Upload 上传流量
func (this *TrafficStatManager) Upload() error { func (this *TrafficStatManager) Upload() error {
var config = this.configFunc() var regionId int64
if config == nil { nodeConfig, _ := nodeconfigs.SharedNodeConfig()
return nil if nodeConfig != nil {
regionId = nodeConfig.RegionId
} }
client, err := rpc.SharedRPC() client, err := rpc.SharedRPC()
@@ -154,10 +167,14 @@ func (this *TrafficStatManager) Upload() error {
} }
this.locker.Lock() this.locker.Lock()
var itemMap = this.itemMap var itemMap = this.itemMap
var domainMap = this.domainsMap var domainMap = this.domainsMap
// reset
this.itemMap = map[string]*TrafficItem{} this.itemMap = map[string]*TrafficItem{}
this.domainsMap = map[string]*TrafficItem{} this.domainsMap = map[string]*TrafficItem{}
this.locker.Unlock() this.locker.Unlock()
// 服务统计 // 服务统计
@@ -174,7 +191,7 @@ func (this *TrafficStatManager) Upload() error {
pbServerStats = append(pbServerStats, &pb.ServerDailyStat{ pbServerStats = append(pbServerStats, &pb.ServerDailyStat{
ServerId: serverId, ServerId: serverId,
RegionId: config.RegionId, NodeRegionId: regionId,
Bytes: item.Bytes, Bytes: item.Bytes,
CachedBytes: item.CachedBytes, CachedBytes: item.CachedBytes,
CountRequests: item.CountRequests, CountRequests: item.CountRequests,
@@ -186,9 +203,6 @@ func (this *TrafficStatManager) Upload() error {
CreatedAt: timestamp, CreatedAt: timestamp,
}) })
} }
if len(pbServerStats) == 0 {
return nil
}
// 域名统计 // 域名统计
var pbDomainStats = []*pb.UploadServerDailyStatsRequest_DomainStat{} var pbDomainStats = []*pb.UploadServerDailyStatsRequest_DomainStat{}
@@ -210,9 +224,40 @@ func (this *TrafficStatManager) Upload() error {
}) })
} }
// 历史未提交记录
if len(this.pbItems) > 0 || len(this.pbDomainItems) > 0 {
var expiredAt = time.Now().Unix() - 1200 // 只保留20分钟
for _, item := range this.pbItems {
if item.CreatedAt > expiredAt {
pbServerStats = append(pbServerStats, item)
}
}
this.pbItems = nil
for _, item := range this.pbDomainItems {
if item.CreatedAt > expiredAt {
pbDomainStats = append(pbDomainStats, item)
}
}
this.pbDomainItems = nil
}
if len(pbServerStats) == 0 && len(pbDomainStats) == 0 {
return nil
}
_, err = client.ServerDailyStatRPC.UploadServerDailyStats(client.Context(), &pb.UploadServerDailyStatsRequest{ _, err = client.ServerDailyStatRPC.UploadServerDailyStats(client.Context(), &pb.UploadServerDailyStatsRequest{
Stats: pbServerStats, Stats: pbServerStats,
DomainStats: pbDomainStats, DomainStats: pbDomainStats,
}) })
return err if err != nil {
// 加回历史记录
this.pbItems = pbServerStats
this.pbDomainItems = pbDomainStats
return err
}
return nil
} }

View File

@@ -73,6 +73,15 @@ func (this *UserAgentParser) Parse(userAgent string) (result UserAgentParserResu
result.BrowserName, result.BrowserVersion = this.parser.Browser() result.BrowserName, result.BrowserVersion = this.parser.Browser()
result.IsMobile = this.parser.Mobile() result.IsMobile = this.parser.Mobile()
// 忽略特殊字符
if len(result.BrowserName) > 0 {
for _, r := range result.BrowserName {
if r == '$' || r == '"' || r == '\'' || r == '<' || r == '>' || r == ')' {
return
}
}
}
if this.cacheCursor == 0 { if this.cacheCursor == 0 {
this.cacheMap1[userAgent] = result this.cacheMap1[userAgent] = result
if len(this.cacheMap1) >= this.maxCacheItems { if len(this.cacheMap1) >= this.maxCacheItems {

View File

@@ -2,47 +2,39 @@
package utils package utils
import "bytes" import (
"bytes"
"sync"
)
var SharedBufferPool = NewBufferPool()
// BufferPool pool for get byte slice // BufferPool pool for get byte slice
type BufferPool struct { type BufferPool struct {
c chan *bytes.Buffer rawPool *sync.Pool
} }
// NewBufferPool 创建新对象 // NewBufferPool 创建新对象
func NewBufferPool(maxSize int) *BufferPool { func NewBufferPool() *BufferPool {
if maxSize <= 0 { var pool = &BufferPool{}
maxSize = 1024 pool.rawPool = &sync.Pool{
} New: func() any {
pool := &BufferPool{ return &bytes.Buffer{}
c: make(chan *bytes.Buffer, maxSize), },
} }
return pool return pool
} }
// Get 获取一个新的Buffer // Get 获取一个新的Buffer
func (this *BufferPool) Get() (b *bytes.Buffer) { func (this *BufferPool) Get() (b *bytes.Buffer) {
select { var buffer = this.rawPool.Get().(*bytes.Buffer)
case b = <-this.c: if buffer.Len() > 0 {
b.Reset() buffer.Reset()
default:
b = &bytes.Buffer{}
} }
return return buffer
} }
// Put 放回一个使用过的byte slice // Put 放回一个使用过的byte slice
func (this *BufferPool) Put(b *bytes.Buffer) { func (this *BufferPool) Put(b *bytes.Buffer) {
b.Reset() this.rawPool.Put(b)
select {
case this.c <- b:
default:
// 已达最大容量,则抛弃
}
}
// Size 当前的数量
func (this *BufferPool) Size() int {
return len(this.c)
} }

View File

@@ -0,0 +1,47 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package utils_test
import (
"bytes"
"github.com/TeaOSLab/EdgeNode/internal/utils"
"strings"
"testing"
)
func TestNewBufferPool(t *testing.T) {
var pool = utils.NewBufferPool()
var b = pool.Get()
b.WriteString("Hello, World")
t.Log(b.String())
pool.Put(b)
t.Log(b.String())
b = pool.Get()
t.Log(b.String())
}
func BenchmarkNewBufferPool1(b *testing.B) {
var data = []byte(strings.Repeat("Hello", 1024))
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
var buffer = &bytes.Buffer{}
buffer.Write(data)
}
})
}
func BenchmarkNewBufferPool2(b *testing.B) {
var pool = utils.NewBufferPool()
var data = []byte(strings.Repeat("Hello", 1024))
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
var buffer = pool.Get()
buffer.Write(data)
pool.Put(buffer)
}
})
}

View File

@@ -39,6 +39,7 @@ func init() {
} }
type ClockManager struct { type ClockManager struct {
lastFailAt int64
} }
func NewClockManager() *ClockManager { func NewClockManager() *ClockManager {
@@ -51,7 +52,13 @@ func (this *ClockManager) Start() {
for range ticker.C { for range ticker.C {
err := this.Sync() err := this.Sync()
if err != nil { if err != nil {
remotelogs.Warn("CLOCK", "sync clock failed: "+err.Error()) var currentTimestamp = time.Now().Unix()
// 每天只提醒一次错误
if currentTimestamp-this.lastFailAt > 86400 {
remotelogs.Warn("CLOCK", "sync clock failed: "+err.Error())
this.lastFailAt = currentTimestamp
}
} }
} }
} }
@@ -72,6 +79,18 @@ func (this *ClockManager) Sync() error {
return nil return nil
} }
// check chrony
if config.CheckChrony {
chronycExe, err := exec.LookPath("chronyc")
if err == nil && len(chronycExe) > 0 {
var chronyCmd = executils.NewTimeoutCmd(3*time.Second, chronycExe, "tracking")
err = chronyCmd.Run()
if err == nil {
return nil
}
}
}
var server = config.Server var server = config.Server
if len(server) == 0 { if len(server) == 0 {
server = "pool.ntp.org" server = "pool.ntp.org"
@@ -118,7 +137,7 @@ func (this *ClockManager) syncNtpdate(ntpdate string, server string) error {
return nil return nil
} }
// 参考自https://medium.com/learning-the-go-programming-language/lets-make-an-ntp-client-in-go-287c4b9a969f // ReadServer 参考自https://medium.com/learning-the-go-programming-language/lets-make-an-ntp-client-in-go-287c4b9a969f
func (this *ClockManager) ReadServer(server string) (time.Time, error) { func (this *ClockManager) ReadServer(server string) (time.Time, error) {
conn, err := net.Dial("udp", server+":123") conn, err := net.Dial("udp", server+":123")
if err != nil { if err != nil {

View File

@@ -20,6 +20,9 @@ func NewDB(rawDB *sql.DB) *DB {
rawDB: rawDB, rawDB: rawDB,
} }
events.OnKey(events.EventQuit, fmt.Sprintf("db_%p", db), func() {
_ = rawDB.Close()
})
events.OnKey(events.EventTerminated, fmt.Sprintf("db_%p", db), func() { events.OnKey(events.EventTerminated, fmt.Sprintf("db_%p", db), func() {
_ = rawDB.Close() _ = rawDB.Close()
}) })

View File

@@ -7,7 +7,7 @@ import (
"github.com/iwind/TeaGo/maps" "github.com/iwind/TeaGo/maps"
) )
func MapToObject(m maps.Map, ptr interface{}) error { func MapToObject(m maps.Map, ptr any) error {
if m == nil { if m == nil {
return nil return nil
} }
@@ -18,7 +18,7 @@ func MapToObject(m maps.Map, ptr interface{}) error {
return json.Unmarshal(mJSON, ptr) return json.Unmarshal(mJSON, ptr)
} }
func ObjectToMap(ptr interface{}) (maps.Map, error) { func ObjectToMap(ptr any) (maps.Map, error) {
if ptr == nil { if ptr == nil {
return maps.Map{}, nil return maps.Map{}, nil
} }
@@ -33,3 +33,12 @@ func ObjectToMap(ptr interface{}) (maps.Map, error) {
} }
return result, nil return result, nil
} }
func Copy(destPtr any, srcPtr any) error {
data, err := json.Marshal(srcPtr)
if err != nil {
return err
}
err = json.Unmarshal(data, destPtr)
return err
}

View File

@@ -3,11 +3,12 @@
package jsonutils package jsonutils
import ( import (
"bytes"
"encoding/json" "encoding/json"
"testing" "testing"
) )
func PrintT(obj interface{}, t *testing.T) { func PrintT(obj any, t *testing.T) {
data, err := json.MarshalIndent(obj, "", " ") data, err := json.MarshalIndent(obj, "", " ")
if err != nil { if err != nil {
t.Log(err) t.Log(err)
@@ -15,3 +16,17 @@ func PrintT(obj interface{}, t *testing.T) {
t.Log(string(data)) t.Log(string(data))
} }
} }
func Equal(obj1 any, obj2 any) bool {
data1, err := json.Marshal(obj1)
if err != nil {
return false
}
data2, err := json.Marshal(obj2)
if err != nil {
return false
}
return bytes.Equal(data1, data2)
}

View File

@@ -0,0 +1,26 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package jsonutils_test
import (
"github.com/TeaOSLab/EdgeNode/internal/utils/jsonutils"
"github.com/iwind/TeaGo/assert"
"github.com/iwind/TeaGo/maps"
"testing"
)
func TestEqual(t *testing.T) {
var a = assert.NewAssertion(t)
{
var m1 = maps.Map{"a": 1, "b2": true}
var m2 = maps.Map{"b2": true, "a": 1}
a.IsTrue(jsonutils.Equal(m1, m2))
}
{
var m1 = maps.Map{"a": 1, "b2": true, "c": nil}
var m2 = maps.Map{"b2": true, "a": 1}
a.IsFalse(jsonutils.Equal(m1, m2))
}
}

View File

@@ -67,3 +67,10 @@ func TestRange_ComposeContentRangeHeader(t *testing.T) {
var r = rangeutils.NewRange(1, 100) var r = rangeutils.NewRange(1, 100)
t.Log(r.ComposeContentRangeHeader("1000")) t.Log(r.ComposeContentRangeHeader("1000"))
} }
func TestRange_SetLength(t *testing.T) {
var r = rangeutils.NewRange(1, 100)
t.Log(r)
t.Log(r.SetLength(1024))
}

View File

@@ -0,0 +1,135 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package readers
import (
"errors"
"github.com/iwind/TeaGo/types"
"io"
"sync"
)
type concurrentSubReader struct {
main *ConcurrentReaderList
index int
}
func (this *concurrentSubReader) Read(p []byte) (n int, err error) {
n, err = this.main.readIndex(p, this.index)
this.index++
return
}
func (this *concurrentSubReader) Close() error {
this.main.removeSubReader(this)
err := this.main.Close()
if err != nil {
return err
}
return nil
}
// ConcurrentReaderList
// TODO 动态调整 pieces = pieces[minPieceIndex:] 以节约内存
type ConcurrentReaderList struct {
locker sync.RWMutex
readLocker sync.Mutex
mainReader io.ReadCloser
subReaderMap map[*concurrentSubReader]bool
pieces [][]byte
lastErr error
}
func NewConcurrentReaderList(mainReader io.ReadCloser) *ConcurrentReaderList {
return &ConcurrentReaderList{
mainReader: mainReader,
subReaderMap: map[*concurrentSubReader]bool{},
}
}
func (this *ConcurrentReaderList) NewReader() io.ReadCloser {
var subReader = &concurrentSubReader{
main: this,
}
this.locker.Lock()
this.subReaderMap[subReader] = true
this.locker.Unlock()
return subReader
}
func (this *ConcurrentReaderList) read(p []byte) (n int, err error) {
n, err = this.mainReader.Read(p)
this.lastErr = err
if n > 0 {
var piece = make([]byte, n)
copy(piece, p[:n])
this.locker.Lock()
this.pieces = append(this.pieces, piece)
this.locker.Unlock()
}
return
}
func (this *ConcurrentReaderList) readIndex(p []byte, index int) (n int, err error) {
// 如果已经有数据
this.locker.RLock()
var countPieces = len(this.pieces)
if index < countPieces {
var piece = this.pieces[index]
this.locker.RUnlock()
var pn = len(piece)
if len(p) < pn {
err = errors.New("invalid buffer length '" + types.String(len(p)) + "' vs '" + types.String(len(piece)) + "'")
return
}
n = pn
copy(p, piece)
return
}
this.locker.RUnlock()
if this.lastErr != nil {
return 0, this.lastErr
}
// 如果没有数据,则读取之
this.readLocker.Lock()
// 再次检查数据是否已更新
this.locker.RLock()
if len(this.pieces) > countPieces || this.lastErr != nil {
this.locker.RUnlock()
this.readLocker.Unlock()
return this.readIndex(p, index)
}
this.locker.RUnlock()
// 从原始Reader中读取
n, err = this.read(p)
this.readLocker.Unlock()
if n > 0 {
// 重新尝试
return this.readIndex(p, index)
}
return
}
func (this *ConcurrentReaderList) removeSubReader(subReader *concurrentSubReader) {
this.locker.Lock()
delete(this.subReaderMap, subReader)
this.locker.Unlock()
}
func (this *ConcurrentReaderList) Close() error {
this.locker.Lock()
if len(this.subReaderMap) == 0 {
this.locker.Unlock()
return this.mainReader.Close()
}
this.locker.Unlock()
return nil
}

View File

@@ -0,0 +1,81 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package readers_test
import (
"bytes"
"github.com/TeaOSLab/EdgeNode/internal/utils/readers"
"io"
"sync"
"testing"
"time"
)
type testReader struct {
t *testing.T
rawReader io.Reader
}
func (this *testReader) Read(p []byte) (n int, err error) {
time.Sleep(1 * time.Second) // 延迟
return this.rawReader.Read(p)
}
func (this *testReader) Close() error {
this.t.Log("close")
return nil
}
func TestNewConcurrentReader(t *testing.T) {
var originBuffer = &bytes.Buffer{}
originBuffer.Write([]byte("0123456789_hello_world"))
var originLength = originBuffer.Len()
var concurrentReader = readers.NewConcurrentReaderList(&testReader{
t: t,
rawReader: originBuffer,
})
var threads = 32
var wg = &sync.WaitGroup{}
wg.Add(threads)
var locker = &sync.Mutex{}
var m = map[int][]byte{} // i => []byte
for i := 0; i < threads; i++ {
go func(i int) {
defer wg.Done()
var reader = concurrentReader.NewReader()
var buf = make([]byte, 4)
for {
n, err := reader.Read(buf)
if n > 0 {
locker.Lock()
m[i] = append(m[i], buf[:n]...)
locker.Unlock()
//t.Log(i, string(buf[:n]))
}
if err != nil {
if err == io.EOF {
break
}
t.Log("ERROR:", err)
}
}
_ = reader.Close()
}(i)
}
wg.Wait()
for i, b := range m {
if len(b) != originLength {
t.Fatal("ERROR:", i, string(b))
}
t.Log(i, string(b))
}
}

View File

@@ -0,0 +1,28 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package readers
import (
"io"
"log"
)
type PrintReader struct {
rawReader io.Reader
tag string
}
func NewPrintReader(rawReader io.Reader, tag string) io.Reader {
return &PrintReader{
rawReader: rawReader,
tag: tag,
}
}
func (this *PrintReader) Read(p []byte) (n int, err error) {
n, err = this.rawReader.Read(p)
if n > 0 {
log.Println("[" + this.tag + "]" + string(p[:n]))
}
return
}

View File

@@ -114,7 +114,8 @@ func (this *ServiceManager) installSystemdService(systemd, exePath string, args
var shortName = teaconst.SystemdServiceName var shortName = teaconst.SystemdServiceName
var longName = "GoEdge Node" // TODO 将来可以修改 var longName = "GoEdge Node" // TODO 将来可以修改
var desc = `# Provides: ` + shortName + ` var desc = `### BEGIN INIT INFO
# Provides: ` + shortName + `
# Required-Start: $all # Required-Start: $all
# Required-Stop: # Required-Stop:
# Default-Start: 2 3 4 5 # Default-Start: 2 3 4 5

View File

@@ -42,8 +42,8 @@ func ToValidUTF8string(v string) string {
return strings.ToValidUTF8(v, "") return strings.ToValidUTF8(v, "")
} }
// ContainsSameStrings 检查两个字符串slice内容是否一致 // EqualStrings 检查两个字符串slice内容是否一致
func ContainsSameStrings(s1 []string, s2 []string) bool { func EqualStrings(s1 []string, s2 []string) bool {
if len(s1) != len(s2) { if len(s1) != len(s2) {
return false return false
} }

View File

@@ -59,9 +59,9 @@ func TestFormatAddressList(t *testing.T) {
func TestContainsSameStrings(t *testing.T) { func TestContainsSameStrings(t *testing.T) {
var a = assert.NewAssertion(t) var a = assert.NewAssertion(t)
a.IsFalse(utils.ContainsSameStrings([]string{"a"}, []string{"b"})) a.IsFalse(utils.EqualStrings([]string{"a"}, []string{"b"}))
a.IsFalse(utils.ContainsSameStrings([]string{"a", "b"}, []string{"b"})) a.IsFalse(utils.EqualStrings([]string{"a", "b"}, []string{"b"}))
a.IsFalse(utils.ContainsSameStrings([]string{"a", "b"}, []string{"a", "b", "c"})) a.IsFalse(utils.EqualStrings([]string{"a", "b"}, []string{"a", "b", "c"}))
a.IsTrue(utils.ContainsSameStrings([]string{"a", "b"}, []string{"a", "b"})) a.IsTrue(utils.EqualStrings([]string{"a", "b"}, []string{"a", "b"}))
a.IsTrue(utils.ContainsSameStrings([]string{"a", "b"}, []string{"b", "a"})) a.IsTrue(utils.EqualStrings([]string{"a", "b"}, []string{"b", "a"}))
} }

View File

@@ -0,0 +1,28 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package writers
import (
"io"
"log"
)
type PrintWriter struct {
rawWriter io.Writer
tag string
}
func NewPrintWriter(rawWriter io.Writer, tag string) io.Writer {
return &PrintWriter{
rawWriter: rawWriter,
tag: tag,
}
}
func (this *PrintWriter) Write(p []byte) (n int, err error) {
n, err = this.rawWriter.Write(p)
if n > 0 {
log.Println("[" + this.tag + "]" + string(p[:n]))
}
return
}

View File

@@ -6,6 +6,7 @@ import (
) )
type Checkpoint struct { type Checkpoint struct {
priority int
} }
func (this *Checkpoint) Init() { func (this *Checkpoint) Init() {
@@ -36,6 +37,14 @@ func (this *Checkpoint) Stop() {
} }
func (this *Checkpoint) SetPriority(priority int) {
this.priority = priority
}
func (this *Checkpoint) Priority() int {
return this.priority
}
func (this *Checkpoint) RequestBodyIsEmpty(req requests.Request) bool { func (this *Checkpoint) RequestBodyIsEmpty(req requests.Request) bool {
if req.WAFRaw().ContentLength == 0 { if req.WAFRaw().ContentLength == 0 {
return true return true

View File

@@ -7,4 +7,5 @@ type CheckpointDefinition struct {
Prefix string Prefix string
HasParams bool // has sub params HasParams bool // has sub params
Instance CheckpointInterface Instance CheckpointInterface
Priority int
} }

View File

@@ -33,4 +33,10 @@ type CheckpointInterface interface {
// Stop stop // Stop stop
Stop() Stop()
// SetPriority set priority
SetPriority(priority int)
// get priority
Priority() int
} }

View File

@@ -0,0 +1,22 @@
package checkpoints
import (
"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
"github.com/iwind/TeaGo/maps"
)
type RequestCNAMECheckpoint struct {
Checkpoint
}
func (this *RequestCNAMECheckpoint) RequestValue(req requests.Request, param string, options maps.Map, ruleId int64) (value interface{}, hasRequestBody bool, sysErr error, userErr error) {
value = req.Format("${cname}")
return
}
func (this *RequestCNAMECheckpoint) ResponseValue(req requests.Request, resp *requests.Response, param string, options maps.Map, ruleId int64) (value interface{}, hasRequestBody bool, sysErr error, userErr error) {
if this.IsRequest() {
return this.RequestValue(req, param, options, ruleId)
}
return
}

View File

@@ -0,0 +1,26 @@
package checkpoints
import (
"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
"github.com/iwind/TeaGo/maps"
)
type RequestIsCNAMECheckpoint struct {
Checkpoint
}
func (this *RequestIsCNAMECheckpoint) RequestValue(req requests.Request, param string, options maps.Map, ruleId int64) (value interface{}, hasRequestBody bool, sysErr error, userErr error) {
if req.Format("${cname}") == req.Format("${host}") {
value = 1
} else {
value = 0
}
return
}
func (this *RequestIsCNAMECheckpoint) ResponseValue(req requests.Request, resp *requests.Response, param string, options maps.Map, ruleId int64) (value interface{}, hasRequestBody bool, sysErr error, userErr error) {
if this.IsRequest() {
return this.RequestValue(req, param, options, ruleId)
}
return
}

View File

@@ -41,19 +41,45 @@ func (this *RequestRefererBlockCheckpoint) RequestValue(req requests.Request, pa
return return
} }
var domains = options.GetSlice("allowDomains") // allow domains
var domainStrings = []string{} var allowDomains = options.GetSlice("allowDomains")
for _, domain := range domains { var allowDomainStrings = []string{}
domainStrings = append(domainStrings, types.String(domain)) for _, domain := range allowDomains {
allowDomainStrings = append(allowDomainStrings, types.String(domain))
} }
if len(domainStrings) == 0 { // deny domains
var denyDomains = options.GetSlice("denyDomains")
var denyDomainStrings = []string{}
for _, domain := range denyDomains {
denyDomainStrings = append(denyDomainStrings, types.String(domain))
}
if len(allowDomainStrings) == 0 {
if len(denyDomainStrings) > 0 {
if configutils.MatchDomains(denyDomainStrings, host) {
value = 0
} else {
value = 1
}
return
}
value = 0 value = 0
return return
} }
if configutils.MatchDomains(domainStrings, host) { if configutils.MatchDomains(allowDomainStrings, host) {
if len(denyDomainStrings) > 0 {
if configutils.MatchDomains(denyDomainStrings, host) {
value = 0
} else {
value = 1
}
return
}
value = 1 value = 1
return
} else { } else {
value = 0 value = 0
} }

View File

@@ -8,6 +8,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "通用Header比如Cache-Control、Accept之类的长度限制防止缓冲区溢出攻击", Description: "通用Header比如Cache-Control、Accept之类的长度限制防止缓冲区溢出攻击",
HasParams: false, HasParams: false,
Instance: new(RequestGeneralHeaderLengthCheckpoint), Instance: new(RequestGeneralHeaderLengthCheckpoint),
Priority: 100,
}, },
{ {
Name: "客户端地址IP", Name: "客户端地址IP",
@@ -15,6 +16,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "试图通过分析X-Forwarded-For等Header获取的客户端地址比如192.168.1.100", Description: "试图通过分析X-Forwarded-For等Header获取的客户端地址比如192.168.1.100",
HasParams: false, HasParams: false,
Instance: new(RequestRemoteAddrCheckpoint), Instance: new(RequestRemoteAddrCheckpoint),
Priority: 100,
}, },
{ {
Name: "客户端源地址IP", Name: "客户端源地址IP",
@@ -22,6 +24,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "直接连接的客户端地址比如192.168.1.100", Description: "直接连接的客户端地址比如192.168.1.100",
HasParams: false, HasParams: false,
Instance: new(RequestRawRemoteAddrCheckpoint), Instance: new(RequestRawRemoteAddrCheckpoint),
Priority: 100,
}, },
{ {
Name: "客户端端口", Name: "客户端端口",
@@ -29,6 +32,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "直接连接的客户端地址端口", Description: "直接连接的客户端地址端口",
HasParams: false, HasParams: false,
Instance: new(RequestRemotePortCheckpoint), Instance: new(RequestRemotePortCheckpoint),
Priority: 100,
}, },
{ {
Name: "客户端用户名", Name: "客户端用户名",
@@ -36,6 +40,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "通过BasicAuth登录的客户端用户名", Description: "通过BasicAuth登录的客户端用户名",
HasParams: false, HasParams: false,
Instance: new(RequestRemoteUserCheckpoint), Instance: new(RequestRemoteUserCheckpoint),
Priority: 100,
}, },
{ {
Name: "请求URI", Name: "请求URI",
@@ -43,6 +48,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "包含URL参数的请求URI类似于 /hello/world?lang=go", Description: "包含URL参数的请求URI类似于 /hello/world?lang=go",
HasParams: false, HasParams: false,
Instance: new(RequestURICheckpoint), Instance: new(RequestURICheckpoint),
Priority: 100,
}, },
{ {
Name: "请求路径", Name: "请求路径",
@@ -50,6 +56,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "不包含URL参数的请求路径类似于 /hello/world", Description: "不包含URL参数的请求路径类似于 /hello/world",
HasParams: false, HasParams: false,
Instance: new(RequestPathCheckpoint), Instance: new(RequestPathCheckpoint),
Priority: 100,
}, },
{ {
Name: "请求URL", Name: "请求URL",
@@ -57,6 +64,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "完整的请求URL包含协议、域名、请求路径、参数等类似于 https://example.com/hello?name=lily", Description: "完整的请求URL包含协议、域名、请求路径、参数等类似于 https://example.com/hello?name=lily",
HasParams: false, HasParams: false,
Instance: new(RequestURLCheckpoint), Instance: new(RequestURLCheckpoint),
Priority: 100,
}, },
{ {
Name: "请求内容长度", Name: "请求内容长度",
@@ -64,6 +72,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "请求Header中的Content-Length", Description: "请求Header中的Content-Length",
HasParams: false, HasParams: false,
Instance: new(RequestLengthCheckpoint), Instance: new(RequestLengthCheckpoint),
Priority: 100,
}, },
{ {
Name: "请求体内容", Name: "请求体内容",
@@ -71,6 +80,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "通常在POST或者PUT等操作时会附带请求体最大限制32M", Description: "通常在POST或者PUT等操作时会附带请求体最大限制32M",
HasParams: false, HasParams: false,
Instance: new(RequestBodyCheckpoint), Instance: new(RequestBodyCheckpoint),
Priority: 5,
}, },
{ {
Name: "请求URI和请求体组合", Name: "请求URI和请求体组合",
@@ -78,6 +88,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "${requestURI}和${requestBody}组合", Description: "${requestURI}和${requestBody}组合",
HasParams: false, HasParams: false,
Instance: new(RequestAllCheckpoint), Instance: new(RequestAllCheckpoint),
Priority: 5,
}, },
{ {
Name: "请求表单参数", Name: "请求表单参数",
@@ -85,6 +96,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "获取POST或者其他方法发送的表单参数最大请求体限制32M", Description: "获取POST或者其他方法发送的表单参数最大请求体限制32M",
HasParams: true, HasParams: true,
Instance: new(RequestFormArgCheckpoint), Instance: new(RequestFormArgCheckpoint),
Priority: 5,
}, },
{ {
Name: "上传文件", Name: "上传文件",
@@ -92,6 +104,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "获取POST上传的文件信息最大请求体限制32M", Description: "获取POST上传的文件信息最大请求体限制32M",
HasParams: true, HasParams: true,
Instance: new(RequestUploadCheckpoint), Instance: new(RequestUploadCheckpoint),
Priority: 20,
}, },
{ {
Name: "请求JSON参数", Name: "请求JSON参数",
@@ -99,6 +112,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "获取POST或者其他方法发送的JSON最大请求体限制32M使用点.)符号表示多级数据", Description: "获取POST或者其他方法发送的JSON最大请求体限制32M使用点.)符号表示多级数据",
HasParams: true, HasParams: true,
Instance: new(RequestJSONArgCheckpoint), Instance: new(RequestJSONArgCheckpoint),
Priority: 5,
}, },
{ {
Name: "请求方法", Name: "请求方法",
@@ -106,6 +120,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "比如GET、POST", Description: "比如GET、POST",
HasParams: false, HasParams: false,
Instance: new(RequestMethodCheckpoint), Instance: new(RequestMethodCheckpoint),
Priority: 100,
}, },
{ {
Name: "请求协议", Name: "请求协议",
@@ -113,6 +128,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "比如http或https", Description: "比如http或https",
HasParams: false, HasParams: false,
Instance: new(RequestSchemeCheckpoint), Instance: new(RequestSchemeCheckpoint),
Priority: 100,
}, },
{ {
Name: "HTTP协议版本", Name: "HTTP协议版本",
@@ -120,6 +136,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "比如HTTP/1.1", Description: "比如HTTP/1.1",
HasParams: false, HasParams: false,
Instance: new(RequestProtoCheckpoint), Instance: new(RequestProtoCheckpoint),
Priority: 100,
}, },
{ {
Name: "主机名", Name: "主机名",
@@ -127,6 +144,23 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "比如teaos.cn", Description: "比如teaos.cn",
HasParams: false, HasParams: false,
Instance: new(RequestHostCheckpoint), Instance: new(RequestHostCheckpoint),
Priority: 100,
},
{
Name: "CNAME",
Prefix: "cname",
Description: "当前网站服务CNAME比如38b48e4f.goedge.cn",
HasParams: false,
Instance: new(RequestCNAMECheckpoint),
Priority: 100,
},
{
Name: "是否为CNAME",
Prefix: "isCNAME",
Description: "是否为CNAME值为1或0",
HasParams: false,
Instance: new(RequestIsCNAMECheckpoint),
Priority: 100,
}, },
{ {
Name: "请求来源URL", Name: "请求来源URL",
@@ -134,6 +168,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "请求Header中的Referer值", Description: "请求Header中的Referer值",
HasParams: false, HasParams: false,
Instance: new(RequestRefererCheckpoint), Instance: new(RequestRefererCheckpoint),
Priority: 100,
}, },
{ {
Name: "客户端信息", Name: "客户端信息",
@@ -141,6 +176,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "比如Mozilla/5.0 AppleWebKit/537.36 (KHTML, like Gecko) Chrome/73.0.3683.103", Description: "比如Mozilla/5.0 AppleWebKit/537.36 (KHTML, like Gecko) Chrome/73.0.3683.103",
HasParams: false, HasParams: false,
Instance: new(RequestUserAgentCheckpoint), Instance: new(RequestUserAgentCheckpoint),
Priority: 100,
}, },
{ {
Name: "内容类型", Name: "内容类型",
@@ -148,6 +184,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "请求Header的Content-Type", Description: "请求Header的Content-Type",
HasParams: false, HasParams: false,
Instance: new(RequestContentTypeCheckpoint), Instance: new(RequestContentTypeCheckpoint),
Priority: 100,
}, },
{ {
Name: "所有cookie组合字符串", Name: "所有cookie组合字符串",
@@ -155,6 +192,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "比如sid=IxZVPFhE&city=beijing&uid=18237", Description: "比如sid=IxZVPFhE&city=beijing&uid=18237",
HasParams: false, HasParams: false,
Instance: new(RequestCookiesCheckpoint), Instance: new(RequestCookiesCheckpoint),
Priority: 100,
}, },
{ {
Name: "单个cookie值", Name: "单个cookie值",
@@ -162,6 +200,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "单个cookie值", Description: "单个cookie值",
HasParams: true, HasParams: true,
Instance: new(RequestCookieCheckpoint), Instance: new(RequestCookieCheckpoint),
Priority: 100,
}, },
{ {
Name: "所有URL参数组合", Name: "所有URL参数组合",
@@ -169,6 +208,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "比如name=lu&age=20", Description: "比如name=lu&age=20",
HasParams: false, HasParams: false,
Instance: new(RequestArgsCheckpoint), Instance: new(RequestArgsCheckpoint),
Priority: 100,
}, },
{ {
Name: "单个URL参数值", Name: "单个URL参数值",
@@ -176,6 +216,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "单个URL参数值", Description: "单个URL参数值",
HasParams: true, HasParams: true,
Instance: new(RequestArgCheckpoint), Instance: new(RequestArgCheckpoint),
Priority: 100,
}, },
{ {
Name: "所有Header信息", Name: "所有Header信息",
@@ -183,6 +224,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "使用\\n隔开的Header信息字符串", Description: "使用\\n隔开的Header信息字符串",
HasParams: false, HasParams: false,
Instance: new(RequestHeadersCheckpoint), Instance: new(RequestHeadersCheckpoint),
Priority: 100,
}, },
{ {
Name: "单个Header值", Name: "单个Header值",
@@ -190,6 +232,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "单个Header值", Description: "单个Header值",
HasParams: true, HasParams: true,
Instance: new(RequestHeaderCheckpoint), Instance: new(RequestHeaderCheckpoint),
Priority: 100,
}, },
{ {
Name: "国家/地区名称", Name: "国家/地区名称",
@@ -197,6 +240,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "国家/地区名称", Description: "国家/地区名称",
HasParams: false, HasParams: false,
Instance: new(RequestGeoCountryNameCheckpoint), Instance: new(RequestGeoCountryNameCheckpoint),
Priority: 90,
}, },
{ {
Name: "省份名称", Name: "省份名称",
@@ -204,6 +248,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "中国省份名称", Description: "中国省份名称",
HasParams: false, HasParams: false,
Instance: new(RequestGeoProvinceNameCheckpoint), Instance: new(RequestGeoProvinceNameCheckpoint),
Priority: 90,
}, },
{ {
Name: "城市名称", Name: "城市名称",
@@ -211,6 +256,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "中国城市名称", Description: "中国城市名称",
HasParams: false, HasParams: false,
Instance: new(RequestGeoCityNameCheckpoint), Instance: new(RequestGeoCityNameCheckpoint),
Priority: 90,
}, },
{ {
Name: "ISP名称", Name: "ISP名称",
@@ -218,6 +264,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "ISP名称", Description: "ISP名称",
HasParams: false, HasParams: false,
Instance: new(RequestISPNameCheckpoint), Instance: new(RequestISPNameCheckpoint),
Priority: 90,
}, },
{ {
Name: "CC统计", Name: "CC统计",
@@ -225,6 +272,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "统计某段时间段内的请求信息", Description: "统计某段时间段内的请求信息",
HasParams: true, HasParams: true,
Instance: new(CCCheckpoint), Instance: new(CCCheckpoint),
Priority: 10,
}, },
{ {
Name: "CC统计", Name: "CC统计",
@@ -232,6 +280,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "统计某段时间段内的请求信息", Description: "统计某段时间段内的请求信息",
HasParams: true, HasParams: true,
Instance: new(CC2Checkpoint), Instance: new(CC2Checkpoint),
Priority: 10,
}, },
{ {
Name: "防盗链", Name: "防盗链",
@@ -239,6 +288,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "阻止一些域名访问引用本站资源", Description: "阻止一些域名访问引用本站资源",
HasParams: true, HasParams: true,
Instance: new(RequestRefererBlockCheckpoint), Instance: new(RequestRefererBlockCheckpoint),
Priority: 20,
}, },
{ {
Name: "通用响应Header长度限制", Name: "通用响应Header长度限制",
@@ -246,6 +296,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "通用Header比如Cache-Control、Accept之类的长度限制防止缓冲区溢出攻击", Description: "通用Header比如Cache-Control、Accept之类的长度限制防止缓冲区溢出攻击",
HasParams: false, HasParams: false,
Instance: new(ResponseGeneralHeaderLengthCheckpoint), Instance: new(ResponseGeneralHeaderLengthCheckpoint),
Priority: 100,
}, },
{ {
Name: "响应状态码", Name: "响应状态码",
@@ -253,6 +304,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "响应状态码比如200、404、500", Description: "响应状态码比如200、404、500",
HasParams: false, HasParams: false,
Instance: new(ResponseStatusCheckpoint), Instance: new(ResponseStatusCheckpoint),
Priority: 100,
}, },
{ {
Name: "响应Header", Name: "响应Header",
@@ -260,6 +312,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "响应Header值", Description: "响应Header值",
HasParams: true, HasParams: true,
Instance: new(ResponseHeaderCheckpoint), Instance: new(ResponseHeaderCheckpoint),
Priority: 100,
}, },
{ {
Name: "响应内容", Name: "响应内容",
@@ -267,6 +320,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "响应内容字符串", Description: "响应内容字符串",
HasParams: false, HasParams: false,
Instance: new(ResponseBodyCheckpoint), Instance: new(ResponseBodyCheckpoint),
Priority: 5,
}, },
{ {
Name: "响应内容长度", Name: "响应内容长度",
@@ -274,6 +328,7 @@ var AllCheckpoints = []*CheckpointDefinition{
Description: "响应内容长度通过响应的Header Content-Length获取", Description: "响应内容长度通过响应的Header Content-Length获取",
HasParams: false, HasParams: false,
Instance: new(ResponseBytesSentCheckpoint), Instance: new(ResponseBytesSentCheckpoint),
Priority: 100,
}, },
} }
@@ -281,6 +336,7 @@ var AllCheckpoints = []*CheckpointDefinition{
func FindCheckpoint(prefix string) CheckpointInterface { func FindCheckpoint(prefix string) CheckpointInterface {
for _, def := range AllCheckpoints { for _, def := range AllCheckpoints {
if def.Prefix == prefix { if def.Prefix == prefix {
def.Instance.SetPriority(def.Priority)
return def.Instance return def.Instance
} }
} }

View File

@@ -35,6 +35,7 @@ type Rule struct {
Value string `yaml:"value" json:"value"` // compared value Value string `yaml:"value" json:"value"` // compared value
IsCaseInsensitive bool `yaml:"isCaseInsensitive" json:"isCaseInsensitive"` IsCaseInsensitive bool `yaml:"isCaseInsensitive" json:"isCaseInsensitive"`
CheckpointOptions map[string]interface{} `yaml:"checkpointOptions" json:"checkpointOptions"` CheckpointOptions map[string]interface{} `yaml:"checkpointOptions" json:"checkpointOptions"`
Priority int `yaml:"priority" json:"priority"`
checkpointFinder func(prefix string) checkpoints.CheckpointInterface checkpointFinder func(prefix string) checkpoints.CheckpointInterface
@@ -132,9 +133,9 @@ func (this *Rule) Init() error {
} }
if singleParamRegexp.MatchString(this.Param) { if singleParamRegexp.MatchString(this.Param) {
param := this.Param[2 : len(this.Param)-1] var param = this.Param[2 : len(this.Param)-1]
pieces := strings.SplitN(param, ".", 2) var pieces = strings.SplitN(param, ".", 2)
prefix := pieces[0] var prefix = pieces[0]
if len(pieces) == 1 { if len(pieces) == 1 {
this.singleParam = "" this.singleParam = ""
} else { } else {
@@ -142,18 +143,20 @@ func (this *Rule) Init() error {
} }
if this.checkpointFinder != nil { if this.checkpointFinder != nil {
checkpoint := this.checkpointFinder(prefix) var checkpoint = this.checkpointFinder(prefix)
if checkpoint == nil { if checkpoint == nil {
return errors.New("no check point '" + prefix + "' found") return errors.New("no check point '" + prefix + "' found")
} }
this.singleCheckpoint = checkpoint this.singleCheckpoint = checkpoint
this.Priority = checkpoint.Priority()
} else { } else {
checkpoint := checkpoints.FindCheckpoint(prefix) var checkpoint = checkpoints.FindCheckpoint(prefix)
if checkpoint == nil { if checkpoint == nil {
return errors.New("no check point '" + prefix + "' found") return errors.New("no check point '" + prefix + "' found")
} }
checkpoint.Init() checkpoint.Init()
this.singleCheckpoint = checkpoint this.singleCheckpoint = checkpoint
this.Priority = checkpoint.Priority()
} }
return nil return nil
@@ -162,22 +165,24 @@ func (this *Rule) Init() error {
this.multipleCheckpoints = map[string]checkpoints.CheckpointInterface{} this.multipleCheckpoints = map[string]checkpoints.CheckpointInterface{}
var err error = nil var err error = nil
configutils.ParseVariables(this.Param, func(varName string) (value string) { configutils.ParseVariables(this.Param, func(varName string) (value string) {
pieces := strings.SplitN(varName, ".", 2) var pieces = strings.SplitN(varName, ".", 2)
prefix := pieces[0] var prefix = pieces[0]
if this.checkpointFinder != nil { if this.checkpointFinder != nil {
checkpoint := this.checkpointFinder(prefix) var checkpoint = this.checkpointFinder(prefix)
if checkpoint == nil { if checkpoint == nil {
err = errors.New("no check point '" + prefix + "' found") err = errors.New("no check point '" + prefix + "' found")
} else { } else {
this.multipleCheckpoints[prefix] = checkpoint this.multipleCheckpoints[prefix] = checkpoint
this.Priority = checkpoint.Priority()
} }
} else { } else {
checkpoint := checkpoints.FindCheckpoint(prefix) var checkpoint = checkpoints.FindCheckpoint(prefix)
if checkpoint == nil { if checkpoint == nil {
err = errors.New("no check point '" + prefix + "' found") err = errors.New("no check point '" + prefix + "' found")
} else { } else {
checkpoint.Init() checkpoint.Init()
this.multipleCheckpoints[prefix] = checkpoint this.multipleCheckpoints[prefix] = checkpoint
this.Priority = checkpoint.Priority()
} }
} }
return "" return ""

View File

@@ -52,6 +52,11 @@ func (this *RuleSet) Init(waf *WAF) error {
return errors.New("init rule '" + rule.Param + " " + rule.Operator + " " + types.String(rule.Value) + "' failed: " + err.Error()) return errors.New("init rule '" + rule.Param + " " + rule.Operator + " " + types.String(rule.Value) + "' failed: " + err.Error())
} }
} }
// sort by priority
sort.Slice(this.Rules, func(i, j int) bool {
return this.Rules[i].Priority > this.Rules[j].Priority
})
} }
// action codes // action codes

View File

@@ -73,6 +73,7 @@ func (this *WAF) Init() (resultErrors []error) {
for _, def := range checkpoints.AllCheckpoints { for _, def := range checkpoints.AllCheckpoints {
instance := reflect.New(reflect.Indirect(reflect.ValueOf(def.Instance)).Type()).Interface().(checkpoints.CheckpointInterface) instance := reflect.New(reflect.Indirect(reflect.ValueOf(def.Instance)).Type()).Interface().(checkpoints.CheckpointInterface)
instance.Init() instance.Init()
instance.SetPriority(def.Priority)
this.checkpointsMap[def.Prefix] = instance this.checkpointsMap[def.Prefix] = instance
} }