path: root/cmd/utils/cmd.go

// Copyright 2014 The go-ethereum Authors
// This file is part of go-ethereum.
//
// go-ethereum is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// go-ethereum is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with go-ethereum. If not, see <http://www.gnu.org/licenses/>.

// Package utils contains internal helper functions for go-ethereum commands.
package utils

import (
    "compress/gzip"
    "fmt"
    "io"
    "os"
    "os/signal"
    "runtime"
    "strings"

    "github.com/ethereum/go-ethereum/core"
    "github.com/ethereum/go-ethereum/core/types"
    "github.com/ethereum/go-ethereum/internal/debug"
    "github.com/ethereum/go-ethereum/log"
    "github.com/ethereum/go-ethereum/node"
    "github.com/ethereum/go-ethereum/rlp"
)

const (
    importBatchSize = 2500
)

// Fatalf formats a message to standard error and exits the program.
// The message is also printed to standard output if standard error
// is redirected to a different file.
func Fatalf(format string, args ...interface{}) {
    w := io.MultiWriter(os.Stdout, os.Stderr)
    if runtime.GOOS == "windows" {
        // The SameFile check below doesn't work on Windows.
        // stdout is unlikely to get redirected though, so just print there.
        w = os.Stdout
    } else {
        outf, _ := os.Stdout.Stat()
        errf, _ := os.Stderr.Stat()
        if outf != nil && errf != nil && os.SameFile(outf, errf) {
            w = os.Stderr
        }
    }
    fmt.Fprintf(w, "Fatal: "+format+"\n", args...)
    os.Exit(1)
}

func StartNode(stack *node.Node) {
    if err := stack.Start(); err != nil {
        Fatalf("Error starting protocol stack: %v", err)
    }
    go func() {
        sigc := make(chan os.Signal, 1)
        signal.Notify(sigc, os.Interrupt)
        defer signal.Stop(sigc)
        <-sigc
        log.Info("Got interrupt, shutting down...")
        go stack.Stop()
        for i := 10; i > 0; i-- {
            <-sigc
            if i > 1 {
                log.Warn("Already shutting down, interrupt more to panic.", "times", i-1)
            }
        }
        debug.Exit() // ensure trace and CPU profile data is flushed.
        debug.LoudPanic("boom")
    }()
}

func ImportChain(chain *core.BlockChain, fn string) error {
    // Watch for Ctrl-C while the import is running.
    // If a signal is received, the import will stop at the next batch.
    interrupt := make(chan os.Signal, 1)
    stop := make(chan struct{})
    signal.Notify(interrupt, os.Interrupt)
    defer signal.Stop(interrupt)
    defer close(interrupt)
    go func() {
        if _, ok := <-interrupt; ok {
            log.Info("Interrupted during import, stopping at next batch")
        }
        close(stop)
    }()
    checkInterrupt := func() bool {
        select {
        case <-stop:
            return true
        default:
            return false
        }
    }

    log.Info("Importing blockchain", "file", fn)
    fh, err := os.Open(fn)
    if err != nil {
        return err
    }
    defer fh.Close()

    var reader io.Reader = fh
    if strings.HasSuffix(fn, ".gz") {
        if reader, err = gzip.NewReader(reader); err != nil {
            return err
        }
    }

    stream := rlp.NewStream(reader, 0)

    // Run actual the import.
    blocks := make(types.Blocks, importBatchSize)
    n := 0
    for batch := 0; ; batch++ {
        // Load a batch of RLP blocks.
        if checkInterrupt() {
            return fmt.Errorf("interrupted")
        }
        i := 0
        for ; i < importBatchSize; i++ {
            var b types.Block
            if err := stream.Decode(&b); err == io.EOF {
                break
            } else if err != nil {
                return fmt.Errorf("at block %d: %v", n, err)
            }
            // don't import first block
            if b.NumberU64() == 0 {
                i--
                continue
            }
            blocks[i] = &b
            n++
        }
        if i == 0 {
            break
        }
        // Import the batch.
        if checkInterrupt() {
            return fmt.Errorf("interrupted")
        }
        if hasAllBlocks(chain, blocks[:i]) {
            log.Info("Skipping batch as all blocks present", "batch", batch, "first", blocks[0].Hash(), "last", blocks[i-1].Hash())
            continue
        }

        if _, err := chain.InsertChain(blocks[:i]); err != nil {
            return fmt.Errorf("invalid block %d: %v", n, err)
        }
    }
    return nil
}

func hasAllBlocks(chain *core.BlockChain, bs []*types.Block) bool {
    for _, b := range bs {
        if !chain.HasBlock(b.Hash()) {
            return false
        }
    }
    return true
}

func ExportChain(blockchain *core.BlockChain, fn string) error {
    log.Info("Exporting blockchain", "file", fn)
    fh, err := os.OpenFile(fn, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
    if err != nil {
        return err
    }
    defer fh.Close()

    var writer io.Writer = fh
    if strings.HasSuffix(fn, ".gz") {
        writer = gzip.NewWriter(writer)
        defer writer.(*gzip.Writer).Close()
    }

    if err := blockchain.Export(writer); err != nil {
        return err
    }
    log.Info("Exported blockchain", "file", fn)

    return nil
}

func ExportAppendChain(blockchain *core.BlockChain, fn string, first uint64, last uint64) error {
    log.Info("Exporting blockchain", "file", fn)
    // TODO verify mode perms
    fh, err := os.OpenFile(fn, os.O_CREATE|os.O_APPEND|os.O_WRONLY, os.ModePerm)
    if err != nil {
        return err
    }
    defer fh.Close()

    var writer io.Writer = fh
    if strings.HasSuffix(fn, ".gz") {
        writer = gzip.NewWriter(writer)
        defer writer.(*gzip.Writer).Close()
    }

    if err := blockchain.ExportN(writer, first, last); err != nil {
        return err
    }
    log.Info("Exported blockchain to", "file", fn)
    return nil
}