feat: support private policy (#5)

* feat: support private policy

pkg/paymasterclient/client.go

@@ -16,7 +16,7 @@ type Client interface {
 	// IsSponsorable checks if a transaction is sponsorable
 	IsSponsorable(ctx context.Context, tx TransactionArgs) (*IsSponsorableResponse, error)
 	// SendRawTransaction sends a raw transaction to the connected domain
-	SendRawTransaction(ctx context.Context, input hexutil.Bytes) (common.Hash, error)
+	SendRawTransaction(ctx context.Context, input hexutil.Bytes, opts *TransactionOptions) (common.Hash, error)
 	// GetGaslessTransactionByHash returns a gasless transaction by hash
 	GetGaslessTransactionByHash(ctx context.Context, txHash common.Hash) (userTx *TransactionResponse, err error)
 
@@ -31,16 +31,31 @@ type Client interface {
 }
 
 type client struct {
-	c *rpc.Client
+	c                 *rpc.Client
+	PrivatePolicyUUID *string
 }
 
+// New creates a new Client with the given URL and options.
+// The URL is typically in the format of https://bsc-megafuel.nodereal.io/
 func New(ctx context.Context, url string, options ...rpc.ClientOption) (Client, error) {
 	c, err := rpc.DialOptions(ctx, url, options...)
 	if err != nil {
 		return nil, err
 	}
 
-	return &client{c}, nil
+	return &client{c, nil}, nil
+}
+
+// NewPrivatePaymaster creates a new Client with private policy functionality.
+// The URL for this function should be in the format:
+// https://open-platform-ap.nodereal.io/{$apikey}/megafuel
+func NewPrivatePaymaster(ctx context.Context, url, privatePolicyUUID string, options ...rpc.ClientOption) (Client, error) {
+	c, err := rpc.DialOptions(ctx, url, options...)
+	if err != nil {
+		return nil, err
+	}
+
+	return &client{c, &privatePolicyUUID}, nil
 }
 
 func (c *client) ChainID(ctx context.Context) (*big.Int, error) {
@@ -54,19 +69,37 @@ func (c *client) ChainID(ctx context.Context) (*big.Int, error) {
 
 func (c *client) IsSponsorable(ctx context.Context, tx TransactionArgs) (*IsSponsorableResponse, error) {
 	var result IsSponsorableResponse
+
+	if c.PrivatePolicyUUID != nil {
+		c.c.SetHeader("X-MegaFuel-Policy-Uuid", *c.PrivatePolicyUUID)
+	}
+
 	err := c.c.CallContext(ctx, &result, "pm_isSponsorable", tx)
 	if err != nil {
 		return nil, err
 	}
+
 	return &result, nil
 }
 
-func (c *client) SendRawTransaction(ctx context.Context, input hexutil.Bytes) (common.Hash, error) {
+func (c *client) SendRawTransaction(ctx context.Context, input hexutil.Bytes, opts *TransactionOptions) (common.Hash, error) {
 	var result common.Hash
+
+	if opts != nil {
+		if opts.UserAgent != "" {
+			c.c.SetHeader("User-Agent", opts.UserAgent)
+		}
+	}
+
+	if c.PrivatePolicyUUID != nil {
+		c.c.SetHeader("X-MegaFuel-Policy-Uuid", *c.PrivatePolicyUUID)
+	}
+
 	err := c.c.CallContext(ctx, &result, "eth_sendRawTransaction", input)
 	if err != nil {
 		return common.Hash{}, err
 	}
+
 	return result, nil
 }
 

pkg/paymasterclient/types.go

@@ -16,6 +16,12 @@ type TransactionArgs struct {
 	Data  *hexutil.Bytes  `json:"data"`
 }
 
+// TransactionOptions defines the options for the SendRawTransaction method.
+type TransactionOptions struct {
+	// UserAgent is an optional field to set a custom User-Agent header for the request.
+	UserAgent string
+}
+
 type IsSponsorableResponse struct {
 	Sponsorable    bool   `json:"sponsorable"`              // Sponsorable is a mandatory field, bool value, indicating if a given tx is able to sponsor.
 	SponsorName    string `json:"sponsorName,omitempty"`    // SponsorName is an optional field, string value, shows the name of the policy sponsor.

pkg/sponsorclient/client.go

@@ -33,7 +33,6 @@ func New(ctx context.Context, url string, options ...rpc.ClientOption) (Client,
 	if err != nil {
 		return nil, err
 	}
-
 	return &client{c}, nil
 }
 

test/paymaster_test.go

@@ -22,34 +22,64 @@ import (
 	"github.com/node-real/megafuel-go-sdk/pkg/paymasterclient"
 )
 
-const PAYMASTER_URL = "https://bsc-megafuel-testnet.nodereal.io/97"
+const (
+	PAYMASTER_URL  = "https://bsc-megafuel-testnet.nodereal.io/"
+	PRIVATE_POLICY = "90f1ba4c-1f93-4759-b8a9-da4d59c668b4"
+)
 
 var log = logrus.New()
 
-// paymasterSetup initializes a paymaster client using the environment variable.
-func paymasterSetup(t *testing.T) (*ethclient.Client, paymasterclient.Client, string, error) {
+// setupCommon contains the common setup logic for both paymaster types
+func setupCommon(t *testing.T) (*ethclient.Client, string, string, error) {
 	t.Helper()
 
 	key := os.Getenv("OPEN_PLATFORM_PRIVATE_KEY")
 	if key == "" {
-		log.Fatal("Environment variable OPEN_PLATFORM_PRIVATE_KEY is not set")
+		return nil, "", "", fmt.Errorf("environment variable OPEN_PLATFORM_PRIVATE_KEY is not set")
 	}
 
 	yourPrivateKey := os.Getenv("YOUR_PRIVATE_KEY")
 	if yourPrivateKey == "" {
-		log.Fatal("Environment variable YOUR_PRIVATE_KEY is not set")
+		return nil, "", "", fmt.Errorf("environment variable YOUR_PRIVATE_KEY is not set")
 	}
 
 	// Connect to an Ethereum node (for transaction assembly)
 	client, err := ethclient.Dial(fmt.Sprintf("https://bsc-testnet.nodereal.io/v1/%s", key))
 	if err != nil {
-		log.Fatalf("Failed to connect to the Ethereum network: %v", err)
+		return nil, "", "", fmt.Errorf("failed to connect to the Ethereum network: %v", err)
+	}
+
+	return client, key, yourPrivateKey, nil
+}
+
+// paymasterSetup initializes a standard paymaster client using the environment variable.
+func paymasterSetup(t *testing.T) (*ethclient.Client, paymasterclient.Client, string, error) {
+	client, _, yourPrivateKey, err := setupCommon(t)
+	if err != nil {
+		return nil, nil, "", err
 	}
 
 	// Create a PaymasterClient (for transaction sending)
 	paymasterClient, err := paymasterclient.New(context.Background(), PAYMASTER_URL)
 	if err != nil {
-		log.Fatalf("Failed to create PaymasterClient: %v", err)
+		return nil, nil, "", fmt.Errorf("failed to create PaymasterClient: %v", err)
+	}
+
+	return client, paymasterClient, yourPrivateKey, nil
+}
+
+// privatePaymasterSetup initializes a private paymaster client using the environment variable.
+func privatePaymasterSetup(t *testing.T) (*ethclient.Client, paymasterclient.Client, string, error) {
+	client, key, yourPrivateKey, err := setupCommon(t)
+	if err != nil {
+		return nil, nil, "", err
+	}
+
+	sponsorURL := fmt.Sprintf("https://open-platform-ap.nodereal.io/%s/megafuel-testnet/97", key)
+	// Create a Private PaymasterClient (for transaction sending)
+	paymasterClient, err := paymasterclient.NewPrivatePaymaster(context.Background(), sponsorURL, PRIVATE_POLICY)
+	if err != nil {
+		return nil, nil, "", fmt.Errorf("failed to create Private PaymasterClient: %v", err)
 	}
 
 	return client, paymasterClient, yourPrivateKey, nil
@@ -117,7 +147,7 @@ func TestPaymasterAPI(t *testing.T) {
 	require.NoError(t, err, "Failed to marshal transaction")
 
 	// Send the signed transaction and check for successful submission.
-	paymasterTx, err := paymasterClient.SendRawTransaction(context.Background(), txInput)
+	paymasterTx, err := paymasterClient.SendRawTransaction(context.Background(), txInput, nil)
 	require.NoError(t, err, "Failed to send sponsorable transaction")
 	log.Infof("Sponsorable transaction sent: %s", signedTx.Hash())
 	log.Info("Waiting for transaction confirmation")
@@ -126,7 +156,7 @@ func TestPaymasterAPI(t *testing.T) {
 	// Check the Paymaster client's chain ID for consistency.
 	payMasterChainID, err := paymasterClient.ChainID(context.Background())
 	require.NoError(t, err, "failed to get paymaster chain id")
-	assert.Equal(t, payMasterChainID, "0x61")
+	assert.Equal(t, payMasterChainID.String(), "97")
 
 	// Retrieve and verify the transaction details by its hash.
 	txResp, err := paymasterClient.GetGaslessTransactionByHash(context.Background(), paymasterTx)
@@ -153,5 +183,104 @@ func TestPaymasterAPI(t *testing.T) {
 	blockNumber := rpc.PendingBlockNumber
 	count, err := paymasterClient.GetTransactionCount(context.Background(), common.HexToAddress(RECIPIENT_ADDRESS), rpc.BlockNumberOrHash{BlockNumber: &blockNumber})
 	require.NoError(t, err, "failed to GetTransactionCount")
-	assert.Greater(t, *count, hexutil.Uint64(0))
+	assert.Greater(t, count, hexutil.Uint64(0))
+}
+
+// TestPaymasterAPI tests the critical functionalities related to the Paymaster API.
+func TestPrivatePolicyGaslessTransaction(t *testing.T) {
+	// Setup Ethereum client and Paymaster client. Ensure no errors during the setup.
+	client, sponsorClient, yourPrivateKey, err := privatePaymasterSetup(t)
+	require.NoError(t, err, "failed to set up paymaster")
+
+	// Convert the private key from hex string to ECDSA format and check for errors.
+	privateKey, err := crypto.HexToECDSA(yourPrivateKey)
+	require.NoError(t, err, "Failed to load private key")
+
+	// Extract the public key from the private key and assert type casting to ECDSA.
+	publicKey := privateKey.Public()
+	publicKeyECDSA, ok := publicKey.(*ecdsa.PublicKey)
+	if !ok {
+		log.Fatal("Error casting public key to ECDSA")
+	}
+	fromAddress := crypto.PubkeyToAddress(*publicKeyECDSA)
+
+	// Fetch the current nonce for the account to ensure the transaction can be processed sequentially.
+	nonce, err := client.PendingNonceAt(context.Background(), fromAddress)
+	require.NoError(t, err, "Failed to get nonce")
+
+	// Define the recipient Ethereum address.
+	toAddress := common.HexToAddress(RECIPIENT_ADDRESS)
+
+	// Construct a new Ethereum transaction.
+	tx := types.NewTx(&types.LegacyTx{
+		Nonce:    nonce,
+		GasPrice: big.NewInt(0),
+		Gas:      21000,
+		To:       &toAddress,
+		Value:    big.NewInt(0),
+	})
+
+	// Prepare a transaction argument for checking if it's sponsorable.
+	gasLimit := tx.Gas()
+
+	privatePolicySponsorableTx := paymasterclient.TransactionArgs{
+		To:    &toAddress,
+		From:  fromAddress,
+		Value: (*hexutil.Big)(big.NewInt(0)),
+		Gas:   (*hexutil.Uint64)(&gasLimit),
+		Data:  &hexutil.Bytes{},
+	}
+
+	privatePolicySponsorableInfo, err := sponsorClient.IsSponsorable(context.Background(), privatePolicySponsorableTx)
+	require.NoError(t, err, "Error checking sponsorable private policy status")
+	require.True(t, privatePolicySponsorableInfo.Sponsorable)
+
+	// Retrieve the blockchain ID to ensure that the transaction is signed correctly.
+	chainID, err := client.ChainID(context.Background())
+	require.NoError(t, err, "Failed to get chain ID")
+
+	// Sign the transaction using the provided private key and the current chain ID.
+	signedTx, err := types.SignTx(tx, types.NewEIP155Signer(chainID), privateKey)
+	require.NoError(t, err, "Failed to sign transaction")
+
+	// Marshal the signed transaction into a binary format for transmission.
+	txInput, err := signedTx.MarshalBinary()
+	require.NoError(t, err, "Failed to marshal transaction")
+
+	transaction, err := sponsorClient.SendRawTransaction(context.Background(), txInput, &paymasterclient.TransactionOptions{UserAgent: "Test User Agent"})
+	require.NoError(t, err, "Failed to send sponsorable private policy transaction")
+	log.Infof("Sponsorable private policy transaction sent: %s", signedTx.Hash())
+	time.Sleep(10 * time.Second) // Consider replacing with a non-blocking wait or event-driven notification.
+
+	// Check the Paymaster client's chain ID for consistency.
+	payMasterChainID, err := sponsorClient.ChainID(context.Background())
+	require.NoError(t, err, "failed to get paymaster chain id")
+	assert.Equal(t, payMasterChainID.String(), "97")
+
+	// Retrieve and verify the transaction details by its hash.
+	txResp, err := sponsorClient.GetGaslessTransactionByHash(context.Background(), transaction)
+	require.NoError(t, err, "failed to GetGaslessTransactionByHash")
+	assert.Equal(t, txResp.TxHash.String(), transaction.String())
+
+	// Check for the related transaction bundle based on the UUID.
+	bundleUuid := txResp.BundleUUID
+	sponsorTx, err := sponsorClient.GetSponsorTxByBundleUUID(context.Background(), bundleUuid)
+	require.NoError(t, err)
+
+	// Retrieve the full bundle using the UUID and verify its existence.
+	bundle, err := sponsorClient.GetBundleByUUID(context.Background(), bundleUuid)
+	require.NoError(t, err)
+
+	// Further validate the bundle by fetching the transaction via its hash.
+	sponsorTx, err = sponsorClient.GetSponsorTxByTxHash(context.Background(), sponsorTx.TxHash)
+	require.NoError(t, err)
+
+	// Log the UUID of the bundle for reference.
+	log.Infof("Bundle UUID: %s", bundle.BundleUUID)
+
+	// Obtain and verify the transaction count for the recipient address.
+	blockNumber := rpc.PendingBlockNumber
+	count, err := sponsorClient.GetTransactionCount(context.Background(), common.HexToAddress(RECIPIENT_ADDRESS), rpc.BlockNumberOrHash{BlockNumber: &blockNumber})
+	require.NoError(t, err, "failed to GetTransactionCount")
+	assert.Greater(t, count, hexutil.Uint64(0))
 }

test/sponsor_test.go

@@ -2,6 +2,7 @@ package test
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"testing"
 
@@ -25,7 +26,7 @@ func sponsorSetup(t *testing.T) (sponsorclient.Client, error) {
 	if key == "" {
 		log.Fatal("Environment variable OPEN_PLATFORM_PRIVATE_KEY is not set")
 	}
-	return sponsorclient.New(context.Background(), "https://open-platform-ap.nodereal.io/"+key+"/megafuel-testnet")
+	return sponsorclient.New(context.Background(), fmt.Sprintf("https://open-platform-ap.nodereal.io/%s/megafuel-testnet", key))
 }
 
 // TestSponsorAPI conducts several whitelist operations.