Merge pull request #103 from OdyseeTeam/firewall

merge high level firewall measures

.github/workflows/pipeline.yml

@@ -24,7 +24,8 @@ jobs:
       - name: Check running containers
         run: docker ps -a
 
-      - uses: actions/setup-go@v3
+      - name: Set up Go 1.20
+        uses: actions/setup-go@v3
         with:
           go-version: '1.20.x'
         id: go

cmd/main.go

@@ -2,7 +2,7 @@ package cmd
 
 import (
 	"fmt"
-	"io/ioutil"
+	"io"
 	"net/http"
 	"os"
 	"time"
@@ -260,7 +260,7 @@ func initPubkey() {
 	if err != nil {
 		l.Fatal(err)
 	}
-	rawKey, err := ioutil.ReadAll(r.Body)
+	rawKey, err := io.ReadAll(r.Body)
 	if err != nil {
 		l.Fatal(err)
 	}

firewall/firewall.go

@@ -0,0 +1,54 @@
+package firewall
+
+import (
+	"errors"
+	"sync"
+	"time"
+
+	"github.com/bluele/gcache"
+)
+
+var WindowSize = 120 * time.Second
+
+const MaxStringsPerIp = 6
+
+var resourcesForIPCache = gcache.New(1000).Simple().Build()
+var whitelist = map[string]bool{
+	"51.210.0.171": true,
+}
+
+func CheckAndRateLimitIp(ip string, endpoint string) (bool, int) {
+	if ip == "" {
+		return false, 0
+	}
+	if whitelist[ip] {
+		return false, 0
+	}
+	resources, err := resourcesForIPCache.Get(ip)
+	if errors.Is(err, gcache.KeyNotFoundError) {
+		tokensMap := &sync.Map{}
+		tokensMap.Store(endpoint, time.Now())
+		err := resourcesForIPCache.SetWithExpire(ip, tokensMap, WindowSize*10)
+		if err != nil {
+			return false, 1
+		}
+		return false, 1
+	}
+	tokensForIP, _ := resources.(*sync.Map)
+	currentTime := time.Now()
+	tokensForIP.Store(endpoint, currentTime)
+	resourcesCount := 0
+	flagged := false
+	tokensForIP.Range(func(k, v interface{}) bool {
+		if currentTime.Sub(v.(time.Time)) > WindowSize {
+			tokensForIP.Delete(k)
+			return true
+		}
+		resourcesCount++
+		if !flagged && resourcesCount > MaxStringsPerIp {
+			flagged = true
+		}
+		return true
+	})
+	return flagged, resourcesCount
+}

firewall/firewall_test.go

@@ -0,0 +1,31 @@
+package firewall
+
+import (
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCheckIPAccess(t *testing.T) {
+	ip := "192.168.0.1"
+	endpoint := "/api/v1/example"
+	WindowSize = 7 * time.Second
+	// Test the first five accesses for an IP don't exceed the limit
+	for i := 1; i <= 6; i++ {
+		result, _ := CheckAndRateLimitIp(ip, endpoint+strconv.Itoa(i))
+		assert.False(t, result, "Expected result to be false, got %v for endpoint %s", result, endpoint+strconv.Itoa(i))
+	}
+
+	// Test the sixth access for an IP exceeds the limit
+	result, _ := CheckAndRateLimitIp(ip, endpoint+"7")
+	assert.True(t, result, "Expected result to be true, got %v for endpoint %s", result, endpoint+"7")
+
+	// Wait for the window size to elapse
+	time.Sleep(WindowSize)
+
+	// Test the access for an IP after the window size elapses doesn't exceed the limit
+	result, _ = CheckAndRateLimitIp(ip, endpoint+"7")
+	assert.False(t, result, "Expected result to be false, got %v for endpoint %s", result, endpoint+"7")
+}

init.sql

@@ -0,0 +1,14 @@
+use godycdn;
+CREATE TABLE `object`
+(
+    `id`               bigint unsigned                  NOT NULL AUTO_INCREMENT,
+    `hash`             char(64) COLLATE utf8_unicode_ci NOT NULL,
+    `is_stored`        tinyint(1)                       NOT NULL DEFAULT '0',
+    `length`           bigint unsigned                           DEFAULT NULL,
+    `last_accessed_at` timestamp                        NULL     DEFAULT NULL,
+    PRIMARY KEY (`id`),
+    UNIQUE KEY `id` (`id`),
+    UNIQUE KEY `hash_idx` (`hash`),
+    KEY `last_accessed_idx` (`last_accessed_at`),
+    KEY `is_stored_idx` (`is_stored`)
+);

internal/iapi/api.go

@@ -2,7 +2,7 @@ package iapi
 
 import (
 	"encoding/json"
-	"io/ioutil"
+	"io"
 	"net/http"
 	"time"
 
@@ -43,7 +43,7 @@ func GetBlockedContent() (map[string]bool, error) {
 	if res.StatusCode != 200 {
 		return nil, errors.Err("unexpected status code %d", res.StatusCode)
 	}
-	body, err := ioutil.ReadAll(res.Body)
+	body, err := io.ReadAll(res.Body)
 	if err != nil {
 		return nil, errors.Err(err)
 	}

pkg/paid/handlers_test.go

@@ -1,7 +1,7 @@
 package paid
 
 import (
-	"io/ioutil"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -18,7 +18,7 @@ func TestHandlePublicKeyRequest(t *testing.T) {
 	HandlePublicKeyRequest(rr, r)
 
 	response := rr.Result()
-	key, err := ioutil.ReadAll(response.Body)
+	key, err := io.ReadAll(response.Body)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, response.StatusCode)
 	assert.NotZero(t, key)