mirror of
				https://github.com/TeaOSLab/EdgeNode.git
				synced 2025-11-04 07:40:56 +08:00 
			
		
		
		
	优化WAF文件上传处理
This commit is contained in:
		@@ -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, ",")
 | 
			
		||||
	}
 | 
			
		||||
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user