update at 2019-11-25 17:23:04 by ehlxr

main.go

@@ -2,65 +2,42 @@ package main
 
 import (
 	"context"
-	"encoding/base64"
 	"fmt"
 	dt "github.com/JetBlink/dingtalk-notify-go-sdk"
-	"github.com/jessevdk/go-flags"
+	"github.com/ehlxr/ddgo/pkg"
 	"io"
 	"net/http"
 	"os"
 	"os/signal"
 	"strings"
+	"time"
 	log "unknwon.dev/clog/v2"
 )
 
 var (
-	AppName   string
-	Version   string
-	BuildTime string
-	GitCommit string
-	GoVersion string
-
-	versionTpl = `%s
-Name: %s
-Version: %s
-BuildTime: %s
-GitCommit: %s
-GoVersion: %s
-
-`
-	bannerBase64 = "DQogX19fXyAgX19fXyAgICBfX18gIF9fX19fIA0KKCAgXyBcKCAgXyBcICAvIF9fKSggIF8gICkNCiApKF8pICkpKF8pICkoIChfLS4gKShfKSggDQooX19fXy8oX19fXy8gIFxfX18vKF9fX19fKQ0K"
-
-	opts struct {
-		Addr    string `short:"a" long:"addr" default:":80" env:"ADDR" description:"Addr to listen on for HTTP server"`
-		Version bool   `short:"v" long:"version" description:"Show version info"`
-		Robot   robot  `group:"DingTalk Robot Options" namespace:"robot" env-namespace:"ROBOT" `
-	}
-
 	dingTalk *dt.Robot
+	limiter  *pkg.LimiterServer
 )
 
-type robot struct {
-	Token     string   `short:"t" long:"token" env:"TOKEN" description:"DingTalk robot access token" required:"true"`
-	Secret    string   `short:"s" long:"secret" env:"SECRET" description:"DingTalk robot secret"`
-	AtMobiles []string `short:"m" long:"at-mobiles" env:"AT_MOBILES" env-delim:"," description:"The mobile of the person will be at"`
-	IsAtAll   bool     `short:"e" long:"at-all" env:"AT_ALL" description:"Whether at everyone"`
-}
-
 func init() {
 	initLog()
 }
 
 func main() {
-	parseArg()
+	pkg.ParseArg()
 
-	dingTalk = dt.NewRobot(opts.Robot.Token, opts.Robot.Secret)
+	dingTalk = dt.NewRobot(pkg.Opts.Robot.Token, pkg.Opts.Robot.Secret)
+	limiter = pkg.NewLimiterServer(1*time.Minute, 20)
 
+	start()
+}
+
+func start() {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", requestHandle)
 
 	server := &http.Server{
-		Addr:    opts.Addr,
+		Addr:    pkg.Opts.Addr,
 		Handler: mux,
 	}
 
@@ -74,7 +51,7 @@ func main() {
 		}
 	}()
 
