Refactor HDKey For BTC/BCH: Strategy Pattern Approach

by Editorial Team 54 views
Iklan Headers

Hey guys! Today, we're diving deep into a crucial refactoring project focused on separating the Bitcoin (BTC) and Bitcoin Cash (BCH) implementations within our HDKey functionality. Currently, the code in internal/infrastructure/wallet/key/hd_wallet.go handles logic for multiple coin types (BTC, BCH, ETH, XRP) using switch statements based on coinTypeCode. This approach introduces significant risks, as modifications intended for one coin could inadvertently affect others. Let's break down why this separation is so important and how we plan to achieve it using the Strategy Pattern.

The Critical Differences Between Bitcoin (BTC) and Bitcoin Cash (BCH)

When dealing with cryptocurrencies, it's vital to recognize that even closely related coins like BTC and BCH have fundamental differences that necessitate distinct handling. Bitcoin (BTC) and Bitcoin Cash (BCH), while sharing a common origin, have diverged significantly in their features and functionalities. These differences aren't just minor tweaks; they represent core architectural variations that impact how transactions are processed, how addresses are generated, and how the overall security model operates. For instance, BTC has embraced technologies like SegWit and Taproot, which are absent in BCH. Understanding these nuances is crucial for maintaining the integrity and reliability of our cryptocurrency wallet infrastructure. This section highlights the key distinctions that drive the need for a clear separation in our codebase.

Feature Bitcoin (BTC) Bitcoin Cash (BCH)
SegWit ✅ Activated 2017 ❌ Not supported
Taproot ✅ Activated 2021 ❌ Not supported
Schnorr Signatures ✅ BIP340 ❌ Not available
MuSig2 ✅ BIP327 ❌ Not available
Address Format Bech32/Bech32m CashAddr (P2PKH only)
Transaction Malleability ✅ Fixed by SegWit ⚠️ Still present

Why These Differences Matter

The implications of these differences are profound. Bitcoin (BTC), with its support for SegWit and Taproot, boasts a more sophisticated address structure, including P2PKH, P2SH-SegWit, Bech32, and Taproot formats. In contrast, Bitcoin Cash (BCH) is limited to P2PKH addresses in the CashAddr format. This divergence means that the logic for key generation and address derivation must be tailored to each coin. Specifically, BTC's key generation process incorporates SegWit/Taproot-specific steps, which are entirely irrelevant and potentially harmful for BCH. Imagine trying to fit a square peg into a round hole – that's the kind of problem we're trying to avoid by separating these implementations.

The Risks of Shared Code

The current implementation, where HDKey contains coin-specific logic embedded within switch statements, poses a significant risk. For example, the createKeysWithIndex() function (lines 291-398) and the getP2SHSegWitAddr() function (lines 532-583) both rely heavily on coinTypeCode to determine the appropriate behavior. This creates a situation where a seemingly innocuous change to BTC's Taproot address generation could inadvertently break BCH functionality. The intermingling of code paths makes it easy to overlook coin-specific edge cases, leading to potential bugs and vulnerabilities. We need to ensure that changes to one coin's logic do not have unintended consequences for other coins. The goal is to create a robust and maintainable system where each coin is treated as a distinct entity with its own set of rules and behaviors.

The Strategy Pattern: Our Solution for Clean Separation

To address these challenges, we're adopting the Strategy Pattern, a design pattern that allows us to encapsulate different algorithms or strategies behind a common interface. This will enable us to cleanly separate the BTC and BCH implementations, ensuring that each coin is handled independently and safely. By leveraging the Strategy Pattern, we can achieve a more modular, testable, and maintainable codebase.

Phase 1: Defining the CoinKeyStrategy Interface

First, we need to define an interface that outlines the common operations required for key generation and address derivation for each coin. This interface, which we'll call CoinKeyStrategy, will serve as the contract that all coin-specific implementations must adhere to.

// internal/application/ports/wallet/coin_key_strategy.go
type CoinKeyStrategy interface {
    // GenerateWalletKey creates a WalletKey from a private key
    GenerateWalletKey(privKey *btcec.PrivateKey) (*domainKey.WalletKey, error)
    
    // CoinTypeCode returns the coin type this strategy handles
    CoinTypeCode() domainCoin.CoinTypeCode
    
    // SupportedAddressTypes returns the address types this coin supports
    SupportedAddressTypes() []domainAddress.AddrType
}

This interface defines three key methods:

  • GenerateWalletKey: This method takes a private key as input and generates a WalletKey object, which encapsulates the necessary information for managing the key.
  • CoinTypeCode: This method returns the coin type code associated with the strategy, allowing us to identify which coin the strategy is responsible for.
  • SupportedAddressTypes: This method returns a list of address types supported by the coin, providing information about the available address formats.

Phase 2: Creating Coin-Specific Implementations

Next, we'll create concrete implementations of the CoinKeyStrategy interface for each coin we need to support. This will involve creating separate files for BTC, BCH, ETH, and XRP, each containing the specific logic for handling key generation and address derivation for that coin.

