优化WAF文件上传处理

internal/waf/checkpoints/request_upload.go

@@ -5,12 +5,18 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
 	"github.com/iwind/TeaGo/lists"
 	"github.com/iwind/TeaGo/maps"
+	"github.com/iwind/TeaGo/types"
 	"io"
 	"net/http"
 	"path/filepath"
+	"regexp"
+	"strconv"
 	"strings"
 )
 
+var multipartHeaderRegexp = regexp.MustCompile(`(?i)(?:^|\r\n)--+\w+\r\n((([\w-]+: .+)\r\n)+)`)
+var multipartHeaderContentRangeRegexp = regexp.MustCompile(`/(\d+)`)
+
 // RequestUploadCheckpoint ${requestUpload.arg}
 type RequestUploadCheckpoint struct {
 	Checkpoint
@@ -36,6 +42,16 @@ func (this *RequestUploadCheckpoint) RequestValue(req requests.Request, param st
 	}
 
 	hasRequestBody = true
+
+	var requestContentLength = req.WAFRaw().ContentLength
+
+	var fields []string
+	var minSize int64
+	var maxSize int64
+	var names []string
+	var extensions []string
+
+	if requestContentLength <= req.WAFMaxRequestSize() { // full read
 		if req.WAFRaw().MultipartForm == nil {
 			var bodyData = req.WAFGetCacheBody()
 			if len(bodyData) == 0 {
@@ -49,10 +65,43 @@ func (this *RequestUploadCheckpoint) RequestValue(req requests.Request, param st
 				req.WAFSetCacheBody(data)
 				defer req.WAFRestoreBody(data)
 			}
-		oldBody := req.WAFRaw().Body
+			var oldBody = req.WAFRaw().Body
 			req.WAFRaw().Body = io.NopCloser(bytes.NewBuffer(bodyData))
-
 			err := req.WAFRaw().ParseMultipartForm(req.WAFMaxRequestSize())
+			if err == nil {
+				for field, files := range req.WAFRaw().MultipartForm.File {
+					if param == "field" {
+						fields = append(fields, field)
+					} else if param == "minSize" {
+						for _, file := range files {
+							if minSize == 0 || minSize > file.Size {
+								minSize = file.Size
+							}
+						}
+					} else if param == "maxSize" {
+						for _, file := range files {
+							if maxSize < file.Size {
+								maxSize = file.Size
+							}
+						}
+					} else if param == "name" {
+						for _, file := range files {
+							if !lists.ContainsString(names, file.Filename) {
+								names = append(names, file.Filename)
+							}
+						}
+					} else if param == "ext" {
+						for _, file := range files {
+							if len(file.Filename) > 0 {
+								exit := strings.ToLower(filepath.Ext(file.Filename))
+								if !lists.ContainsString(extensions, exit) {
+									extensions = append(extensions, exit)
+								}
+							}
+						}
+					}
+				}
+			}
 
 			// 还原
 			req.WAFRaw().Body = oldBody
@@ -66,55 +115,98 @@ func (this *RequestUploadCheckpoint) RequestValue(req requests.Request, param st
 				return
 			}
 		}
+	} else { // read first part
+		var bodyData = req.WAFGetCacheBody()
+		if len(bodyData) == 0 {
+			data, err := req.WAFReadBody(req.WAFMaxRequestSize())
+			if err != nil {
+				sysErr = err
+				return
+			}
+
+			bodyData = data
+			req.WAFSetCacheBody(data)
+			defer req.WAFRestoreBody(data)
+		}
+
+		var subMatches = multipartHeaderRegexp.FindAllSubmatch(bodyData, -1)
+		for _, subMatch := range subMatches {
+			var headers = bytes.Split(subMatch[1], []byte{'\r', '\n'})
+			var partContentLength int64 = -1
+			for _, header := range headers {
+				if len(header) > 2 {
+					var kv = bytes.SplitN(header, []byte{':'}, 2)
+					if len(kv) == 2 {
+						var k = kv[0]
+						var v = kv[1]
+						k = []byte("Content-Range")   // TODO
+						v = []byte("bytes 0-10/1024") // TODO
+						switch string(bytes.ToLower(k)) {
+						case "content-disposition":
+							var props = bytes.Split(v, []byte{';', ' '})
+							for _, prop := range props {
+								var propKV = bytes.SplitN(prop, []byte{'='}, 2)
+								if len(propKV) == 2 {
+									var propValue = string(propKV[1])
+									switch string(propKV[0]) {
+									case "name":
+										if param == "field" {
+											propValue, _ = strconv.Unquote(propValue)
+											fields = append(fields, propValue)
+										}
+									case "filename":
+										if param == "name" {
+											propValue, _ = strconv.Unquote(propValue)
+											names = append(names, propValue)
+										} else if param == "ext" {
+											propValue, _ = strconv.Unquote(propValue)
+											extensions = append(extensions, strings.ToLower(filepath.Ext(propValue)))
+										}
+									}
+								}
+							}
+						case "content-range":
+							if partContentLength <= 0 {
+								var contentRange = multipartHeaderContentRangeRegexp.FindSubmatch(v)
+								if len(contentRange) >= 2 {
+									partContentLength = types.Int64(string(contentRange[1]))
+								}
+							}
+						case "content-length":
+							if partContentLength <= 0 {
+								partContentLength = types.Int64(string(v))
+							}
+						}
+					}
+				}
+			}
+
+			// minSize & maxSize
+			if partContentLength > 0 {
+				if param == "minSize" && (minSize == 0 /** not set yet **/ || partContentLength < minSize) {
+					minSize = partContentLength
+				} else if param == "maxSize" && partContentLength > maxSize {
+					maxSize = partContentLength
+				}
+			}
+		}
+	}
 
 	if param == "field" { // field
-		fields := []string{}
-		for field := range req.WAFRaw().MultipartForm.File {
-			fields = append(fields, field)
-		}
 		value = strings.Join(fields, ",")
 	} else if param == "minSize" { // minSize
-		minSize := int64(0)
+		if minSize == 0 && requestContentLength > 0 {
-		for _, files := range req.WAFRaw().MultipartForm.File {
+			minSize = requestContentLength
-			for _, file := range files {
-				if minSize == 0 || minSize > file.Size {
-					minSize = file.Size
-				}
-			}
 		}
 		value = minSize
 	} else if param == "maxSize" { // maxSize
-		maxSize := int64(0)
+		if maxSize == 0 && requestContentLength > 0 {
-		for _, files := range req.WAFRaw().MultipartForm.File {
+			maxSize = requestContentLength
-			for _, file := range files {
-				if maxSize < file.Size {
-					maxSize = file.Size
-				}
-			}
 		}
 		value = maxSize
 	} else if param == "name" { // name
-		names := []string{}
-		for _, files := range req.WAFRaw().MultipartForm.File {
-			for _, file := range files {
-				if !lists.ContainsString(names, file.Filename) {
-					names = append(names, file.Filename)
-				}
-			}
-		}
 		value = strings.Join(names, ",")
 	} else if param == "ext" { // ext
-		extensions := []string{}
-		for _, files := range req.WAFRaw().MultipartForm.File {
-			for _, file := range files {
-				if len(file.Filename) > 0 {
-					exit := strings.ToLower(filepath.Ext(file.Filename))
-					if !lists.ContainsString(extensions, exit) {
-						extensions = append(extensions, exit)
-					}
-				}
-			}
-		}
 		value = strings.Join(extensions, ",")
 	}