From b55080468451cccd9429b90192ce860c543bd4cd Mon Sep 17 00:00:00 2001 From: qydysky Date: Mon, 21 Oct 2024 02:11:15 +0800 Subject: [PATCH] 1 --- README.md | 3 +++ config.go | 55 ++++++++++++++++++++++++++++++++++++++++++++++++++++--- main.go | 42 ++++++++++++++++++++++-------------------- 3 files changed, 77 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index ec7b72e..fb5f2cf 100755 --- a/README.md +++ b/README.md @@ -21,6 +21,9 @@ config: - *addr*: string 监听端口 例：`0.0.0.0:8081` - *matchRule*: string 匹配规则 `prefix`：当未匹配到时，返回最近的/匹配， `all`：当未匹配到时，返回404 - *copyBlocks*: int 转发的块数量，默认`1000` +- *retryBlocks*: {} 重试 + - *sizeB*: int 重试的块大小，默认`1000000` + - *num*: int 重试的块数量，默认`1000`，为`-1`时停用重试 - *tls*: {} 启用tls - *pub*: string 公钥pem路径 - *key*: string 私钥pem路径 diff --git a/config.go b/config.go index 8bd47a0..3d1f521 100755 --- a/config.go +++ b/config.go @@ -1,11 +1,13 @@ package front import ( + "bytes" "context" "crypto/tls" "encoding/pem" "errors" "fmt" + "io" "net" "net/http" "os" @@ -13,6 +15,7 @@ import ( "sync" "time" + "github.com/dustin/go-humanize" "github.com/qydysky/front/dealer" filiter "github.com/qydysky/front/filiter" component2 "github.com/qydysky/part/component2" @@ -28,9 +31,14 @@ type Config struct { Pub string `json:"pub"` Key string `json:"key"` } `json:"tls"` - MatchRule string `json:"matchRule"` - CopyBlocks int `json:"copyBlocks"` - BlocksI pslice.BlocksI[byte] `json:"-"` + RetryBlocks struct { + SizeB int `json:"sizeB"` + Num int `json:"num"` + } `json:"retryBlocks"` + RetryBlocksI pslice.BlocksI[byte] `json:"-"` + MatchRule string `json:"matchRule"` + CopyBlocks int `json:"copyBlocks"` + BlocksI pslice.BlocksI[byte] `json:"-"` routeP pweb.WebPath routeMap sync.Map `json:"-"` @@ -71,6 +79,15 @@ func (t *Config) Run(ctx context.Context, logger Logger) { } t.BlocksI = pslice.NewBlocks[byte](16*1024, t.CopyBlocks) } + if t.RetryBlocks.SizeB == 0 { + t.RetryBlocks.SizeB = humanize.MByte + } + if t.RetryBlocks.Num == 0 { + t.RetryBlocks.Num = 1000 + } + if t.RetryBlocks.SizeB > 0 && t.RetryBlocks.Num > 0 { + t.RetryBlocksI = pslice.NewBlocks[byte](t.RetryBlocks.SizeB, t.RetryBlocks.Num) + } syncWeb := pweb.NewSyncMap(&httpSer, &t.routeP, matchfunc) defer syncWeb.Shutdown() @@ -168,6 +185,27 @@ func (t *Config) SwapSign(ctx context.Context, logger Logger) { Deal(ctx context.Context, w http.ResponseWriter, r *http.Request, routePath string, chosenBack *Back, logger Logger, blocksi pslice.BlocksI[byte]) error } + // repack + var reBuf []byte + if t.RetryBlocksI != nil && r.Body != nil { + var putBack func() + var e error + reBuf, putBack, e = t.RetryBlocksI.Get() + if e != nil { + logger.Warn(`W:`, fmt.Sprintf(logFormat, r.RemoteAddr, route.config.Addr, routePath, "Err", ErrReqReBodyFail)) + w.Header().Add(header+"Error", ErrReqReBodyFail.Error()) + w.WriteHeader(http.StatusServiceUnavailable) + return + } + defer putBack() + if n, _ := r.Body.Read(reBuf); n == cap(reBuf) { + logger.Warn(`W:`, fmt.Sprintf(logFormat, r.RemoteAddr, route.config.Addr, routePath, "Err", ErrReqReBodyOverflow)) + w.Header().Add(header+"Error", ErrReqReBodyOverflow.Error()) + w.WriteHeader(http.StatusServiceUnavailable) + return + } + } + for i := 0; i < len(backIs); i++ { if !backIs[i].IsLive() { continue @@ -177,6 +215,10 @@ func (t *Config) SwapSign(ctx context.Context, logger Logger) { backIs[i].lastChosenT = time.Now() backIs[i].lock.Unlock() + if len(reBuf) != 0 { + r.Body = io.NopCloser(bytes.NewBuffer(reBuf)) + } + if !strings.Contains(backIs[i].To, "://") { e = component2.Get[reqDealer]("local").Deal(r.Context(), w, r, routePath, backIs[i], logger, t.BlocksI) } else if strings.ToLower((r.Header.Get("Upgrade"))) == "websocket" { @@ -303,6 +345,13 @@ func (t *Route) SwapSign(logger Logger) { func (t *Route) FiliterBackByRequest(r *http.Request) []*Back { var backLink []*Back for i := 0; i < len(t.Backs); i++ { + if ok, e := t.Backs[i].getFiliterReqUri().Match(r); ok && e == nil { + t.Backs[i].route = t + for k := uint(0); k < t.Backs[i].Weight; k++ { + backLink = append(backLink, &t.Backs[i]) + } + } + if ok, e := t.Backs[i].getFiliterReqHeader().Match(r.Header); ok && e == nil { t.Backs[i].route = t for k := uint(0); k < t.Backs[i].Weight; k++ { diff --git a/main.go b/main.go index a7ffa5d..3c57b3a 100755 --- a/main.go +++ b/main.go @@ -179,24 +179,26 @@ func copyHeader(s, t http.Header, app []dealer.HeaderDealer) { var cookie = fmt.Sprintf("%p", &struct{}{}) var header = "X-Front-" var ( - ErrRedirect = errors.New("ErrRedirect") - ErrNoHttp = errors.New("ErrNoHttp") - ErrNoWs = errors.New("ErrNoWs") - ErrCopy = errors.New("ErrCopy") - ErrReqCreFail = errors.New("ErrReqCreFail") - ErrReqDoFail = errors.New("ErrReqDoFail") - ErrResDoFail = errors.New("ErrResDoFail") - ErrResFail = errors.New("ErrResFail") - ErrResTO = errors.New("ErrResTO") - ErrUriTooLong = errors.New("ErrUriTooLong") - ErrPatherCheckFail = errors.New("ErrPatherCheckFail") - ErrHeaderCheckFail = errors.New("ErrHeaderCheckFail") - ErrBodyCheckFail = errors.New("ErrBodyCheckFail") - ErrAllBacksFail = errors.New("ErrAllBacksFail") - ErrBackFail = errors.New("ErrBackFail") - ErrNoRoute = errors.New("ErrNoRoute") - ErrDealReqUri = errors.New("ErrDealReqUri") - ErrDealReqHeader = errors.New("ErrDealReqHeader") - ErrDealResHeader = errors.New("ErrDealResHeader") - ErrCerVerify = errors.New("ErrCerVerify") + ErrRedirect = errors.New("ErrRedirect") + ErrNoHttp = errors.New("ErrNoHttp") + ErrNoWs = errors.New("ErrNoWs") + ErrCopy = errors.New("ErrCopy") + ErrReqReBodyFail = errors.New("ErrReqReBodyFail") + ErrReqReBodyOverflow = errors.New("ErrReqReBodyOverflow") + ErrReqCreFail = errors.New("ErrReqCreFail") + ErrReqDoFail = errors.New("ErrReqDoFail") + ErrResDoFail = errors.New("ErrResDoFail") + ErrResFail = errors.New("ErrResFail") + ErrResTO = errors.New("ErrResTO") + ErrUriTooLong = errors.New("ErrUriTooLong") + ErrPatherCheckFail = errors.New("ErrPatherCheckFail") + ErrHeaderCheckFail = errors.New("ErrHeaderCheckFail") + ErrBodyCheckFail = errors.New("ErrBodyCheckFail") + ErrAllBacksFail = errors.New("ErrAllBacksFail") + ErrBackFail = errors.New("ErrBackFail") + ErrNoRoute = errors.New("ErrNoRoute") + ErrDealReqUri = errors.New("ErrDealReqUri") + ErrDealReqHeader = errors.New("ErrDealReqHeader") + ErrDealResHeader = errors.New("ErrDealResHeader") + ErrCerVerify = errors.New("ErrCerVerify") ) -- 2.39.2