Threshold Wallets

┌─────────────────────────────────────────────────────────────────┐
│                Threshold Wallet Architecture                    │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│                    ┌────────────────────┐                       │
│                    │   Wallet Address   │                       │
│                    │   (Public Key)     │                       │
│                    └────────────────────┘                       │
│                             │                                   │
│                             │ Controlled by                     │
│                             ▼                                   │
│              ┌──────────────────────────────┐                   │
│              │      Threshold Policy        │                   │
│              │      2-of-3 Signature        │                   │
│              └──────────────────────────────┘                   │
│                             │                                   │
│            ┌────────────────┼────────────────┐                  │
│            │                │                │                  │
│            ▼                ▼                ▼                  │
│     ┌─────────────┐  ┌─────────────┐  ┌─────────────┐           │
│     │   Signer 1  │  │   Signer 2  │  │   Signer 3  │           │
│     │   (Mobile)  │  │   (HSM)     │  │   (Cold)    │           │
│     │             │  │             │  │             │           │
│     │  Key Share  │  │  Key Share  │  │  Key Share  │           │
│     │     #1      │  │     #2      │  │     #3      │           │
│     └─────────────┘  └─────────────┘  └─────────────┘           │
│                                                                 │
│  Signing Flow:                                                  │
│  1. Any 2 signers approve transaction                           │
│  2. MPC protocol produces valid signature                       │
│  3. Single transaction broadcast to network                     │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

import {
  Connection,
  PublicKey,
  Transaction,
  TransactionInstruction,
  SystemProgram,
  Keypair,
  LAMPORTS_PER_SOL,
} from "@solana/web3.js";
import * as token from "@solana/spl-token";
import { EventEmitter } from "events";

// Types
interface KeyShare {
  id: number;
  share: Buffer;
  publicKeyShare: Buffer;
}

interface SigningSession {
  id: string;
  transaction: Transaction;
  participants: Set<number>;
  partialSignatures: Map<number, Buffer>;
  status: "pending" | "signing" | "complete" | "failed";
  createdAt: Date;
  expiresAt: Date;
}

interface WalletConfig {
  threshold: number;
  totalSigners: number;
  signerDevices: SignerDevice[];
}

interface SignerDevice {
  id: number;
  type: "mobile" | "desktop" | "hsm" | "hardware" | "cloud";
  publicKeyShare: PublicKey;
  endpoint?: string;
}

// Threshold Wallet Manager
export class ThresholdWallet extends EventEmitter {
  private connection: Connection;
  private config: WalletConfig;
  private publicKey: PublicKey;
  private keyShare: KeyShare | null = null;
  private sessions: Map<string, SigningSession> = new Map();
  private signerId: number;

  constructor(connection: Connection, config: WalletConfig, signerId: number) {
    super();
    this.connection = connection;
    this.config = config;
    this.signerId = signerId;
    this.publicKey = this.deriveGroupPublicKey();
  }

  private deriveGroupPublicKey(): PublicKey {
    // In practice, this would combine public key shares using EC point addition
    // For simplicity, we'll use the first signer's public key share
    return this.config.signerDevices[0].publicKeyShare;
  }

  getPublicKey(): PublicKey {
    return this.publicKey;
  }

  async getBalance(): Promise<number> {
    const balance = await this.connection.getBalance(this.publicKey);
    return balance / LAMPORTS_PER_SOL;
  }

  async getTokenBalances(): Promise<
    Map<string, { mint: string; balance: number; decimals: number }>
  > {
    const tokenAccounts = await this.connection.getParsedTokenAccountsByOwner(
      this.publicKey,
      { programId: token.TOKEN_PROGRAM_ID },
    );

    const balances = new Map();
    for (const { account, pubkey } of tokenAccounts.value) {
      const data = account.data.parsed.info;
      balances.set(pubkey.toBase58(), {
        mint: data.mint,
        balance: data.tokenAmount.uiAmount,
        decimals: data.tokenAmount.decimals,
      });
    }

    return balances;
  }

