重构对HTTP请求的处理方法：缓存、压缩、WebP、限速

internal/caches/reader_file.go

@@ -278,17 +278,27 @@ func (this *FileReader) Read(buf []byte) (n int, err error) {
 	}()
 
 	// 直接返回从Header中剩余的
-	if this.bodyBufLen > 0 && len(buf) >= this.bodyBufLen {
-		copy(buf, this.bodyBuf)
-		isOk = true
-		n = this.bodyBufLen
+	if this.bodyBufLen > 0 {
+		var bufLen = len(buf)
+		if bufLen < this.bodyBufLen {
+			this.bodyBufLen -= bufLen
+			copy(buf, this.bodyBuf[:bufLen])
+			this.bodyBuf = this.bodyBuf[bufLen:]
 
-		if this.bodySize <= int64(this.bodyBufLen) {
-			err = io.EOF
-			return
+			n = bufLen
+		} else {
+			copy(buf, this.bodyBuf)
+			this.bodyBuf = nil
+
+			if this.bodySize <= int64(this.bodyBufLen) {
+				err = io.EOF
+			}
+
+			n = this.bodyBufLen
+			this.bodyBufLen = 0
 		}
 
-		this.bodyBufLen = 0
+		isOk = true
 		return
 	}
 

internal/nodes/http_request.go

@@ -315,12 +315,12 @@ func (this *HTTPRequest) doEnd() {
 	// TODO 增加Header统计，考虑从Conn中读取
 	if this.ReqServer != nil {
 		if this.isCached {
-			stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.sentBodyBytes, this.writer.sentBodyBytes, 1, 1, 0, 0, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
+			stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes(), this.writer.SentBodyBytes(), 1, 1, 0, 0, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
 		} else {
 			if this.isAttack {
-				stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.sentBodyBytes, 0, 1, 0, 1, this.writer.sentBodyBytes, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
+				stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes(), 0, 1, 0, 1, this.writer.SentBodyBytes(), this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
 			} else {
-				stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.sentBodyBytes, 0, 1, 0, 0, 0, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
+				stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes(), 0, 1, 0, 0, 0, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
 			}
 		}
 	}

internal/nodes/http_request_cache.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
 	"github.com/TeaOSLab/EdgeNode/internal/caches"
-	"github.com/TeaOSLab/EdgeNode/internal/compressions"
 	"github.com/TeaOSLab/EdgeNode/internal/goman"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/TeaOSLab/EdgeNode/internal/rpc"
@@ -162,11 +161,15 @@ func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
 	var err error
 
 	// 是否优先检查WebP
+	var isWebP = false
 	if this.web.WebP != nil &&
 		this.web.WebP.IsOn &&
 		this.web.WebP.MatchRequest(filepath.Ext(this.Path()), this.Format) &&
 		this.web.WebP.MatchAccept(this.requestHeader("Accept")) {
 		reader, _ = storage.OpenReader(key+webpSuffix, useStale)
+		if reader != nil {
+			isWebP = true
+		}
 	}
 
 	// 检查正常的文件
@@ -189,8 +192,11 @@ func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
 			return
 		}
 	}
+
 	defer func() {
-		_ = reader.Close()
+		if !this.writer.DelayRead() {
+			_ = reader.Close()
+		}
 	}()
 
 	if useStale {
@@ -257,7 +263,11 @@ func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
 	var eTag = ""
 	var lastModifiedAt = reader.LastModified()
 	if lastModifiedAt > 0 {
-		eTag = "\"" + strconv.FormatInt(lastModifiedAt, 10) + "\""
+		if isWebP {
+			eTag = "\"" + strconv.FormatInt(lastModifiedAt, 10) + "_webp" + "\""
+		} else {
+			eTag = "\"" + strconv.FormatInt(lastModifiedAt, 10) + "\""
+		}
 		respHeader.Del("Etag")
 		respHeader["ETag"] = []string{eTag}
 	}
@@ -439,25 +449,11 @@ func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
 				return true
 			}
 		} else { // 没有Range
-			var body io.Reader = reader
-			var contentEncoding = this.writer.Header().Get("Content-Encoding")
-			if len(contentEncoding) > 0 && !httpAcceptEncoding(this.RawReq.Header.Get("Accept-Encoding"), contentEncoding) {
-				decompressReader, err := compressions.NewReader(body, contentEncoding)
-				if err == nil {
-					body = decompressReader
-					defer func() {
-						_ = decompressReader.Close()
-					}()
-
-					this.writer.Header().Del("Content-Encoding")
-					this.writer.Header().Del("Content-Length")
-				}
-			}
-
-			this.writer.PrepareCompression(reader.BodySize())
+			var resp = &http.Response{Body: reader}
+			this.writer.Prepare(resp, reader.BodySize(), reader.Status(), false)
 			this.writer.WriteHeader(reader.Status())
 
-			_, err = io.CopyBuffer(this.writer, body, buf)
+			_, err = io.CopyBuffer(this.writer, resp.Body, buf)
 			if err == io.EOF {
 				err = nil
 			}

internal/nodes/http_request_fastcgi.go

@@ -190,7 +190,7 @@ func (this *HTTPRequest) doFastcgi() (shouldStop bool) {
 	this.processResponseHeaders(resp.StatusCode)
 
 	// 准备
-	this.writer.Prepare(resp.ContentLength, resp.StatusCode)
+	this.writer.Prepare(resp, resp.ContentLength, resp.StatusCode, true)
 
 	// 设置响应代码
 	this.writer.WriteHeader(resp.StatusCode)

internal/nodes/http_request_page.go

@@ -61,11 +61,11 @@ func (this *HTTPRequest) doPage(status int) (shouldStop bool) {
 					if page.NewStatus > 0 {
 						// 自定义响应Headers
 						this.processResponseHeaders(page.NewStatus)
-						this.writer.Prepare(stat.Size(), page.NewStatus)
+						this.writer.Prepare(nil, stat.Size(), page.NewStatus, true)
 						this.writer.WriteHeader(page.NewStatus)
 					} else {
 						this.processResponseHeaders(status)
-						this.writer.Prepare(stat.Size(), status)
+						this.writer.Prepare(nil, stat.Size(), status, true)
 						this.writer.WriteHeader(status)
 					}
 					buf := utils.BytePool1k.Get()
@@ -100,11 +100,11 @@ func (this *HTTPRequest) doPage(status int) (shouldStop bool) {
 				if page.NewStatus > 0 {
 					// 自定义响应Headers
 					this.processResponseHeaders(page.NewStatus)
-					this.writer.Prepare(int64(len(content)), page.NewStatus)
+					this.writer.Prepare(nil, int64(len(content)), page.NewStatus, true)
 					this.writer.WriteHeader(page.NewStatus)
 				} else {
 					this.processResponseHeaders(status)
-					this.writer.Prepare(int64(len(content)), status)
+					this.writer.Prepare(nil, int64(len(content)), status, true)
 					this.writer.WriteHeader(status)
 				}
 

internal/nodes/http_request_reverse_proxy.go

@@ -285,10 +285,10 @@ func (this *HTTPRequest) doReverseProxy() {
 	this.processResponseHeaders(resp.StatusCode)
 
 	// 是否需要刷新
-	shouldAutoFlush := this.reverseProxy.AutoFlush || this.RawReq.Header.Get("Accept") == "text/event-stream"
+	var shouldAutoFlush = this.reverseProxy.AutoFlush || this.RawReq.Header.Get("Accept") == "text/event-stream"
 
 	// 准备
-	delayHeaders := this.writer.Prepare(resp.ContentLength, resp.StatusCode)
+	var delayHeaders = this.writer.Prepare(resp, resp.ContentLength, resp.StatusCode, true)
 
 	// 设置响应代码
 	if !delayHeaders {

internal/nodes/http_request_root.go

@@ -302,7 +302,7 @@ func (this *HTTPRequest) doRoot() (isBreak bool) {
 		this.cacheRef = nil // 不支持缓存
 	}
 
-	this.writer.Prepare(fileSize, http.StatusOK)
+	this.writer.Prepare(nil, fileSize, http.StatusOK, true)
 
 	pool := this.bytePool(fileSize)
 	buf := pool.Get()

internal/nodes/http_request_url.go

@@ -54,9 +54,9 @@ func (this *HTTPRequest) doURL(method string, url string, host string, statusCod
 	}
 	this.writer.AddHeaders(resp.Header)
 	if statusCode <= 0 {
-		this.writer.Prepare(resp.ContentLength, resp.StatusCode)
+		this.writer.Prepare(resp, resp.ContentLength, resp.StatusCode, true)
 	} else {
-		this.writer.Prepare(resp.ContentLength, statusCode)
+		this.writer.Prepare(resp, resp.ContentLength, statusCode, true)
 	}
 
 	// 设置响应代码

internal/nodes/http_writer_rate.go

@@ -1,102 +0,0 @@
-// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
-
-package nodes
-
-import (
-	"bufio"
-	"github.com/iwind/TeaGo/types"
-	"net"
-	"net/http"
-	"time"
-)
-
-// HTTPRateWriter 限速写入
-type HTTPRateWriter struct {
-	parentWriter http.ResponseWriter
-
-	rateBytes int
-	lastBytes int
-	timeCost  time.Duration
-}
-
-func NewHTTPRateWriter(writer http.ResponseWriter, rateBytes int64) http.ResponseWriter {
-	return &HTTPRateWriter{
-		parentWriter: writer,
-		rateBytes:    types.Int(rateBytes),
-	}
-}
-
-func (this *HTTPRateWriter) Header() http.Header {
-	return this.parentWriter.Header()
-}
-
-func (this *HTTPRateWriter) Write(data []byte) (int, error) {
-	if len(data) == 0 {
-		return 0, nil
-	}
-
-	var left = this.rateBytes - this.lastBytes
-
-	if left <= 0 {
-		if this.timeCost > 0 && this.timeCost < 1*time.Second {
-			time.Sleep(1*time.Second - this.timeCost)
-		}
-
-		this.lastBytes = 0
-		this.timeCost = 0
-		return this.Write(data)
-	}
-
-	var n = len(data)
-
-	// n <= left
-	if n <= left {
-		this.lastBytes += n
-
-		var before = time.Now()
-		defer func() {
-			this.timeCost += time.Since(before)
-		}()
-		return this.parentWriter.Write(data)
-	}
-
-	// n > left
-	var before = time.Now()
-	result, err := this.parentWriter.Write(data[:left])
-	this.timeCost += time.Since(before)
-
-	if err != nil {
-		return result, err
-	}
-	this.lastBytes += left
-
-	return this.Write(data[left:])
-}
-
-func (this *HTTPRateWriter) WriteHeader(statusCode int) {
-	this.parentWriter.WriteHeader(statusCode)
-}
-
-// Hijack Hijack
-func (this *HTTPRateWriter) Hijack() (conn net.Conn, buf *bufio.ReadWriter, err error) {
-	if this.parentWriter == nil {
-		return
-	}
-	hijack, ok := this.parentWriter.(http.Hijacker)
-	if ok {
-		return hijack.Hijack()
-	}
-	return
-}
-
-// Flush Flush
-func (this *HTTPRateWriter) Flush() {
-	if this.parentWriter == nil {
-		return
-	}
-	flusher, ok := this.parentWriter.(http.Flusher)
-	if ok {
-		flusher.Flush()
-		return
-	}
-}

internal/utils/readers/bytes_counter_reader.go

@@ -0,0 +1,26 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package readers
+
+import "io"
+
+type BytesCounterReader struct {
+	rawReader io.Reader
+	count     int64
+}
+
+func NewBytesCounterReader(rawReader io.Reader) *BytesCounterReader {
+	return &BytesCounterReader{
+		rawReader: rawReader,
+	}
+}
+
+func (this *BytesCounterReader) Read(p []byte) (n int, err error) {
+	n, err = this.rawReader.Read(p)
+	this.count += int64(n)
+	return
+}
+
+func (this *BytesCounterReader) TotalBytes() int64 {
+	return this.count
+}

internal/utils/readers/filter_reader.go

@@ -0,0 +1,34 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package readers
+
+import "io"
+
+type FilterFunc = func(p []byte, err error) error
+
+type FilterReader struct {
+	rawReader io.Reader
+	filters   []FilterFunc
+}
+
+func NewFilterReader(rawReader io.Reader) *FilterReader {
+	return &FilterReader{
+		rawReader: rawReader,
+	}
+}
+
+func (this *FilterReader) Add(filter FilterFunc) {
+	this.filters = append(this.filters, filter)
+}
+
+func (this *FilterReader) Read(p []byte) (n int, err error) {
+	n, err = this.rawReader.Read(p)
+	for _, filter := range this.filters {
+		filterErr := filter(p[:n], err)
+		if filterErr != nil {
+			err = filterErr
+			return
+		}
+	}
+	return
+}

internal/utils/readers/filter_reader_test.go

@@ -0,0 +1,41 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package readers_test
+
+import (
+	"bytes"
+	"errors"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/readers"
+	"testing"
+)
+
+func TestNewFilterReader(t *testing.T) {
+	var reader = readers.NewFilterReader(bytes.NewBufferString("0123456789"))
+	reader.Add(func(p []byte, err error) error {
+		t.Log("filter1:", string(p), err)
+		return nil
+	})
+	reader.Add(func(p []byte, err error) error {
+		t.Log("filter2:", string(p), err)
+		if string(p) == "345" {
+			return errors.New("end")
+		}
+		return nil
+	})
+	reader.Add(func(p []byte, err error) error {
+		t.Log("filter3:", string(p), err)
+		return nil
+	})
+
+	var buf = make([]byte, 3)
+	for {
+		n, err := reader.Read(buf)
+		if n > 0 {
+			t.Log(string(buf[:n]))
+		}
+		if err != nil {
+			t.Log(err)
+			break
+		}
+	}
+}

internal/utils/readers/tee_reader.go

@@ -0,0 +1,52 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package readers
+
+import (
+	"io"
+)
+
+type TeeReader struct {
+	r io.Reader
+	w io.Writer
+
+	onFail func(err error)
+	onEOF  func()
+}
+
+func NewTeeReader(reader io.Reader, writer io.Writer) *TeeReader {
+	return &TeeReader{
+		r: reader,
+		w: writer,
+	}
+}
+
+func (this *TeeReader) Read(p []byte) (n int, err error) {
+	n, err = this.r.Read(p)
+	if n > 0 {
+		_, wErr := this.w.Write(p[:n])
+		if err == nil && wErr != nil {
+			err = wErr
+		}
+	}
+	if err != nil {
+		if err == io.EOF {
+			if this.onEOF != nil {
+				this.onEOF()
+			}
+		} else {
+			if this.onFail != nil {
+				this.onFail(err)
+			}
+		}
+	}
+	return
+}
+
+func (this *TeeReader) OnFail(onFail func(err error)) {
+	this.onFail = onFail
+}
+
+func (this *TeeReader) OnEOF(onEOF func()) {
+	this.onEOF = onEOF
+}

internal/utils/readers/tee_reader_closer.go

@@ -0,0 +1,58 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package readers
+
+import "io"
+
+type TeeReaderCloser struct {
+	r io.Reader
+	w io.Writer
+
+	onFail func(err error)
+	onEOF  func()
+}
+
+func NewTeeReaderCloser(reader io.Reader, writer io.Writer) *TeeReaderCloser {
+	return &TeeReaderCloser{
+		r: reader,
+		w: writer,
+	}
+}
+
+func (this *TeeReaderCloser) Read(p []byte) (n int, err error) {
+	n, err = this.r.Read(p)
+	if n > 0 {
+		_, wErr := this.w.Write(p[:n])
+		if err == nil && wErr != nil {
+			err = wErr
+		}
+	}
+	if err != nil {
+		if err == io.EOF {
+			if this.onEOF != nil {
+				this.onEOF()
+			}
+		} else {
+			if this.onFail != nil {
+				this.onFail(err)
+			}
+		}
+	}
+	return
+}
+
+func (this *TeeReaderCloser) Close() error {
+	r, ok := this.r.(io.Closer)
+	if ok {
+		return r.Close()
+	}
+	return nil
+}
+
+func (this *TeeReaderCloser) OnFail(onFail func(err error)) {
+	this.onFail = onFail
+}
+
+func (this *TeeReaderCloser) OnEOF(onEOF func()) {
+	this.onEOF = onEOF
+}

internal/utils/writers/bytes_counter_writer.go

@@ -0,0 +1,28 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package writers
+
+import "io"
+
+type BytesCounterWriter struct {
+	writer io.Writer
+	count  int64
+}
+
+func NewBytesCounterWriter(rawWriter io.Writer) *BytesCounterWriter {
+	return &BytesCounterWriter{writer: rawWriter}
+}
+
+func (this *BytesCounterWriter) Write(p []byte) (n int, err error) {
+	n, err = this.writer.Write(p)
+	this.count += int64(n)
+	return
+}
+
+func (this *BytesCounterWriter) Close() error {
+	return nil
+}
+
+func (this *BytesCounterWriter) TotalBytes() int64 {
+	return this.count
+}

internal/utils/writers/rate_limit_writer.go

@@ -0,0 +1,87 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package writers
+
+import (
+	"github.com/iwind/TeaGo/types"
+	"io"
+	"time"
+)
+
+// RateLimitWriter 限速写入
+type RateLimitWriter struct {
+	rawWriter io.WriteCloser
+
+	rateBytes int
+
+	written int
+	before  time.Time
+}
+
+func NewRateLimitWriter(rawWriter io.WriteCloser, rateBytes int64) io.WriteCloser {
+	return &RateLimitWriter{
+		rawWriter: rawWriter,
+		rateBytes: types.Int(rateBytes),
+		before:    time.Now(),
+	}
+}
+
+func (this *RateLimitWriter) Write(p []byte) (n int, err error) {
+	if this.rateBytes <= 0 {
+		return this.write(p)
+	}
+
+	var size = len(p)
+	if size == 0 {
+		return 0, nil
+	}
+
+	if size <= this.rateBytes {
+		return this.write(p)
+	}
+
+	for {
+		size = len(p)
+
+		var limit = this.rateBytes
+		if limit > size {
+			limit = size
+		}
+		n1, wErr := this.write(p[:limit])
+		n += n1
+		if wErr != nil {
+			return n, wErr
+		}
+
+		if size > limit {
+			p = p[limit:]
+		} else {
+			break
+		}
+	}
+
+	return
+}
+
+func (this *RateLimitWriter) Close() error {
+	return this.rawWriter.Close()
+}
+
+func (this *RateLimitWriter) write(p []byte) (n int, err error) {
+	n, err = this.rawWriter.Write(p)
+
+	if err == nil {
+		this.written += n
+
+		if this.written >= this.rateBytes {
+			var duration = 1*time.Second - time.Now().Sub(this.before)
+			if duration > 0 {
+				time.Sleep(duration)
+			}
+			this.before = time.Now()
+			this.written = 0
+		}
+	}
+
+	return
+}

internal/utils/writers/rate_limit_writer_test.go

@@ -0,0 +1,41 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package writers
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestSleep(t *testing.T) {
+	var count = 2000
+	var wg = sync.WaitGroup{}
+	wg.Add(count)
+	var before = time.Now()
+	for i := 0; i < count; i++ {
+		go func() {
+			defer wg.Done()
+			time.Sleep(1 * time.Second)
+		}()
+	}
+	wg.Wait()
+	t.Log(time.Since(before).Seconds()*1000, "ms")
+}
+
+func TestTimeout(t *testing.T) {
+	var count = 2000
+	var wg = sync.WaitGroup{}
+	wg.Add(count)
+	var before = time.Now()
+	for i := 0; i < count; i++ {
+		go func() {
+			defer wg.Done()
+
+			var timeout = time.NewTimer(1 * time.Second)
+			<-timeout.C
+		}()
+	}
+	wg.Wait()
+	t.Log(time.Since(before).Seconds()*1000, "ms")
+}

internal/utils/writers/tee_writer_closer.go

@@ -0,0 +1,51 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package writers
+
+import "io"
+
+type TeeWriterCloser struct {
+	primaryW   io.WriteCloser
+	secondaryW io.WriteCloser
+
+	onFail func(err error)
+}
+
+func NewTeeWriterCloser(primaryW io.WriteCloser, secondaryW io.WriteCloser) *TeeWriterCloser {
+	return &TeeWriterCloser{
+		primaryW:   primaryW,
+		secondaryW: secondaryW,
+	}
+}
+
+func (this *TeeWriterCloser) Write(p []byte) (n int, err error) {
+	{
+		n, err = this.primaryW.Write(p)
+
+		if err != nil {
+			if this.onFail != nil {
+				this.onFail(err)
+			}
+		}
+	}
+
+	{
+		_, err2 := this.secondaryW.Write(p)
+		if err2 != nil {
+			if this.onFail != nil {
+				this.onFail(err2)
+			}
+		}
+	}
+
+	return
+}
+
+func (this *TeeWriterCloser) Close() error {
+	// 这里不关闭secondary
+	return this.primaryW.Close()
+}
+
+func (this *TeeWriterCloser) OnFail(onFail func(err error)) {
+	this.onFail = onFail
+}