修复连接快速关闭时数据可能无法完整传输的问题

go.mod

@@ -8,8 +8,9 @@ require (
 	github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d // indirect
 	github.com/TeaOSLab/EdgeCommon v0.0.0-00010101000000-000000000000
 	github.com/andybalholm/brotli v1.0.3
+	github.com/biessek/golang-ico v0.0.0-20180326222316-d348d9ea4670
 	github.com/cespare/xxhash v1.1.0
-	github.com/chai2010/webp v1.1.0
+	github.com/chai2010/webp v1.1.0 // indirect
 	github.com/dchest/captcha v0.0.0-20200903113550-03f5f0333e1f
 	github.com/dop251/goja v0.0.0-20210804101310-32956a348b49
 	github.com/go-ole/go-ole v1.2.4 // indirect
@@ -18,9 +19,13 @@ require (
 	github.com/iwind/TeaGo v0.0.0-20210628135026-38575a4ab060
 	github.com/iwind/gofcgi v0.0.0-20210528023741-a92711d45f11
 	github.com/iwind/gosock v0.0.0-20210722083328-12b2d66abec3
+	github.com/iwind/gowebp v0.0.0-20211029040624-7331ecc78ed8
+	github.com/jsummers/gobmp v0.0.0-20151104160322-e2ba15ffa76e // indirect
 	github.com/lionsoul2014/ip2region v2.2.0-release+incompatible
 	github.com/mattn/go-sqlite3 v2.0.3+incompatible
+	github.com/miekg/dns v1.1.43
 	github.com/mssola/user_agent v0.5.2
+	github.com/pires/go-proxyproto v0.6.1
 	github.com/shirou/gopsutil v3.21.5+incompatible
 	github.com/tklauser/go-sysconf v0.3.6 // indirect
 	golang.org/x/image v0.0.0-20190802002840-cff245a6509b

go.sum

@@ -10,6 +10,8 @@ github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d/go.mod h1:3eOhrU
 github.com/andybalholm/brotli v1.0.3 h1:fpcw+r1N1h0Poc1F/pHbW40cUm/lMEQslZtCkBQ0UnM=
 github.com/andybalholm/brotli v1.0.3/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
 github.com/benbjohnson/clock v1.0.3/go.mod h1:bGMdMPoPVvcYyt1gHDf4J2KE153Yf9BuiUKYMaxlTDM=
+github.com/biessek/golang-ico v0.0.0-20180326222316-d348d9ea4670 h1:FQPKKjDhzG0T4ew6dm6MGrXb4PRAi8ZmTuYuxcF62BM=
+github.com/biessek/golang-ico v0.0.0-20180326222316-d348d9ea4670/go.mod h1:iRWAFbKXMMkVQyxZ1PfGlkBr1TjATx1zy2MRprV7A3Q=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
@@ -80,7 +82,11 @@ github.com/iwind/gofcgi v0.0.0-20210528023741-a92711d45f11 h1:DaQjoWZhLNxjhIXedV
 github.com/iwind/gofcgi v0.0.0-20210528023741-a92711d45f11/go.mod h1:JtbX20untAjUVjZs1ZBtq80f5rJWvwtQNRL6EnuYRnY=
 github.com/iwind/gosock v0.0.0-20210722083328-12b2d66abec3 h1:aBSonas7vFcgTj9u96/bWGILGv1ZbUSTLiOzcI1ZT6c=
 github.com/iwind/gosock v0.0.0-20210722083328-12b2d66abec3/go.mod h1:H5Q7SXwbx3a97ecJkaS2sD77gspzE7HFUafBO0peEyA=
+github.com/iwind/gowebp v0.0.0-20211029040624-7331ecc78ed8 h1:AojsHz9Es9B3He2MQQxeRq3TyD//o9huxUo7r1wh44g=
+github.com/iwind/gowebp v0.0.0-20211029040624-7331ecc78ed8/go.mod h1:QJBY2txYhLMzwLO29iB5ujDJ3s3V7DsZ582nw4Ss+tM=
 github.com/json-iterator/go v1.1.10/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
+github.com/jsummers/gobmp v0.0.0-20151104160322-e2ba15ffa76e h1:LvL4XsI70QxOGHed6yhQtAU34Kx3Qq2wwBzGFKY8zKk=
+github.com/jsummers/gobmp v0.0.0-20151104160322-e2ba15ffa76e/go.mod h1:kLgvv7o6UM+0QSf0QjAse3wReFDsb9qbZJdfexWlrQw=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
@@ -92,6 +98,8 @@ github.com/lionsoul2014/ip2region v2.2.0-release+incompatible h1:1qp9iks+69h7IGL
 github.com/lionsoul2014/ip2region v2.2.0-release+incompatible/go.mod h1:+ZBN7PBoh5gG6/y0ZQ85vJDBe21WnfbRrQQwTfliJJI=
 github.com/mattn/go-sqlite3 v2.0.3+incompatible h1:gXHsfypPkaMZrKbD5209QV9jbUTJKjyR5WD3HYQSd+U=
 github.com/mattn/go-sqlite3 v2.0.3+incompatible/go.mod h1:FPy6KqzDD04eiIsT53CuJW3U88zkxoIYsOqkbpncsNc=
+github.com/miekg/dns v1.1.43 h1:JKfpVSCB84vrAmHzyrsxB5NAr5kLoMXZArPSw7Qlgyg=
+github.com/miekg/dns v1.1.43/go.mod h1:+evo5L0630/F6ca/Z9+GAqzhjGyn8/c+TBaOyfEl0V4=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/mssola/user_agent v0.5.2 h1:CZkTUahjL1+OcZ5zv3kZr8QiJ8jy2H08vZIEkBeRbxo=
@@ -106,6 +114,8 @@ github.com/onsi/gomega v1.7.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1Cpa
 github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
 github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo=
 github.com/opentracing/opentracing-go v1.1.1-0.20190913142402-a7454ce5950e/go.mod h1:UkNAQd3GIcIGf0SeVgPpRdFStlNbqXla1AfSYxPUl2o=
+github.com/pires/go-proxyproto v0.6.1 h1:EBupykFmo22SDjv4fQVQd2J9NOoLPmyZA/15ldOGkPw=
+github.com/pires/go-proxyproto v0.6.1/go.mod h1:Odh9VFOZJCf9G8cLW5o435Xf1J95Jw9Gw5rnCjcwzAY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
@@ -157,6 +167,7 @@ golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAG
 golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=

internal/apps/app_cmd.go

@@ -184,6 +184,8 @@ func (this *AppCmd) runStart() {
 		return
 	}
 
+	_ = os.Setenv("EdgeBackground", "on")
+
 	cmd := exec.Command(os.Args[0])
 	err := cmd.Start()
 	if err != nil {

internal/apps/log_writer.go

@@ -6,6 +6,10 @@ import (
 	"github.com/iwind/TeaGo/logs"
 	"github.com/iwind/TeaGo/utils/time"
 	"log"
+	"os"
+	"path/filepath"
+	"runtime"
+	"strconv"
 )
 
 type LogWriter struct {
@@ -34,7 +38,24 @@ func (this *LogWriter) Init() {
 }
 
 func (this *LogWriter) Write(message string) {
-	log.Println(message)
+	// 文件和行号
+	var callDepth = 2
+	var file string
+	var line int
+	var ok bool
+	_, file, line, ok = runtime.Caller(callDepth)
+	if ok {
+		file = filepath.Base(file)
+	}
+
+	backgroundEnv, _ := os.LookupEnv("EdgeBackground")
+	if backgroundEnv != "on" {
+		if len(file) > 0 {
+			log.Println(message + " (" + file + ":" + strconv.Itoa(line) + ")")
+		} else {
+			log.Println(message)
+		}
+	}
 
 	if this.fileAppender != nil {
 		_, err := this.fileAppender.AppendString(timeutil.Format("Y/m/d H:i:s ") + message + "\n")

internal/caches/item.go

@@ -27,7 +27,7 @@ func (this *Item) IsExpired() bool {
 }
 
 func (this *Item) TotalSize() int64 {
-	return this.Size() + this.MetaSize + int64(len(this.Key)) + 64
+	return this.Size() + this.MetaSize + int64(len(this.Key)) + int64(len(this.Host)) + 64
 }
 
 func (this *Item) Size() int64 {

internal/caches/storage_memory_test.go

@@ -1,9 +1,12 @@
 package caches
 
 import (
+	"bytes"
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
 	"github.com/iwind/TeaGo/logs"
 	"github.com/iwind/TeaGo/rands"
+	"runtime"
+	"runtime/debug"
 	"strconv"
 	"testing"
 	"time"
@@ -265,3 +268,38 @@ func TestMemoryStorage_Locker(t *testing.T) {
 	}
 	t.Log("ok")
 }
+
+func TestMemoryStorage_Stop(t *testing.T) {
+	var stat1 = &runtime.MemStats{}
+	runtime.ReadMemStats(stat1)
+
+	var m = map[uint64]*MemoryItem{}
+	for i := 0; i < 1_000_000; i++ {
+		m[uint64(i)] = &MemoryItem{
+			HeaderValue: []byte("Hello, World"),
+			BodyValue:   bytes.Repeat([]byte("Hello"), 1024),
+		}
+	}
+
+	m = map[uint64]*MemoryItem{}
+
+	var before = time.Now()
+	//runtime.GC()
+	debug.FreeOSMemory()
+	/**go func() {
+		time.Sleep(10 * time.Second)
+		runtime.GC()
+	}()**/
+	t.Log(time.Since(before).Seconds()*1000, "ms")
+
+	var stat2 = &runtime.MemStats{}
+	runtime.ReadMemStats(stat2)
+
+	if stat2.HeapInuse > stat1.HeapInuse {
+		t.Log(stat2.HeapInuse, stat1.HeapInuse, (stat2.HeapInuse-stat1.HeapInuse)/1024/1024, "MB")
+	} else {
+		t.Log("0 MB")
+	}
+
+	t.Log(len(m))
+}

internal/compressions/reader.go

@@ -0,0 +1,8 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package compressions
+
+type Reader interface {
+	Read(p []byte) (n int, err error)
+	Close() error
+}

internal/compressions/reader_brotli.go

@@ -0,0 +1,24 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package compressions
+
+import (
+	"github.com/andybalholm/brotli"
+	"io"
+)
+
+type BrotliReader struct {
+	reader *brotli.Reader
+}
+
+func NewBrotliReader(reader io.Reader) (Reader, error) {
+	return &BrotliReader{reader: brotli.NewReader(reader)}, nil
+}
+
+func (this *BrotliReader) Read(p []byte) (n int, err error) {
+	return this.reader.Read(p)
+}
+
+func (this *BrotliReader) Close() error {
+	return nil
+}

internal/compressions/reader_deflate.go

@@ -0,0 +1,24 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package compressions
+
+import (
+	"compress/flate"
+	"io"
+)
+
+type DeflateReader struct {
+	reader io.ReadCloser
+}
+
+func NewDeflateReader(reader io.Reader) (Reader, error) {
+	return &DeflateReader{reader: flate.NewReader(reader)}, nil
+}
+
+func (this *DeflateReader) Read(p []byte) (n int, err error) {
+	return this.reader.Read(p)
+}
+
+func (this *DeflateReader) Close() error {
+	return this.reader.Close()
+}

internal/compressions/reader_gzip.go

@@ -0,0 +1,30 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package compressions
+
+import (
+	"compress/gzip"
+	"io"
+)
+
+type GzipReader struct {
+	reader *gzip.Reader
+}
+
+func NewGzipReader(reader io.Reader) (Reader, error) {
+	r, err := gzip.NewReader(reader)
+	if err != nil {
+		return nil, err
+	}
+	return &GzipReader{
+		reader: r,
+	}, nil
+}
+
+func (this *GzipReader) Read(p []byte) (n int, err error) {
+	return this.reader.Read(p)
+}
+
+func (this *GzipReader) Close() error {
+	return this.reader.Close()
+}

internal/compressions/utils.go

@@ -8,6 +8,31 @@ import (
 	"io"
 )
 
+type ContentEncoding = string
+
+const (
+	ContentEncodingBr      ContentEncoding = "br"
+	ContentEncodingGzip    ContentEncoding = "gzip"
+	ContentEncodingDeflate ContentEncoding = "deflate"
+)
+
+var ErrNotSupportedContentEncoding = errors.New("not supported content encoding")
+
+// NewReader 获取Reader
+func NewReader(reader io.Reader, contentEncoding ContentEncoding) (Reader, error) {
+	switch contentEncoding {
+	case ContentEncodingBr:
+		return NewBrotliReader(reader)
+	case ContentEncodingGzip:
+		return NewGzipReader(reader)
+	case ContentEncodingDeflate:
+		return NewDeflateReader(reader)
+	}
+	return nil, ErrNotSupportedContentEncoding
+}
+
+// NewWriter 获取Writer
+// TODO 考虑重用Writer
 func NewWriter(writer io.Writer, compressType serverconfigs.HTTPCompressionType, level int) (Writer, error) {
 	switch compressType {
 	case serverconfigs.HTTPCompressionTypeGzip:

internal/compressions/writer_encoding.go

@@ -0,0 +1,51 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package compressions
+
+import (
+	"bytes"
+	"io"
+)
+
+type EncodingWriter struct {
+	contentEncoding ContentEncoding
+	writer          Writer
+	buf             *bytes.Buffer
+}
+
+func NewEncodingWriter(contentEncoding ContentEncoding, writer Writer) (Writer, error) {
+	return &EncodingWriter{
+		contentEncoding: contentEncoding,
+		writer:          writer,
+		buf:             &bytes.Buffer{},
+	}, nil
+}
+
+func (this *EncodingWriter) Write(p []byte) (int, error) {
+	return this.buf.Write(p)
+}
+
+func (this *EncodingWriter) Flush() error {
+	return this.writer.Flush()
+}
+
+func (this *EncodingWriter) Close() error {
+	reader, err := NewReader(this.buf, this.contentEncoding)
+	if err != nil {
+		_ = this.writer.Close()
+		return err
+	}
+	_, err = io.Copy(this.writer, reader)
+	if err != nil {
+		_ = reader.Close()
+		_ = this.writer.Close()
+		return err
+	}
+
+	_ = reader.Close()
+	return this.writer.Close()
+}
+
+func (this *EncodingWriter) Level() int {
+	return this.writer.Level()
+}

internal/compressions/writer_encoding_test.go

@@ -0,0 +1,47 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package compressions
+
+import (
+	"bytes"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
+	"testing"
+)
+
+func TestNewEncodingWriter(t *testing.T) {
+	var buf = &bytes.Buffer{}
+
+	subWriter, err := NewWriter(buf, serverconfigs.HTTPCompressionTypeGzip, 5)
+	if err != nil {
+		t.Fatal(err)
+	}
+	writer, err := NewEncodingWriter(ContentEncodingGzip, subWriter)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	gzipBuf := &bytes.Buffer{}
+	gzipWriter, err := NewGzipWriter(gzipBuf, 5)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = gzipWriter.Write([]byte("Hello"))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = gzipWriter.Write([]byte("World"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	_ = gzipWriter.Close()
+
+	_, err = writer.Write(gzipBuf.Bytes())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_ = writer.Close()
+
+	t.Log(buf.String())
+}

internal/const/const.go

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

internal/nodes/client_conn.go

@@ -42,17 +42,25 @@ func init() {
 	})
 }
 
-// TrafficConn 用于统计流量的连接
+// ClientConn 客户端连接
-type TrafficConn struct {
+type ClientConn struct {
 	rawConn  net.Conn
 	isClosed bool
 }
 
-func NewTrafficConn(conn net.Conn) net.Conn {
+func NewClientConn(conn net.Conn, quickClose bool) net.Conn {
-	return &TrafficConn{rawConn: conn}
+	if quickClose {
+		tcpConn, ok := conn.(*net.TCPConn)
+		if ok {
+			// TODO 可以设置此值
+			_ = tcpConn.SetLinger(3)
+		}
+	}
+
+	return &ClientConn{rawConn: conn}
 }
 
-func (this *TrafficConn) Read(b []byte) (n int, err error) {
+func (this *ClientConn) Read(b []byte) (n int, err error) {
 	n, err = this.rawConn.Read(b)
 	if n > 0 {
 		atomic.AddUint64(&inTrafficBytes, uint64(n))
@@ -60,7 +68,7 @@ func (this *TrafficConn) Read(b []byte) (n int, err error) {
 	return
 }
 
-func (this *TrafficConn) Write(b []byte) (n int, err error) {
+func (this *ClientConn) Write(b []byte) (n int, err error) {
 	n, err = this.rawConn.Write(b)
 	if n > 0 {
 		atomic.AddUint64(&outTrafficBytes, uint64(n))
@@ -68,31 +76,31 @@ func (this *TrafficConn) Write(b []byte) (n int, err error) {
 	return
 }
 
-func (this *TrafficConn) Close() error {
+func (this *ClientConn) Close() error {
 	this.isClosed = true
 	return this.rawConn.Close()
 }
 
-func (this *TrafficConn) LocalAddr() net.Addr {
+func (this *ClientConn) LocalAddr() net.Addr {
 	return this.rawConn.LocalAddr()
 }
 
-func (this *TrafficConn) RemoteAddr() net.Addr {
+func (this *ClientConn) RemoteAddr() net.Addr {
 	return this.rawConn.RemoteAddr()
 }
 
-func (this *TrafficConn) SetDeadline(t time.Time) error {
+func (this *ClientConn) SetDeadline(t time.Time) error {
 	return this.rawConn.SetDeadline(t)
 }
 
-func (this *TrafficConn) SetReadDeadline(t time.Time) error {
+func (this *ClientConn) SetReadDeadline(t time.Time) error {
 	return this.rawConn.SetReadDeadline(t)
 }
 
-func (this *TrafficConn) SetWriteDeadline(t time.Time) error {
+func (this *ClientConn) SetWriteDeadline(t time.Time) error {
 	return this.rawConn.SetWriteDeadline(t)
 }
 
-func (this *TrafficConn) IsClosed() bool {
+func (this *ClientConn) IsClosed() bool {
 	return this.isClosed
 }

internal/nodes/client_conn_utils.go

@@ -0,0 +1,22 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"net"
+)
+
+// 判断客户端连接是否已关闭
+func isClientConnClosed(conn net.Conn) bool {
+	if conn == nil {
+		return true
+	}
+	clientConn, ok := conn.(*ClientConn)
+	if ok {
+		return clientConn.IsClosed()
+	}
+
+	// TODO 解决tls.Conn无法获取底层连接对象的问题
+
+	return false
+}

internal/nodes/client_listener.go

@@ -0,0 +1,54 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
+	"github.com/TeaOSLab/EdgeNode/internal/waf"
+	"net"
+)
+
+// ClientListener 客户端网络监听
+type ClientListener struct {
+	rawListener net.Listener
+	quickClose  bool
+}
+
+func NewClientListener(listener net.Listener, quickClose bool) net.Listener {
+	return &ClientListener{
+		rawListener: listener,
+		quickClose:  quickClose,
+	}
+}
+
+func (this *ClientListener) Accept() (net.Conn, error) {
+	conn, err := this.rawListener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	// 是否在WAF名单中
+	ip, _, err := net.SplitHostPort(conn.RemoteAddr().String())
+	if err == nil {
+		if !waf.SharedIPWhiteList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, ip) &&
+			waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, ip) {
+
+			tcpConn, ok := conn.(*net.TCPConn)
+			if ok {
+				_ = tcpConn.SetLinger(0)
+			}
+
+			_ = conn.Close()
+			return this.Accept()
+		}
+	}
+
+	return NewClientConn(conn, this.quickClose), nil
+}
+
+func (this *ClientListener) Close() error {
+	return this.rawListener.Close()
+}
+
+func (this *ClientListener) Addr() net.Addr {
+	return this.rawListener.Addr()
+}

internal/nodes/http_client_pool.go

@@ -6,10 +6,12 @@ import (
 	"errors"
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+	"github.com/pires/go-proxyproto"
 	"net"
 	"net/http"
 	"runtime"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 )
@@ -37,7 +39,7 @@ func NewHTTPClientPool() *HTTPClientPool {
 }
 
 // Client 根据地址获取客户端
-func (this *HTTPClientPool) Client(req *http.Request, origin *serverconfigs.OriginConfig, originAddr string) (rawClient *http.Client, err error) {
+func (this *HTTPClientPool) Client(req *HTTPRequest, origin *serverconfigs.OriginConfig, originAddr string, proxyProtocol *serverconfigs.ProxyProtocolConfig) (rawClient *http.Client, err error) {
 	if origin.Addr == nil {
 		return nil, errors.New("origin addr should not be empty (originId:" + strconv.FormatInt(origin.Id, 10) + ")")
 	}
@@ -73,11 +75,11 @@ func (this *HTTPClientPool) Client(req *http.Request, origin *serverconfigs.Orig
 		numberCPU = 8
 	}
 	if maxConnections <= 0 {
-		maxConnections = numberCPU * 8
+		maxConnections = numberCPU * 32
 	}
 
 	if idleConns <= 0 {
-		idleConns = numberCPU * 4
+		idleConns = numberCPU * 8
 	}
 	//logs.Println("[ORIGIN]max connections:", maxConnections)
 
@@ -105,7 +107,7 @@ func (this *HTTPClientPool) Client(req *http.Request, origin *serverconfigs.Orig
 				for i := 1; i <= retries; i++ {
 					port := int(toaConfig.RandLocalPort())
 					// TODO 思考是否支持X-Real-IP/X-Forwarded-IP
-					err := sharedTOAManager.SendMsg("add:" + strconv.Itoa(port) + ":" + req.RemoteAddr)
+					err := sharedTOAManager.SendMsg("add:" + strconv.Itoa(port) + ":" + req.requestRemoteAddr(true))
 					if err != nil {
 						remotelogs.Error("TOA", "add failed: "+err.Error())
 					} else {
@@ -126,10 +128,43 @@ func (this *HTTPClientPool) Client(req *http.Request, origin *serverconfigs.Orig
 			}
 
 			// 普通的连接
-			return (&net.Dialer{
+			conn, err := (&net.Dialer{
 				Timeout:   connectionTimeout,
 				KeepAlive: 1 * time.Minute,
 			}).DialContext(ctx, network, originAddr)
+			if err != nil {
+				return nil, err
+			}
+
+			if proxyProtocol != nil && proxyProtocol.IsOn && (proxyProtocol.Version == serverconfigs.ProxyProtocolVersion1 || proxyProtocol.Version == serverconfigs.ProxyProtocolVersion2) {
+				var remoteAddr = req.requestRemoteAddr(true)
+				var transportProtocol = proxyproto.TCPv4
+				if strings.Contains(remoteAddr, ":") {
+					transportProtocol = proxyproto.TCPv6
+				}
+				var destAddr = conn.RemoteAddr()
+				var reqConn = req.RawReq.Context().Value(HTTPConnContextKey)
+				if reqConn != nil {
+					destAddr = reqConn.(net.Conn).LocalAddr()
+				}
+				header := proxyproto.Header{
+					Version:           byte(proxyProtocol.Version),
+					Command:           proxyproto.PROXY,
+					TransportProtocol: transportProtocol,
+					SourceAddr: &net.TCPAddr{
+						IP:   net.ParseIP(remoteAddr),
+						Port: req.requestRemotePort(),
+					},
+					DestinationAddr: destAddr,
+				}
+				_, err = header.WriteTo(conn)
+				if err != nil {
+					_ = conn.Close()
+					return nil, err
+				}
+			}
+
+			return conn, nil
 		},
 		MaxIdleConns:          0,
 		MaxIdleConnsPerHost:   idleConns,

internal/nodes/http_client_pool_test.go

@@ -21,14 +21,14 @@ func TestHTTPClientPool_Client(t *testing.T) {
 			t.Fatal(err)
 		}
 		{
-			client, err := pool.Client(nil, origin, origin.Addr.PickAddress())
+			client, err := pool.Client(nil, origin, origin.Addr.PickAddress(), nil)
 			if err != nil {
 				t.Fatal(err)
 			}
 			t.Log("client:", client)
 		}
 		for i := 0; i < 10; i++ {
-			client, err := pool.Client(nil, origin, origin.Addr.PickAddress())
+			client, err := pool.Client(nil, origin, origin.Addr.PickAddress(), nil)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -53,7 +53,7 @@ func TestHTTPClientPool_cleanClients(t *testing.T) {
 
 	for i := 0; i < 10; i++ {
 		t.Log("get", i)
-		_, _ = pool.Client(nil, origin, origin.Addr.PickAddress())
+		_, _ = pool.Client(nil, origin, origin.Addr.PickAddress(), nil)
 		time.Sleep(1 * time.Second)
 	}
 }
@@ -73,6 +73,6 @@ func BenchmarkHTTPClientPool_Client(b *testing.B) {
 
 	pool := NewHTTPClientPool()
 	for i := 0; i < b.N; i++ {
-		_, _ = pool.Client(nil, origin, origin.Addr.PickAddress())
+		_, _ = pool.Client(nil, origin, origin.Addr.PickAddress(), nil)
 	}
 }

internal/nodes/http_request.go

@@ -126,6 +126,24 @@ func (this *HTTPRequest) Do() {
 		return
 	}
 
+	// 特殊URL处理
+	if len(this.rawURI) > 1 && this.rawURI[1] == '.' {
+		// ACME
+		// TODO 需要配置是否启用ACME检测
+		if strings.HasPrefix(this.rawURI, "/.well-known/acme-challenge/") {
+			this.doACME()
+			this.doEnd()
+			return
+		}
+	}
+
+	// 带宽限制
+	if this.Server.BandwidthLimit != nil && this.Server.BandwidthLimit.IsOn && !this.Server.BandwidthLimit.IsEmpty() && this.Server.BandwidthLimitStatus != nil && this.Server.BandwidthLimitStatus.IsValid() {
+		this.doBandwidthLimit()
+		this.doEnd()
+		return
+	}
+
 	// WAF
 	if this.web.FirewallRef != nil && this.web.FirewallRef.IsOn {
 		if this.doWAFRequest() {
@@ -167,12 +185,10 @@ func (this *HTTPRequest) Do() {
 
 // 开始调用
 func (this *HTTPRequest) doBegin() {
-	// 特殊URL处理
+	// 处理健康检查
-	if len(this.rawURI) > 1 && this.rawURI[1] == '.' {
+	var healthCheckKey = this.RawReq.Header.Get(serverconfigs.HealthCheckHeaderName)
-		// ACME
+	if len(healthCheckKey) > 0 {
-		// TODO 需要配置是否启用ACME检测
+		if this.doHealthCheck(healthCheckKey) {
-		if strings.HasPrefix(this.rawURI, "/.well-known/acme-challenge/") {
-			this.doACME()
 			return
 		}
 	}

internal/nodes/http_request_bandwidth_limit.go

@@ -0,0 +1,24 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
+)
+
+// 带宽限制
+func (this *HTTPRequest) doBandwidthLimit() {
+	var config = this.Server.BandwidthLimit
+
+	this.tags = append(this.tags, "bandwidth")
+
+	var statusCode = 509
+	this.processResponseHeaders(statusCode)
+
+	this.writer.WriteHeader(statusCode)
+	if len(config.NoticePageBody) != 0 {
+		_, _ = this.writer.WriteString(config.NoticePageBody)
+	} else {
+		_, _ = this.writer.WriteString(serverconfigs.DefaultBandwidthLimitNoticePageBody)
+	}
+}

internal/nodes/http_request_cache.go

@@ -3,8 +3,10 @@ package nodes
 import (
 	"bytes"
 	"errors"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
 	"github.com/TeaOSLab/EdgeNode/internal/caches"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+	"github.com/TeaOSLab/EdgeNode/internal/rpc"
 	"net/http"
 	"path/filepath"
 	"strconv"
@@ -109,6 +111,32 @@ func (this *HTTPRequest) doCacheRead() (shouldStop bool) {
 		return
 	}
 
+	// 判断是否在Purge
+	if this.web.Cache.PurgeIsOn && strings.ToUpper(this.RawReq.Method) == "PURGE" && this.RawReq.Header.Get("X-Edge-Purge-Key") == this.web.Cache.PurgeKey {
+		err := storage.Delete(key)
+		if err != nil {
+			remotelogs.Error("HTTP_REQUEST_CACHE", "purge failed: "+err.Error())
+		}
+
+		go func() {
+			rpcClient, err := rpc.SharedRPC()
+			if err == nil {
+				for _, rpcServerService := range rpcClient.ServerRPCList() {
+					_, err = rpcServerService.PurgeServerCache(rpcClient.Context(), &pb.PurgeServerCacheRequest{
+						Domains:  []string{this.Host},
+						Keys:     []string{key},
+						Prefixes: nil,
+					})
+					if err != nil {
+						remotelogs.Error("HTTP_REQUEST_CACHE", "purge failed: "+err.Error())
+					}
+				}
+			}
+		}()
+
+		return true
+	}
+
 	buf := bytePool32k.Get()
 	defer func() {
 		bytePool32k.Put(buf)

internal/nodes/http_request_health_check.go

@@ -0,0 +1,27 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/nodeutils"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+)
+
+// 健康检查
+func (this *HTTPRequest) doHealthCheck(key string) (stop bool) {
+	this.tags = append(this.tags, "healthCheck")
+
+	this.RawReq.Header.Del(serverconfigs.HealthCheckHeaderName)
+
+	data, err := nodeutils.DecryptData(sharedNodeConfig.NodeId, sharedNodeConfig.Secret, key)
+	if err != nil {
+		remotelogs.Error("HTTP_REQUEST_HEALTH_CHECK", "decode key failed: "+err.Error())
+		return
+	}
+
+	if data.GetBool("onlyBasicRequest") {
+		return true
+	}
+	return
+}

internal/nodes/http_request_page.go

@@ -40,13 +40,28 @@ func (this *HTTPRequest) doPage(status int) (shouldStop bool) {
 						return true
 					}
 
+					stat, err := fp.Stat()
+					if err != nil {
+						logs.Error(err)
+						msg := "404 could not read page content: '" + page.URL + "'"
+
+						this.writer.WriteHeader(http.StatusNotFound)
+						_, err := this.writer.Write([]byte(msg))
+						if err != nil {
+							logs.Error(err)
+						}
+						return true
+					}
+
 					// 修改状态码
 					if page.NewStatus > 0 {
 						// 自定义响应Headers
 						this.processResponseHeaders(page.NewStatus)
+						this.writer.Prepare(stat.Size(), page.NewStatus)
 						this.writer.WriteHeader(page.NewStatus)
 					} else {
 						this.processResponseHeaders(status)
+						this.writer.Prepare(stat.Size(), status)
 						this.writer.WriteHeader(status)
 					}
 					buf := bytePool1k.Get()
@@ -69,17 +84,21 @@ func (this *HTTPRequest) doPage(status int) (shouldStop bool) {
 
 				return true
 			} else if page.BodyType == shared.BodyTypeHTML {
+				var content = this.Format(page.Body)
+
 				// 修改状态码
 				if page.NewStatus > 0 {
 					// 自定义响应Headers
 					this.processResponseHeaders(page.NewStatus)
+					this.writer.Prepare(int64(len(content)), page.NewStatus)
 					this.writer.WriteHeader(page.NewStatus)
 				} else {
 					this.processResponseHeaders(status)
+					this.writer.Prepare(int64(len(content)), status)
 					this.writer.WriteHeader(status)
 				}
 
-				_, err := this.writer.WriteString(this.Format(page.Body))
+				_, err := this.writer.WriteString(content)
 				if err != nil {
 					if !this.canIgnore(err) {
 						remotelogs.Warn("HTTP_REQUEST_PAGE", "write to client failed: "+err.Error())

internal/nodes/http_request_reverse_proxy.go

@@ -145,7 +145,7 @@ func (this *HTTPRequest) doReverseProxy() {
 	}
 
 	// 获取请求客户端
-	client, err := SharedHTTPClientPool.Client(this.RawReq, origin, originAddr)
+	client, err := SharedHTTPClientPool.Client(this, origin, originAddr, this.reverseProxy.ProxyProtocol)
 	if err != nil {
 		remotelogs.Error("HTTP_REQUEST_REVERSE_PROXY", err.Error())
 		this.write50x(err, http.StatusBadGateway)
@@ -249,10 +249,12 @@ func (this *HTTPRequest) doReverseProxy() {
 	shouldAutoFlush := this.reverseProxy.AutoFlush || this.RawReq.Header.Get("Accept") == "text/event-stream"
 
 	// 准备
-	this.writer.Prepare(resp.ContentLength, resp.StatusCode)
+	delayHeaders := this.writer.Prepare(resp.ContentLength, resp.StatusCode)
 
 	// 设置响应代码
-	this.writer.WriteHeader(resp.StatusCode)
+	if !delayHeaders {
+		this.writer.WriteHeader(resp.StatusCode)
+	}
 
 	// 输出到客户端
 	pool := this.bytePool(resp.ContentLength)

internal/nodes/http_request_waf.go

@@ -20,12 +20,23 @@ func (this *HTTPRequest) doWAFRequest() (blocked bool) {
 	// 当前连接是否已关闭
 	var conn = this.RawReq.Context().Value(HTTPConnContextKey)
 	if conn != nil {
-		trafficConn, ok := conn.(*TrafficConn)
+		if isClientConnClosed(conn.(net.Conn)) {
-		if ok && trafficConn.IsClosed() {
+			this.disableLog = true
 			return true
 		}
 	}
 
+	// 检查是否在临时黑名单中
+	var remoteAddr = this.WAFRemoteIP()
+	if waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeService, this.Server.Id, remoteAddr) || waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, remoteAddr) {
+		this.disableLog = true
+		if conn != nil {
+			_ = conn.(net.Conn).Close()
+		}
+
+		return true
+	}
+
 	// 当前服务的独立设置
 	if this.web.FirewallPolicy != nil && this.web.FirewallPolicy.IsOn {
 		blocked, breakChecking := this.checkWAFRequest(this.web.FirewallPolicy)

internal/nodes/http_request_websocket.go

@@ -56,7 +56,7 @@ func (this *HTTPRequest) doWebsocket() {
 	}
 
 	clientConn, _, err := this.writer.Hijack()
-	if err != nil {
+	if err != nil || clientConn == nil {
 		this.write50x(err, http.StatusInternalServerError)
 		return
 	}

internal/nodes/http_writer.go

@@ -3,17 +3,22 @@ package nodes
 import (
 	"bufio"
 	"bytes"
+	"compress/flate"
+	"compress/gzip"
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
 	"github.com/TeaOSLab/EdgeNode/internal/caches"
 	"github.com/TeaOSLab/EdgeNode/internal/compressions"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
-	"github.com/chai2010/webp"
+	"github.com/andybalholm/brotli"
+	_ "github.com/biessek/golang-ico"
 	"github.com/iwind/TeaGo/lists"
 	"github.com/iwind/TeaGo/types"
+	"github.com/iwind/gowebp"
 	_ "golang.org/x/image/bmp"
 	_ "golang.org/x/image/webp"
 	"image"
+	"image/gif"
 	_ "image/gif"
 	_ "image/jpeg"
 	_ "image/png"
@@ -39,9 +44,11 @@ type HTTPWriter struct {
 
 	size int64
 
-	webpIsEncoding bool
+	webpIsEncoding        bool
-	webpBuffer     *bytes.Buffer
+	webpBuffer            *bytes.Buffer
-	webpIsWriting  bool
+	webpIsWriting         bool
+	webpOriginContentType string
+	webpOriginEncoding    string // gzip
 
 	compressionConfig *serverconfigs.HTTPCompressionConfig
 	compressionWriter compressions.Writer
@@ -92,12 +99,16 @@ func (this *HTTPWriter) SetCompression(config *serverconfigs.HTTPCompressionConf
 
 // Prepare 准备输出
 // 缓存不调用此函数
-func (this *HTTPWriter) Prepare(size int64, status int) {
+func (this *HTTPWriter) Prepare(size int64, status int) (delayHeaders bool) {
 	this.size = size
 	this.statusCode = status
 
 	if status == http.StatusOK {
 		this.prepareWebP(size)
+
+		if this.webpIsEncoding {
+			delayHeaders = true
+		}
 	}
 
 	this.prepareCache(size)
@@ -106,6 +117,8 @@ func (this *HTTPWriter) Prepare(size int64, status int) {
 	if !this.webpIsEncoding {
 		this.PrepareCompression(size)
 	}
+
+	return
 }
 
 // Raw 包装前的原始的Writer
@@ -261,11 +274,61 @@ func (this *HTTPWriter) Close() {
 	// webp writer
 	if this.isOk && this.webpIsEncoding {
 		var bufferLen = int64(this.webpBuffer.Len())
-		atomic.AddInt64(&webpTotalBufferSize, bufferLen*8)
+		atomic.AddInt64(&webpTotalBufferSize, bufferLen*4)
 
-		imageData, _, err := image.Decode(this.webpBuffer)
+		// 需要把字节读取出来做备份，防止在image.Decode()过程中丢失
+		var imageBytes = this.webpBuffer.Bytes()
+		var imageData image.Image
+		var gifImage *gif.GIF
+		var isGif = strings.Contains(this.webpOriginContentType, "image/gif")
+
+		var err error
+		if this.webpOriginEncoding == "gzip" {
+			this.Header().Del("Content-Encoding")
+			var reader *gzip.Reader
+			reader, err = gzip.NewReader(this.webpBuffer)
+			if err == nil {
+				defer func() {
+					_ = reader.Close()
+				}()
+				if isGif {
+					gifImage, err = gif.DecodeAll(reader)
+				} else {
+					imageData, _, err = image.Decode(reader)
+				}
+			}
+		} else if this.webpOriginEncoding == "deflate" {
+			this.Header().Del("Content-Encoding")
+			var reader io.ReadCloser
+			reader = flate.NewReader(this.webpBuffer)
+			defer func() {
+				_ = reader.Close()
+			}()
+			if isGif {
+				gifImage, err = gif.DecodeAll(reader)
+			} else {
+				imageData, _, err = image.Decode(reader)
+			}
+		} else if this.webpOriginEncoding == "br" {
+			this.Header().Del("Content-Encoding")
+			var reader *brotli.Reader
+			reader = brotli.NewReader(this.webpBuffer)
+			if isGif {
+				gifImage, err = gif.DecodeAll(reader)
+			} else {
+				imageData, _, err = image.Decode(reader)
+			}
+		} else {
+			if isGif {
+				gifImage, err = gif.DecodeAll(this.webpBuffer)
+			} else {
+				imageData, _, err = image.Decode(this.webpBuffer)
+			}
+		}
 		if err != nil {
-			_, _ = io.Copy(this.writer, this.webpBuffer)
+			this.Header().Set("Content-Type", this.webpOriginContentType)
+			this.WriteHeader(http.StatusOK)
+			_, _ = this.writer.Write(imageBytes)
 
 			// 处理缓存
 			if this.cacheWriter != nil {
@@ -279,16 +342,47 @@ func (this *HTTPWriter) Close() {
 			}
 			this.webpIsWriting = true
 
-			err = webp.Encode(this, imageData, &webp.Options{
+			if imageData != nil {
-				Lossless: false,
+				err = gowebp.Encode(this, imageData, &gowebp.Options{
-				Quality:  f,
+					Lossless: false,
-				Exact:    true,
+					Quality:  f,
-			})
+					Exact:    true,
+				})
+			} else if gifImage != nil {
+				anim := gowebp.NewWebpAnimation(gifImage.Config.Width, gifImage.Config.Height, gifImage.LoopCount)
+				anim.WebPAnimEncoderOptions.SetKmin(9)
+				anim.WebPAnimEncoderOptions.SetKmax(17)
+				defer anim.ReleaseMemory()
+				webpConfig := gowebp.NewWebpConfig()
+				//webpConfig.SetLossless(1)
+				webpConfig.SetQuality(f)
+
+				timeline := 0
+
+				for i, img := range gifImage.Image {
+					err = anim.AddFrame(img, timeline, webpConfig)
+					if err != nil {
+						break
+					}
+					timeline += gifImage.Delay[i] * 10
+				}
+				if err == nil {
+					err = anim.AddFrame(nil, timeline, webpConfig)
+
+					if err == nil {
+						err = anim.Encode(this)
+					}
+				}
+			}
 			if err != nil {
 				if !this.req.canIgnore(err) {
 					remotelogs.Error("HTTP_WRITER", "encode webp failed: "+err.Error())
 				}
 
+				this.Header().Set("Content-Type", this.webpOriginContentType)
+				this.WriteHeader(http.StatusOK)
+				_, _ = this.writer.Write(imageBytes)
+
 				// 处理缓存
 				if this.cacheWriter != nil {
 					_ = this.cacheWriter.Discard()
@@ -297,7 +391,7 @@ func (this *HTTPWriter) Close() {
 			}
 		}
 
-		atomic.AddInt64(&webpTotalBufferSize, -bufferLen*8)
+		atomic.AddInt64(&webpTotalBufferSize, -bufferLen*4)
 		this.webpBuffer.Reset()
 	}
 
@@ -368,9 +462,19 @@ func (this *HTTPWriter) prepareWebP(size int64) {
 		this.req.web.WebP.IsOn &&
 		this.req.web.WebP.MatchResponse(this.Header().Get("Content-Type"), size, filepath.Ext(this.req.requestPath()), this.req.Format) &&
 		this.req.web.WebP.MatchAccept(this.req.requestHeader("Accept")) &&
-		len(this.writer.Header().Get("Content-Encoding")) == 0 &&
 		atomic.LoadInt64(&webpTotalBufferSize) < webpMaxBufferSize {
+
+		var contentEncoding = this.writer.Header().Get("Content-Encoding")
+		switch contentEncoding {
+		case "gzip", "deflate", "br":
+			this.webpOriginEncoding = contentEncoding
+		case "": // 空
+		default:
+			return
+		}
+
 		this.webpIsEncoding = true
+		this.webpOriginContentType = this.Header().Get("Content-Type")
 		this.webpBuffer = webpBufferPool.Get()
 
 		this.Header().Del("Content-Length")
@@ -387,7 +491,8 @@ func (this *HTTPWriter) PrepareCompression(size int64) {
 	}
 
 	// 如果已经有编码则不处理
-	if len(this.writer.Header().Get("Content-Encoding")) > 0 {
+	var contentEncoding = this.writer.Header().Get("Content-Encoding")
+	if len(contentEncoding) > 0 && (!this.compressionConfig.DecompressData || !lists.ContainsString([]string{"gzip", "deflate", "br"}, contentEncoding)) {
 		return
 	}
 
@@ -401,6 +506,12 @@ func (this *HTTPWriter) PrepareCompression(size int64) {
 	if !ok {
 		return
 	}
+
+	// 压缩前后如果编码一致，则不处理
+	if compressionEncoding == contentEncoding {
+		return
+	}
+
 	this.compressionType = compressionType
 
 	// compression writer
@@ -411,6 +522,15 @@ func (this *HTTPWriter) PrepareCompression(size int64) {
 		return
 	}
 
+	// convert between encodings
+	if len(contentEncoding) > 0 {
+		this.compressionWriter, err = compressions.NewEncodingWriter(contentEncoding, this.compressionWriter)
+		if err != nil {
+			remotelogs.Error("HTTP_WRITER", err.Error())
+			return
+		}
+	}
+
 	// body copy
 	if this.bodyCopying {
 		this.compressionBodyBuffer = bytes.NewBuffer([]byte{})

internal/nodes/listener.go

@@ -59,7 +59,7 @@ func (this *Listener) listenTCP() error {
 	if err != nil {
 		return err
 	}
-	netListener = NewTrafficListener(netListener)
+	netListener = NewClientListener(netListener, protocol.IsHTTPFamily() || protocol.IsHTTPSFamily())
 	events.On(events.EventQuit, func() {
 		remotelogs.Println("LISTENER", "quit "+this.group.FullAddr())
 		_ = netListener.Close()

internal/nodes/listener_base.go

@@ -21,19 +21,19 @@ type BaseListener struct {
 	countActiveConnections int64 // 当前活跃的连接数
 }
 
-// 初始化
+// Init 初始化
 func (this *BaseListener) Init() {
 	this.namedServers = map[string]*NamedServer{}
 }
 
-// 清除既有配置
+// Reset 清除既有配置
 func (this *BaseListener) Reset() {
 	this.namedServersLocker.Lock()
 	this.namedServers = map[string]*NamedServer{}
 	this.namedServersLocker.Unlock()
 }
 
-// 获取当前活跃连接数
+// CountActiveListeners 获取当前活跃连接数
 func (this *BaseListener) CountActiveListeners() int {
 	return types.Int(this.countActiveConnections)
 }
@@ -253,3 +253,38 @@ func (this *BaseListener) findNamedServerMatched(name string) (serverConfig *ser
 
 	return nil, name
 }
+
+// 使用CNAME来查找服务
+// TODO 防止单IP随机生成域名攻击
+func (this *BaseListener) findServerWithCname(domain string) *serverconfigs.ServerConfig {
+	if !sharedNodeConfig.SupportCNAME {
+		return nil
+	}
+
+	var realName = sharedCNAMEManager.Lookup(domain)
+	if len(realName) == 0 {
+		return nil
+	}
+
+	this.serversLocker.Lock()
+	defer this.serversLocker.Unlock()
+
+	group := this.Group
+	if group == nil {
+		return nil
+	}
+
+	currentServers := group.Servers
+	countServers := len(currentServers)
+	if countServers == 0 {
+		return nil
+	}
+
+	for _, server := range currentServers {
+		if server.SupportCNAME && lists.ContainsString(server.AliasServerNames, realName) {
+			return server
+		}
+	}
+
+	return nil
+}

internal/nodes/listener_http.go

@@ -157,33 +157,38 @@ func (this *HTTPListener) handleHTTP(rawWriter http.ResponseWriter, rawReq *http
 
 	server, serverName := this.findNamedServer(domain)
 	if server == nil {
-		// 严格匹配域名模式下，我们拒绝用户访问
+		server = this.findServerWithCname(domain)
-		if sharedNodeConfig.GlobalConfig != nil && sharedNodeConfig.GlobalConfig.HTTPAll.MatchDomainStrictly {
+		if server == nil {
-			httpAllConfig := sharedNodeConfig.GlobalConfig.HTTPAll
+			// 严格匹配域名模式下，我们拒绝用户访问
-			mismatchAction := httpAllConfig.DomainMismatchAction
+			if sharedNodeConfig.GlobalConfig != nil && sharedNodeConfig.GlobalConfig.HTTPAll.MatchDomainStrictly {
-			if mismatchAction != nil && mismatchAction.Code == "page" {
+				httpAllConfig := sharedNodeConfig.GlobalConfig.HTTPAll
-				if mismatchAction.Options != nil {
+				mismatchAction := httpAllConfig.DomainMismatchAction
-					rawWriter.Header().Set("Content-Type", "text/html; charset=utf-8")
+				if mismatchAction != nil && mismatchAction.Code == "page" {
-					rawWriter.WriteHeader(mismatchAction.Options.GetInt("statusCode"))
+					if mismatchAction.Options != nil {
-					_, _ = rawWriter.Write([]byte(mismatchAction.Options.GetString("contentHTML")))
+						rawWriter.Header().Set("Content-Type", "text/html; charset=utf-8")
+						rawWriter.WriteHeader(mismatchAction.Options.GetInt("statusCode"))
+						_, _ = rawWriter.Write([]byte(mismatchAction.Options.GetString("contentHTML")))
+					} else {
+						http.Error(rawWriter, "404 page not found: '"+rawReq.URL.String()+"'", http.StatusNotFound)
+					}
+					return
 				} else {
-					http.Error(rawWriter, "404 page not found: '"+rawReq.URL.String()+"'", http.StatusNotFound)
+					hijacker, ok := rawWriter.(http.Hijacker)
-				}
+					if ok {
-				return
+						conn, _, _ := hijacker.Hijack()
-			} else {
+						if conn != nil {
-				hijacker, ok := rawWriter.(http.Hijacker)
+							_ = conn.Close()
-				if ok {
+							return
-					conn, _, _ := hijacker.Hijack()
+						}
-					if conn != nil {
-						_ = conn.Close()
-						return
 					}
 				}
 			}
-		}
 
-		http.Error(rawWriter, "404 page not found: '"+rawReq.URL.String()+"'", http.StatusNotFound)
+			http.Error(rawWriter, "404 page not found: '"+rawReq.URL.String()+"'", http.StatusNotFound)
-		return
+			return
+		} else {
+			serverName = domain
+		}
 	}
 
 	// 包装新请求对象

internal/nodes/listener_tcp.go

@@ -6,7 +6,9 @@ import (
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/TeaOSLab/EdgeNode/internal/stats"
+	"github.com/pires/go-proxyproto"
 	"net"
+	"strings"
 	"sync/atomic"
 )
 
@@ -83,6 +85,31 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 		_ = originConn.Close()
 	}
 
+	// PROXY Protocol
+	if firstServer.ReverseProxy != nil &&
+		firstServer.ReverseProxy.ProxyProtocol != nil &&
+		firstServer.ReverseProxy.ProxyProtocol.IsOn &&
+		(firstServer.ReverseProxy.ProxyProtocol.Version == serverconfigs.ProxyProtocolVersion1 || firstServer.ReverseProxy.ProxyProtocol.Version == serverconfigs.ProxyProtocolVersion2) {
+		var remoteAddr = conn.RemoteAddr()
+		var transportProtocol = proxyproto.TCPv4
+		if strings.Contains(remoteAddr.String(), "[") {
+			transportProtocol = proxyproto.TCPv6
+		}
+		header := proxyproto.Header{
+			Version:           byte(firstServer.ReverseProxy.ProxyProtocol.Version),
+			Command:           proxyproto.PROXY,
+			TransportProtocol: transportProtocol,
+			SourceAddr:        remoteAddr,
+			DestinationAddr:   conn.LocalAddr(),
+		}
+		_, err = header.WriteTo(originConn)
+		if err != nil {
+			closer()
+			return err
+		}
+	}
+
+	// 从源站读取
 	go func() {
 		originBuffer := bytePool32k.Get()
 		defer func() {
@@ -107,6 +134,7 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 		}
 	}()
 
+	// 从客户端读取
 	clientBuffer := bytePool32k.Get()
 	defer func() {
 		bytePool32k.Put(clientBuffer)

internal/nodes/listener_test.go

@@ -8,7 +8,7 @@ import (
 func TestListener_Listen(t *testing.T) {
 	listener := NewListener()
 
-	group := serverconfigs.NewServerGroup("http://:1234")
+	group := serverconfigs.NewServerAddressGroup("https://:1234")
 
 	listener.Reload(group)
 	err := listener.Listen()

internal/nodes/listener_udp.go

@@ -6,7 +6,9 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/TeaOSLab/EdgeNode/internal/stats"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
+	"github.com/pires/go-proxyproto"
 	"net"
+	"strings"
 	"sync"
 	"time"
 )
@@ -19,6 +21,8 @@ type UDPListener struct {
 	connMap    map[string]*UDPConn
 	connLocker sync.Mutex
 	connTicker *utils.Ticker
+
+	reverseProxy *serverconfigs.ReverseProxyConfig
 }
 
 func (this *UDPListener) Serve() error {
@@ -26,8 +30,9 @@ func (this *UDPListener) Serve() error {
 	if firstServer == nil {
 		return errors.New("no server available")
 	}
-	if firstServer.ReverseProxy == nil {
+	this.reverseProxy = firstServer.ReverseProxy
-		return errors.New("no ReverseProxy configured for the server")
+	if this.reverseProxy == nil {
+		return errors.New("no ReverseProxy configured for the server '" + firstServer.Name + "'")
 	}
 
 	this.connMap = map[string]*UDPConn{}
@@ -50,7 +55,7 @@ func (this *UDPListener) Serve() error {
 				ok = false
 			}
 			if !ok {
-				originConn, err := this.connectOrigin(firstServer.ReverseProxy, "")
+				originConn, err := this.connectOrigin(this.reverseProxy, addr)
 				if err != nil {
 					remotelogs.Error("UDP_LISTENER", "unable to connect to origin server: "+err.Error())
 					continue
@@ -87,9 +92,16 @@ func (this *UDPListener) Close() error {
 func (this *UDPListener) Reload(group *serverconfigs.ServerAddressGroup) {
 	this.Group = group
 	this.Reset()
+
+	// 重置配置
+	firstServer := this.Group.FirstServer()
+	if firstServer == nil {
+		return
+	}
+	this.reverseProxy = firstServer.ReverseProxy
 }
 
-func (this *UDPListener) connectOrigin(reverseProxy *serverconfigs.ReverseProxyConfig, remoteAddr string) (conn net.Conn, err error) {
+func (this *UDPListener) connectOrigin(reverseProxy *serverconfigs.ReverseProxyConfig, remoteAddr net.Addr) (conn net.Conn, err error) {
 	if reverseProxy == nil {
 		return nil, errors.New("no reverse proxy config")
 	}
@@ -100,11 +112,34 @@ func (this *UDPListener) connectOrigin(reverseProxy *serverconfigs.ReverseProxyC
 		if origin == nil {
 			continue
 		}
-		conn, err = OriginConnect(origin, remoteAddr)
+		conn, err = OriginConnect(origin, remoteAddr.String())
 		if err != nil {
 			remotelogs.Error("UDP_LISTENER", "unable to connect origin: "+origin.Addr.Host+":"+origin.Addr.PortRange+": "+err.Error())
 			continue
 		} else {
+			// PROXY Protocol
+			if reverseProxy != nil &&
+				reverseProxy.ProxyProtocol != nil &&
+				reverseProxy.ProxyProtocol.IsOn &&
+				(reverseProxy.ProxyProtocol.Version == serverconfigs.ProxyProtocolVersion1 || reverseProxy.ProxyProtocol.Version == serverconfigs.ProxyProtocolVersion2) {
+				var transportProtocol = proxyproto.UDPv4
+				if strings.Contains(remoteAddr.String(), "[") {
+					transportProtocol = proxyproto.UDPv6
+				}
+				header := proxyproto.Header{
+					Version:           byte(reverseProxy.ProxyProtocol.Version),
+					Command:           proxyproto.PROXY,
+					TransportProtocol: transportProtocol,
+					SourceAddr:        remoteAddr,
+					DestinationAddr:   this.Listener.LocalAddr(),
+				}
+				_, err = header.WriteTo(conn)
+				if err != nil {
+					_ = conn.Close()
+					return nil, err
+				}
+			}
+
 			return
 		}
 	}

internal/nodes/node.go

@@ -75,6 +75,9 @@ func (this *Node) Start() {
 		DaemonPid = os.Getppid()
 	}
 
+	// 处理异常
+	this.handlePanic()
+
 	// 启动事件
 	events.Notify(events.EventStart)
 
@@ -178,6 +181,7 @@ func (this *Node) Daemon() {
 
 				// 可以标记当前是从守护进程启动的
 				_ = os.Setenv("EdgeDaemon", "on")
+				_ = os.Setenv("EdgeBackground", "on")
 
 				cmd := exec.Command(exe)
 				err = cmd.Start()

internal/nodes/node_panic.go

@@ -0,0 +1,46 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+//go:build !arm64
+// +build !arm64
+
+package nodes
+
+import (
+	"bytes"
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/logs"
+	"io/ioutil"
+	"os"
+	"syscall"
+)
+
+// 处理异常
+func (this *Node) handlePanic() {
+	// 如果是在前台运行就直接返回
+	backgroundEnv, _ := os.LookupEnv("EdgeBackground")
+	if backgroundEnv != "on" {
+		return
+	}
+
+	var panicFile = Tea.Root + "/logs/panic.log"
+
+	// 分析panic
+	data, err := ioutil.ReadFile(panicFile)
+	if err == nil {
+		var index = bytes.Index(data, []byte("panic:"))
+		if index >= 0 {
+			remotelogs.Error("NODE", "系统错误，请上报给开发者: "+string(data[index:]))
+		}
+	}
+
+	fp, err := os.OpenFile(panicFile, os.O_CREATE|os.O_TRUNC|os.O_WRONLY|os.O_APPEND, 0777)
+	if err != nil {
+		logs.Println("NODE", "open 'panic.log' failed: "+err.Error())
+		return
+	}
+	err = syscall.Dup2(int(fp.Fd()), int(os.Stderr.Fd()))
+	if err != nil {
+		logs.Println("NODE", "write to 'panic.log' failed: "+err.Error())
+	}
+}
+

internal/nodes/node_panic_arm64.go

@@ -0,0 +1,10 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+//go:build arm64
+// +build arm64
+
+package nodes
+
+// 处理异常
+func (this *Node) handlePanic() {
+	return
+}

internal/nodes/server_cname_manager.go

@@ -0,0 +1,48 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/ttlcache"
+	"github.com/TeaOSLab/EdgeNode/internal/utils"
+	"github.com/iwind/TeaGo/types"
+	"strings"
+	"sync"
+	"time"
+)
+
+var sharedCNAMEManager = NewServerCNAMEManager()
+
+// ServerCNAMEManager 服务CNAME管理
+// TODO 需要自动更新缓存里的记录
+type ServerCNAMEManager struct {
+	ttlCache *ttlcache.Cache
+
+	locker sync.Mutex
+}
+
+func NewServerCNAMEManager() *ServerCNAMEManager {
+	return &ServerCNAMEManager{
+		ttlCache: ttlcache.NewCache(),
+	}
+}
+
+func (this *ServerCNAMEManager) Lookup(domain string) string {
+	if len(domain) == 0 {
+		return ""
+	}
+
+	var item = this.ttlCache.Read(domain)
+	if item != nil {
+		return types.String(item.Value)
+	}
+
+	cname, _ := utils.LookupCNAME(domain)
+	if len(cname) > 0 {
+		cname = strings.TrimSuffix(cname, ".")
+	}
+
+	this.ttlCache.Write(domain, cname, time.Now().Unix()+600)
+
+	return cname
+}

internal/nodes/server_cname_manager_test.go

@@ -0,0 +1,19 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"testing"
+	"time"
+)
+
+func TestServerCNameManager_Lookup(t *testing.T) {
+	var cnameManager = NewServerCNAMEManager()
+	t.Log(cnameManager.Lookup("www.yun4s.cn"))
+
+	var before = time.Now()
+	defer func() {
+		t.Log(time.Since(before).Seconds()*1000, "ms")
+	}()
+	t.Log(cnameManager.Lookup("www.yun4s.cn"))
+}

internal/nodes/timezone.go

@@ -0,0 +1,35 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/events"
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+	"time"
+)
+
+func init() {
+	// 管理时区
+	var lastTimeZone = ""
+
+	events.On(events.EventReload, func() {
+		if sharedNodeConfig != nil {
+			var timeZone = sharedNodeConfig.TimeZone
+			if len(timeZone) == 0 {
+				timeZone = "Asia/Shanghai"
+			}
+
+			if lastTimeZone != timeZone {
+				location, err := time.LoadLocation(timeZone)
+				if err != nil {
+					remotelogs.Error("TIMEZONE", "change time zone failed: "+err.Error())
+					return
+				}
+
+				remotelogs.Println("TIMEZONE", "change time zone to '"+timeZone+"'")
+				time.Local = location
+				lastTimeZone = timeZone
+			}
+		}
+	})
+}

internal/nodes/traffic_listener.go

@@ -1,44 +0,0 @@
-// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
-
-package nodes
-
-import (
-	"github.com/TeaOSLab/EdgeNode/internal/waf"
-	"net"
-)
-
-// TrafficListener 用于统计流量的网络监听
-type TrafficListener struct {
-	rawListener net.Listener
-}
-
-func NewTrafficListener(listener net.Listener) net.Listener {
-	return &TrafficListener{rawListener: listener}
-}
-
-func (this *TrafficListener) Accept() (net.Conn, error) {
-	conn, err := this.rawListener.Accept()
-	if err != nil {
-		return nil, err
-	}
-	// 是否在WAF名单中
-	ip, _, err := net.SplitHostPort(conn.RemoteAddr().String())
-	if err == nil {
-		if !waf.SharedIPWhiteList.Contains(waf.IPTypeAll, ip) && waf.SharedIPBlackList.Contains(waf.IPTypeAll, ip) {
-			defer func() {
-				_ = conn.Close()
-			}()
-			return conn, nil
-		}
-	}
-
-	return NewTrafficConn(conn), nil
-}
-
-func (this *TrafficListener) Close() error {
-	return this.rawListener.Close()
-}
-
-func (this *TrafficListener) Addr() net.Addr {
-	return this.rawListener.Addr()
-}

internal/rpc/rpc_client.go

@@ -105,6 +105,18 @@ func (this *RPCClient) ServerRPC() pb.ServerServiceClient {
 	return pb.NewServerServiceClient(this.pickConn())
 }
 
+func (this *RPCClient) ServerRPCList() []pb.ServerServiceClient {
+	this.locker.Lock()
+	defer this.locker.Unlock()
+
+	var clients = []pb.ServerServiceClient{}
+	for _, conn := range this.conns {
+		clients = append(clients, pb.NewServerServiceClient(conn))
+	}
+
+	return clients
+}
+
 func (this *RPCClient) ServerDailyStatRPC() pb.ServerDailyStatServiceClient {
 	return pb.NewServerDailyStatServiceClient(this.pickConn())
 }

internal/utils/lookup.go

@@ -0,0 +1,35 @@
+package utils
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
+	"github.com/miekg/dns"
+)
+
+// LookupCNAME 获取CNAME
+func LookupCNAME(host string) (string, error) {
+	config, err := dns.ClientConfigFromFile("/etc/resolv.conf")
+	if err != nil {
+		return "", err
+	}
+
+	c := new(dns.Client)
+	m := new(dns.Msg)
+
+	m.SetQuestion(host+".", dns.TypeCNAME)
+	m.RecursionDesired = true
+
+	var lastErr error
+	for _, serverAddr := range config.Servers {
+		r, _, err := c.Exchange(m, configutils.QuoteIP(serverAddr)+":"+config.Port)
+		if err != nil {
+			lastErr = err
+			continue
+		}
+		if len(r.Answer) == 0 {
+			continue
+		}
+
+		return r.Answer[0].(*dns.CNAME).Target, nil
+	}
+	return "", lastErr
+}

internal/utils/lookup_test.go

@@ -0,0 +1,9 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package utils
+
+import "testing"
+
+func TestLookupCNAME(t *testing.T) {
+	t.Log(LookupCNAME("www.yun4s.cn"))
+}

internal/utils/net.go

@@ -1,3 +1,6 @@
+//go:build !freebsd
+// +build !freebsd
+
 package utils
 
 import (
@@ -7,7 +10,7 @@ import (
 	"syscall"
 )
 
-// 监听可重用的端口
+// ListenReuseAddr 监听可重用的端口
 func ListenReuseAddr(network string, addr string) (net.Listener, error) {
 	config := &net.ListenConfig{
 		Control: func(network, address string, c syscall.RawConn) error {

internal/waf/action_base.go

@@ -2,7 +2,9 @@
 
 package waf
 
-import "net/http"
+import (
+	"net/http"
+)
 
 type BaseAction struct {
 }
@@ -13,7 +15,7 @@ func (this *BaseAction) CloseConn(writer http.ResponseWriter) error {
 	hijack, ok := writer.(http.Hijacker)
 	if ok {
 		conn, _, err := hijack.Hijack()
-		if err == nil {
+		if err == nil && conn != nil {
 			return conn.Close()
 		}
 	}

internal/waf/action_block.go

@@ -24,6 +24,7 @@ type BlockAction struct {
 	Body       string `yaml:"body" json:"body"` // supports HTML
 	URL        string `yaml:"url" json:"url"`
 	Timeout    int32  `yaml:"timeout" json:"timeout"`
+	Scope      string `yaml:"scope" json:"scope"`
 }
 
 func (this *BlockAction) Init(waf *WAF) error {
@@ -62,7 +63,7 @@ func (this *BlockAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, reque
 	if timeout <= 0 {
 		timeout = 60 // 默认封锁60秒
 	}
-	SharedIPBlackList.Add(IPTypeAll, request.WAFRemoteIP(), time.Now().Unix()+int64(timeout))
+	SharedIPBlackList.Add(IPTypeAll, this.Scope, request.WAFServerId(), request.WAFRemoteIP(), time.Now().Unix()+int64(timeout))
 
 	if writer != nil {
 		// close the connection

internal/waf/action_captcha.go

@@ -23,6 +23,7 @@ type CaptchaAction struct {
 	Life           int32  `yaml:"life" json:"life"`
 	Language       string `yaml:"language" json:"language"`             // 语言，zh-CN, en-US ...
 	AddToWhiteList bool   `yaml:"addToWhiteList" json:"addToWhiteList"` // 是否加入到白名单
+	Scope          string `yaml:"scope" json:"scope"`
 }
 
 func (this *CaptchaAction) Init(waf *WAF) error {
@@ -43,7 +44,7 @@ func (this *CaptchaAction) WillChange() bool {
 
 func (this *CaptchaAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) (allow bool) {
 	// 是否在白名单中
-	if SharedIPWhiteList.Contains("set:"+set.Id, request.WAFRemoteIP()) {
+	if SharedIPWhiteList.Contains("set:"+set.Id, this.Scope, request.WAFServerId(), request.WAFRemoteIP()) {
 		return true
 	}
 

internal/waf/action_get_302.go

@@ -20,7 +20,8 @@ const (
 type Get302Action struct {
 	BaseAction
 
-	Life int32 `yaml:"life" json:"life"`
+	Life  int32  `yaml:"life" json:"life"`
+	Scope string `yaml:"scope" json:"scope"`
 }
 
 func (this *Get302Action) Init(waf *WAF) error {
@@ -46,7 +47,7 @@ func (this *Get302Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, requ
 	}
 
 	// 是否已经在白名单中
-	if SharedIPWhiteList.Contains("set:"+set.Id, request.WAFRemoteIP()) {
+	if SharedIPWhiteList.Contains("set:"+set.Id, this.Scope, request.WAFServerId(), request.WAFRemoteIP()) {
 		return true
 	}
 
@@ -54,6 +55,7 @@ func (this *Get302Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, requ
 		"url":       request.WAFRaw().URL.String(),
 		"timestamp": time.Now().Unix(),
 		"life":      this.Life,
+		"scope":     this.Scope,
 		"setId":     set.Id,
 	}
 	info, err := utils.SimpleEncryptMap(m)

internal/waf/action_page.go

@@ -0,0 +1,41 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package waf
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
+	"net/http"
+)
+
+type PageAction struct {
+	BaseAction
+
+	Status int    `yaml:"status" json:"status"`
+	Body   string `yaml:"body" json:"body"`
+}
+
+func (this *PageAction) Init(waf *WAF) error {
+	return nil
+}
+
+func (this *PageAction) Code() string {
+	return ActionPage
+}
+
+func (this *PageAction) IsAttack() bool {
+	return false
+}
+
+// WillChange determine if the action will change the request
+func (this *PageAction) WillChange() bool {
+	return true
+}
+
+// Perform the action
+func (this *PageAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) (allow bool) {
+	writer.Header().Set("Content-Type", "text/html; charset=utf-8")
+	writer.WriteHeader(this.Status)
+	_, _ = writer.Write([]byte(request.Format(this.Body)))
+
+	return false
+}

internal/waf/action_post_307.go

@@ -10,7 +10,8 @@ import (
 )
 
 type Post307Action struct {
-	Life int32 `yaml:"life" json:"life"`
+	Life  int32  `yaml:"life" json:"life"`
+	Scope string `yaml:"scope" json:"scope"`
 
 	BaseAction
 }
@@ -40,7 +41,7 @@ func (this *Post307Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req
 	}
 
 	// 是否已经在白名单中
-	if SharedIPWhiteList.Contains("set:"+set.Id, request.WAFRemoteIP()) {
+	if SharedIPWhiteList.Contains("set:"+set.Id, this.Scope, request.WAFServerId(), request.WAFRemoteIP()) {
 		return true
 	}
 
@@ -54,7 +55,7 @@ func (this *Post307Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req
 				life = 600 // 默认10分钟
 			}
 			var setId = m.GetString("setId")
-			SharedIPWhiteList.Add("set:"+setId, request.WAFRemoteIP(), time.Now().Unix()+life)
+			SharedIPWhiteList.Add("set:"+setId, this.Scope, request.WAFServerId(), request.WAFRemoteIP(), time.Now().Unix()+life)
 			return true
 		}
 	}
@@ -62,6 +63,7 @@ func (this *Post307Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req
 	var m = maps.Map{
 		"timestamp": time.Now().Unix(),
 		"life":      this.Life,
+		"scope":     this.Scope,
 		"setId":     set.Id,
 		"remoteIP":  request.WAFRemoteIP(),
 	}

internal/waf/action_record_ip.go

@@ -58,6 +58,7 @@ type RecordIPAction struct {
 	IPListId int64  `yaml:"ipListId" json:"ipListId"`
 	Level    string `yaml:"level" json:"level"`
 	Timeout  int32  `yaml:"timeout" json:"timeout"`
+	Scope    string `yaml:"scope" json:"scope"`
 }
 
 func (this *RecordIPAction) Init(waf *WAF) error {
@@ -78,11 +79,10 @@ func (this *RecordIPAction) WillChange() bool {
 
 func (this *RecordIPAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) (allow bool) {
 	// 是否在本地白名单中
-	if SharedIPWhiteList.Contains("set:"+set.Id, set.Id) {
+	if SharedIPWhiteList.Contains("set:"+set.Id, this.Scope, request.WAFServerId(), request.WAFRemoteIP()) {
 		return true
 	}
 
-	// 先加入本地的黑名单
 	timeout := this.Timeout
 	if timeout <= 0 {
 		timeout = 86400 // 1天
@@ -90,16 +90,15 @@ func (this *RecordIPAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, re
 	expiredAt := time.Now().Unix() + int64(timeout)
 
 	if this.Type == "black" {
-		_ = this.CloseConn(writer)
+		writer.WriteHeader(http.StatusForbidden)
 
-		SharedIPBlackList.Add(IPTypeAll, request.WAFRemoteIP(), expiredAt)
+		request.WAFClose()
+
+		// 先加入本地的黑名单
+		SharedIPBlackList.Add(IPTypeAll, this.Scope, request.WAFServerId(), request.WAFRemoteIP(), expiredAt)
 	} else {
 		// 加入本地白名单
-		timeout := this.Timeout
+		SharedIPWhiteList.Add("set:"+set.Id, this.Scope, request.WAFServerId(), request.WAFRemoteIP(), expiredAt)
-		if timeout <= 0 {
-			timeout = 86400 // 1天
-		}
-		SharedIPWhiteList.Add("set:"+set.Id, request.WAFRemoteIP(), expiredAt)
 	}
 
 	// 上报

internal/waf/action_types.go

@@ -13,6 +13,7 @@ const (
 	ActionPost307  ActionString = "post_307"  // 针对POST的307重定向认证
 	ActionRecordIP ActionString = "record_ip" // 记录IP
 	ActionTag      ActionString = "tag"       // 标签
+	ActionPage     ActionString = "page"      // 显示网页
 	ActionAllow    ActionString = "allow"     // allow
 	ActionGoGroup  ActionString = "go_group"  // go to next rule group
 	ActionGoSet    ActionString = "go_set"    // go to next rule set
@@ -73,6 +74,12 @@ var AllActions = []*ActionDefinition{
 		Instance: new(TagAction),
 		Type:     reflect.TypeOf(new(TagAction)).Elem(),
 	},
+	{
+		Name:     "显示页面",
+		Code:     ActionPage,
+		Instance: new(PageAction),
+		Type:     reflect.TypeOf(new(PageAction)).Elem(),
+	},
 	{
 		Name:     "跳到下一个规则分组",
 		Code:     ActionGoGroup,

internal/waf/captcha_validator.go

@@ -143,7 +143,7 @@ func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, setId int64,
 			}
 
 			// 加入到白名单
-			SharedIPWhiteList.Add("set:"+strconv.FormatInt(setId, 10), request.WAFRemoteIP(), time.Now().Unix()+int64(life)) // TODO
+			SharedIPWhiteList.Add("set:"+strconv.FormatInt(setId, 10), actionConfig.Scope, request.WAFServerId(), request.WAFRemoteIP(), time.Now().Unix()+int64(life))
 
 			http.Redirect(writer, request.WAFRaw(), originURL, http.StatusSeeOther)
 

internal/waf/checkpoints/request_referer_block.go

@@ -0,0 +1,66 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package checkpoints
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
+	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
+	"github.com/iwind/TeaGo/maps"
+	"github.com/iwind/TeaGo/types"
+	"net/url"
+)
+
+// RequestRefererBlockCheckpoint 防盗链
+type RequestRefererBlockCheckpoint struct {
+	Checkpoint
+}
+
+// RequestValue 计算checkpoint值
+// 选项：allowEmpty, allowSameDomain, allowDomains
+func (this *RequestRefererBlockCheckpoint) RequestValue(req requests.Request, param string, options maps.Map) (value interface{}, sysErr error, userErr error) {
+	var referer = req.WAFRaw().Referer()
+
+	if len(referer) == 0 {
+		if options.GetBool("allowEmpty") {
+			value = 1
+			return
+		}
+		value = 0
+		return
+	}
+
+	u, err := url.Parse(referer)
+	if err != nil {
+		value = 0
+		return
+	}
+	var host = u.Host
+
+	if options.GetBool("allowSameDomain") && host == req.WAFRaw().Host {
+		value = 1
+		return
+	}
+
+	var domains = options.GetSlice("allowDomains")
+	var domainStrings = []string{}
+	for _, domain := range domains {
+		domainStrings = append(domainStrings, types.String(domain))
+	}
+
+	if len(domainStrings) == 0 {
+		value = 0
+		return
+	}
+
+	if configutils.MatchDomains(domainStrings, host) {
+		value = 1
+	} else {
+		value = 0
+	}
+
+	return
+}
+
+func (this *RequestRefererBlockCheckpoint) ResponseValue(req requests.Request, resp *requests.Response, param string, options maps.Map) (value interface{}, sysErr error, userErr error) {
+	return
+}

internal/waf/checkpoints/utils.go

@@ -198,6 +198,13 @@ var AllCheckpoints = []*CheckpointDefinition{
 		HasParams:   true,
 		Instance:    new(CC2Checkpoint),
 	},
+	{
+		Name:        "防盗链",
+		Prefix:      "refererBlock",
+		Description: "阻止一些域名访问引用本站资源",
+		HasParams:   true,
+		Instance:    new(RequestRefererBlockCheckpoint),
+	},
 	{
 		Name:        "通用响应Header长度限制",
 		Prefix:      "responseGeneralHeaderLength",

internal/waf/get302_validator.go

@@ -44,7 +44,7 @@ func (this *Get302Validator) Run(request requests.Request, writer http.ResponseW
 		life = 600 // 默认10分钟
 	}
 	setId := m.GetString("setId")
-	SharedIPWhiteList.Add("set:"+setId, request.WAFRemoteIP(), time.Now().Unix()+life)
+	SharedIPWhiteList.Add("set:"+setId, m.GetString("scope"), request.WAFServerId(), request.WAFRemoteIP(), time.Now().Unix()+life)
 
 	// 返回原始URL
 	var url = m.GetString("url")

internal/waf/ip_list.go

@@ -3,7 +3,9 @@
 package waf
 
 import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
 	"github.com/TeaOSLab/EdgeNode/internal/utils/expires"
+	"github.com/iwind/TeaGo/types"
 	"sync"
 	"sync/atomic"
 )
@@ -43,8 +45,15 @@ func NewIPList() *IPList {
 }
 
 // Add 添加IP
-func (this *IPList) Add(ipType string, ip string, expiresAt int64) {
+func (this *IPList) Add(ipType string, scope firewallconfigs.FirewallScope, serverId int64, ip string, expiresAt int64) {
-	ip = ip + "@" + ipType
+	switch scope {
+	case firewallconfigs.FirewallScopeGlobal:
+		ip = "*@" + ip + "@" + ipType
+	case firewallconfigs.FirewallScopeService:
+		ip = types.String(serverId) + "@" + ip + "@" + ipType
+	default:
+		ip = "*@" + ip + "@" + ipType
+	}
 
 	var id = this.nextId()
 	this.expireList.Add(id, expiresAt)
@@ -55,8 +64,15 @@ func (this *IPList) Add(ipType string, ip string, expiresAt int64) {
 }
 
 // Contains 判断是否有某个IP
-func (this *IPList) Contains(ipType string, ip string) bool {
+func (this *IPList) Contains(ipType string, scope firewallconfigs.FirewallScope, serverId int64, ip string) bool {
-	ip = ip + "@" + ipType
+	switch scope {
+	case firewallconfigs.FirewallScopeGlobal:
+		ip = "*@" + ip + "@" + ipType
+	case firewallconfigs.FirewallScopeService:
+		ip = types.String(serverId) + "@" + ip + "@" + ipType
+	default:
+		ip = "*@" + ip + "@" + ipType
+	}
 
 	this.locker.RLock()
 	defer this.locker.RUnlock()

internal/waf/ip_list_test.go

@@ -3,6 +3,7 @@
 package waf
 
 import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
 	"github.com/iwind/TeaGo/assert"
 	"github.com/iwind/TeaGo/logs"
 	"runtime"
@@ -13,11 +14,11 @@ import (
 
 func TestNewIPList(t *testing.T) {
 	list := NewIPList()
-	list.Add(IPTypeAll, "127.0.0.1", time.Now().Unix())
+	list.Add(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.1", time.Now().Unix())
-	list.Add(IPTypeAll, "127.0.0.2", time.Now().Unix()+1)
+	list.Add(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.2", time.Now().Unix()+1)
-	list.Add(IPTypeAll, "127.0.0.1", time.Now().Unix()+2)
+	list.Add(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.1", time.Now().Unix()+2)
-	list.Add(IPTypeAll, "127.0.0.3", time.Now().Unix()+3)
+	list.Add(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.3", time.Now().Unix()+3)
-	list.Add(IPTypeAll, "127.0.0.10", time.Now().Unix()+10)
+	list.Add(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "127.0.0.10", time.Now().Unix()+10)
 
 	var ticker = time.NewTicker(1 * time.Second)
 	for range ticker.C {
@@ -36,10 +37,10 @@ func TestIPList_Contains(t *testing.T) {
 	list := NewIPList()
 
 	for i := 0; i < 1_0000; i++ {
-		list.Add(IPTypeAll, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
+		list.Add(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
 	}
-	a.IsTrue(list.Contains(IPTypeAll, "192.168.1.100"))
+	a.IsTrue(list.Contains(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1.100"))
-	a.IsFalse(list.Contains(IPTypeAll, "192.168.2.100"))
+	a.IsFalse(list.Contains(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.2.100"))
 }
 
 func BenchmarkIPList_Add(b *testing.B) {
@@ -47,7 +48,7 @@ func BenchmarkIPList_Add(b *testing.B) {
 
 	list := NewIPList()
 	for i := 0; i < b.N; i++ {
-		list.Add(IPTypeAll, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
+		list.Add(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
 	}
 	b.Log(len(list.ipMap))
 }
@@ -58,10 +59,10 @@ func BenchmarkIPList_Has(b *testing.B) {
 	list := NewIPList()
 
 	for i := 0; i < 1_0000; i++ {
-		list.Add(IPTypeAll, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
+		list.Add(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1."+strconv.Itoa(i), time.Now().Unix()+3600)
 	}
 
 	for i := 0; i < b.N; i++ {
-		list.Contains(IPTypeAll, "192.168.1.100")
+		list.Contains(IPTypeAll, firewallconfigs.FirewallScopeGlobal, 1, "192.168.1.100")
 	}
 }

internal/waf/rule.go

@@ -23,7 +23,7 @@ import (
 
 var singleParamRegexp = regexp.MustCompile("^\\${[\\w.-]+}$")
 
-// rule
+// Rule
 type Rule struct {
 	Description       string                 `yaml:"description" json:"description"`
 	Param             string                 `yaml:"param" json:"param"` // such as ${arg.name} or ${args}, can be composite as ${arg.firstName}${arg.lastName}

internal/waf/rule_set.go

@@ -67,13 +67,16 @@ func (this *RuleSet) Init(waf *WAF) error {
 		instance := FindActionInstance(action.Code, action.Options)
 		if instance == nil {
 			remotelogs.Error("WAF_RULE_SET", "can not find instance for action '"+action.Code+"'")
-		} else {
+			continue
-			this.actionInstances = append(this.actionInstances, instance)
 		}
+
 		err := instance.Init(waf)
 		if err != nil {
 			remotelogs.Error("WAF_RULE_SET", "init action '"+action.Code+"' failed: "+err.Error())
+			continue
 		}
+
+		this.actionInstances = append(this.actionInstances, instance)
 	}
 
 	// sort actions