synapsecns/sanguine

package forker

import (
    "bufio"
    "context"
    "crypto/rand"
    "encoding/base32"
    // embed is used for anvil binaries.
    _ "embed"
    "fmt"
    "github.com/ethereum/go-ethereum/ethclient"
    "github.com/ipfs/go-log"
    "github.com/phayes/freeport"
    "golang.org/x/sync/errgroup"
    "io"
    "k8s.io/apimachinery/pkg/util/wait"
    "os"
    "os/exec"
    "runtime"
    "strconv"
    "time"
)

//go:embed anvilbin/anvil_darwin_amd64_bin
var anvilDarwinAmd64Bin []byte

//go:embed anvilbin/anvil_darwin_arm64_bin
var anvilDarwinArm64Bin []byte

//go:embed anvilbin/anvil_linux_amd64_bin
var anvilLinuxAmd64Bin []byte

func getRandomString(length int) (string, error) {
    randomBytes := make([]byte, 32)
    _, err := rand.Read(randomBytes)
    if err != nil {
        return "", fmt.Errorf("rand read failed to generate random bytes: %w", err)
    }
    return base32.StdEncoding.EncodeToString(randomBytes)[:length], nil
}

func getForkCommandWaiter(cmd *exec.Cmd) func() error {
    return func() error {
        err := cmd.Wait()
        if err != nil {
            return fmt.Errorf("could not start forked node: %w", err)
        }
        return nil
    }
}

func getForkClientConnectionHandler(client *ethclient.Client, processPort int, chainID uint64, cancel context.CancelFunc) func(ctx context.Context) {
    return func(ctx context.Context) {
        var err error
        client, err = ethclient.DialContext(ctx, fmt.Sprintf("http://localhost:%d", processPort))
        if err != nil {
            logger.Errorf("failed to dial forked node: %s", err)
            return
        }

        rpcChainID, err := client.ChainID(ctx)
        if err != nil {
            return
        }

        if rpcChainID.Uint64() == chainID {
            cancel()
        }
    }
}

func getForkClientHandler(ctx context.Context, cmd *exec.Cmd, chainID uint64, processPort int, clientFunc func(client *ethclient.Client)) func() error {
    return func() error {
        defer cmd.Process.Kill() //nolint: errcheck
        startupCtx, cancel := context.WithCancel(ctx)
        defer cancel()

        var client *ethclient.Client
        wait.UntilWithContext(startupCtx, getForkClientConnectionHandler(client, processPort, chainID, cancel), time.Millisecond*10)
        clientFunc(client)
        return nil
    }
}

func getAnvilFilePayload() ([]byte, error) {
    var filePayload []byte
    osType := runtime.GOOS
    archType := runtime.GOARCH
    osTypeWithArch := osType + "_" + archType
    switch osTypeWithArch {
    case "darwin_amd64":
        filePayload = anvilDarwinAmd64Bin
    case "darwin_arm64":
        filePayload = anvilDarwinArm64Bin
    case "linux_amd64":
        filePayload = anvilLinuxAmd64Bin
    default:
        return []byte{}, fmt.Errorf("unsupported arch %s for os %s", archType, osType)
    }
    return filePayload, nil
}

// Fork will fork the given evm blockchain for the purposes of testing.
func Fork(ctx context.Context, rpcURL string, chainID uint64, clientFunc func(client *ethclient.Client)) error {
    processPort, err := freeport.GetFreePort()
    if err != nil {
        return fmt.Errorf("failed to get free port: %w", err)
    }

    now := time.Now().UnixNano()
    randStr, err := getRandomString(32)
    if err != nil {
        return fmt.Errorf("failed to generate random string")
    }

    embeddedFileName := "./anvil_embedded_" + strconv.FormatInt(now, 10) + "_" + randStr

    filePayload, err := getAnvilFilePayload()
    if err != nil {
        return fmt.Errorf("failed to get anvil file binary: %w", err)
    }

    err = os.WriteFile(embeddedFileName, filePayload, 0755) //nolint: gosec
    if err != nil {
        return fmt.Errorf("failed to write %s: %w", embeddedFileName, err)
    }
    defer os.Remove(embeddedFileName) //nolint: errcheck

    cmd := exec.Command(embeddedFileName, "--fork-url", rpcURL, "--chain-id", strconv.Itoa(int(chainID)), "--port", strconv.Itoa(processPort)) //nolint: gosec
    cmd.Env = os.Environ()

    stderr, err := cmd.StderrPipe()
    if err != nil {
        return fmt.Errorf("failed to get stderr pipe: %w", err)
    }

    stdout, err := cmd.StdoutPipe()
    if err != nil {
        return fmt.Errorf("failed to get stdout pipe: %w", err)
    }

    go outputLogs(stderr, log.LevelError)
    go outputLogs(stdout, log.LevelInfo)

    err = cmd.Start()
    if err != nil {
        return fmt.Errorf("could not start forked node: %w", err)
    }

    g, ctx := errgroup.WithContext(ctx)
    g.Go(getForkCommandWaiter(cmd))
    g.Go(getForkClientHandler(ctx, cmd, chainID, processPort, clientFunc))

    err = g.Wait()
    if err != nil {
        logger.Errorf("waitgroup exited: %s", err)
    }

    return nil
}

func outputLogs(reader io.ReadCloser, logLevel log.LogLevel) {
    s := bufio.NewScanner(reader)

    for s.Scan() {
        // TODO: make this exhautive
        switch logLevel {
        case log.LevelInfo:
            logger.Infof("%s", s.Text())
        case log.LevelError:
            logger.Errorf("%s", s.Text())
        }
    }
}