Skip to content

Commit

Permalink
Added support for setting the parameter size accepted by the interfac…
Browse files Browse the repository at this point in the history
…e and custom timeout and maxbytes in API file (#1713)

* Added support for setting the parameter size accepted by the interface

* support custom timeout and maxbytes in API file

* support timeout used unit

* remove goctl maxBytes
  • Loading branch information
jiang4869 committed Mar 31, 2022
1 parent 500bd87 commit 321dc2d
Show file tree
Hide file tree
Showing 5 changed files with 74 additions and 2 deletions.
11 changes: 10 additions & 1 deletion rest/engine.go
Expand Up @@ -94,7 +94,7 @@ func (ng *engine) bindRoute(fr featuredRoutes, router httpx.Router, metrics *sta
handler.TimeoutHandler(ng.checkedTimeout(fr.timeout)),
handler.RecoverHandler,
handler.MetricHandler(metrics),
handler.MaxBytesHandler(ng.conf.MaxBytes),
handler.MaxBytesHandler(ng.checkedMaxBytes(fr.maxBytes)),
handler.GunzipHandler,
)
chain = ng.appendAuthHandler(fr, chain, verifier)
Expand Down Expand Up @@ -127,6 +127,15 @@ func (ng *engine) checkedTimeout(timeout time.Duration) time.Duration {
return time.Duration(ng.conf.Timeout) * time.Millisecond
}

func (ng *engine) checkedMaxBytes(bytes int64) int64 {

if bytes > 0 {
return bytes
}

return ng.conf.MaxBytes
}

func (ng *engine) createMetrics() *stat.Metrics {
var metrics *stat.Metrics

Expand Down
35 changes: 35 additions & 0 deletions rest/engine_test.go
Expand Up @@ -194,6 +194,41 @@ func TestEngine_checkedTimeout(t *testing.T) {
}
}

func TestEngine_checkedMaxBytes(t *testing.T) {
tests := []struct {
name string
maxBytes int64
expect int64
}{
{
name: "not set",
expect: 1000,
},
{
name: "less",
maxBytes: 500,
expect: 500,
},
{
name: "equal",
maxBytes: 1000,
expect: 1000,
},
{
name: "more",
maxBytes: 1500,
expect: 1500,
},
}

ng := newEngine(RestConf{
MaxBytes: 1000,
})
for _, test := range tests {
assert.Equal(t, test.expect, ng.checkedMaxBytes(test.maxBytes))
}
}

func TestEngine_notFoundHandler(t *testing.T) {
logx.Disable()

Expand Down
7 changes: 7 additions & 0 deletions rest/server.go
Expand Up @@ -223,6 +223,13 @@ func WithTimeout(timeout time.Duration) RouteOption {
}
}

// WithMaxBytes returns a RouteOption to set maxBytes with given value.
func WithMaxBytes(maxBytes int64) RouteOption {
return func(r *featuredRoutes) {
r.maxBytes = maxBytes
}
}

// WithTLSConfig returns a RunOption that with given tls config.
func WithTLSConfig(cfg *tls.Config) RunOption {
return func(svr *Server) {
Expand Down
1 change: 1 addition & 0 deletions rest/types.go
Expand Up @@ -36,5 +36,6 @@ type (
jwt jwtSetting
signature signatureSetting
routes []Route
maxBytes int64
}
)
22 changes: 21 additions & 1 deletion tools/goctl/api/gogen/genroutes.go
Expand Up @@ -7,6 +7,7 @@ import (
"sort"
"strings"
"text/template"
"time"

"github.com/zeromicro/go-zero/core/collection"
"github.com/zeromicro/go-zero/tools/goctl/api/spec"
Expand Down Expand Up @@ -34,7 +35,7 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
`
routesAdditionTemplate = `
server.AddRoutes(
{{.routes}} {{.jwt}}{{.signature}} {{.prefix}}
{{.routes}} {{.jwt}}{{.signature}} {{.prefix}} {{.timeout}}
)
`
)
Expand All @@ -57,6 +58,8 @@ type (
jwtEnabled bool
signatureEnabled bool
authName string
timeout string
timeoutEnable bool
middlewares []string
prefix string
jwtTrans string
Expand Down Expand Up @@ -110,6 +113,15 @@ func genRoutes(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error
rest.WithPrefix("%s"),`, g.prefix)
}

var timeout string
if g.timeoutEnable {
duration, err := time.ParseDuration(g.timeout)
if err != nil {
panic(err)
}
timeout = fmt.Sprintf("rest.WithTimeout(%d),", duration)
}

var routes string
if len(g.middlewares) > 0 {
gbuilder.WriteString("\n}...,")
Expand All @@ -130,6 +142,7 @@ rest.WithPrefix("%s"),`, g.prefix)
"jwt": jwt,
"signature": signature,
"prefix": prefix,
"timeout": timeout,
}); err != nil {
return err
}
Expand Down Expand Up @@ -205,6 +218,13 @@ func getRoutes(api *spec.ApiSpec) ([]group, error) {
})
}

timeout := g.GetAnnotation("timeout")

if len(timeout) > 0 {
groupedRoutes.timeoutEnable = true
groupedRoutes.timeout = timeout
}

jwt := g.GetAnnotation("jwt")
if len(jwt) > 0 {
groupedRoutes.authName = jwt
Expand Down

0 comments on commit 321dc2d

Please sign in to comment.