internal/infrastructure/wallet/key/strategy/
├── interface.go       # CoinKeyStrategy interface (if not in ports)
├── btc_strategy.go    # BTC: P2PKH, P2SH-SegWit, Bech32, Taproot
├── bch_strategy.go    # BCH: P2PKH only (CashAddr format)
├── eth_strategy.go    # ETH: Ethereum addresses
└── xrp_strategy.go    # XRP: XRP addresses
  • BTCKeyStrategy: This implementation will handle the generation of P2PKH, P2SH-SegWit, Bech32, and Taproot addresses for Bitcoin.
  • BCHKeyStrategy: This implementation will handle the generation of P2PKH (CashAddr) addresses for Bitcoin Cash, explicitly setting SegWit fields to empty to avoid any potential conflicts.
  • ETHKeyStrategy: This implementation will handle the generation of Ethereum addresses.
  • XRPKeyStrategy: This implementation will handle the generation of XRP addresses.

Key Separation Points

The key to this separation is ensuring that each strategy is responsible for its own coin-specific logic. For example, the BTCKeyStrategy will include the necessary code for generating SegWit and Taproot addresses, while the BCHKeyStrategy will explicitly exclude this code. This will prevent any accidental mixing of logic and ensure that each coin is handled correctly.

Phase 3: Refactoring HDKey

With the strategy implementations in place, we can now refactor the HDKey struct to leverage the Strategy Pattern. This will involve removing all switch statements based on coinTypeCode and instead accepting a CoinKeyStrategy as a dependency.

  1. HDKey becomes coin-agnostic (shared BIP32 derivation logic only).
  2. HDKey accepts a CoinKeyStrategy for coin-specific operations.
  3. Remove all switch statements on coinTypeCode from HDKey.

This will make the HDKey struct much cleaner and easier to maintain, as it will no longer be responsible for handling coin-specific logic. Instead, it will delegate these responsibilities to the appropriate CoinKeyStrategy implementation.

Phase 4: Updating the Factory

Finally, we need to update the factory that creates HDKey instances to ensure that the correct CoinKeyStrategy is used for each coin. This will involve modifying the factory to create the appropriate strategy based on the coinTypeCode and passing it to the HDKey constructor.

  1. Factory creates appropriate strategy based on coinTypeCode.
  2. BIP generators remain unchanged externally (backward compatible).

This will ensure that the correct strategy is used for each coin, maintaining backward compatibility and ensuring that existing functionality continues to work as expected.

Benefits of the Strategy Pattern

By adopting the Strategy Pattern, we can achieve a number of significant benefits:

Aspect Before After
Separation Mixed logic Clean separation per coin
Safety Risk of cross-coin impact Changes isolated to coin strategy
BTC Features Mixed with BCH checks Pure BTC with SegWit/Taproot
BCH Features Conditional exclusions Pure BCH without SegWit code
Testability Hard to test per coin Easy to unit test each strategy
Extensibility Add switch case Add new strategy file
  • Improved Separation: The Strategy Pattern enforces a clear separation of concerns, ensuring that each coin is handled independently and safely.
  • Increased Safety: By isolating coin-specific logic, we reduce the risk of cross-coin impact, preventing unintended consequences from changes to one coin affecting others.
  • Enhanced Testability: With each strategy implemented as a separate unit, it becomes much easier to write unit tests and verify the correctness of each implementation.
  • Greater Extensibility: Adding support for new coins becomes a simple matter of creating a new strategy implementation, without the need to modify existing code.

References

  • BCH documentation: docs/crypto/bch/README.md (BTC vs BCH comparison)
  • Existing pattern: internal/infrastructure/api/btc/bch/ (embedding + override)

Acceptance Criteria

To ensure that the refactoring is successful, we've defined a set of acceptance criteria that must be met:

  • [ ] Define CoinKeyStrategy interface in ports layer
  • [ ] Implement BTCKeyStrategy with all address types (P2PKH, P2SH-SegWit, Bech32, Taproot)
  • [ ] Implement BCHKeyStrategy with P2PKH (CashAddr format) only
  • [ ] Implement ETHKeyStrategy for Ethereum addresses
  • [ ] Implement XRPKeyStrategy for XRP addresses
  • [ ] Refactor HDKey to use strategy pattern (remove coin switch statements)
  • [ ] Update Factory to create appropriate strategy
  • [ ] All existing tests pass (make gotest)
  • [ ] Verification: make go-lint && make tidy && make check-build && make gotest
  • [ ] No functional changes (same input/output behavior)

Test Plan

To thoroughly test the refactoring, we'll follow a comprehensive test plan:

  • [ ] Existing hd_wallet_test.go tests pass unchanged
  • [ ] Existing hd_wallet_consistency_test.go tests pass
  • [ ] Add unit tests for each strategy implementation
  • [ ] Verify generated addresses match before/after refactoring
  • [ ] Verify BCH strategy does NOT generate SegWit/Taproot addresses
  • [ ] Verify BTC strategy generates all 4 address types

Risks and Mitigation

The risk associated with this refactoring is considered low, as it primarily involves reorganizing existing code without introducing new functionality. However, to mitigate any potential issues, we'll ensure comprehensive test coverage before and after the refactoring. This will allow us to quickly identify and address any regressions or unexpected behavior.

  • Low: This is pure refactoring with no external behavior change
  • Mitigation: Comprehensive test coverage before refactoring

By following this approach, we can confidently refactor the HDKey implementation to cleanly separate BTC and BCH logic, resulting in a more robust, maintainable, and testable codebase. Let's get to work!