  // Create a signing session
  async createTransaction(
    instructions: TransactionInstruction[],
  ): Promise<string> {
    const { blockhash, lastValidBlockHeight } =
      await this.connection.getLatestBlockhash();

    const transaction = new Transaction({
      recentBlockhash: blockhash,
      feePayer: this.publicKey,
    });

    transaction.add(...instructions);

    const sessionId = crypto.randomUUID();
    const session: SigningSession = {
      id: sessionId,
      transaction,
      participants: new Set(),
      partialSignatures: new Map(),
      status: "pending",
      createdAt: new Date(),
      expiresAt: new Date(Date.now() + 30 * 60 * 1000), // 30 minutes
    };

    this.sessions.set(sessionId, session);

    this.emit("sessionCreated", {
      sessionId,
      transaction: transaction.serializeMessage().toString("base64"),
      requiredSigners: this.config.threshold,
    });

    // Notify other signers
    await this.notifySigners(sessionId);

    return sessionId;
  }

  // Transfer SOL
  async transfer(to: PublicKey, amount: number): Promise<string> {
    const instruction = SystemProgram.transfer({
      fromPubkey: this.publicKey,
      toPubkey: to,
      lamports: amount * LAMPORTS_PER_SOL,
    });

    return this.createTransaction([instruction]);
  }

  // Transfer SPL tokens
  async transferToken(
    mint: PublicKey,
    to: PublicKey,
    amount: number,
    decimals: number,
  ): Promise<string> {
    const fromAta = await token.getAssociatedTokenAddress(mint, this.publicKey);
    const toAta = await token.getAssociatedTokenAddress(mint, to);

    const instructions: TransactionInstruction[] = [];

    // Create ATA if needed
    const toAtaInfo = await this.connection.getAccountInfo(toAta);
    if (!toAtaInfo) {
      instructions.push(
        token.createAssociatedTokenAccountInstruction(
          this.publicKey, // payer
          toAta,
          to,
          mint,
        ),
      );
    }

    instructions.push(
      token.createTransferInstruction(
        fromAta,
        toAta,
        this.publicKey,
        amount * Math.pow(10, decimals),
      ),
    );

    return this.createTransaction(instructions);
  }

  // Approve a session (called by each signer)
  async approveSession(sessionId: string): Promise<void> {
    const session = this.sessions.get(sessionId);
    if (!session) {
      throw new Error("Session not found");
    }

    if (session.status !== "pending" && session.status !== "signing") {
      throw new Error(`Session is ${session.status}`);
    }

    if (new Date() > session.expiresAt) {
      session.status = "failed";
      throw new Error("Session expired");
    }

    session.participants.add(this.signerId);
    session.status = "signing";

    // Generate partial signature
    const partialSig = await this.generatePartialSignature(
      session.transaction,
      this.signerId,
    );
    session.partialSignatures.set(this.signerId, partialSig);

    this.emit("approved", { sessionId, signerId: this.signerId });

    // Check if we have enough signatures
    if (session.partialSignatures.size >= this.config.threshold) {
      await this.finalizeSignature(session);
    }
  }

  // Reject a session
  async rejectSession(sessionId: string, reason?: string): Promise<void> {
    const session = this.sessions.get(sessionId);
    if (!session) {
      throw new Error("Session not found");
    }

    session.status = "failed";

    this.emit("rejected", { sessionId, signerId: this.signerId, reason });
  }

  private async generatePartialSignature(
    transaction: Transaction,
    signerId: number,
  ): Promise<Buffer> {
    if (!this.keyShare) {
      throw new Error("No key share available");
    }

    const message = transaction.serializeMessage();

    // In a real implementation, this would use the MPC signing protocol
    // For demonstration, we'll create a partial signature
    const partialSig = await this.mpcPartialSign(message, signerId);

    return partialSig;
  }

