From 92ee4fe2bdf786f02901e82585881d2e7a492426 Mon Sep 17 00:00:00 2001
From: Romain de Laage <romain.delaage@rdelaage.ovh>
Date: Sun, 18 Apr 2021 10:43:12 +0200
Subject: [PATCH] Make rate limite configurable, little entrypoint fix

---
 config.json |  3 ++-
 main.go     |  3 ++-
 server.go   |  8 ++++----
 start.sh    | 12 +++++++++++-
 util.go     |  5 +++--
 5 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/config.json b/config.json
index 932573d..f9c2e64 100644
--- a/config.json
+++ b/config.json
@@ -4,5 +4,6 @@
 	"key_path": "certs/key.rsa",
 	"base_url": "https://mamot.fr",
 	"title": "MastoGem",
-	"home_message": "Welcome on MastoGem, a Mastodon proxy for Gemini.\nYou can view the last 20 toots of a Mastodon account by providing its id, for example:\n=> gemini://localhost/310515 My Mastodon account"
+	"home_message": "Welcome on MastoGem, a Mastodon proxy for Gemini.\nYou can view the last 20 toots of a Mastodon account by providing its id, for example:\n=> gemini://localhost/310515 My Mastodon account",
+	"rate_limit": 45
 }
diff --git a/main.go b/main.go
index be18f87..6f0fc71 100644
--- a/main.go
+++ b/main.go
@@ -38,6 +38,7 @@ type Config struct {
 	BaseURL     string `json:"base_url"`
 	Title       string `json:"title"`
 	HomeMessage string `json:"home_message"`
+	RateLimit   int    `json:"rate_limit"`
 }
 
 type Account struct {
@@ -76,5 +77,5 @@ func main() {
 	log.Println("Server successfully started")
 	log.Println("Server is listening at " + config.Listen)
 
-	serve(listener, config.BaseURL, config.Title, config.HomeMessage)
+	serve(listener, config.BaseURL, config.Title, config.HomeMessage, config.RateLimit)
 }
diff --git a/server.go b/server.go
index 7e5bb5b..6a08c2a 100644
--- a/server.go
+++ b/server.go
@@ -48,14 +48,14 @@ func listen(address, certFile, keyFile string) net.Listener {
 	return listener
 }
 
-func serve(listener net.Listener, baseURL, title, home_message string) {
+func serve(listener net.Listener, baseURL, title, home_message string, rateLimit int) {
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
 			log.Println(err)
 		}
 
-		go handleConn(conn.(*tls.Conn), baseURL, title, home_message)
+		go handleConn(conn.(*tls.Conn), baseURL, title, home_message, rateLimit)
 	}
 }
 
@@ -87,10 +87,10 @@ func getPath(conn *tls.Conn) (string, string, error) {
 	return parsedURL.Path, parsedURL.RawQuery, nil
 }
 
-func handleConn(conn *tls.Conn, baseURL, title, home_message string) {
+func handleConn(conn *tls.Conn, baseURL, title, home_message string, rateLimit int) {
 	defer conn.Close()
 
-	if !rateIsOk(rateMap, strings.Split(conn.RemoteAddr().String(), ":")[0], 60) {
+	if !rateIsOk(rateMap, strings.Split(conn.RemoteAddr().String(), ":")[0], rateLimit) {
 		log.Printf("Too many requests for %s\n", conn.RemoteAddr().String())
 		_, err := fmt.Fprintf(conn, "44 60\r\n")
 		if err != nil {
diff --git a/start.sh b/start.sh
index 8780af4..7ba27df 100644
--- a/start.sh
+++ b/start.sh
@@ -15,18 +15,27 @@ fi
 if [ ! -f /key.rsa ]
 then
 	echo "You must bind a private key at /key.rsa"
+	exit 1
 fi
 
 if [ -z "$TITLE" ]
 then
+	echo "Using default title"
 	TITLE=MastoGem
 fi
 
 if [ -z "$HOME_MESSAGE" ]
 then
+	echo "Using default home message"
 	HOME_MESSAGE="Welcome on MastoGem, a Mastodon proxy for Gemini !"
 fi
 
+if [ -z "$RATE_LIMIT" ]
+then
+	echo "Using default rate limit"
+	RATE_LIMIT=45
+fi
+
 cat << EOF > /config.json
 {
 	"listen": "0.0.0.0:1965",
@@ -34,7 +43,8 @@ cat << EOF > /config.json
 	"key_path": "/key.rsa",
 	"base_url": "$MASTODON_BASE_URL",
 	"title": "$TITLE",
-	"home_message": "$HOME_MESSAGE"
+	"home_message": "$HOME_MESSAGE",
+	"rate_limit": $RATE_LIMIT
 }
 EOF
 
diff --git a/util.go b/util.go
index 98df018..96befb6 100644
--- a/util.go
+++ b/util.go
@@ -40,6 +40,7 @@ func getConfig() Config {
 			BaseURL: "https://mamot.fr",
 			Title: "MastoGem",
 			HomeMessage: "Welcome on MastoGem, this is a Mastodon proxy for Gemini. You can view the last 20 toots of a Mastodon account by providing its ID.",
+			RateLimit: 45,
 		}
 
 		return config
@@ -109,11 +110,11 @@ func formatBlog(toot Blog) string {
 func rateIsOk(tab map[string]Rate, remoteIP string, limit int) bool {
 	elmt, ok := tab[remoteIP]
 	if ok == false {
-		tab[remoteIP] = Rate{time.Now(), 0}
+		tab[remoteIP] = Rate{time.Now(), 1}
 		return true
 	} else {
 		if time.Since(elmt.Date).Minutes() >= 1 {
-			tab[remoteIP] = Rate{time.Now(), 0}
+			tab[remoteIP] = Rate{time.Now(), 1}
 			return true
 		} else {
 			if elmt.Count < limit {