Trezor integration in Go

trezor.go

package main

import (
	"bytes"
	"encoding/hex"
	"flag"
	"fmt"
	"math/big"
	"os"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/golang/protobuf/proto"
	"github.com/kr/pretty"
	"github.com/rendaw/go-trezor"
	"github.com/rendaw/go-trezor/messages"
)

var (
	nonce    = flag.Int("nonce", -1, "Current nonce")
	gasPrice = flag.Float64("gasPrice", 0, "Gas price in Gwei")
	to       = flag.String("to", "", "Ethereum destination address")
)

func main() {
	flag.Parse()

	// Check if *nonce has the sentinel value, which would mean that the user
	// probably didn't set the flag.
	if *nonce == -1 {
		fmt.Fprintln(os.Stderr, "--nonce is required")
		os.Exit(1)
	}

	// Ensure the user didn't provide a negative nonce.
	if *nonce < 0 {
		fmt.Fprintln(os.Stderr, "--nonce must not be negative")
		os.Exit(1)
	}

	// Ensure gas price is positive.
	if *gasPrice <= 0 {
		fmt.Fprintln(os.Stderr, "--gasPrice must be set to a positive value")
		os.Exit(1)
	}

	if *to == "" {
		fmt.Fprintln(os.Stderr, "--to is required")
		os.Exit(1)
	}

	toAddress := common.HexToAddress(*to)

	devices, err := trezor.Enumerate()
	check(err, "failed to enumerate Trezor devices")

	if len(devices) == 0 {
		fmt.Println("No Trezor devices found.")
		return
	}

	device := devices[0]

	// Convenience function for getting proto.Message values from device.Read()
	unmarshal := func(msgType messages.MessageType, msgBytes []byte, err error) (proto.Message, error) {
		if err != nil {
			return nil, err
		}
		var msg proto.Message
		switch msgType {
		case messages.MessageType_MessageType_Failure:
			msg = &messages.Failure{}
		case messages.MessageType_MessageType_PinMatrixRequest:
			msg = &messages.PinMatrixRequest{}
		case messages.MessageType_MessageType_Features:
			msg = &messages.Features{}
		case messages.MessageType_MessageType_EthereumAddress:
			msg = &messages.EthereumAddress{}
		case messages.MessageType_MessageType_ButtonRequest:
			msg = &messages.ButtonRequest{}
		case messages.MessageType_MessageType_EthereumTxRequest:
			msg = &messages.EthereumTxRequest{}
		default:
			return nil, fmt.Errorf("unhandled message type: %v", msgType)
		}
		err = proto.Unmarshal(msgBytes, msg)
		if err != nil {
			return nil, err
		}
		return msg, nil
	}

	mustReadMessage := func(description string) proto.Message {
		msg, err := unmarshal(device.Read())
		check(err, "failed to read "+description)
		return msg
	}

	// Initialize
	{
		check(device.Open(), "opening trezor device")
		defer device.Close()

		check(device.Write(&messages.Initialize{}), "initializing trezor")
		msg := mustReadMessage("initialize response")
		if _, ok := msg.(*messages.Features); !ok {
			fmt.Fprintf(os.Stderr, "Failed to initialize Trezor device. Received response %v.\n")
			os.Exit(1)
		}
	}

	var hdPath []uint32
	{
		var hardened uint32 = 1 << 31
		hdPath = []uint32{
			44 | hardened,
			60 | hardened,
			0 | hardened,
			0,
			0,
		}
	}

	// Get a response from the device, approving PIN requests
	// and button requests first if necessary.
	getResponse := func(description string) proto.Message {
		for {
			msg := mustReadMessage(description)
			switch msg := msg.(type) {
			case *messages.PinMatrixRequest:
				fmt.Println("requested pin")
				ok, pin, err := doPin("enter the device pin")
				if !ok {
					check(device.Write(&messages.Cancel{}), "sending 'cancel' message to Trezor")
					fmt.Println("cancelled")
					os.Exit(1)
				}
				check(err, "getting pin")
				check(device.Write(&messages.PinMatrixAck{
					Pin: proto.String(pin),
				}), "writing PinMatrixAck")
			case *messages.ButtonRequest:
				check(device.Write(&messages.ButtonAck{}), "sending button ack to Trezor")
			default:
				return msg
			}
		}
	}

	// Request ethereum address
	if true {
		check(device.Write(&messages.EthereumGetAddress{
			AddressN: hdPath,
		}), "writing EthereumGetAddress message")

		msg := getResponse("EthereumGetAddress response")
		if msg, ok := msg.(*messages.EthereumAddress); ok {
			fmt.Println("0x" + hex.EncodeToString(msg.Address))
		} else {
			fmt.Println("unhandled message type", msg)
			return
		}
	}

	// Send some ETH. Construct a transaction here and sign it.
	var (
		amount          = big.NewInt(4e16)
		gasLimit uint64 = 21000
		gasPrice        = new(big.Int).SetUint64(uint64(*gasPrice * 1e9))
		data     []byte

		r, s []byte
		v    uint32
	)
	{
		uintToBytes := func(u uint64) []byte {
			return new(big.Int).SetUint64(u).Bytes()
		}
		initialChunk := data
		if len(initialChunk) > 1024 {
			initialChunk = initialChunk[:1024]
		}
		check(device.Write(&messages.EthereumSignTx{
			AddressN:         hdPath,
			Nonce:            uintToBytes(uint64(*nonce)),
			GasPrice:         gasPrice.Bytes(),
			GasLimit:         uintToBytes(gasLimit),
			To:               toAddress.Bytes(),
			Value:            amount.Bytes(),
			DataInitialChunk: initialChunk,
			DataLength:       proto.Uint32(uint32(len(data))),
			ChainId:          proto.Uint32(1),
		}), "sending EthereumSignTx message")

		msg := getResponse("EthereumSignTx result")
		if msg, ok := msg.(*messages.EthereumTxRequest); !ok {
			fmt.Fprintln(os.Stderr, "unexpected message after EthereumSignTx request:", msg)
			os.Exit(1)
		} else {
			r = []byte(string(msg.SignatureR))
			s = []byte(string(msg.SignatureS))
			v = *msg.SignatureV
		}
		pretty.Println(msg)
	}

	// Serialize transaction to send online.
	{
		tx, err := types.NewTransaction(
			uint64(*nonce),
			toAddress,
			amount,
			gasLimit,
			gasPrice,
			data,
		).WithSignature(
			types.NewEIP155Signer(common.Big1),
			append(append(r, s...), byte(v-37)),
		)
		check(err, "applying signature to transaction")

		buf := new(bytes.Buffer)
		check(tx.EncodeRLP(buf), "RLP-encoding transaction")
		fmt.Println("0x" + hex.EncodeToString(buf.Bytes()))
	}
}

func check(err error, msg string) {
	if err != nil {
		fmt.Fprintln(os.Stderr, msg+":", err)
		os.Exit(1)
	}
}

This was written before I knew about go-ethereum's usbwallet package. You should probably use that package instead of doing something like this.