  private async mpcPartialSign(
    message: Buffer,
    signerId: number,
  ): Promise<Buffer> {
    // This is a placeholder for actual MPC signing
    // In production, use a library like tss-lib or implement FROST

    // Simulated partial signature
    const crypto = require("crypto");
    const hash = crypto
      .createHash("sha256")
      .update(message)
      .update(Buffer.from([signerId]))
      .digest();

    return hash;
  }

  private async finalizeSignature(session: SigningSession): Promise<void> {
    // Combine partial signatures using MPC aggregation
    const partialSigs = Array.from(session.partialSignatures.values());
    const combinedSignature = await this.aggregateSignatures(partialSigs);

    // Add signature to transaction
    session.transaction.addSignature(this.publicKey, combinedSignature);

    // Send transaction
    try {
      const signature = await this.connection.sendRawTransaction(
        session.transaction.serialize(),
      );

      const confirmation = await this.connection.confirmTransaction(signature);

      session.status = "complete";

      this.emit("complete", {
        sessionId: session.id,
        signature,
        confirmation,
      });
    } catch (error) {
      session.status = "failed";
      this.emit("failed", { sessionId: session.id, error });
      throw error;
    }
  }

  private async aggregateSignatures(partialSigs: Buffer[]): Promise<Buffer> {
    // In a real implementation, this would use Lagrange interpolation
    // to combine partial signatures into a valid EdDSA signature

    // Placeholder: XOR all partial signatures
    const result = Buffer.alloc(64);
    for (const sig of partialSigs) {
      for (let i = 0; i < Math.min(sig.length, result.length); i++) {
        result[i] ^= sig[i];
      }
    }

    return result;
  }

  private async notifySigners(sessionId: string): Promise<void> {
    const session = this.sessions.get(sessionId);
    if (!session) return;

    // Notify all signer devices
    for (const device of this.config.signerDevices) {
      if (device.endpoint) {
        try {
          await fetch(`${device.endpoint}/notify`, {
            method: "POST",
            headers: { "Content-Type": "application/json" },
            body: JSON.stringify({
              sessionId,
              transaction: session.transaction
                .serializeMessage()
                .toString("base64"),
              requiredSigners: this.config.threshold,
            }),
          });
        } catch (error) {
          console.error(`Failed to notify device ${device.id}:`, error);
        }
      }
    }
  }

  // Get session status
  getSession(sessionId: string): SigningSession | undefined {
    return this.sessions.get(sessionId);
  }

  // List pending sessions
  getPendingSessions(): SigningSession[] {
    return Array.from(this.sessions.values()).filter(
      (s) => s.status === "pending" || s.status === "signing",
    );
  }
}

import express from "express";
import { EventEmitter } from "events";

interface SignerConfig {
  id: number;
  keyShare: Buffer;
  publicKey: PublicKey;
  allowedAddresses?: PublicKey[];
  maxAmount?: number;
  requireConfirmation: boolean;
}

export class SignerService extends EventEmitter {
  private app: express.Application;
  private config: SignerConfig;
  private pendingApprovals: Map<string, () => Promise<void>> = new Map();

  constructor(config: SignerConfig) {
    super();
    this.config = config;
    this.app = express();
    this.app.use(express.json());
    this.setupRoutes();
  }

