实现基础的206 partial content缓存

internal/utils/readers/filter_reader.go

@@ -1,34 +0,0 @@
-// 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/handlers.go

@@ -0,0 +1,3 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package readers

internal/utils/readers/reader_base.go

@@ -0,0 +1,6 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package readers
+
+type BaseReader struct {
+}

internal/utils/readers/reader_closer_byte_ranges.go

@@ -0,0 +1,157 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package readers
+
+import (
+	"bytes"
+	"github.com/iwind/TeaGo/types"
+	"io"
+	"mime/multipart"
+	"net/textproto"
+	"regexp"
+	"strings"
+)
+
+type OnPartReadHandler func(start int64, end int64, total int64, data []byte, header textproto.MIMEHeader)
+
+var contentRangeRegexp = regexp.MustCompile(`^(\d+)-(\d+)/(\d+|\*)`)
+
+type ByteRangesReaderCloser struct {
+	BaseReader
+
+	rawReader io.ReadCloser
+	boundary  string
+
+	mReader *multipart.Reader
+	part    *multipart.Part
+
+	buf   *bytes.Buffer
+	isEOF bool
+
+	onPartReadHandler OnPartReadHandler
+	rangeStart        int64
+	rangeEnd          int64
+	total             int64
+
+	isStarted bool
+	nl        string
+}
+
+func NewByteRangesReaderCloser(reader io.ReadCloser, boundary string) *ByteRangesReaderCloser {
+	return &ByteRangesReaderCloser{
+		rawReader: reader,
+		mReader:   multipart.NewReader(reader, boundary),
+		boundary:  boundary,
+		buf:       &bytes.Buffer{},
+		nl:        "\r\n",
+	}
+}
+
+func (this *ByteRangesReaderCloser) Read(p []byte) (n int, err error) {
+	n, err = this.read(p)
+	return
+}
+
+func (this *ByteRangesReaderCloser) Close() error {
+	return this.rawReader.Close()
+}
+
+func (this *ByteRangesReaderCloser) OnPartRead(handler OnPartReadHandler) {
+	this.onPartReadHandler = handler
+}
+
+func (this *ByteRangesReaderCloser) read(p []byte) (n int, err error) {
+	// read from buffer
+	n, err = this.buf.Read(p)
+	if !this.isEOF {
+		err = nil
+	}
+	if n > 0 {
+		return
+	}
+	if this.isEOF {
+		return
+	}
+
+	if this.part == nil {
+		part, partErr := this.mReader.NextPart()
+		if partErr != nil {
+			if partErr == io.EOF {
+				this.buf.WriteString(this.nl + "--" + this.boundary + "--" + this.nl)
+				this.isEOF = true
+				n, _ = this.buf.Read(p)
+				return
+			}
+
+			return 0, partErr
+		}
+
+		if !this.isStarted {
+			this.isStarted = true
+			this.buf.WriteString("--" + this.boundary + this.nl)
+		} else {
+			this.buf.WriteString(this.nl + "--" + this.boundary + this.nl)
+		}
+
+		// Headers
+		var hasRange = false
+		for k, v := range part.Header {
+			for _, v1 := range v {
+				this.buf.WriteString(k + ": " + v1 + this.nl)
+
+				// parse range
+				if k == "Content-Range" {
+					var bytesPrefix = "bytes "
+					if strings.HasPrefix(v1, bytesPrefix) {
+						var r = v1[len(bytesPrefix):]
+						var matches = contentRangeRegexp.FindStringSubmatch(r)
+						if len(matches) > 2 {
+							var start = types.Int64(matches[1])
+							var end = types.Int64(matches[2])
+							var total int64 = 0
+							if matches[3] != "*" {
+								total = types.Int64(matches[3])
+							}
+							if start <= end {
+								hasRange = true
+								this.rangeStart = start
+								this.rangeEnd = end
+								this.total = total
+							}
+						}
+					}
+				}
+			}
+		}
+
+		if !hasRange {
+			this.rangeStart = -1
+			this.rangeEnd = -1
+		}
+
+		this.buf.WriteString(this.nl)
+		this.part = part
+
+		n, _ = this.buf.Read(p)
+		return
+	}
+
+	n, err = this.part.Read(p)
+
+	if this.onPartReadHandler != nil && n > 0 && this.rangeStart >= 0 && this.rangeEnd >= 0 {
+		this.onPartReadHandler(this.rangeStart, this.rangeEnd, this.total, p[:n], this.part.Header)
+		this.rangeStart += int64(n)
+	}
+
+	if err == io.EOF {
+		this.part = nil
+		err = nil
+
+		// 如果没有读取到内容，则直接跳到下一个Part
+		if n == 0 {
+			return this.read(p)
+		}
+	}
+
+	return
+}

