implement total claimable limit

Author: Your Name

contribs/gnofaucet/cooldown.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"time"
 
@@ -12,52 +13,71 @@ import (
 // CooldownLimiter limits a specific user to one claim per cooldown period
 // this limiter keeps track of which keys are on cooldown using a badger database (written to a local file)
 type CooldownLimiter struct {
-	redis        *redis.Client
-	cooldownTime time.Duration
+	redis             *redis.Client
+	cooldownTime      time.Duration
+	maxlifeTimeAmount int64
 }
 
 // NewCooldownLimiter initializes a Cooldown Limiter with a given duration
-func NewCooldownLimiter(cooldown time.Duration, redis *redis.Client) *CooldownLimiter {
+func NewCooldownLimiter(cooldown time.Duration, redis *redis.Client, maxlifeTimeAmount int64) *CooldownLimiter {
 	return &CooldownLimiter{
-		redis:        redis,
-		cooldownTime: cooldown,
+		redis:             redis,
+		cooldownTime:      cooldown,
+		maxlifeTimeAmount: maxlifeTimeAmount,
 	}
 }
 
 // CheckCooldown checks if a key can make a claim or if it is still within the cooldown period
+// also checks that the user will not exceed the max lifetime allowed amount
 // Returns true if the key is not on cooldown, and marks the key as on cooldown
 // Returns false if the key is on cooldown or if an error occurs
-func (rl *CooldownLimiter) CheckCooldown(ctx context.Context, key string) (bool, error) {
-	isOnCooldown, err := rl.isOnCooldown(ctx, key)
+func (rl *CooldownLimiter) CheckCooldown(ctx context.Context, key string, amountClaimed int64) (bool, error) {
+	claimData, err := rl.getClaimsData(ctx, key)
 	if err != nil {
 		return false, fmt.Errorf("unable to check if key is on cooldown, %w", err)
 	}
-	if isOnCooldown {
-		return false, nil // Deny claim if within cooldown period
+	// Deny claim if within cooldown period
+	if claimData.LastClaimed.Add(rl.cooldownTime).After(time.Now()) {
+		return false, nil
+	}
+	// check that user will not exceed max lifetime allowed amount
+	if claimData.TotalClaimed+amountClaimed > rl.maxlifeTimeAmount {
+		return false, nil
 	}
 
-	return true, rl.markOnCooldown(ctx, key)
+	return true, rl.declareClaimedValue(ctx, key, amountClaimed, claimData)
 }
 
-func (rl *CooldownLimiter) isOnCooldown(ctx context.Context, key string) (bool, error) {
-	_, err := rl.redis.Get(ctx, key).Result()
+func (rl *CooldownLimiter) getClaimsData(ctx context.Context, key string) (*claimData, error) {
+	storedData, err := rl.redis.Get(ctx, key).Result()
 	if err != nil {
-		// Since we use redis's TTL feature to manage cooldown periods,
-		// an error redis.Nil simply indicates that the key is not on cooldown.
+		// Here we return an empty claimData because is the first time the user is making a claim
+		// the total amount claimed is 0 and the lastClaimed is the default time value
 		if errors.Is(err, redis.Nil) {
-			return false, nil
+			return &claimData{}, nil
 		}
 		// Any other unexpected error is returned.
-		return false, err
+		return nil, err
 	}
 
-	// Key found: it is on cooldown
-	return true, nil
+	claimData := &claimData{}
+	err = json.Unmarshal([]byte(storedData), claimData)
+	return claimData, err
 }
 
-func (rl *CooldownLimiter) markOnCooldown(ctx context.Context, key string) error {
-	// The value set here does not matter, as we only rely on
-	// redis's TTL feature to check if a key is still on cooldown
-	return rl.redis.Set(ctx, key, "claimed", rl.cooldownTime).Err()
+func (rl *CooldownLimiter) declareClaimedValue(ctx context.Context, key string, amountClaimed int64, currentData *claimData) error {
+	currentData.LastClaimed = time.Now()
+	currentData.TotalClaimed += amountClaimed
+
+	data, err := json.Marshal(currentData)
+	if err != nil {
+		return fmt.Errorf("unable to marshal claim data, %w", err)
+	}
+
+	return rl.redis.Set(ctx, key, data, 0).Err()
+}
 
+type claimData struct {
+	LastClaimed  time.Time
+	TotalClaimed int64
 }

contribs/gnofaucet/cooldown_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"math"
 	"testing"
 	"time"
 
@@ -11,33 +12,34 @@ import (
 )
 
 func TestCooldownLimiter(t *testing.T) {
+	var tenGnots int64 = 10_000_000
 	redisServer := miniredis.RunT(t)
 	rdb := redis.NewClient(&redis.Options{
 		Addr: redisServer.Addr(),
 	})
 
 	cooldownDuration := time.Second
-	limiter := NewCooldownLimiter(cooldownDuration, rdb)
+	limiter := NewCooldownLimiter(cooldownDuration, rdb, math.MaxInt64)
 	ctx := context.Background()
 	user := "testUser"
 
 	// First check should be allowed
-	allowed, err := limiter.CheckCooldown(ctx, user)
+	allowed, err := limiter.CheckCooldown(ctx, user, tenGnots)
 	require.NoError(t, err)
 
 	if !allowed {
 		t.Errorf("Expected first CheckCooldown to return true, but got false")
 	}
 
-	allowed, err = limiter.CheckCooldown(ctx, user)
+	allowed, err = limiter.CheckCooldown(ctx, user, tenGnots)
 	require.NoError(t, err)
 	// Second check immediately should be denied
 	if allowed {
 		t.Errorf("Expected second CheckCooldown to return false, but got true")
 	}
 
 	require.Eventually(t, func() bool {
-		allowed, err := limiter.CheckCooldown(ctx, user)
+		allowed, err := limiter.CheckCooldown(ctx, user, tenGnots)
 		return err == nil && !allowed
 	}, 2*cooldownDuration, 10*time.Millisecond, "Expected CheckCooldown to return true after cooldown period")
 }

contribs/gnofaucet/gh.go

@@ -1,9 +1,11 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"strings"
 
@@ -25,14 +27,6 @@ func getGithubMiddleware(clientID, secret string, coolDownLimiter *CooldownLimit
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(
 			func(w http.ResponseWriter, r *http.Request) {
-				// github Oauth flow is enabled
-				if secret == "" || clientID == "" {
-					// Continue with serving the faucet request
-					next.ServeHTTP(w, r)
-
-					return
-				}
-
 				// Extracts the authorization code returned by the GitHub OAuth flow.
 				//
 				// When a user successfully authenticates via GitHub OAuth, GitHub redirects them
@@ -52,8 +46,14 @@ func getGithubMiddleware(clientID, secret string, coolDownLimiter *CooldownLimit
 					return
 				}
 
+				claimAmount, err := getClaimAmount(r)
+				if err != nil {
+					http.Error(w, err.Error(), http.StatusBadRequest)
+					return
+				}
+
 				// Just check if given account have asked for faucet before the cooldown period
-				allowedToClaim, err := coolDownLimiter.CheckCooldown(r.Context(), user.GetLogin())
+				allowedToClaim, err := coolDownLimiter.CheckCooldown(r.Context(), user.GetLogin(), claimAmount)
 				if err != nil {
 					http.Error(w, err.Error(), http.StatusInternalServerError)
 					return
@@ -71,6 +71,25 @@ func getGithubMiddleware(clientID, secret string, coolDownLimiter *CooldownLimit
 	}
 }
 
+type request struct {
+	Amount int64 `json:"amount"`
+}
+
+func getClaimAmount(r *http.Request) (int64, error) {
+	body, err := io.ReadAll(r.Body)
+	if err != nil {
+		return 0, err
+	}
+
+	var data request
+	err = json.Unmarshal(body, &data)
+	if err != nil {
+		return 0, err
+	}
+	r.Body = io.NopCloser(bytes.NewBuffer(body))
+	return data.Amount, nil
+}
+
 type gitHubTokenResponse struct {
 	AccessToken string `json:"access_token"`
 }

contribs/gnofaucet/gh_test.go

@@ -1,8 +1,11 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"errors"
+	"fmt"
+	"math"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -25,10 +28,11 @@ func mockExchangeCodeForToken(ctx context.Context, secret, clientID, code string
 func TestGitHubMiddleware(t *testing.T) {
 	cooldown := 2 * time.Minute
 	exchangeCodeForUser = mockExchangeCodeForToken
-	t.Run("Midleware without credentials", func(t *testing.T) {
-		middleware := getGithubMiddleware("", "", getCooldownLimiter(t, cooldown))
-		// Test missing clientID and secret, middleware does nothing
-		req := httptest.NewRequest("GET", "http://localhost", nil)
+	var tenGnots int64 = 10000000
+	claimBody := fmt.Sprintf(`{"amount": %d}`, tenGnots)
+	t.Run("request without code", func(t *testing.T) {
+		middleware := getGithubMiddleware("mockClientID", "mockSecret", getCooldownLimiter(t, cooldown, math.MaxInt64))
+		req := httptest.NewRequest("GET", "http://localhost?code=", bytes.NewBufferString(claimBody))
 		rec := httptest.NewRecorder()
 
 		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -37,13 +41,14 @@ func TestGitHubMiddleware(t *testing.T) {
 
 		handler.ServeHTTP(rec, req)
 
-		if rec.Code != http.StatusOK {
-			t.Errorf("Expected status OK, got %d", rec.Code)
+		if rec.Code != http.StatusBadRequest {
+			t.Errorf("Expected status BadRequest, got %d", rec.Code)
 		}
 	})
-	t.Run("request without code", func(t *testing.T) {
-		middleware := getGithubMiddleware("mockClientID", "mockSecret", getCooldownLimiter(t, cooldown))
-		req := httptest.NewRequest("GET", "http://localhost?code=", nil)
+
+	t.Run("request invalid code", func(t *testing.T) {
+		middleware := getGithubMiddleware("mockClientID", "mockSecret", getCooldownLimiter(t, cooldown, math.MaxInt64))
+		req := httptest.NewRequest("GET", "http://localhost?code=invalid", bytes.NewBufferString(claimBody))
 		rec := httptest.NewRecorder()
 
 		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -57,9 +62,9 @@ func TestGitHubMiddleware(t *testing.T) {
 		}
 	})
 
-	t.Run("request invalid code", func(t *testing.T) {
-		middleware := getGithubMiddleware("mockClientID", "mockSecret", getCooldownLimiter(t, cooldown))
-		req := httptest.NewRequest("GET", "http://localhost?code=invalid", nil)
+	t.Run("OK", func(t *testing.T) {
+		middleware := getGithubMiddleware("mockClientID", "mockSecret", getCooldownLimiter(t, cooldown, math.MaxInt64))
+		req := httptest.NewRequest("GET", "http://localhost?code=valid", bytes.NewBufferString(claimBody))
 		rec := httptest.NewRecorder()
 
 		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -68,14 +73,14 @@ func TestGitHubMiddleware(t *testing.T) {
 
 		handler.ServeHTTP(rec, req)
 
-		if rec.Code != http.StatusBadRequest {
-			t.Errorf("Expected status BadRequest, got %d", rec.Code)
+		if rec.Code != http.StatusOK {
+			t.Errorf("Expected status OK, got %d", rec.Code)
 		}
 	})
 
-	t.Run("OK", func(t *testing.T) {
-		middleware := getGithubMiddleware("mockClientID", "mockSecret", getCooldownLimiter(t, cooldown))
-		req := httptest.NewRequest("GET", "http://localhost?code=valid", nil)
+	t.Run("Cooldown active", func(t *testing.T) {
+		middleware := getGithubMiddleware("mockClientID", "mockSecret", getCooldownLimiter(t, cooldown, math.MaxInt64))
+		req := httptest.NewRequest("GET", "http://localhost?code=valid", bytes.NewBufferString(claimBody))
 		rec := httptest.NewRecorder()
 
 		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -87,41 +92,66 @@ func TestGitHubMiddleware(t *testing.T) {
 		if rec.Code != http.StatusOK {
 			t.Errorf("Expected status OK, got %d", rec.Code)
 		}
+
+		req = httptest.NewRequest("GET", "http://localhost?code=valid", bytes.NewBufferString(claimBody))
+		rec = httptest.NewRecorder()
+
+		handler.ServeHTTP(rec, req)
+		if rec.Code != http.StatusTooManyRequests {
+			t.Errorf("Expected status TooManyRequest, got %d", rec.Code)
+		}
 	})
 
-	t.Run("Cooldown active", func(t *testing.T) {
-		middleware := getGithubMiddleware("mockClientID", "mockSecret", getCooldownLimiter(t, cooldown))
-		req := httptest.NewRequest("GET", "http://localhost?code=valid", nil)
+	t.Run("User exceeded lifetime limit", func(t *testing.T) {
+		cooldown = time.Millisecond
+		// Max lifetime amount is 20 Gnots so we should be able to make 2 claims
+		middleware := getGithubMiddleware("mockClientID", "mockSecret", getCooldownLimiter(t, cooldown, tenGnots*2))
+		req := httptest.NewRequest("GET", "http://localhost?code=valid", bytes.NewBufferString(claimBody))
 		rec := httptest.NewRecorder()
 
 		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(http.StatusOK)
 		}))
 
 		handler.ServeHTTP(rec, req)
+		// First claim ok
+		if rec.Code != http.StatusOK {
+			t.Errorf("Expected status OK, got %d", rec.Code)
+		}
+		// Wait 2 times the cooldown
+		time.Sleep(2 * cooldown)
+
+		req = httptest.NewRequest("GET", "http://localhost?code=valid", bytes.NewBufferString(claimBody))
+		rec = httptest.NewRecorder()
 
+		handler.ServeHTTP(rec, req)
+		//Second claim should also be ok
 		if rec.Code != http.StatusOK {
 			t.Errorf("Expected status OK, got %d", rec.Code)
 		}
 
-		req = httptest.NewRequest("GET", "http://localhost?code=valid", nil)
+		// Third one should fail
+		time.Sleep(2 * cooldown)
+
+		req = httptest.NewRequest("GET", "http://localhost?code=valid", bytes.NewBufferString(claimBody))
 		rec = httptest.NewRecorder()
 
 		handler.ServeHTTP(rec, req)
+		//Second claim should also be ok
 		if rec.Code != http.StatusTooManyRequests {
-			t.Errorf("Expected status TooManyRequest, got %d", rec.Code)
+			t.Errorf("Expected status OK, got %d", rec.Code)
 		}
 	})
 }
 
-func getCooldownLimiter(t *testing.T, duration time.Duration) *CooldownLimiter {
+func getCooldownLimiter(t *testing.T, duration time.Duration, maxlifeTimeAmount int64) *CooldownLimiter {
 	t.Helper()
 	redisServer := miniredis.RunT(t)
 	rdb := redis.NewClient(&redis.Options{
 		Addr: redisServer.Addr(),
 	})
 
-	limiter := NewCooldownLimiter(duration, rdb)
+	limiter := NewCooldownLimiter(duration, rdb, maxlifeTimeAmount)
 
 	return limiter
 }