Merge pull request #8 from fatedier/dev

bump version to v0.2.0

CHANGELOG.md

@@ -1,3 +1,8 @@
+### v0.2.0
+
+- Support bandwidth limit on fftw.
+- Support limit the max traffic fftw can used every day.
+
 ### v0.1.0
 
 - Support upload and download process bar.

CHANGELOG_zh.md

@@ -1,3 +1,8 @@
+### v0.2.0
+
+- fftw 支持限速。
+- fftw 支持限制每天使用的流量，超过限制后会自动从服务器端注销，第二天恢复。
+
 ### v0.1.0
 
 - 支持上传下载进度条显示。

cmd/fftw/root.go

@@ -20,6 +20,8 @@ func init() {
 	rootCmd.PersistentFlags().StringVarP(&options.ServerAddr, "server_addr", "s", version.DefaultServerAddr(), "remote fft server address")
 	rootCmd.PersistentFlags().StringVarP(&options.BindAddr, "bind_addr", "b", "0.0.0.0:7778", "bind address")
 	rootCmd.PersistentFlags().StringVarP(&options.AdvicePublicIP, "advice_public_ip", "p", "", "fft worker's advice public ip")
+	rootCmd.PersistentFlags().IntVarP(&options.RateKB, "rate", "", 4096, "max bandwidth fftw will provide, unit is KB, default is 4096KB and min value is 50KB")
+	rootCmd.PersistentFlags().IntVarP(&options.MaxTrafficMBPerDay, "max_traffic_per_day", "", 0, "max traffic fftw can use every day, 0 means no limit, unit is MB, default is 0MB and min value is 128MB")
 
 	rootCmd.PersistentFlags().StringVarP(&options.LogFile, "log_file", "", "console", "log file path")
 	rootCmd.PersistentFlags().StringVarP(&options.LogLevel, "log_level", "", "info", "log level")

go.mod

@@ -12,5 +12,6 @@ require (
 	github.com/mattn/go-runewidth v0.0.4 // indirect
 	github.com/spf13/cobra v0.0.3
 	github.com/spf13/pflag v1.0.3 // indirect
+	golang.org/x/time v0.0.0-20190308202827-9d24e82272b4
 	gopkg.in/cheggaaa/pb.v1 v1.0.28 // indirect
 )

go.sum

@@ -30,5 +30,7 @@ golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a h1:1BGLXjeY4akVXGgbC9HugT3Jv
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223 h1:DH4skfRX4EBpamg7iV4ZlCpblAHI6s6TDM39bFZumv8=
 golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/time v0.0.0-20190308202827-9d24e82272b4 h1:SvFZT6jyqRaOeXpc5h/JSfZenJ2O330aBsf7JfSUXmQ=
+golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 gopkg.in/cheggaaa/pb.v1 v1.0.28 h1:n1tBJnnK2r7g9OW2btFH91V92STTUevLXYFb8gy9EMk=
 gopkg.in/cheggaaa/pb.v1 v1.0.28/go.mod h1:V/YB90LKu/1FcN3WVnfiiE5oMCibMjukxqG/qStrOgw=

pkg/io/limit.go

@@ -0,0 +1,33 @@
+package io
+
+import (
+	"context"
+	"io"
+
+	"golang.org/x/time/rate"
+)
+
+type RateReader struct {
+	underlying io.Reader
+	limiter    *rate.Limiter
+}
+
+func NewRateReader(r io.Reader, limiter *rate.Limiter) *RateReader {
+	return &RateReader{
+		underlying: r,
+		limiter:    limiter,
+	}
+}
+
+func (rr *RateReader) Read(p []byte) (n int, err error) {
+	n, err = rr.underlying.Read(p)
+	if err != nil {
+		return
+	}
+
+	err = rr.limiter.WaitN(context.Background(), n)
+	if err != nil {
+		return
+	}
+	return
+}