diff --git a/internal/waf/checkpoints/request_upload.go b/internal/waf/checkpoints/request_upload.go index db1cd80..af0a2f0 100644 --- a/internal/waf/checkpoints/request_upload.go +++ b/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,7 +42,80 @@ func (this *RequestUploadCheckpoint) RequestValue(req requests.Request, param st } hasRequestBody = true - if req.WAFRaw().MultipartForm == nil { + + 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 { + data, err := req.WAFReadBody(req.WAFMaxRequestSize()) + if err != nil { + sysErr = err + return + } + + bodyData = data + req.WAFSetCacheBody(data) + defer req.WAFRestoreBody(data) + } + 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 + + if err != nil { + userErr = err + return + } + + if req.WAFRaw().MultipartForm == nil { + return + } + } + } else { // read first part var bodyData = req.WAFGetCacheBody() if len(bodyData) == 0 { data, err := req.WAFReadBody(req.WAFMaxRequestSize()) @@ -49,72 +128,85 @@ func (this *RequestUploadCheckpoint) RequestValue(req requests.Request, param st req.WAFSetCacheBody(data) defer req.WAFRestoreBody(data) } - oldBody := req.WAFRaw().Body - req.WAFRaw().Body = io.NopCloser(bytes.NewBuffer(bodyData)) - err := req.WAFRaw().ParseMultipartForm(req.WAFMaxRequestSize()) + 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)) + } + } + } + } + } - // 还原 - req.WAFRaw().Body = oldBody - - if err != nil { - userErr = err - return - } - - if req.WAFRaw().MultipartForm == nil { - return + // 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) - for _, files := range req.WAFRaw().MultipartForm.File { - for _, file := range files { - if minSize == 0 || minSize > file.Size { - minSize = file.Size - } - } + if minSize == 0 && requestContentLength > 0 { + minSize = requestContentLength } value = minSize } else if param == "maxSize" { // maxSize - maxSize := int64(0) - for _, files := range req.WAFRaw().MultipartForm.File { - for _, file := range files { - if maxSize < file.Size { - maxSize = file.Size - } - } + if maxSize == 0 && requestContentLength > 0 { + maxSize = requestContentLength } 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, ",") }