internal/utils/readers/reader_closer_byte_ranges_test.go

@@ -0,0 +1,52 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package readers_test
+
+import (
+	"bytes"
+	"fmt"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/readers"
+	"io"
+	"io/ioutil"
+	"net/textproto"
+	"testing"
+)
+
+func TestNewByteRangesReader(t *testing.T) {
+	var boundary = "7143cd51d2ee12a1"
+	var dashBoundary = "--" + boundary
+	var b = bytes.NewReader([]byte(dashBoundary + "\r\nContent-Range: bytes 0-4/36\r\nContent-Type: text/plain\r\n\r\n01234\r\n" + dashBoundary + "\r\nContent-Range: bytes 5-9/36\r\nContent-Type: text/plain\r\n\r\n56789\r\n--" + boundary + "\r\nContent-Range: bytes 10-12/36\r\nContent-Type: text/plain\r\n\r\nabc\r\n" + dashBoundary + "--\r\n"))
+
+	var reader = readers.NewByteRangesReaderCloser(ioutil.NopCloser(b), boundary)
+	var p = make([]byte, 16)
+	for {
+		n, err := reader.Read(p)
+		if n > 0 {
+			fmt.Print(string(p[:n]))
+		}
+		if err != nil {
+			if err != io.EOF {
+				t.Fatal(err)
+			}
+			break
+		}
+	}
+}
+
+func TestByteRangesReader_OnPartRead(t *testing.T) {
+	var boundary = "7143cd51d2ee12a1"
+	var dashBoundary = "--" + boundary
+	var b = bytes.NewReader([]byte(dashBoundary + "\r\nContent-Range: bytes 0-4/36\r\nContent-Type: text/plain\r\n\r\n01234\r\n" + dashBoundary + "\r\nContent-Range: bytes 5-9/36\r\nContent-Type: text/plain\r\n\r\n56789\r\n--" + boundary + "\r\nContent-Range: bytes 10-12/36\r\nContent-Type: text/plain\r\n\r\nabc\r\n" + dashBoundary + "--\r\n"))
+
+	var reader = readers.NewByteRangesReaderCloser(ioutil.NopCloser(b), boundary)
+	reader.OnPartRead(func(start int64, end int64, total int64, data []byte, header textproto.MIMEHeader) {
+		t.Log(start, "-", end, "/", total, string(data))
+	})
+	var p = make([]byte, 3)
+	for {
+		_, err := reader.Read(p)
+		if err != nil {
+			break
+		}
+	}
+}

internal/utils/readers/reader_closer_filter.go

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

internal/utils/readers/reader_closer_filter_test.go

@@ -10,7 +10,7 @@ import (
 )
 
 func TestNewFilterReader(t *testing.T) {
-	var reader = readers.NewFilterReader(bytes.NewBufferString("0123456789"))
+	var reader = readers.NewFilterReaderCloser(bytes.NewBufferString("0123456789"))
 	reader.Add(func(p []byte, err error) error {
 		t.Log("filter1:", string(p), err)
 		return nil