From c47e76cd3d0878680e64842a6d420efb7f6a34d3 Mon Sep 17 00:00:00 2001
From: NielsAD <niels.a.d@gmail.com>
Date: Tue, 11 Sep 2018 14:04:51 +0200
Subject: [PATCH] Add rate limiter

---
 .gitmodules                     |  6 +++---
 main.go                         | 23 ++++++++++++++++-------
 vendor/github.com/pkg/errors    |  1 +
 vendor/github.com/ulule/limiter |  1 +
 4 files changed, 21 insertions(+), 10 deletions(-)
 create mode 160000 vendor/github.com/pkg/errors
 create mode 160000 vendor/github.com/ulule/limiter

diff --git a/.gitmodules b/.gitmodules
index 0d64eae..7d5b80b 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -7,6 +7,6 @@
 [submodule "vendor/github.com/karrick/godirwalk"]
 	path = vendor/github.com/karrick/godirwalk
 	url = https://github.com/karrick/godirwalk.git
-[submodule "github.com/mattn/go-sqlite3"]
-	path = github.com/mattn/go-sqlite3
-	url = https://github.com/mattn/go-sqlite3.git
+[submodule "vendor/github.com/ulule/limiter"]
+	path = vendor/github.com/ulule/limiter
+	url = https://github.com/ulule/limiter.git
diff --git a/main.go b/main.go
index 2108ae6..e05390b 100644
--- a/main.go
+++ b/main.go
@@ -18,14 +18,18 @@ import (
 	"time"
 
 	_ "github.com/mattn/go-sqlite3"
+	"github.com/ulule/limiter"
+	"github.com/ulule/limiter/drivers/middleware/stdlib"
+	"github.com/ulule/limiter/drivers/store/memory"
 )
 
 var (
-	addr    = flag.String("a", ":80", "TCP network address to listen for connections")
-	db      = flag.String("d", ":memory:", "Database location")
-	dir     = flag.String("r", ".", "Root directory to serve")
-	refresh = flag.String("i", "1h", "Refresh interval")
-	cached  = flag.Bool("cached", false, "Serve everything from cache (rather than search/recursive queries only)")
+	addr      = flag.String("a", ":80", "TCP network address to listen for connections")
+	db        = flag.String("d", ":memory:", "Database location")
+	dir       = flag.String("r", ".", "Root directory to serve")
+	refresh   = flag.String("i", "1h", "Refresh interval")
+	forwarded = flag.Bool("forwarded", false, "Trust X-Real-IP and X-Forwarded-For headers")
+	cached    = flag.Bool("cached", false, "Serve everything from cache (rather than search/recursive queries only)")
 )
 
 var logOut = log.New(os.Stdout, "", 0)
@@ -78,9 +82,14 @@ func main() {
 		}
 	})
 
+	limit := stdlib.NewMiddleware(
+		limiter.New(memory.NewStore(), limiter.Rate{Period: 1 * time.Second, Limit: 5}),
+		stdlib.WithForwardHeader(*forwarded),
+	)
+
 	srv := &http.Server{Addr: *addr}
-	http.Handle("/idx/", logRequest(http.StripPrefix("/idx/", fs)))
-	http.Handle("/dl/", logRequest(http.StripPrefix("/dl/", nodir(http.FileServer(http.Dir(fs.Root))))))
+	http.Handle("/idx/", limit.Handler(logRequest(http.StripPrefix("/idx/", fs))))
+	http.Handle("/dl/", limit.Handler(logRequest(http.StripPrefix("/dl/", nodir(http.FileServer(http.Dir(fs.Root)))))))
 	http.Handle("/sitemap.txt", http.HandlerFunc(fs.Sitemap))
 	http.Handle("/", pub)
 
diff --git a/vendor/github.com/pkg/errors b/vendor/github.com/pkg/errors
new file mode 160000
index 0000000..c059e47
--- /dev/null
+++ b/vendor/github.com/pkg/errors
@@ -0,0 +1 @@
+Subproject commit c059e472caf75dbe73903f6521a20abac245b17f
diff --git a/vendor/github.com/ulule/limiter b/vendor/github.com/ulule/limiter
new file mode 160000
index 0000000..1f43cf9
--- /dev/null
+++ b/vendor/github.com/ulule/limiter
@@ -0,0 +1 @@
+Subproject commit 1f43cf9e49e56c9b502d08637597c2cc60aeeec5