Merge pull request #42 from afr1ka/main

Deleted cache env vars to simplify APIFW configuration process

cmd/api-firewall/internal/handlers/check.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/sirupsen/logrus"
 	"github.com/valyala/fasthttp"
-
 	"github.com/wallarm/api-firewall/internal/platform/proxy"
 	"github.com/wallarm/api-firewall/internal/platform/web"
 )

cmd/api-firewall/internal/handlers/openapi.go

@@ -13,7 +13,6 @@ import (
 	"github.com/sirupsen/logrus"
 	"github.com/valyala/fasthttp"
 	"github.com/valyala/fastjson"
-
 	"github.com/wallarm/api-firewall/internal/config"
 	"github.com/wallarm/api-firewall/internal/platform/oauth2"
 	"github.com/wallarm/api-firewall/internal/platform/openapi3"

cmd/api-firewall/internal/handlers/routes.go

@@ -2,7 +2,6 @@ package handlers
 
 import (
 	"crypto/rsa"
-	"github.com/wallarm/api-firewall/internal/platform/denylist"
 	"io/ioutil"
 	"net/url"
 	"os"
@@ -14,9 +13,9 @@ import (
 	"github.com/sirupsen/logrus"
 	"github.com/valyala/fasthttp"
 	"github.com/valyala/fastjson"
-
 	"github.com/wallarm/api-firewall/internal/config"
 	"github.com/wallarm/api-firewall/internal/mid"
+	"github.com/wallarm/api-firewall/internal/platform/denylist"
 	woauth2 "github.com/wallarm/api-firewall/internal/platform/oauth2"
 	"github.com/wallarm/api-firewall/internal/platform/openapi3"
 	"github.com/wallarm/api-firewall/internal/platform/proxy"

cmd/api-firewall/tests/main_test.go

@@ -4,10 +4,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"github.com/golang/mock/gomock"
-	"github.com/sirupsen/logrus"
-	"github.com/valyala/fasthttp"
-	"github.com/wallarm/api-firewall/internal/platform/denylist"
 	"io"
 	"net"
 	"net/url"
@@ -17,8 +13,12 @@ import (
 	"testing"
 	"time"
 
+	"github.com/golang/mock/gomock"
+	"github.com/sirupsen/logrus"
+	"github.com/valyala/fasthttp"
 	"github.com/wallarm/api-firewall/cmd/api-firewall/internal/handlers"
 	"github.com/wallarm/api-firewall/internal/config"
+	"github.com/wallarm/api-firewall/internal/platform/denylist"
 	"github.com/wallarm/api-firewall/internal/platform/openapi3"
 	"github.com/wallarm/api-firewall/internal/platform/router"
 	"github.com/wallarm/api-firewall/internal/platform/tests"
@@ -324,12 +324,6 @@ func (s *ServiceTests) testBlockMode(t *testing.T) {
 
 func (s *ServiceTests) testDenylist(t *testing.T) {
 
-	cacheCfg := config.Cache{
-		NumCounters: 100000000,
-		MaxCost:     2147483648,
-		BufferItems: 64,
-	}
-
 	tokensCfg := config.Token{
 		CookieName: testDeniedCookieName,
 		HeaderName: "",
@@ -346,8 +340,7 @@ func (s *ServiceTests) testDenylist(t *testing.T) {
 		},
 		Denylist: struct {
 			Tokens config.Token
-			Cache  config.Cache
-		}{Tokens: tokensCfg, Cache: cacheCfg},
+		}{Tokens: tokensCfg},
 	}
 
 	logger := logrus.New()

demo/docker-compose/docker-compose.yml

@@ -20,10 +20,6 @@ services:
       APIFW_DENYLIST_TOKENS_COOKIE_NAME: "test"
       APIFW_DENYLIST_TOKENS_HEADER_NAME: ""
       APIFW_DENYLIST_TOKENS_TRIM_BEARER_PREFIX: "true"
-      # Denylist: Cache
-      APIFW_DENYLIST_CACHE_NUM_COUNTERS: "10000000"
-      APIFW_DENYLIST_CACHE_MAX_COST: "2147483648"
-      APIFW_DENYLIST_CACHE_BUFFER_ITEMS: "64"
     volumes:
       - ./volumes/api-firewall:/opt/resources:ro
     ports:

internal/config/config.go

@@ -30,12 +30,6 @@ type JWT struct {
 	SecretKey          string `conf:""`
 }
 
-type Cache struct {
-	NumCounters int64 `conf:"default:100000000"`
-	MaxCost     int64 `conf:"default:2147483648"`
-	BufferItems int64 `conf:"default:64"`
-}
-
 type Token struct {
 	CookieName       string `conf:""`
 	HeaderName       string `conf:""`
@@ -45,7 +39,6 @@ type Token struct {
 
 type Denylist struct {
 	Tokens Token
-	Cache  Cache
 }
 
 type Introspection struct {

internal/platform/denylist/denylist.go

@@ -4,82 +4,101 @@ import (
 	"bufio"
 	"io"
 	"os"
+	"strings"
 
 	"github.com/dgraph-io/ristretto"
 	"github.com/sirupsen/logrus"
 	"github.com/wallarm/api-firewall/internal/config"
 )
 
+const (
+	BufferItems = 64
+	ElementCost = 1
+)
+
 type DeniedTokens struct {
 	Cache       *ristretto.Cache
-	ElementsNum int
+	ElementsNum int64
 }
 
 func New(cfg *config.APIFWConfiguration, logger *logrus.Logger) (*DeniedTokens, error) {
 
-	cache, err := ristretto.NewCache(&ristretto.Config{
-		NumCounters: cfg.Denylist.Cache.NumCounters,
-		MaxCost:     cfg.Denylist.Cache.MaxCost,
-		BufferItems: cfg.Denylist.Cache.BufferItems,
-	})
+	if cfg.Denylist.Tokens.File == "" {
+		return nil, nil
+	}
+
+	var totalEntries int64
+	var totalCacheCapacity int64
+
+	// open tokens storage
+	f, err := os.Open(cfg.Denylist.Tokens.File)
 	if err != nil {
 		return nil, err
 	}
 
-	totalEntries := 0
+	// count non-empty entries and total cache capacity in bytes
+	c := bufio.NewScanner(f)
+	for c.Scan() {
+		if c.Text() != "" {
+			totalCacheCapacity += int64(len(c.Text()))
+			totalEntries += 1
+		}
+	}
+	err = c.Err()
+	if err != nil {
+		return nil, err
+	}
 
-	// Loading tokens to the cache
-	if cfg.Denylist.Tokens.File != "" {
+	// go to the beginning of the storage file
+	if _, err = f.Seek(0, io.SeekStart); err != nil {
+		return nil, err
+	}
 
-		f, err := os.Open(cfg.Denylist.Tokens.File)
-		if err != nil {
-			return nil, err
-		}
+	logger.Debugf("Denylist: total entries (lines) found in the file: %d", totalEntries)
 
-		// count non-empty entries
-		c := bufio.NewScanner(f)
-		for c.Scan() {
-			if c.Text() != "" {
-				totalEntries += 1
-			}
-		}
-		err = c.Err()
-		if err != nil {
-			return nil, err
-		}
+	// max cost = total bytes found in the storage + 5%
+	maxCost := totalCacheCapacity + totalCacheCapacity/20
 
-		if _, err = f.Seek(0, io.SeekStart); err != nil {
-			return nil, err
-		}
+	logger.Debugf("Denylist: cache capacity: %d bytes", maxCost)
+
+	cache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: maxCost * 10, // recommended value
+		MaxCost:     maxCost,
+		BufferItems: BufferItems,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	var numOfElements int64
+	totalEntries10P := totalEntries / 10
 
-		logger.Debugf("Denylist: total entries (lines) found in the file: %d", totalEntries)
-
-		totalEntries10P := totalEntries / 10
-		numOfElements := 0
-		current10P := 0
-		s := bufio.NewScanner(f)
-		for s.Scan() {
-			if s.Text() != "" {
-				if ok := cache.Set(s.Text(), nil, 1); ok {
-					numOfElements += 1
-					if numOfElements%totalEntries10P == 0 {
-						current10P += 10
-						logger.Debugf("Denylist: loaded %d perecents of tokens. Total elements in the cache: %d", current10P, numOfElements)
-					}
-				} else {
-					logger.Errorf("Denylist: can't add the token to the cache: %s", s.Text())
+	// 10% counter
+	counter10P := 0
+
+	// tokens loading to the cache
+	s := bufio.NewScanner(f)
+	for s.Scan() {
+		if s.Text() != "" {
+			if ok := cache.Set(strings.TrimSpace(s.Text()), nil, ElementCost); ok {
+				numOfElements += 1
+				if numOfElements%totalEntries10P == 0 {
+					counter10P += 10
+					logger.Debugf("Denylist: loaded %d perecents of tokens. Total elements in the cache: %d", counter10P, numOfElements)
 				}
-				cache.Wait()
+			} else {
+				logger.Errorf("Denylist: can't add the token to the cache: %s", s.Text())
 			}
+			cache.Wait()
 		}
-		err = s.Err()
-		if err != nil {
-			return nil, err
-		}
+	}
+	err = s.Err()
+	if err != nil {
+		return nil, err
+	}
 
-		if err := f.Close(); err != nil {
-			return nil, err
-		}
+	if err := f.Close(); err != nil {
+		return nil, err
 	}
 
 	return &DeniedTokens{Cache: cache, ElementsNum: totalEntries}, nil