Browse Source

Add Retry-After header

Volodymyr Tkach 2 years ago
parent
commit
3e6211976d
2 changed files with 16 additions and 0 deletions
  1. 1 0
      utils/http/servlimit/servlimit.go
  2. 15 0
      utils/http/servlimit/servlimit_test.go

+ 1 - 0
utils/http/servlimit/servlimit.go

@@ -37,6 +37,7 @@ func ReqPerSecond(handler http.Handler, requests int) http.Handler {
 
 
 		// Restrict access
 		// Restrict access
 		if reqs >= requests {
 		if reqs >= requests {
+			w.Header().Set("Retry-After", "1")
 			w.WriteHeader(429)
 			w.WriteHeader(429)
 			if _, err := w.Write([]byte("Too Many Requests\n")); err != nil {
 			if _, err := w.Write([]byte("Too Many Requests\n")); err != nil {
 				log.Printf("%s\n", err.Error())
 				log.Printf("%s\n", err.Error())

+ 15 - 0
utils/http/servlimit/servlimit_test.go

@@ -42,6 +42,7 @@ var _ = Describe("servlimit", func() {
 			defer resp.Body.Close()
 			defer resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
+			Expect(resp.Header.Get("Retry-After")).To(Equal(""))
 
 
 			body, err := io.ReadAll(resp.Body)
 			body, err := io.ReadAll(resp.Body)
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
@@ -54,6 +55,7 @@ var _ = Describe("servlimit", func() {
 			resp.Body.Close()
 			resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
+			Expect(resp.Header.Get("Retry-After")).To(Equal(""))
 
 
 			time.Sleep(1 * time.Second)
 			time.Sleep(1 * time.Second)
 
 
@@ -62,6 +64,7 @@ var _ = Describe("servlimit", func() {
 			defer resp.Body.Close()
 			defer resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
+			Expect(resp.Header.Get("Retry-After")).To(Equal(""))
 
 
 			body, err := io.ReadAll(resp.Body)
 			body, err := io.ReadAll(resp.Body)
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
@@ -74,12 +77,14 @@ var _ = Describe("servlimit", func() {
 			resp.Body.Close()
 			resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
+			Expect(resp.Header.Get("Retry-After")).To(Equal(""))
 
 
 			resp, err = client.Get(srv.URL + "/")
 			resp, err = client.Get(srv.URL + "/")
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
 			defer resp.Body.Close()
 			defer resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
+			Expect(resp.Header.Get("Retry-After")).To(Equal("1"))
 
 
 			body, err := io.ReadAll(resp.Body)
 			body, err := io.ReadAll(resp.Body)
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
@@ -92,30 +97,35 @@ var _ = Describe("servlimit", func() {
 			resp.Body.Close()
 			resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
+			Expect(resp.Header.Get("Retry-After")).To(Equal(""))
 
 
 			resp, err = client.Get(srv.URL + "/")
 			resp, err = client.Get(srv.URL + "/")
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
 			resp.Body.Close()
 			resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
+			Expect(resp.Header.Get("Retry-After")).To(Equal("1"))
 
 
 			resp, err = client.Get(srv.URL + "/")
 			resp, err = client.Get(srv.URL + "/")
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
 			resp.Body.Close()
 			resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
+			Expect(resp.Header.Get("Retry-After")).To(Equal("1"))
 
 
 			resp, err = client.Get(srv.URL + "/")
 			resp, err = client.Get(srv.URL + "/")
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
 			resp.Body.Close()
 			resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
+			Expect(resp.Header.Get("Retry-After")).To(Equal("1"))
 
 
 			resp, err = client.Get(srv.URL + "/")
 			resp, err = client.Get(srv.URL + "/")
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
 			defer resp.Body.Close()
 			defer resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
+			Expect(resp.Header.Get("Retry-After")).To(Equal("1"))
 
 
 			body, err := io.ReadAll(resp.Body)
 			body, err := io.ReadAll(resp.Body)
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
@@ -134,6 +144,7 @@ var _ = Describe("servlimit", func() {
 			defer resp.Body.Close()
 			defer resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
+			Expect(resp.Header.Get("Retry-After")).To(Equal(""))
 
 
 			body, err := io.ReadAll(resp.Body)
 			body, err := io.ReadAll(resp.Body)
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
@@ -150,24 +161,28 @@ var _ = Describe("servlimit", func() {
 			resp.Body.Close()
 			resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
+			Expect(resp.Header.Get("Retry-After")).To(Equal(""))
 
 
 			resp, err = client.Get(srv.URL + "/")
 			resp, err = client.Get(srv.URL + "/")
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
 			resp.Body.Close()
 			resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
+			Expect(resp.Header.Get("Retry-After")).To(Equal(""))
 
 
 			resp, err = client.Get(srv.URL + "/")
 			resp, err = client.Get(srv.URL + "/")
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
 			resp.Body.Close()
 			resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
 			Expect(resp.StatusCode).To(Equal(http.StatusOK))
+			Expect(resp.Header.Get("Retry-After")).To(Equal(""))
 
 
 			resp, err = client.Get(srv.URL + "/")
 			resp, err = client.Get(srv.URL + "/")
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())
 			defer resp.Body.Close()
 			defer resp.Body.Close()
 
 
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
 			Expect(resp.StatusCode).To(Equal(http.StatusTooManyRequests))
+			Expect(resp.Header.Get("Retry-After")).To(Equal("1"))
 
 
 			body, err := io.ReadAll(resp.Body)
 			body, err := io.ReadAll(resp.Body)
 			Expect(err).To(Succeed())
 			Expect(err).To(Succeed())