-	log.Info("Starting HTTP server on http://%s", opts.Addr)
+	log.Info("Starting HTTP server on http://%s", pkg.Opts.Addr)
 	if err := server.ListenAndServe(); err != nil {
 		if err == http.ErrServerClosed {
 			log.Info("Server closed under request")
@@ -103,33 +80,6 @@ func initLog() {
 	}
 }
 
-func parseArg() {
-	parser := flags.NewParser(&opts, flags.HelpFlag|flags.PassDoubleDash)
-	parser.NamespaceDelimiter = "-"
-
-	if AppName != "" {
-		parser.Name = AppName
-	}
-
-	if _, err := parser.Parse(); err != nil {
-		if opts.Version {
-			printVersion()
-			os.Exit(0)
-		}
-
-		if flagsErr, ok := err.(*flags.Error); ok && flagsErr.Type == flags.ErrHelp {
-			_, _ = fmt.Fprintln(os.Stdout, err)
-			os.Exit(0)
-		}
-
-		_, _ = fmt.Fprintln(os.Stderr, err)
-
-		parser.WriteHelp(os.Stderr)
-
-		os.Exit(1)
-	}
-}
-
 func requestHandle(w http.ResponseWriter, r *http.Request) {
 	err := r.ParseForm()
 	if err != nil {
@@ -147,13 +97,19 @@ func requestHandle(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	ats := opts.Robot.AtMobiles
+	ats := pkg.Opts.Robot.AtMobiles
 	at := r.Form.Get("at")
 	if at != "" {
 		ats = append(ats, strings.Split(at, ",")...)
 	}
 
-	err = dingTalk.SendTextMessage(content, ats, opts.Robot.IsAtAll)
+	app := r.Form.Get("app")
+	if app != "" {
+		content = fmt.Sprintf("%s\n%s", app, content)
+	}
+
+	if limiter.IsAvailable() {
+		err = dingTalk.SendTextMessage(content, ats, pkg.Opts.Robot.IsAtAll)
 		if err != nil {
 			log.Error("%+v", err)
 			_, _ = fmt.Fprintln(w, err)
@@ -161,11 +117,13 @@ func requestHandle(w http.ResponseWriter, r *http.Request) {
 		}
 
 		log.Info("send message <%s> success", content)
-	_, _ = io.WriteString(w, "send message success")
-}
+		_, _ = io.WriteString(w, "success")
+		return
+	} else {
+		log.Error("dingTalk 1 m allow send 20 msg. msg %v discarded.",
+			content)
 
-// printVersion Print out version information
-func printVersion() {
-	banner, _ := base64.StdEncoding.DecodeString(bannerBase64)
-	fmt.Printf(versionTpl, banner, AppName, Version, BuildTime, GitCommit, GoVersion)
+		_, _ = io.WriteString(w, "dingTalk 1 m allow send 20 msg. msg discarded.")
+		return
+	}
 }

pkg/limiter.go

@@ -0,0 +1,42 @@
+package pkg
+
+import (
+	"sync"
+	"time"
+)
+
+type LimiterServer struct {
+	interval time.Duration
+	maxCount int
+	sync.Mutex
+	reqCount  int
+	startTime time.Time
+}
+
+func NewLimiterServer(i time.Duration, c int) *LimiterServer {
+	return &LimiterServer{
+		interval: i,
+		maxCount: c,
+	}
+}
+
+func (limiter *LimiterServer) IsAvailable() bool {
+	limiter.Lock()
+	defer limiter.Unlock()
+
+	if limiter.startTime.IsZero() ||
+		limiter.startTime.Add(limiter.interval).Before(time.Now()) {
+		limiter.reqCount = 1
+		limiter.startTime = time.Now()
+
+		return true
+	}
+
+	if limiter.reqCount < limiter.maxCount {
+		limiter.reqCount += 1
+
+		return true
+	}
+
+	return false
+}

pkg/limiter_test.go

@@ -0,0 +1,18 @@
+package pkg
+
+import (
+	"testing"
+	"time"
+)
+
+func TestLimiter(t *testing.T) {
+	limiter := NewLimiterServer(1*time.Second, 5)
+
+	for {
+		if limiter.IsAvailable() {
+			t.Log("hello...", limiter.reqCount)
+		} else {
+			return
+		}
+	}
+}

pkg/parser.go

@@ -0,0 +1,49 @@
+package pkg
+
+import (
+	"fmt"
+	"github.com/jessevdk/go-flags"
+	"os"
+)
+
+var (
+	Opts struct {
+		Addr    string `short:"a" long:"addr" default:":80" env:"ADDR" description:"Addr to listen on for HTTP server"`
+		Version bool   `short:"v" long:"version" description:"Show version info"`
+		Robot   Robot  `group:"DingTalk Robot Options" namespace:"robot" env-namespace:"ROBOT" `
+	}
+)
+
+type Robot struct {
+	Token     string   `short:"t" long:"token" env:"TOKEN" description:"DingTalk robot access token" required:"true"`
+	Secret    string   `short:"s" long:"secret" env:"SECRET" description:"DingTalk robot secret"`
+	AtMobiles []string `short:"m" long:"at-mobiles" env:"AT_MOBILES" env-delim:"," description:"The mobile of the person will be at"`
+	IsAtAll   bool     `short:"e" long:"at-all" env:"AT_ALL" description:"Whether at everyone"`
+}
+
+func ParseArg() {
+	parser := flags.NewParser(&Opts, flags.HelpFlag|flags.PassDoubleDash)
+	parser.NamespaceDelimiter = "-"
+
+	if AppName != "" {
+		parser.Name = AppName
+	}
+
+	if _, err := parser.Parse(); err != nil {
+		if Opts.Version {
+			PrintVersion()
+			os.Exit(0)
+		}
+
+		if flagsErr, ok := err.(*flags.Error); ok && flagsErr.Type == flags.ErrHelp {
+			_, _ = fmt.Fprintln(os.Stdout, err)
+			os.Exit(0)
+		}
+
+		_, _ = fmt.Fprintln(os.Stderr, err)
+
+		parser.WriteHelp(os.Stderr)
+
+		os.Exit(1)
+	}
+}

pkg/version.go

@@ -0,0 +1,30 @@
+package pkg
+
+import (
+	"encoding/base64"
+	"fmt"
+)
+
+var (
+	AppName   string
+	Version   string
+	BuildTime string
+	GitCommit string
+	GoVersion string
+
+	versionTpl = `%s
+Name: %s
+Version: %s
+BuildTime: %s
+GitCommit: %s
+GoVersion: %s
+
+`
+	bannerBase64 = "DQogX19fXyAgX19fXyAgICBfX18gIF9fX19fIA0KKCAgXyBcKCAgXyBcICAvIF9fKSggIF8gICkNCiApKF8pICkpKF8pICkoIChfLS4gKShfKSggDQooX19fXy8oX19fXy8gIFxfX18vKF9fX19fKQ0K"
+)
+
+// PrintVersion Print out version information
+func PrintVersion() {
+	banner, _ := base64.StdEncoding.DecodeString(bannerBase64)
+	fmt.Printf(versionTpl, banner, AppName, Version, BuildTime, GitCommit, GoVersion)
+}

version

@@ -1 +1 @@
-v0.0.3
+v0.0.4