Add ratelimiter

cmd/web/BUILD

@@ -14,6 +14,7 @@ go_library(
         "//internal/util:date",
         "//internal/util:io",
         "//internal/web:mdiningserver",
+        "//internal/web:ratelimiter",
         "@com_github_anders617_mdining_proto//proto:mdining_go_proto",
         "@com_github_golang_glog//:go_default_library",
         "@com_github_golang_protobuf//proto:go_default_library",

cmd/web/main.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/MichiganDiningAPI/api/analytics/analyticsclient"
 	"github.com/MichiganDiningAPI/internal/web/mdiningserver"
+	"github.com/MichiganDiningAPI/internal/web/ratelimiter"
 	pb "github.com/anders617/mdining-proto/proto/mdining"
 	"github.com/golang/glog"
 	"github.com/grpc-ecosystem/grpc-gateway/runtime"
@@ -56,9 +57,7 @@ func allowCORS(h http.Handler) http.Handler {
 	})
 }
 
-//
 // Serves GRPC requests
-//
 func serveGRPC(port string, server *mdiningserver.Server) {
 	lis, err := net.Listen("tcp", ":"+port)
 	if err != nil {
@@ -124,6 +123,7 @@ func main() {
 	pb.RegisterMDiningServer(grpcServer, mDiningServer)
 	// Wrap it in a grpcweb handler in order to also serve grpc-web requests
 	wrappedGrpc := grpcweb.WrapServer(grpcServer, grpcweb.WithAllowedRequestHeaders([]string{"*"}))
+	menuRateLimiter := ratelimiter.New()
 	grpcWebHandler := http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
 		if wrappedGrpc.IsGrpcWebRequest(req) {
 			wrappedGrpc.ServeHTTP(resp, req)
@@ -146,6 +146,10 @@ func main() {
 			http.Error(resp, "Unavailable", http.StatusInternalServerError)
 			return
 		}
+		if !menuRateLimiter.ShouldAllow(req) {
+			http.Error(resp, "Please do not abuse this API. Rate limit reached.", http.StatusInternalServerError)
+			return
+		}
 		// Fall back to other servers.
 		mux.ServeHTTP(resp, req)
 	})

internal/web/BUILD

@@ -9,10 +9,22 @@ go_library(
         "//db:dynamoclient",
         "//internal/processing:mdiningprocessing",
         "//internal/util:date",
+        "//internal/web:ratelimiter",
         "@com_github_anders617_mdining_proto//proto:mdining_go_proto",
         "@com_github_golang_glog//:go_default_library",
         "@com_github_google_uuid//:go_default_library",
         "@org_golang_google_grpc//codes:go_default_library",
         "@org_golang_google_grpc//status:go_default_library",
     ],
 )
+
+go_library(
+    name = "ratelimiter",
+    srcs = ["ratelimiter.go"],
+    importpath = "github.com/MichiganDiningAPI/internal/web/ratelimiter",
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_github_golang_glog//:go_default_library",
+        "@com_github_google_uuid//:go_default_library",
+    ],
+)

internal/web/mdiningserver.go

@@ -134,9 +134,7 @@ func (s *Server) fetchFoodStats(wg *sync.WaitGroup) {
 	glog.Infof("QueryFoodStats Success")
 }
 
-//
 // Used for healthcheck to see if data is available.
-//
 func (s *Server) IsAvailable() bool {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
@@ -146,9 +144,7 @@ func (s *Server) IsAvailable() bool {
 	return true
 }
 
-//
 // Handler for GetDiningHalls request
-//
 func (s *Server) GetDiningHalls(ctx context.Context, req *pb.DiningHallsRequest) (*pb.DiningHallsReply, error) {
 	glog.Infof("GetDiningHalls req{%v}", req)
 	s.mu.RLock()
@@ -160,9 +156,7 @@ func (s *Server) GetDiningHalls(ctx context.Context, req *pb.DiningHallsRequest)
 	return &pb.DiningHallsReply{DiningHalls: s.diningHalls.DiningHalls}, nil
 }
 
-//
 // Handler for GetItems request
-//
 func (s *Server) GetItems(ctx context.Context, req *pb.ItemsRequest) (*pb.ItemsReply, error) {
 	glog.Infof("GetItems req{%v}", req)
 	s.mu.RLock()

internal/web/ratelimiter.go

@@ -0,0 +1,54 @@
+package ratelimiter
+
+import (
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/golang/glog"
+)
+
+const maxPerHour = 500
+
+func getIPFromRemoteAddr(remoteAddr string) string {
+	comps := strings.Split(remoteAddr, ":")
+	if len(comps) != 2 {
+		// If the address is not parsable, just return localhost
+		return "localhost"
+	}
+	return comps[0]
+}
+
+type hourRate struct {
+	hour  int
+	count int
+}
+
+type RateLimiter struct {
+	rates map[string]*hourRate
+	mu    sync.RWMutex
+}
+
+func New() *RateLimiter {
+	return &RateLimiter{rates: make(map[string]*hourRate)}
+}
+
+func (r *RateLimiter) ShouldAllow(req *http.Request) bool {
+	if !strings.Contains(req.URL.Path, "menus") {
+		return true
+	}
+	hour := time.Now().Hour()
+	ip := getIPFromRemoteAddr(req.RemoteAddr)
+	r.mu.Lock()
+	rate, ok := r.rates[ip]
+	if !ok || rate.hour != hour {
+		r.rates[ip] = &hourRate{hour, 0}
+		rate = r.rates[ip]
+	}
+	var requestRate = rate.count + 1
+	rate.count = requestRate
+	r.mu.Unlock()
+	glog.Infof("RateLimiter %s {hour: %v, count: %v} Allow: %b", ip, hour, requestRate, requestRate <= maxPerHour)
+	return requestRate <= maxPerHour
+}