Functions: tracking subscriptions

core/scripts/functions/templates/oracle.toml

@@ -43,6 +43,13 @@ maxSecretsSizesList = [10_240, 20_480, 51_200, 102_400, 307_200, 512_000, 1_048_
   updateFrequencySec = 30
   updateTimeoutSec = 10
 
+  [pluginConfig.OnchainSubscriptions]
+  blockConfirmations = 1
+  routerAddress = "{{router_contract_address}}"
+  queryFrequencySec = 30
+  queryTimeoutSec = 10
+  queryRangeSize = 100
+
   [pluginConfig.RateLimiter]
   globalBurst = 5
   globalRPS = 10

core/services/functions/connector_handler.go

@@ -21,31 +21,33 @@ import (
 type functionsConnectorHandler struct {
 	utils.StartStopOnce
 
-	connector   connector.GatewayConnector
-	signerKey   *ecdsa.PrivateKey
-	nodeAddress string
-	storage     s4.Storage
-	allowlist   functions.OnchainAllowlist
-	rateLimiter *hc.RateLimiter
-	lggr        logger.Logger
+	connector     connector.GatewayConnector
+	signerKey     *ecdsa.PrivateKey
+	nodeAddress   string
+	storage       s4.Storage
+	allowlist     functions.OnchainAllowlist
+	subscriptions functions.OnchainSubscriptions
+	rateLimiter   *hc.RateLimiter
+	lggr          logger.Logger
 }
 
 var (
 	_ connector.Signer                  = &functionsConnectorHandler{}
 	_ connector.GatewayConnectorHandler = &functionsConnectorHandler{}
 )
 
-func NewFunctionsConnectorHandler(nodeAddress string, signerKey *ecdsa.PrivateKey, storage s4.Storage, allowlist functions.OnchainAllowlist, rateLimiter *hc.RateLimiter, lggr logger.Logger) (*functionsConnectorHandler, error) {
-	if signerKey == nil || storage == nil || allowlist == nil || rateLimiter == nil {
-		return nil, fmt.Errorf("signerKey, storage, allowlist and rateLimiter must be non-nil")
+func NewFunctionsConnectorHandler(nodeAddress string, signerKey *ecdsa.PrivateKey, storage s4.Storage, allowlist functions.OnchainAllowlist, rateLimiter *hc.RateLimiter, subscriptions functions.OnchainSubscriptions, lggr logger.Logger) (*functionsConnectorHandler, error) {
+	if signerKey == nil || storage == nil || allowlist == nil || rateLimiter == nil || subscriptions == nil {
+		return nil, fmt.Errorf("signerKey, storage, allowlist, rateLimiter and subscriptions must be non-nil")
 	}
 	return &functionsConnectorHandler{
-		nodeAddress: nodeAddress,
-		signerKey:   signerKey,
-		storage:     storage,
-		allowlist:   allowlist,
-		rateLimiter: rateLimiter,
-		lggr:        lggr.Named("FunctionsConnectorHandler"),
+		nodeAddress:   nodeAddress,
+		signerKey:     signerKey,
+		storage:       storage,
+		allowlist:     allowlist,
+		rateLimiter:   rateLimiter,
+		subscriptions: subscriptions,
+		lggr:          lggr.Named("FunctionsConnectorHandler"),
 	}, nil
 }
 
@@ -68,6 +70,10 @@ func (h *functionsConnectorHandler) HandleGatewayMessage(ctx context.Context, ga
 		h.lggr.Errorw("request rate-limited", "id", gatewayId, "address", fromAddr)
 		return
 	}
+	if h.subscriptions.GetSubscription(fromAddr) == nil {
+		h.lggr.Errorw("request is not backed with a valid subscription", "id", gatewayId, "address", fromAddr)
+		return
+	}
 
 	h.lggr.Debugw("handling gateway request", "id", gatewayId, "method", body.Method)
 

core/services/functions/connector_handler_test.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/smartcontractkit/chainlink/v2/core/gethwrappers/functions/generated/functions_router"
 	"github.com/smartcontractkit/chainlink/v2/core/internal/testutils"
 	"github.com/smartcontractkit/chainlink/v2/core/logger"
 	"github.com/smartcontractkit/chainlink/v2/core/services/functions"
@@ -31,10 +32,11 @@ func TestFunctionsConnectorHandler(t *testing.T) {
 	connector := gcmocks.NewGatewayConnector(t)
 	allowlist := gfmocks.NewOnchainAllowlist(t)
 	rateLimiter, err := hc.NewRateLimiter(hc.RateLimiterConfig{GlobalRPS: 100.0, GlobalBurst: 100, PerSenderRPS: 100.0, PerSenderBurst: 100})
+	subscriptions := gfmocks.NewOnchainSubscriptions(t)
 	require.NoError(t, err)
 	allowlist.On("Start", mock.Anything).Return(nil)
 	allowlist.On("Close", mock.Anything).Return(nil)
-	handler, err := functions.NewFunctionsConnectorHandler(addr.Hex(), privateKey, storage, allowlist, rateLimiter, logger)
+	handler, err := functions.NewFunctionsConnectorHandler(addr.Hex(), privateKey, storage, allowlist, rateLimiter, subscriptions, logger)
 	require.NoError(t, err)
 
 	handler.SetConnector(connector)
@@ -73,6 +75,9 @@ func TestFunctionsConnectorHandler(t *testing.T) {
 			}
 			storage.On("List", ctx, addr).Return(snapshot, nil).Once()
 			allowlist.On("Allow", addr).Return(true).Once()
+			subscriptions.On("GetSubscription", mock.Anything).Return(
+				&functions_router.IFunctionsSubscriptionsSubscription{},
+			)
 			connector.On("SendToGateway", ctx, "gw1", mock.Anything).Run(func(args mock.Arguments) {
 				msg, ok := args[2].(*api.Message)
 				require.True(t, ok)
@@ -129,6 +134,9 @@ func TestFunctionsConnectorHandler(t *testing.T) {
 
 			storage.On("Put", ctx, &key, &record, signature).Return(nil).Once()
 			allowlist.On("Allow", addr).Return(true).Once()
+			subscriptions.On("GetSubscription", mock.Anything).Return(
+				&functions_router.IFunctionsSubscriptionsSubscription{},
+			)
 			connector.On("SendToGateway", ctx, "gw1", mock.Anything).Run(func(args mock.Arguments) {
 				msg, ok := args[2].(*api.Message)
 				require.True(t, ok)

core/services/gateway/handlers/functions/subscriptions.go

@@ -0,0 +1,200 @@
+package functions
+
+import (
+	"context"
+	"fmt"
+	"math/big"
+	"sync"
+	"time"
+
+	"github.com/ethereum/go-ethereum/accounts/abi/bind"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/pkg/errors"
+
+	evmclient "github.com/smartcontractkit/chainlink/v2/core/chains/evm/client"
+	"github.com/smartcontractkit/chainlink/v2/core/gethwrappers/functions/generated/functions_router"
+	"github.com/smartcontractkit/chainlink/v2/core/logger"
+	"github.com/smartcontractkit/chainlink/v2/core/services/job"
+	"github.com/smartcontractkit/chainlink/v2/core/utils"
+)
+
+type OnchainSubscriptionsConfig struct {
+	RouterAddress      common.Address `json:"routerAddress"`
+	BlockConfirmations uint           `json:"blockConfirmations"`
+	QueryFrequencySec  uint           `json:"queryFrequencySec"`
+	QueryTimeoutSec    uint           `json:"queryTimeoutSec"`
+	QueryRangeSize     uint           `json:"queryRangeSize"`
+}
+
+// OnchainSubscriptions maintains a mirror of all subscriptions fetched from the blockchain (EVM-only).
+// All methods are thread-safe.
+//
+//go:generate mockery --quiet --name OnchainSubscriptions --output ./mocks/ --case=underscore
+type OnchainSubscriptions interface {
+	job.ServiceCtx
+
+	// GetSubscription returns a subscription for the given user address, or null if not found
+	GetSubscription(common.Address) *functions_router.IFunctionsSubscriptionsSubscription
+}
+
+type onchainSubscriptions struct {
+	utils.StartStopOnce
+
+	config             OnchainSubscriptionsConfig
+	subscriptions      map[common.Address]*functions_router.IFunctionsSubscriptionsSubscription
+	client             evmclient.Client
+	router             *functions_router.FunctionsRouter
+	blockConfirmations *big.Int
+	lggr               logger.Logger
+	closeWait          sync.WaitGroup
+	rwMutex            sync.RWMutex
+	stopCh             utils.StopChan
+}
+
+func NewOnchainSubscriptions(client evmclient.Client, config OnchainSubscriptionsConfig, lggr logger.Logger) (OnchainSubscriptions, error) {
+	if client == nil {
+		return nil, errors.New("client is nil")
+	}
+	if lggr == nil {
+		return nil, errors.New("logger is nil")
+	}
+	router, err := functions_router.NewFunctionsRouter(config.RouterAddress, client)
+	if err != nil {
+		return nil, fmt.Errorf("unexpected error during functions_router.NewFunctionsRouter: %s", err)
+	}
+	return &onchainSubscriptions{
+		config:             config,
+		subscriptions:      make(map[common.Address]*functions_router.IFunctionsSubscriptionsSubscription),
+		client:             client,
+		router:             router,
+		blockConfirmations: big.NewInt(int64(config.BlockConfirmations)),
+		lggr:               lggr.Named("OnchainSubscriptions"),
+		stopCh:             make(utils.StopChan),
+	}, nil
+}
+
+func (s *onchainSubscriptions) Start(ctx context.Context) error {
+	return s.StartOnce("OnchainSubscriptions", func() error {
+		s.lggr.Info("starting onchain subscriptions")
+		if s.config.QueryFrequencySec == 0 {
+			return errors.New("OnchainSubscriptionsConfig.UpdateFrequencySec must be greater than 0")
+		}
+		if s.config.QueryTimeoutSec == 0 {
+			return errors.New("OnchainSubscriptionsConfig.UpdateTimeoutSec must be greater than 0")
+		}
+		if s.config.QueryRangeSize == 0 {
+			return errors.New("OnchainSubscriptionsConfig.QueryRangeSize must be greater than 0")
+		}
+
+		s.closeWait.Add(1)
+		go s.queryLoop()
+
+		return nil
+	})
+}
+
+func (s *onchainSubscriptions) Close() error {
+	return s.StopOnce("OnchainSubscriptions", func() (err error) {
+		s.lggr.Info("closing onchain subscriptions")
+		close(s.stopCh)
+		s.closeWait.Wait()
+		return nil
+	})
+}
+
+func (s *onchainSubscriptions) GetSubscription(address common.Address) *functions_router.IFunctionsSubscriptionsSubscription {
+	s.rwMutex.RLock()
+	defer s.rwMutex.RUnlock()
+	subscription, ok := s.subscriptions[address]
+	if !ok {
+		return nil
+	}
+	return subscription
+}
+
+func (s *onchainSubscriptions) queryLoop() {
+	defer s.closeWait.Done()
+
+	ticker := time.NewTicker(time.Duration(s.config.QueryFrequencySec) * time.Second)
+	defer ticker.Stop()
+
+	var start uint64 = 1
+
+	queryFunc := func() {
+		ctx, cancel := utils.ContextFromChanWithTimeout(s.stopCh, time.Duration(s.config.QueryTimeoutSec)*time.Second)
+		defer cancel()
+
+		latestBlockHeight, err := s.client.LatestBlockHeight(ctx)
+		if err != nil || latestBlockHeight == nil {
+			s.lggr.Errorw("Error calling LatestBlockHeight", "err", err, "latestBlockHeight", latestBlockHeight.Int64())
+			return
+		}
+
+		blockNumber := big.NewInt(0).Sub(latestBlockHeight, s.blockConfirmations)
+
+		count, err := s.getSubscriptionsCount(ctx, blockNumber)
+		if err != nil {
+			s.lggr.Errorw("Error getting subscriptions count", "err", err)
+			return
+		}
+		if count == 0 {
+			s.lggr.Info("Router has no subscriptions yet")
+			return
+		}
+
+		if start > count {
+			start = 1
+		}
+
+		end := start + uint64(s.config.QueryRangeSize)
+		if end > count {
+			end = count
+		}
+		if err := s.querySubscriptionsRange(ctx, blockNumber, start, end); err != nil {
+			s.lggr.Errorw("Error querying subscriptions", "err", err, "start", start, "end", end)
+			return
+		}
+
+		start = end + 1
+	}
+
+	for {
+		select {
+		case <-s.stopCh:
+			return
+		case <-ticker.C:
+			queryFunc()
+		}
+	}
+}
+
+func (s *onchainSubscriptions) querySubscriptionsRange(ctx context.Context, blockNumber *big.Int, start, end uint64) error {
+	subscriptions, err := s.router.GetSubscriptionsInRange(&bind.CallOpts{
+		Pending:     false,
+		BlockNumber: blockNumber,
+		Context:     ctx,
+	}, start, end)
+	if err != nil {
+		return errors.Wrap(err, "unexpected error during functions_router.GetSubscriptionsInRange")
+	}
+
+	s.rwMutex.Lock()
+	defer s.rwMutex.Unlock()
+	for _, subscription := range subscriptions {
+		if subscription.Owner == utils.ZeroAddress {
+			continue
+		}
+		subscription := subscription
+		s.subscriptions[subscription.Owner] = &subscription
+	}
+
+	return nil
+}
+
+func (s *onchainSubscriptions) getSubscriptionsCount(ctx context.Context, blockNumber *big.Int) (uint64, error) {
+	return s.router.GetSubscriptionCount(&bind.CallOpts{
+		Pending:     false,
+		BlockNumber: blockNumber,
+		Context:     ctx,
+	})
+}