From 2213934795d1609aae4ecccbb07882008abbea94 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Fri, 5 Apr 2024 11:45:18 +0800
Subject: [PATCH] =?UTF-8?q?=E7=BC=93=E5=AD=98=E5=86=99=E5=85=A5=E7=BB=93?=
 =?UTF-8?q?=E6=9D=9F=E6=97=B6=E6=A3=80=E6=9F=A5Content-Length=E6=98=AF?=
 =?UTF-8?q?=E5=90=A6=E5=92=8C=E5=AE=9E=E9=99=85=E5=86=85=E5=AE=B9=E9=95=BF?=
 =?UTF-8?q?=E5=BA=A6=E4=B8=80=E8=87=B4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/caches/errors.go        | 15 ++++++++-------
 internal/caches/writer_file.go   |  7 +++++++
 internal/caches/writer_memory.go | 11 +++++++++--
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/internal/caches/errors.go b/internal/caches/errors.go
index 31244d6..2fc49ba 100644
--- a/internal/caches/errors.go
+++ b/internal/caches/errors.go
@@ -6,13 +6,14 @@ import "errors"
 
 // 常用的几个错误
 var (
-	ErrNotFound           = errors.New("cache not found")
-	ErrFileIsWriting      = errors.New("the cache file is updating")
-	ErrInvalidRange       = errors.New("invalid range")
-	ErrEntityTooLarge     = errors.New("entity too large")
-	ErrWritingUnavailable = errors.New("writing unavailable")
-	ErrWritingQueueFull   = errors.New("writing queue full")
-	ErrServerIsBusy       = errors.New("server is busy")
+	ErrNotFound                = errors.New("cache not found")
+	ErrFileIsWriting           = errors.New("the cache file is updating")
+	ErrInvalidRange            = errors.New("invalid range")
+	ErrEntityTooLarge          = errors.New("entity too large")
+	ErrWritingUnavailable      = errors.New("writing unavailable")
+	ErrWritingQueueFull        = errors.New("writing queue full")
+	ErrServerIsBusy            = errors.New("server is busy")
+	ErrUnexpectedContentLength = errors.New("unexpected content length")
 )
 
 // CapacityError 容量错误
diff --git a/internal/caches/writer_file.go b/internal/caches/writer_file.go
index 49c957c..37745fa 100644
--- a/internal/caches/writer_file.go
+++ b/internal/caches/writer_file.go
@@ -136,6 +136,13 @@ func (this *FileWriter) Close() error {
 
 	var path = this.rawWriter.Name()
 
+	// check content length
+	if this.metaBodySize > 0 && this.bodySize != this.metaBodySize {
+		_ = this.rawWriter.Close()
+		_ = os.Remove(path)
+		return ErrUnexpectedContentLength
+	}
+
 	err := this.WriteHeaderLength(types.Int(this.headerSize))
 	if err != nil {
 		fsutils.WriteBegin()
diff --git a/internal/caches/writer_memory.go b/internal/caches/writer_memory.go
index 2910c08..7efa960 100644
--- a/internal/caches/writer_memory.go
+++ b/internal/caches/writer_memory.go
@@ -121,6 +121,14 @@ func (this *MemoryWriter) Close() error {
 		return nil
 	}
 
+	// check content length
+	if this.expectedBodySize > 0 && this.bodySize != this.expectedBodySize {
+		this.storage.locker.Lock()
+		delete(this.storage.valuesMap, this.hash)
+		this.storage.locker.Unlock()
+		return ErrUnexpectedContentLength
+	}
+
 	this.storage.locker.Lock()
 	this.item.IsDone = true
 	var err error
@@ -129,7 +137,7 @@ func (this *MemoryWriter) Close() error {
 			this.storage.valuesMap[this.hash] = this.item
 
 			select {
-			case this.storage.dirtyChan <- types.String(this.bodySize) + "@" +this.key  :
+			case this.storage.dirtyChan <- types.String(this.bodySize) + "@" + this.key:
 				atomic.AddInt64(&this.storage.totalDirtySize, this.bodySize)
 			default:
 				// remove from values map
@@ -158,7 +166,6 @@ func (this *MemoryWriter) Discard() error {
 
 	this.storage.locker.Lock()
 	delete(this.storage.valuesMap, this.hash)
-
 	this.storage.locker.Unlock()
 	return nil
 }