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

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
+}