From 9769e79fe3491746991f5441b85f6d2347a0e16a Mon Sep 17 00:00:00 2001
From: dudulu <wack1001@163.com>
Date: Wed, 18 Mar 2020 18:13:59 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E9=95=BF=E5=BA=A6=E9=AA=8C?=
 =?UTF-8?q?=E8=AF=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/service/service.go |  4 +++-
 pkg/service/service.go | 17 ++++++++++++-----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/cmd/service/service.go b/cmd/service/service.go
index ead973b..e79024a 100644
--- 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
index fbf9618..a4cc0bb 100644
--- 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,