Merge pull request #32 from pyama86/hook-error

Change scope with newLimitError

error.go

@@ -15,7 +15,7 @@ type LimitError struct {
 	lh         *limitHandler
 }
 
-func newLimitError(statusCode int, err error, lh *limitHandler) *LimitError {
+func NewLimitError(statusCode int, err error, lh *limitHandler) *LimitError {
 	return &LimitError{
 		statusCode: statusCode,
 		err:        err,

rl.go

@@ -115,23 +115,23 @@ func (lm *limitMw) Handler(next http.Handler) http.Handler {
 				case <-ctx.Done():
 					// Increment must be called even if the request limit is already exceeded
 					if err := lh.limiter.Increment(lh.key, currWindow); err != nil {
-						return newLimitError(http.StatusInternalServerError, err, lh)
+						return NewLimitError(http.StatusInternalServerError, err, lh)
 					}
 					return nil
 				default:
 				}
 				rate, err := lh.status(now, currWindow)
 				if err != nil {
-					return newLimitError(http.StatusPreconditionRequired, err, lh)
+					return NewLimitError(http.StatusPreconditionRequired, err, lh)
 				}
 				nrate := int(math.Round(rate))
 				if nrate >= lh.reqLimit {
-					return newLimitError(http.StatusTooManyRequests, ErrRateLimitExceeded, lh)
+					return NewLimitError(http.StatusTooManyRequests, ErrRateLimitExceeded, lh)
 				}
 
 				lh.rateLimitRemaining = lh.reqLimit - nrate
 				if err := lh.limiter.Increment(lh.key, currWindow); err != nil {
-					return newLimitError(http.StatusInternalServerError, err, lh)
+					return NewLimitError(http.StatusInternalServerError, err, lh)
 				}
 				return nil
 			})

rl_test.go

@@ -35,7 +35,11 @@ func TestRateLimit(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			r := http.NewServeMux()
 			r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-				_, _ = w.Write([]byte("Hello, world"))
+				_, err := w.Write([]byte("Hello, world"))
+				if err != nil {
+					t.Fatal(err)
+				}
+
 			})
 			m := rl.New(tt.limiters...)
 			ts := httptest.NewServer(m(r))
@@ -87,12 +91,12 @@ func TestRateLimit(t *testing.T) {
 func BenchmarkRL(b *testing.B) { //nostyle:all
 	r := http.NewServeMux()
 	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		_, _ = w.Write([]byte("Hello, world"))
+		_, err := w.Write([]byte("Hello, world"))
+		if err != nil {
+			b.Fatal(err)
+		}
 	})
-	m := rl.New(
-		testutil.NewLimiter(10, httprate.KeyByIP, 0),
-		testutil.NewLimiter(10, testutil.KeyByHost, 0),
-	)
+	m := rl.New(testutil.NewLimiter(10, httprate.KeyByIP, 0), testutil.NewLimiter(10, testutil.KeyByHost, 0))
 	ts := httptest.NewServer(m(r))
 	b.Cleanup(func() {
 		ts.Close()

testutil/limiter.go

@@ -71,11 +71,17 @@ func (l *Limiter) OnRequestLimit(err error) http.HandlerFunc {
 		}
 		le, ok := err.(*rl.LimitError)
 		if !ok {
-			_, _ = w.Write([]byte("Too many requests"))
+			_, err = w.Write([]byte("Too many requests"))
+			if err != nil {
+				return
+			}
 			return
 		}
 		msg := fmt.Sprintf("Too many requests. name: %s, ratelimit: %d req/%s, ratelimit-ramaining: %d, ratelimit-reset: %d", le.LimierName(), le.RequestLimit(), le.WindowLen(), le.RateLimitRemaining(), le.RateLimitReset())
-		_, _ = w.Write([]byte(msg))
+		_, err = w.Write([]byte(msg))
+		if err != nil {
+			return
+		}
 	}
 }