增加长度验证

icowan · Mar 18, 2020 · 9769e79 · 9769e79
1 parent ec0fffa
commit 9769e79
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/cmd/service/service.go b/cmd/service/service.go
@@ -45,6 +45,7 @@ var (
  logLevel = fs.String("log-level", "all", "logging level.")
  devCors = fs.String("dev-cors", "false", "is develop")
  rateBucketNum = fs.Int("rate-bucket", 10, "rate bucket num")
+ maxLength = fs.Int("max-length", -1, "code length")
  err error
 )
 
@@ -64,6 +65,7 @@ func Run() {
  logLevel = envString("LOG_LEVEL", logLevel)
  devCors = envString("DEV_CORS", devCors)
  rateBucketNum = envInt("RATE_BUCKET", rateBucketNum)
+ maxLength = envInt("MAX_LENGTH", maxLength)
 
  logger = logging.SetLogging(logger, logPath, logLevel)
 
@@ -84,7 +86,7 @@ func Run() {
  }
  }
 
- svc := service.New(getServiceMiddleware(logger), logger, repo, *shortUri)
+ svc := service.New(getServiceMiddleware(logger), logger, repo, *shortUri, *maxLength)
  eps := endpoint.New(svc, getEndpointMiddleware(logger))
  g := createService(eps)
  initCancelInterrupt(g)

diff --git a/pkg/service/service.go b/pkg/service/service.go
@@ -30,18 +30,22 @@ type service struct {
  repository Repository
  logger log.Logger
  shortUri string
+ maxLen int
 }
 
-func New(middleware []Middleware, logger log.Logger, repository Repository, shortUri string) Service {
- var svc Service = NewService(logger, repository, shortUri)
+func New(middleware []Middleware, logger log.Logger, repository Repository, shortUri string, maxLength int) Service {
+ var svc Service = NewService(logger, repository, shortUri, maxLength)
  for _, m := range middleware {
  svc = m(svc)
  }
  return svc
 }
 
-func NewService(logger log.Logger, repository Repository, shortUri string) Service {
- return &service{repository: repository, shortUri: shortUri, logger: logger}
+func NewService(logger log.Logger, repository Repository, shortUri string, maxLength int) Service {
+ if maxLength > 9 {
+ maxLength = 9
+ }
+ return &service{repository: repository, shortUri: shortUri, logger: logger, maxLen: maxLength}
 }
 
 func (s *service) Get(ctx context.Context, code string) (redirect *Redirect, err error) {
@@ -52,9 +56,12 @@ func (s *service) Post(ctx context.Context, domain string) (redirect *Redirect,
  now := time.Now()
  now = now.In(time.Local)
  code := shortid.MustGenerate()
-
  // todo 考虑如何处理垃圾数据的问题 得复的url 不同的code
 
+ if s.maxLen > 0 {
+ code = code[:s.maxLen]
+ }
+
  redirect = &Redirect{
  Code: code,
  URL: domain,