Partial Content从源站读取数据时验证本地缓存的Content-MD5是否一致

internal/caches/partial_ranges.go

@@ -12,16 +12,17 @@ import (
 
 // PartialRanges 内容分区范围定义
 type PartialRanges struct {
-	Version  int        `json:"version"`  // 版本号
+	Version    int        `json:"version"`    // 版本号
-	Ranges   [][2]int64 `json:"ranges"`   // 范围
+	Ranges     [][2]int64 `json:"ranges"`     // 范围
-	BodySize int64      `json:"bodySize"` // 总长度
+	BodySize   int64      `json:"bodySize"`   // 总长度
+	ContentMD5 string     `json:"contentMD5"` // 内容md5
 }
 
 // NewPartialRanges 获取新对象
 func NewPartialRanges(expiresAt int64) *PartialRanges {
 	return &PartialRanges{
 		Ranges:  [][2]int64{},
-		Version: 1,
+		Version: 2,
 	}
 }
 
@@ -46,6 +47,8 @@ func NewPartialRangesFromData(data []byte) (*PartialRanges, error) {
 				if commaIndex > 0 {
 					rs.Ranges = append(rs.Ranges, [2]int64{types.Int64(line[colonIndex+1 : commaIndex]), types.Int64(line[commaIndex+1:])})
 				}
+			case "m": // Content-MD5
+				rs.ContentMD5 = string(line[colonIndex+1:])
 			}
 		}
 		data = data[index+1:]
@@ -171,6 +174,9 @@ func (this *PartialRanges) FindRangeAtPosition(position int64) (r rangeutils.Ran
 func (this *PartialRanges) String() string {
 	var s = "v:" + strconv.Itoa(this.Version) + "\n" + // version
 		"b:" + this.formatInt64(this.BodySize) + "\n" // bodySize
+	if len(this.ContentMD5) > 0 {
+		s += "m:" + this.ContentMD5 + "\n" // Content-MD5
+	}
 	for _, r := range this.Ranges {
 		s += "r:" + this.formatInt64(r[0]) + "," + this.formatInt64(r[1]) + "\n" // range
 	}

internal/caches/partial_ranges_test.go

@@ -3,6 +3,8 @@
 package caches_test
 
 import (
+	"crypto/md5"
+	"encoding/base64"
 	"encoding/json"
 	"github.com/TeaOSLab/EdgeNode/internal/caches"
 	"github.com/iwind/TeaGo/assert"
@@ -167,8 +169,13 @@ func TestPartialRanges_Encode_String(t *testing.T) {
 	for i := 0; i < 10; i++ {
 		r.Ranges = append(r.Ranges, [2]int64{int64(i * 100), int64(i*100 + 100)})
 	}
+
+	var sum = md5.Sum([]byte("123456"))
+	r.ContentMD5 = base64.StdEncoding.EncodeToString(sum[:])
+
 	var before = time.Now()
 	var data = r.String()
+	t.Log(data)
 	t.Log(time.Since(before).Seconds()*1000, "ms")
 	t.Log(len(data))
 

internal/caches/reader_file_test.go

@@ -2,6 +2,7 @@ package caches
 
 import (
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
+	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
 	"github.com/iwind/TeaGo/Tea"
 	"os"
 	"testing"
@@ -36,7 +37,7 @@ func TestFileReader(t *testing.T) {
 	defer func() {
 		_ = fp.Close()
 	}()
-	reader := NewFileReader(fp)
+	reader := NewFileReader(fsutils.NewFile(fp, fsutils.FlagRead))
 	err = reader.Init()
 	if err != nil {
 		t.Fatal(err)
@@ -75,7 +76,7 @@ func TestFileReader_ReadHeader(t *testing.T) {
 	defer func() {
 		_ = fp.Close()
 	}()
-	var reader = NewFileReader(fp)
+	var reader = NewFileReader(fsutils.NewFile(fp, fsutils.FlagRead))
 	err = reader.Init()
 	if err != nil {
 		if os.IsNotExist(err) {
@@ -138,7 +139,7 @@ func TestFileReader_Range(t *testing.T) {
 	defer func() {
 		_ = fp.Close()
 	}()
-	reader := NewFileReader(fp)
+	reader := NewFileReader(fsutils.NewFile(fp, fsutils.FlagRead))
 	err = reader.Init()
 	if err != nil {
 		t.Fatal(err)

internal/caches/writer_partial_file.go

@@ -7,6 +7,7 @@ import (
 	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
 	"github.com/iwind/TeaGo/types"
 	"io"
+	"strings"
 	"sync"
 )
 
@@ -183,6 +184,14 @@ func (this *PartialFileWriter) SetBodyLength(bodyLength int64) {
 	this.bodySize = bodyLength
 }
 
+// SetContentMD5 设置内容MD5
+func (this *PartialFileWriter) SetContentMD5(contentMD5 string) {
+	if strings.Contains(contentMD5, "\n") || len(contentMD5) > 128 {
+		return
+	}
+	this.ranges.ContentMD5 = contentMD5
+}
+
 // WriteBodyLength 写入Body长度数据
 func (this *PartialFileWriter) WriteBodyLength(bodyLength int64) error {
 	if this.metaBodySize > 0 && this.metaBodySize == bodyLength {

internal/caches/writer_partial_file_test.go

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

internal/nodes/http_request_cache_partial.go

@@ -42,8 +42,25 @@ func (this *HTTPRequestPartialReader) Read(p []byte) (n int, err error) {
 			err = io.ErrUnexpectedEOF
 			return
 		}
+
 		this.resp = resp
 
+		// 对比Content-MD5
+		partialReader, ok := this.cacheReader.(*caches.PartialFileReader)
+		if ok {
+			if partialReader.Ranges().Version >= 2 && resp.Header.Get("Content-MD5") != partialReader.Ranges().ContentMD5 {
+				err = io.ErrUnexpectedEOF
+
+				var storage = this.req.writer.cacheStorage
+				if storage != nil {
+					_ = storage.Delete(this.req.cacheKey + caches.SuffixPartial)
+				}
+
+				return
+			}
+		}
+
+		// 准备写入
 		this.prepareCacheWriter()
 	}
 

internal/nodes/http_writer.go

@@ -403,6 +403,11 @@ func (this *HTTPWriter) PrepareCache(resp *http.Response, size int64) {
 					return
 				}
 				partialWriter.SetBodyLength(total)
+
+				var contentMD5 = this.rawWriter.Header().Get("Content-MD5")
+				if len(contentMD5) > 0 {
+					partialWriter.SetContentMD5(contentMD5)
+				}
 			}
 			var filterReader = readers.NewFilterReaderCloser(resp.Body)
 			this.cacheIsFinished = true
@@ -462,6 +467,11 @@ func (this *HTTPWriter) PrepareCache(resp *http.Response, size int64) {
 				// 写入total
 				if !writtenTotal && total > 0 {
 					partialWriter.SetBodyLength(total)
+					var contentMD5 = this.rawWriter.Header().Get("Content-MD5")
+					if len(contentMD5) > 0 {
+						partialWriter.SetContentMD5(contentMD5)
+					}
+
 					writtenTotal = true
 				}