  private setupRoutes() {
    // Receive signing request
    this.app.post("/notify", async (req, res) => {
      const { sessionId, transaction, requiredSigners } = req.body;

      // Validate transaction
      const validation = await this.validateTransaction(transaction);
      if (!validation.valid) {
        return res.status(400).json({
          approved: false,
          reason: validation.reason,
        });
      }

      if (this.config.requireConfirmation) {
        // Store for manual approval
        this.pendingApprovals.set(sessionId, async () => {
          await this.signTransaction(sessionId, transaction);
        });

        this.emit("approval_required", {
          sessionId,
          transaction: validation.parsed,
        });

        return res.json({ status: "pending_approval", sessionId });
      }

      // Auto-approve
      await this.signTransaction(sessionId, transaction);
      res.json({ approved: true, sessionId });
    });

    // Manual approval endpoint
    this.app.post("/approve/:sessionId", async (req, res) => {
      const { sessionId } = req.params;
      const approval = this.pendingApprovals.get(sessionId);

      if (!approval) {
        return res.status(404).json({ error: "No pending approval found" });
      }

      try {
        await approval();
        this.pendingApprovals.delete(sessionId);
        res.json({ approved: true });
      } catch (error) {
        res.status(500).json({ error: "Failed to sign" });
      }
    });

    // Reject endpoint
    this.app.post("/reject/:sessionId", async (req, res) => {
      const { sessionId } = req.params;
      const { reason } = req.body;

      this.pendingApprovals.delete(sessionId);

      this.emit("rejected", { sessionId, reason });
      res.json({ rejected: true });
    });

    // Health check
    this.app.get("/health", (req, res) => {
      res.json({ status: "ok", signerId: this.config.id });
    });
  }

  private async validateTransaction(
    transactionBase64: string,
  ): Promise<{ valid: boolean; reason?: string; parsed?: any }> {
    const buffer = Buffer.from(transactionBase64, "base64");

    try {
      const tx = Transaction.from(buffer);

      // Check allowed addresses
      if (this.config.allowedAddresses) {
        for (const ix of tx.instructions) {
          for (const key of ix.keys) {
            if (
              key.isWritable &&
              !this.config.allowedAddresses.some((a) => a.equals(key.pubkey))
            ) {
              return {
                valid: false,
                reason: `Unauthorized address: ${key.pubkey.toBase58()}`,
              };
            }
          }
        }
      }

      // Check amount limits for transfers
      if (this.config.maxAmount) {
        for (const ix of tx.instructions) {
          if (ix.programId.equals(SystemProgram.programId)) {
            // Decode transfer amount
            const data = ix.data;
            if (data.length >= 12) {
              const amount = data.readBigUInt64LE(4);
              if (amount > BigInt(this.config.maxAmount * LAMPORTS_PER_SOL)) {
                return {
                  valid: false,
                  reason: `Amount exceeds limit: ${Number(amount) / LAMPORTS_PER_SOL} SOL`,
                };
              }
            }
          }
        }
      }

      return {
        valid: true,
        parsed: {
          instructions: tx.instructions.length,
          feePayer: tx.feePayer?.toBase58(),
        },
      };
    } catch (error) {
      return { valid: false, reason: "Invalid transaction format" };
    }
  }

  private async signTransaction(
    sessionId: string,
    transactionBase64: string,
  ): Promise<void> {
    // Generate partial signature using key share
    const message = Buffer.from(transactionBase64, "base64");

    // Notify coordinator with partial signature
    this.emit("signed", {
      sessionId,
      signerId: this.config.id,
      // In production, this would be actual partial signature
      partialSignature: Buffer.alloc(32).toString("base64"),
    });
  }

  start(port: number) {
    this.app.listen(port, () => {
      console.log(`Signer ${this.config.id} listening on port ${port}`);
    });
  }
}

import { useState, useEffect } from "react";
import { ThresholdWallet, SigningSession } from "./threshold-wallet";

interface TransactionRequest {
  id: string;
  type: "transfer" | "swap" | "custom";
  amount?: number;
  recipient?: string;
  status: "pending" | "approved" | "complete" | "failed";
  approvals: number;
  requiredApprovals: number;
}

