优化WAF文件上传处理

This commit is contained in:
GoEdgeLab
2023-08-09 17:55:48 +08:00
parent ccef38bda4
commit 1ee50048ac

View File

@@ -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, ",")
}