From 804382e85d876284196d904800263bef4b02109c Mon Sep 17 00:00:00 2001
From: Michael Bumann <hello@michaelbumann.com>
Date: Sun, 25 Aug 2019 17:49:24 +0200
Subject: [PATCH] Add request rate limit option

This allows to configure a request rate limit. By default it is
currently set to 10 requests per second.
---
 invoices_proxy.go | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/invoices_proxy.go b/invoices_proxy.go
index d52e264..f78745b 100644
--- a/invoices_proxy.go
+++ b/invoices_proxy.go
@@ -5,11 +5,26 @@ import (
 	"github.com/bumi/lntip/ln"
 	"github.com/labstack/echo"
 	"github.com/labstack/echo/middleware"
+	"github.com/didip/tollbooth"
+	"github.com/didip/tollbooth/limiter"
 	"log"
 	"net/http"
 	"os"
 )
 
+// move to file
+func LimitMiddleware(lmt *limiter.Limiter) echo.MiddlewareFunc {
+	return func(next echo.HandlerFunc) echo.HandlerFunc {
+		return echo.HandlerFunc(func(c echo.Context) error {
+		httpError := tollbooth.LimitByRequest(lmt, c.Response(), c.Request())
+			if httpError != nil {
+				return c.String(httpError.StatusCode, httpError.Message)
+			}
+			return next(c)
+		})
+	}
+}
+
 var stdOutLogger = log.New(os.Stdout, "", log.LstdFlags)
 
 type Invoice struct {
@@ -24,6 +39,7 @@ func main() {
 	bind := flag.String("bind", ":1323", "Host and port to bind on")
 	staticPath := flag.String("static-path", "", "Path to a static assets directory. Blank to disable serving static files")
 	disableCors := flag.Bool("disable-cors", false, "Disable CORS headers")
+	requestLimit := flag.Float64("request-limit", 10, "Request limit per second")
 
 	flag.Parse()
 
@@ -36,6 +52,11 @@ func main() {
 	}
 	e.Use(middleware.Recover())
 
+	if *requestLimit > 0 {
+		limiter := tollbooth.NewLimiter(*requestLimit, nil)
+		e.Use(LimitMiddleware(limiter))
+	}
+
 	stdOutLogger.Printf("Connection to %s using macaroon %s and cert %s", *address, *macaroonFile, *certFile)
 	lndOptions := ln.LNDoptions{
 		Address:      *address,
@@ -77,5 +98,10 @@ func main() {
 		return c.JSON(http.StatusOK, invoice)
 	})
 
+	// debug test endpoint
+	e.GET("/ping", func(c echo.Context) error {
+		return c.JSON(http.StatusOK, "pong")
+	})
+
 	e.Logger.Fatal(e.Start(*bind))
 }