export function useThresholdWallet(wallet: ThresholdWallet | null) {
  const [balance, setBalance] = useState<number>(0);
  const [pendingTxs, setPendingTxs] = useState<TransactionRequest[]>([]);
  const [loading, setLoading] = useState(false);

  useEffect(() => {
    if (!wallet) return;

    const fetchBalance = async () => {
      const bal = await wallet.getBalance();
      setBalance(bal);
    };

    fetchBalance();
    const interval = setInterval(fetchBalance, 30000);

    // Listen for wallet events
    const handleSessionCreated = (data: any) => {
      setPendingTxs((prev) => [
        ...prev,
        {
          id: data.sessionId,
          type: "transfer",
          status: "pending",
          approvals: 0,
          requiredApprovals: data.requiredSigners,
        },
      ]);
    };

    const handleApproved = (data: any) => {
      setPendingTxs((prev) =>
        prev.map((tx) =>
          tx.id === data.sessionId
            ? { ...tx, approvals: tx.approvals + 1 }
            : tx
        )
      );
    };

    const handleComplete = (data: any) => {
      setPendingTxs((prev) =>
        prev.map((tx) =>
          tx.id === data.sessionId ? { ...tx, status: "complete" } : tx
        )
      );
    };

    wallet.on("sessionCreated", handleSessionCreated);
    wallet.on("approved", handleApproved);
    wallet.on("complete", handleComplete);

    return () => {
      clearInterval(interval);
      wallet.off("sessionCreated", handleSessionCreated);
      wallet.off("approved", handleApproved);
      wallet.off("complete", handleComplete);
    };
  }, [wallet]);

  const sendTransaction = async (to: string, amount: number) => {
    if (!wallet) throw new Error("Wallet not initialized");

    setLoading(true);
    try {
      const sessionId = await wallet.transfer(new PublicKey(to), amount);
      return sessionId;
    } finally {
      setLoading(false);
    }
  };

  const approveTransaction = async (sessionId: string) => {
    if (!wallet) throw new Error("Wallet not initialized");

    setLoading(true);
    try {
      await wallet.approveSession(sessionId);
    } finally {
      setLoading(false);
    }
  };

  return {
    balance,
    pendingTxs,
    loading,
    sendTransaction,
    approveTransaction,
  };
}

// UI Component
export function ThresholdWalletUI({ wallet }: { wallet: ThresholdWallet }) {
  const { balance, pendingTxs, loading, sendTransaction, approveTransaction } =
    useThresholdWallet(wallet);
  const [recipient, setRecipient] = useState("");
  const [amount, setAmount] = useState("");

  const handleSend = async () => {
    await sendTransaction(recipient, parseFloat(amount));
    setRecipient("");
    setAmount("");
  };

  return (
    <div className="p-4 space-y-4">
      <div className="text-2xl font-bold">
        Balance: {balance.toFixed(4)} SOL
      </div>

      <div className="space-y-2">
        <input
          type="text"
          placeholder="Recipient address"
          value={recipient}
          onChange={(e) => setRecipient(e.target.value)}
          className="w-full p-2 border rounded"
        />
        <input
          type="number"
          placeholder="Amount (SOL)"
          value={amount}
          onChange={(e) => setAmount(e.target.value)}
          className="w-full p-2 border rounded"
        />
        <button
          onClick={handleSend}
          disabled={loading || !recipient || !amount}
          className="w-full p-2 bg-blue-500 text-white rounded disabled:opacity-50"
        >
          {loading ? "Processing..." : "Send"}
        </button>
      </div>

      <div className="space-y-2">
        <h3 className="font-bold">Pending Transactions</h3>
        {pendingTxs.map((tx) => (
          <div key={tx.id} className="p-2 border rounded flex justify-between">
            <span>
              {tx.status} ({tx.approvals}/{tx.requiredApprovals})
            </span>
            {tx.status === "pending" && (
              <button
                onClick={() => approveTransaction(tx.id)}
                className="px-2 py-1 bg-green-500 text-white rounded text-sm"
              >
                Approve
              </button>
            )}
          </div>
        ))}
      </div>
    </div>
  );
}