add max balance middleware and tests

contribs/gnofaucet/coins.go

@@ -0,0 +1,54 @@
+package main
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"net/http"
+
+	tm2Client "github.com/gnolang/faucet/client/http"
+	"github.com/gnolang/gno/tm2/pkg/crypto"
+)
+
+func getAccountBalanceMiddleware(tm2Client *tm2Client.Client, maxBalance int64) func(next http.Handler) http.Handler {
+	type request struct {
+		To string `json:"to"`
+	}
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(
+			func(w http.ResponseWriter, r *http.Request) {
+				var data request
+				body, err := io.ReadAll(r.Body)
+				if err != nil {
+					http.Error(w, err.Error(), http.StatusBadRequest)
+					return
+				}
+
+				err = json.Unmarshal(body, &data)
+				r.Body = io.NopCloser(bytes.NewBuffer(body))
+				balance, err := checkAccountBalance(tm2Client, data.To)
+				if err != nil {
+					http.Error(w, err.Error(), http.StatusBadRequest)
+					return
+				}
+				if balance >= maxBalance {
+					http.Error(w, "accounts is already topped up", http.StatusBadRequest)
+					return
+				}
+				next.ServeHTTP(w, r)
+			},
+		)
+	}
+}
+
+var checkAccountBalance = func(tm2Client *tm2Client.Client, walletAddress string) (int64, error) {
+	address, err := crypto.AddressFromString(walletAddress)
+	if err != nil {
+		return 0, err
+	}
+	acc, err := tm2Client.GetAccount(address)
+	if err != nil {
+		return 0, err
+	}
+	return acc.GetCoins().AmountOf("ugnot"), nil
+}

contribs/gnofaucet/coins_test.go

@@ -0,0 +1,82 @@
+package main
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	tm2Client "github.com/gnolang/faucet/client/http"
+	"github.com/stretchr/testify/assert"
+)
+
+func mockedCheckAccountBalance(amount int64, err error) func(tm2Client *tm2Client.Client, walletAddress string) (int64, error) {
+	return func(tm2Client *tm2Client.Client, walletAddress string) (int64, error) {
+		return amount, err
+	}
+}
+
+func TestGetAccountBalanceMiddleware(t *testing.T) {
+	maxBalance := int64(1000)
+
+	tests := []struct {
+		name             string
+		requestBody      map[string]string
+		expectedStatus   int
+		expectedBody     string
+		checkBalanceFunc func(tm2Client *tm2Client.Client, walletAddress string) (int64, error)
+	}{
+		{
+			name:             "Valid address with low balance (should pass)",
+			requestBody:      map[string]string{"to": "valid_address_low_balance"},
+			expectedStatus:   http.StatusOK,
+			expectedBody:     "next handler reached",
+			checkBalanceFunc: mockedCheckAccountBalance(500, nil),
+		},
+		{
+			name:             "Valid address with high balance (should fail)",
+			requestBody:      map[string]string{"To": "valid_address_high_balance"},
+			expectedStatus:   http.StatusBadRequest,
+			expectedBody:     "accounts is already topped up",
+			checkBalanceFunc: mockedCheckAccountBalance(2*maxBalance, nil),
+		},
+		{
+			name:             "Invalid address (should fail)",
+			requestBody:      map[string]string{"To": "invalid_address"},
+			expectedStatus:   http.StatusBadRequest,
+			expectedBody:     "account not found",
+			checkBalanceFunc: mockedCheckAccountBalance(2*maxBalance, errors.New("account not found")),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			checkAccountBalance = tt.checkBalanceFunc
+			// Convert request body to JSON
+			reqBody, _ := json.Marshal(tt.requestBody)
+
+			// Create request
+			req := httptest.NewRequest(http.MethodPost, "/claim", bytes.NewReader(reqBody))
+			req.Header.Set("Content-Type", "application/json")
+
+			// Create ResponseRecorder
+			rr := httptest.NewRecorder()
+
+			// Mock next handler
+			nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusOK)
+				w.Write([]byte("next handler reached"))
+			})
+
+			// Apply middleware
+			handler := getAccountBalanceMiddleware(nil, maxBalance)(nextHandler)
+			handler.ServeHTTP(rr, req)
+
+			// Check response
+			assert.Equal(t, tt.expectedStatus, rr.Code)
+			assert.Contains(t, rr.Body.String(), tt.expectedBody)
+		})
+	}
+}

