Add rate limit

main.go

@@ -17,7 +17,10 @@
 */
 package main
 
-import "log"
+import (
+	"log"
+	"time"
+)
 
 type Blog struct {
 	Id         string  `json:"id"`
@@ -58,8 +61,16 @@ type Tag struct {
 	Name string `json:"name"`
 }
 
+type Rate struct {
+	Date  time.Time
+	Count int
+}
+
+var rateMap map[string]Rate
+
 func main() {
 	config := getConfig()
+	rateMap = make(map[string]Rate)
 
 	listener := listen(config.Listen, config.CertPath, config.KeyPath)
 	log.Println("Server successfully started")

server.go

@@ -90,6 +90,16 @@ func getPath(conn *tls.Conn) (string, string, error) {
 func handleConn(conn *tls.Conn, baseURL, title, home_message string) {
 	defer conn.Close()
 
+	if !rateIsOk(rateMap, strings.Split(conn.RemoteAddr().String(), ":")[0], 60) {
+		log.Printf("Too many requests for %s\n", conn.RemoteAddr().String())
+		_, err := fmt.Fprintf(conn, "44 60\r\n")
+		if err != nil {
+			log.Println("send error: %s", err)
+			return
+		}
+		return
+	}
+
 	path, query, err := getPath(conn)
 	if err != nil {
 		log.Println("get url: %s", err)

util.go

@@ -25,6 +25,7 @@ import (
 	"io/ioutil"
 	"encoding/json"
 	"log"
+	"time"
 )
 
 func getConfig() Config {
@@ -104,3 +105,23 @@ func formatBlog(toot Blog) string {
 
 	return "### Written by " + author + " on " + toot.Date[0:10] + " at " + toot.Date[11:16] + "\n" + content + "\n=> /toot/" + toot.Id + " More informations about this toot"
 }
+
+func rateIsOk(tab map[string]Rate, remoteIP string, limit int) bool {
+	elmt, ok := tab[remoteIP]
+	if ok == false {
+		tab[remoteIP] = Rate{time.Now(), 0}
+		return true
+	} else {
+		if time.Since(elmt.Date).Minutes() >= 1 {
+			tab[remoteIP] = Rate{time.Now(), 0}
+			return true
+		} else {
+			if elmt.Count < limit {
+				tab[remoteIP] = Rate{elmt.Date, elmt.Count + 1}
+				return true
+			} else {
+				return false
+			}
+		}
+	}
+}