contribs/gnofaucet/cooldown_test.go

@@ -0,0 +1,28 @@
+package main
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCooldownLimiter(t *testing.T) {
+	cooldownDuration := time.Second
+	limiter := NewCooldownLimiter(cooldownDuration)
+	user := "testUser"
+
+	// First check should be allowed
+	if !limiter.CheckCooldown(user) {
+		t.Errorf("Expected first CheckCooldown to return true, but got false")
+	}
+
+	// Second check immediately should be denied
+	if limiter.CheckCooldown(user) {
+		t.Errorf("Expected second CheckCooldown to return false, but got true")
+	}
+
+	require.Eventually(t, func() bool {
+		return limiter.CheckCooldown(user)
+	}, 2*cooldownDuration, 10*time.Millisecond, "Expected CheckCooldown to return true after cooldown period")
+}

contribs/gnofaucet/gh.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -15,7 +16,7 @@ func getGithubMiddleware(clientID, secret string, cooldown time.Duration) func(n
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(
 			func(w http.ResponseWriter, r *http.Request) {
-				// Check if the captcha is enabled
+				// github Oauth flow is enabled
 				if secret == "" || clientID == "" {
 					// Continue with serving the faucet request
 					next.ServeHTTP(w, r)
@@ -29,18 +30,12 @@ func getGithubMiddleware(clientID, secret string, cooldown time.Duration) func(n
 					return
 				}
 
-				res, err := exchangeCodeForToken(secret, clientID, code)
+				user, err := exchangeCodeForUser(r.Context(), secret, clientID, code)
 				if err != nil {
 					http.Error(w, err.Error(), http.StatusBadRequest)
 					return
 				}
 
-				client := github.NewClient(http.DefaultClient).WithAuthToken(res.AccessToken)
-				user, _, err := client.Users.Get(r.Context(), "")
-				if err != nil {
-					http.Error(w, err.Error(), http.StatusBadRequest)
-					return
-				}
 				// Just check if given account have asked for faucet before the cooldown period
 				if !coolDownLimiter.CheckCooldown(user.GetLogin()) {
 					http.Error(w, "user is on cooldown", http.StatusTooManyRequests)
@@ -59,7 +54,7 @@ type GitHubTokenResponse struct {
 	AccessToken string `json:"access_token"`
 }
 
-func exchangeCodeForToken(secret, clientID, code string) (*GitHubTokenResponse, error) {
+var exchangeCodeForUser = func(ctx context.Context, secret, clientID, code string) (*github.User, error) {
 	url := "https://github.com/login/oauth/access_token"
 	body := fmt.Sprintf("client_id=%s&client_secret=%s&code=%s", clientID, secret, code)
 	req, err := http.NewRequest("POST", url, strings.NewReader(body))
@@ -84,5 +79,7 @@ func exchangeCodeForToken(secret, clientID, code string) (*GitHubTokenResponse,
 		return nil, fmt.Errorf("unable to exchange code for token")
 	}
 
-	return &tokenResponse, nil
+	ghClient := github.NewClient(http.DefaultClient).WithAuthToken(tokenResponse.AccessToken)
+	user, _, err := ghClient.Users.Get(ctx, "")
+	return user, err
 }

contribs/gnofaucet/gh_test.go

@@ -0,0 +1,126 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/google/go-github/v64/github"
+)
+
+// Mock function for exchangeCodeForToken
+func mockExchangeCodeForToken(ctx context.Context, secret, clientID, code string) (*github.User, error) {
+	login := "mock_login"
+	if code == "valid" {
+		fmt.Println("mockExchangeCodeForToken: valid")
+		return &github.User{Login: &login}, nil
+	}
+	return nil, errors.New("invalid code")
+}
+
+// Mock function for GitHub client
+func mockGetUser(token string) (*github.User, error) {
+	if token == "mock_token" {
+		return &github.User{Login: github.String("testUser")}, nil
+	}
+	return nil, errors.New("invalid token")
+}
+
+func TestGitHubMiddleware(t *testing.T) {
+	cooldown := 2 * time.Minute
+	exchangeCodeForUser = mockExchangeCodeForToken
+	t.Run("Midleware without credentials", func(t *testing.T) {
+		middleware := getGithubMiddleware("", "", cooldown)
+		// Test missing clientID and secret, middleware does nothing
+		req := httptest.NewRequest("GET", "http://localhost", nil)
+		rec := httptest.NewRecorder()
+
+		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		}))
+
+		handler.ServeHTTP(rec, req)
+
+		if rec.Code != http.StatusOK {
+			t.Errorf("Expected status OK, got %d", rec.Code)
+		}
+
+	})
+	t.Run("request without code", func(t *testing.T) {
+		middleware := getGithubMiddleware("mockClientID", "mockSecret", cooldown)
+		req := httptest.NewRequest("GET", "http://localhost?code=", nil)
+		rec := httptest.NewRecorder()
+
+		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		}))
+
+		handler.ServeHTTP(rec, req)
+
+		if rec.Code != http.StatusBadRequest {
+			t.Errorf("Expected status BadRequest, got %d", rec.Code)
+		}
+
+	})
+
+	t.Run("request invalid code", func(t *testing.T) {
+		middleware := getGithubMiddleware("mockClientID", "mockSecret", cooldown)
+		req := httptest.NewRequest("GET", "http://localhost?code=invalid", nil)
+		rec := httptest.NewRecorder()
+
+		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		}))
+
+		handler.ServeHTTP(rec, req)
+
+		if rec.Code != http.StatusBadRequest {
+			t.Errorf("Expected status BadRequest, got %d", rec.Code)
+		}
+	})
+
+	t.Run("OK", func(t *testing.T) {
+		middleware := getGithubMiddleware("mockClientID", "mockSecret", cooldown)
+		req := httptest.NewRequest("GET", "http://localhost?code=valid", nil)
+		rec := httptest.NewRecorder()
+
+		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		}))
+
+		handler.ServeHTTP(rec, req)
+
+		if rec.Code != http.StatusOK {
+			t.Errorf("Expected status OK, got %d", rec.Code)
+		}
+	})
+
+	t.Run("Cooldown active", func(t *testing.T) {
+		middleware := getGithubMiddleware("mockClientID", "mockSecret", cooldown)
+		req := httptest.NewRequest("GET", "http://localhost?code=valid", nil)
+		rec := httptest.NewRecorder()
+
+		handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		}))
+
+		handler.ServeHTTP(rec, req)
+
+		if rec.Code != http.StatusOK {
+			t.Errorf("Expected status OK, got %d", rec.Code)
+		}
+
+		req = httptest.NewRequest("GET", "http://localhost?code=valid", nil)
+		rec = httptest.NewRecorder()
+
+		handler.ServeHTTP(rec, req)
+		if rec.Code != http.StatusTooManyRequests {
+			t.Errorf("Expected status TooManyRequest, got %d", rec.Code)
+		}
+	})
+
+}

contribs/gnofaucet/serve.go

@@ -59,6 +59,7 @@ type serveCfg struct {
 
 	captchaSecret  string
 	ghClientID     string
+	maxBalance     int64
 	ghClientSecret string
 	isBehindProxy  bool
 }
@@ -143,6 +144,13 @@ func (c *serveCfg) RegisterFlags(fs *flag.FlagSet) {
 		"github client id for oauth authentication",
 	)
 
+	fs.Int64Var(
+		&c.maxBalance,
+		"max-balance",
+		10000000, // 10 ugnot
+		"limit of tokens the user can possess to be eligible to claim the faucet",
+	)
+
 	fs.BoolVar(
 		&c.isBehindProxy,
 		"is-behind-proxy",
@@ -212,6 +220,7 @@ func execServe(ctx context.Context, cfg *serveCfg, io commands.IO) error {
 		getIPMiddleware(cfg.isBehindProxy, st),
 		getCaptchaMiddleware(cfg.captchaSecret),
 		getGithubMiddleware(cfg.ghClientID, cfg.ghClientSecret, 1*time.Hour),
+		getAccountBalanceMiddleware(cli, cfg.maxBalance),
 	}
 
 	// Create a new faucet with