<commit_before>package main

import "C"
import (
	"encoding/json"
	"fmt"
	"os"

	"gopkg.in/go-playground/validator.v9"

	"github.com/status-im/status-go/geth/common"
	"github.com/status-im/status-go/geth/log"
	"github.com/status-im/status-go/geth/params"
	"github.com/status-im/status-go/helpers/profiling"
)

//GenerateConfig for status node
//export GenerateConfig
func GenerateConfig(datadir *C.char, networkID C.int, devMode C.int) *C.char {
	config, err := params.NewNodeConfig(C.GoString(datadir), uint64(networkID), devMode == 1)
	if err != nil {
		return makeJSONResponse(err)
	}

	outBytes, err := json.Marshal(&config)
	if err != nil {
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//StartNode - start Status node
//export StartNode
func StartNode(configJSON *C.char) *C.char {
	config, err := params.LoadNodeConfig(C.GoString(configJSON))
	if err != nil {
		return makeJSONResponse(err)
	}

	_, err = statusAPI.StartNodeAsync(config)
	return makeJSONResponse(err)
}

//StopNode - stop status node
//export StopNode
func StopNode() *C.char {
	_, err := statusAPI.StopNodeAsync()
	return makeJSONResponse(err)
}

//ValidateNodeConfig validates config for status node
//export ValidateNodeConfig
func ValidateNodeConfig(configJSON *C.char) *C.char {
	var resp common.APIDetailedResponse

	_, err := params.LoadNodeConfig(C.GoString(configJSON))

	// Convert errors to common.APIDetailedResponse
	switch err := err.(type) {
	case validator.ValidationErrors:
		resp = common.APIDetailedResponse{
			Message:     "validation: validation failed",
			FieldErrors: make([]common.APIFieldError, len(err)),
		}

		for i, ve := range err {
			resp.FieldErrors[i] = common.APIFieldError{
				Parameter: ve.Namespace(),
				Errors: []common.APIError{
					{
						Message: fmt.Sprintf("field validation failed on the '%s' tag", ve.Tag()),
					},
				},
			}
		}
	case error:
		resp = common.APIDetailedResponse{
			Message: fmt.Sprintf("validation: %s", err.Error()),
		}
	case nil:
		resp = common.APIDetailedResponse{
			Status: true,
		}
	}

	respJSON, err := json.Marshal(resp)
	if err != nil {
		return makeJSONResponse(err)
	}

	return C.CString(string(respJSON))
}

//ResetChainData remove chain data from data directory
//export ResetChainData
func ResetChainData() *C.char {
	_, err := statusAPI.ResetChainDataAsync()
	return makeJSONResponse(err)
}

//CallRPC calls status node via rpc
//export CallRPC
func CallRPC(inputJSON *C.char) *C.char {
	outputJSON := statusAPI.CallRPC(C.GoString(inputJSON))
	return C.CString(outputJSON)
}

//CreateAccount is equivalent to creating an account from the command line,
// just modified to handle the function arg passing
//export CreateAccount
func CreateAccount(password *C.char) *C.char {
	address, pubKey, mnemonic, err := statusAPI.CreateAccount(C.GoString(password))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.AccountInfo{
		Address:  address,
		PubKey:   pubKey,
		Mnemonic: mnemonic,
		Error:    errString,
	}
	outBytes, _ := json.Marshal(&out)
	return C.CString(string(outBytes))
}

//CreateChildAccount creates sub-account
//export CreateChildAccount
func CreateChildAccount(parentAddress, password *C.char) *C.char {
	address, pubKey, err := statusAPI.CreateChildAccount(C.GoString(parentAddress), C.GoString(password))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.AccountInfo{
		Address: address,
		PubKey:  pubKey,
		Error:   errString,
	}
	outBytes, _ := json.Marshal(&out)
	return C.CString(string(outBytes))
}

//RecoverAccount re-creates master key using given details
//export RecoverAccount
func RecoverAccount(password, mnemonic *C.char) *C.char {
	address, pubKey, err := statusAPI.RecoverAccount(C.GoString(password), C.GoString(mnemonic))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.AccountInfo{
		Address:  address,
		PubKey:   pubKey,
		Mnemonic: C.GoString(mnemonic),
		Error:    errString,
	}
	outBytes, _ := json.Marshal(&out)
	return C.CString(string(outBytes))
}

//VerifyAccountPassword verifies account password
//export VerifyAccountPassword
func VerifyAccountPassword(keyStoreDir, address, password *C.char) *C.char {
	_, err := statusAPI.VerifyAccountPassword(C.GoString(keyStoreDir), C.GoString(address), C.GoString(password))
	return makeJSONResponse(err)
}

//Login loads a key file (for a given address), tries to decrypt it using the password, to verify ownership
// if verified, purges all the previous identities from Whisper, and injects verified key as shh identity
//export Login
func Login(address, password *C.char) *C.char {
	err := statusAPI.SelectAccount(C.GoString(address), C.GoString(password))
	return makeJSONResponse(err)
}

//Logout is equivalent to clearing whisper identities
//export Logout
func Logout() *C.char {
	err := statusAPI.Logout()
	return makeJSONResponse(err)
}

//CompleteTransaction instructs backend to complete sending of a given transaction
//export CompleteTransaction
func CompleteTransaction(id, password *C.char) *C.char {
	txHash, err := statusAPI.CompleteTransaction(common.QueuedTxID(C.GoString(id)), C.GoString(password))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.CompleteTransactionResult{
		ID:    C.GoString(id),
		Hash:  txHash.Hex(),
		Error: errString,
	}
	outBytes, err := json.Marshal(&out)
	if err != nil {
		log.Error("failed to marshal CompleteTransaction output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//CompleteTransactions instructs backend to complete sending of multiple transactions
//export CompleteTransactions
func CompleteTransactions(ids, password *C.char) *C.char {
	out := common.CompleteTransactionsResult{}
	out.Results = make(map[string]common.CompleteTransactionResult)

	parsedIDs, err := common.ParseJSONArray(C.GoString(ids))
	if err != nil {
		out.Results["none"] = common.CompleteTransactionResult{
			Error: err.Error(),
		}
	} else {
		txIDs := make([]common.QueuedTxID, len(parsedIDs))
		for i, id := range parsedIDs {
			txIDs[i] = common.QueuedTxID(id)
		}

		results := statusAPI.CompleteTransactions(txIDs, C.GoString(password))
		for txID, result := range results {
			txResult := common.CompleteTransactionResult{
				ID:   string(txID),
				Hash: result.Hash.Hex(),
			}
			if result.Error != nil {
				txResult.Error = result.Error.Error()
			}
			out.Results[string(txID)] = txResult
		}
	}

	outBytes, err := json.Marshal(&out)
	if err != nil {
		log.Error("failed to marshal CompleteTransactions output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//DiscardTransaction discards a given transaction from transaction queue
//export DiscardTransaction
func DiscardTransaction(id *C.char) *C.char {
	err := statusAPI.DiscardTransaction(common.QueuedTxID(C.GoString(id)))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.DiscardTransactionResult{
		ID:    C.GoString(id),
		Error: errString,
	}
	outBytes, err := json.Marshal(&out)
	if err != nil {
		log.Error("failed to marshal DiscardTransaction output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//DiscardTransactions discards given multiple transactions from transaction queue
//export DiscardTransactions
func DiscardTransactions(ids *C.char) *C.char {
	out := common.DiscardTransactionsResult{}
	out.Results = make(map[string]common.DiscardTransactionResult)

	parsedIDs, err := common.ParseJSONArray(C.GoString(ids))
	if err != nil {
		out.Results["none"] = common.DiscardTransactionResult{
			Error: err.Error(),
		}
	} else {
		txIDs := make([]common.QueuedTxID, len(parsedIDs))
		for i, id := range parsedIDs {
			txIDs[i] = common.QueuedTxID(id)
		}

		results := statusAPI.DiscardTransactions(txIDs)
		for txID, result := range results {
			txResult := common.DiscardTransactionResult{
				ID: string(txID),
			}
			if result.Error != nil {
				txResult.Error = result.Error.Error()
			}
			out.Results[string(txID)] = txResult
		}
	}

	outBytes, err := json.Marshal(&out)
	if err != nil {
		log.Error("failed to marshal DiscardTransactions output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//InitJail setup initial JavaScript
//export InitJail
func InitJail(js *C.char) {
	statusAPI.JailBaseJS(C.GoString(js))
}

//Parse creates a new jail cell context and executes provided JavaScript code
//export Parse
func Parse(chatID *C.char, js *C.char) *C.char {
	res := statusAPI.JailParse(C.GoString(chatID), C.GoString(js))
	return C.CString(res)
}

//Call executes given JavaScript function
//export Call
func Call(chatID *C.char, path *C.char, params *C.char) *C.char {
	res := statusAPI.JailCall(C.GoString(chatID), C.GoString(path), C.GoString(params))
	return C.CString(res)
}

//StartCPUProfile runs pprof for cpu
//export StartCPUProfile
func StartCPUProfile(dataDir *C.char) *C.char {
	err := profiling.StartCPUProfile(C.GoString(dataDir))
	return makeJSONResponse(err)
}

//StopCPUProfiling stops pprof for cpu
//export StopCPUProfiling
func StopCPUProfiling() *C.char { //nolint: deadcode
	err := profiling.StopCPUProfile()
	return makeJSONResponse(err)
}

//WriteHeapProfile starts pprof for heap
//export WriteHeapProfile
func WriteHeapProfile(dataDir *C.char) *C.char { //nolint: deadcode
	err := profiling.WriteHeapFile(C.GoString(dataDir))
	return makeJSONResponse(err)
}

func makeJSONResponse(err error) *C.char {
	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.APIResponse{
		Error: errString,
	}
	outBytes, _ := json.Marshal(&out)

	return C.CString(string(outBytes))
}

// Notify sends push notification by given token
//export Notify
func Notify(token *C.char) *C.char {
	err := statusAPI.Notify(C.GoString(token))

	res := true
	errString := ""
	if err != nil {
		res = false
		errString = err.Error()
	}

	out := common.NotifyResult{
		Status: res,
		Error:  errString,
	}

	outBytes, err := json.Marshal(&out)
	if err != nil {
		log.Error("failed to marshal Notify output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}
<commit_msg>References for marshal library responces removed<commit_after>package main

import "C"
import (
	"encoding/json"
	"fmt"
	"os"

	"gopkg.in/go-playground/validator.v9"

	"github.com/status-im/status-go/geth/common"
	"github.com/status-im/status-go/geth/log"
	"github.com/status-im/status-go/geth/params"
	"github.com/status-im/status-go/helpers/profiling"
)

//GenerateConfig for status node
//export GenerateConfig
func GenerateConfig(datadir *C.char, networkID C.int, devMode C.int) *C.char {
	config, err := params.NewNodeConfig(C.GoString(datadir), uint64(networkID), devMode == 1)
	if err != nil {
		return makeJSONResponse(err)
	}

	outBytes, err := json.Marshal(config)
	if err != nil {
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//StartNode - start Status node
//export StartNode
func StartNode(configJSON *C.char) *C.char {
	config, err := params.LoadNodeConfig(C.GoString(configJSON))
	if err != nil {
		return makeJSONResponse(err)
	}

	_, err = statusAPI.StartNodeAsync(config)
	return makeJSONResponse(err)
}

//StopNode - stop status node
//export StopNode
func StopNode() *C.char {
	_, err := statusAPI.StopNodeAsync()
	return makeJSONResponse(err)
}

//ValidateNodeConfig validates config for status node
//export ValidateNodeConfig
func ValidateNodeConfig(configJSON *C.char) *C.char {
	var resp common.APIDetailedResponse

	_, err := params.LoadNodeConfig(C.GoString(configJSON))

	// Convert errors to common.APIDetailedResponse
	switch err := err.(type) {
	case validator.ValidationErrors:
		resp = common.APIDetailedResponse{
			Message:     "validation: validation failed",
			FieldErrors: make([]common.APIFieldError, len(err)),
		}

		for i, ve := range err {
			resp.FieldErrors[i] = common.APIFieldError{
				Parameter: ve.Namespace(),
				Errors: []common.APIError{
					{
						Message: fmt.Sprintf("field validation failed on the '%s' tag", ve.Tag()),
					},
				},
			}
		}
	case error:
		resp = common.APIDetailedResponse{
			Message: fmt.Sprintf("validation: %s", err.Error()),
		}
	case nil:
		resp = common.APIDetailedResponse{
			Status: true,
		}
	}

	respJSON, err := json.Marshal(resp)
	if err != nil {
		return makeJSONResponse(err)
	}

	return C.CString(string(respJSON))
}

//ResetChainData remove chain data from data directory
//export ResetChainData
func ResetChainData() *C.char {
	_, err := statusAPI.ResetChainDataAsync()
	return makeJSONResponse(err)
}

//CallRPC calls status node via rpc
//export CallRPC
func CallRPC(inputJSON *C.char) *C.char {
	outputJSON := statusAPI.CallRPC(C.GoString(inputJSON))
	return C.CString(outputJSON)
}

//CreateAccount is equivalent to creating an account from the command line,
// just modified to handle the function arg passing
//export CreateAccount
func CreateAccount(password *C.char) *C.char {
	address, pubKey, mnemonic, err := statusAPI.CreateAccount(C.GoString(password))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.AccountInfo{
		Address:  address,
		PubKey:   pubKey,
		Mnemonic: mnemonic,
		Error:    errString,
	}
	outBytes, _ := json.Marshal(out)
	return C.CString(string(outBytes))
}

//CreateChildAccount creates sub-account
//export CreateChildAccount
func CreateChildAccount(parentAddress, password *C.char) *C.char {
	address, pubKey, err := statusAPI.CreateChildAccount(C.GoString(parentAddress), C.GoString(password))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.AccountInfo{
		Address: address,
		PubKey:  pubKey,
		Error:   errString,
	}
	outBytes, _ := json.Marshal(out)
	return C.CString(string(outBytes))
}

//RecoverAccount re-creates master key using given details
//export RecoverAccount
func RecoverAccount(password, mnemonic *C.char) *C.char {
	address, pubKey, err := statusAPI.RecoverAccount(C.GoString(password), C.GoString(mnemonic))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.AccountInfo{
		Address:  address,
		PubKey:   pubKey,
		Mnemonic: C.GoString(mnemonic),
		Error:    errString,
	}
	outBytes, _ := json.Marshal(out)
	return C.CString(string(outBytes))
}

//VerifyAccountPassword verifies account password
//export VerifyAccountPassword
func VerifyAccountPassword(keyStoreDir, address, password *C.char) *C.char {
	_, err := statusAPI.VerifyAccountPassword(C.GoString(keyStoreDir), C.GoString(address), C.GoString(password))
	return makeJSONResponse(err)
}

//Login loads a key file (for a given address), tries to decrypt it using the password, to verify ownership
// if verified, purges all the previous identities from Whisper, and injects verified key as shh identity
//export Login
func Login(address, password *C.char) *C.char {
	err := statusAPI.SelectAccount(C.GoString(address), C.GoString(password))
	return makeJSONResponse(err)
}

//Logout is equivalent to clearing whisper identities
//export Logout
func Logout() *C.char {
	err := statusAPI.Logout()
	return makeJSONResponse(err)
}

//CompleteTransaction instructs backend to complete sending of a given transaction
//export CompleteTransaction
func CompleteTransaction(id, password *C.char) *C.char {
	txHash, err := statusAPI.CompleteTransaction(common.QueuedTxID(C.GoString(id)), C.GoString(password))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.CompleteTransactionResult{
		ID:    C.GoString(id),
		Hash:  txHash.Hex(),
		Error: errString,
	}
	outBytes, err := json.Marshal(out)
	if err != nil {
		log.Error("failed to marshal CompleteTransaction output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//CompleteTransactions instructs backend to complete sending of multiple transactions
//export CompleteTransactions
func CompleteTransactions(ids, password *C.char) *C.char {
	out := common.CompleteTransactionsResult{}
	out.Results = make(map[string]common.CompleteTransactionResult)

	parsedIDs, err := common.ParseJSONArray(C.GoString(ids))
	if err != nil {
		out.Results["none"] = common.CompleteTransactionResult{
			Error: err.Error(),
		}
	} else {
		txIDs := make([]common.QueuedTxID, len(parsedIDs))
		for i, id := range parsedIDs {
			txIDs[i] = common.QueuedTxID(id)
		}

		results := statusAPI.CompleteTransactions(txIDs, C.GoString(password))
		for txID, result := range results {
			txResult := common.CompleteTransactionResult{
				ID:   string(txID),
				Hash: result.Hash.Hex(),
			}
			if result.Error != nil {
				txResult.Error = result.Error.Error()
			}
			out.Results[string(txID)] = txResult
		}
	}

	outBytes, err := json.Marshal(out)
	if err != nil {
		log.Error("failed to marshal CompleteTransactions output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//DiscardTransaction discards a given transaction from transaction queue
//export DiscardTransaction
func DiscardTransaction(id *C.char) *C.char {
	err := statusAPI.DiscardTransaction(common.QueuedTxID(C.GoString(id)))

	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.DiscardTransactionResult{
		ID:    C.GoString(id),
		Error: errString,
	}
	outBytes, err := json.Marshal(out)
	if err != nil {
		log.Error("failed to marshal DiscardTransaction output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//DiscardTransactions discards given multiple transactions from transaction queue
//export DiscardTransactions
func DiscardTransactions(ids *C.char) *C.char {
	out := common.DiscardTransactionsResult{}
	out.Results = make(map[string]common.DiscardTransactionResult)

	parsedIDs, err := common.ParseJSONArray(C.GoString(ids))
	if err != nil {
		out.Results["none"] = common.DiscardTransactionResult{
			Error: err.Error(),
		}
	} else {
		txIDs := make([]common.QueuedTxID, len(parsedIDs))
		for i, id := range parsedIDs {
			txIDs[i] = common.QueuedTxID(id)
		}

		results := statusAPI.DiscardTransactions(txIDs)
		for txID, result := range results {
			txResult := common.DiscardTransactionResult{
				ID: string(txID),
			}
			if result.Error != nil {
				txResult.Error = result.Error.Error()
			}
			out.Results[string(txID)] = txResult
		}
	}

	outBytes, err := json.Marshal(out)
	if err != nil {
		log.Error("failed to marshal DiscardTransactions output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}

//InitJail setup initial JavaScript
//export InitJail
func InitJail(js *C.char) {
	statusAPI.JailBaseJS(C.GoString(js))
}

//Parse creates a new jail cell context and executes provided JavaScript code
//export Parse
func Parse(chatID *C.char, js *C.char) *C.char {
	res := statusAPI.JailParse(C.GoString(chatID), C.GoString(js))
	return C.CString(res)
}

//Call executes given JavaScript function
//export Call
func Call(chatID *C.char, path *C.char, params *C.char) *C.char {
	res := statusAPI.JailCall(C.GoString(chatID), C.GoString(path), C.GoString(params))
	return C.CString(res)
}

//StartCPUProfile runs pprof for cpu
//export StartCPUProfile
func StartCPUProfile(dataDir *C.char) *C.char {
	err := profiling.StartCPUProfile(C.GoString(dataDir))
	return makeJSONResponse(err)
}

//StopCPUProfiling stops pprof for cpu
//export StopCPUProfiling
func StopCPUProfiling() *C.char { //nolint: deadcode
	err := profiling.StopCPUProfile()
	return makeJSONResponse(err)
}

//WriteHeapProfile starts pprof for heap
//export WriteHeapProfile
func WriteHeapProfile(dataDir *C.char) *C.char { //nolint: deadcode
	err := profiling.WriteHeapFile(C.GoString(dataDir))
	return makeJSONResponse(err)
}

func makeJSONResponse(err error) *C.char {
	errString := ""
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		errString = err.Error()
	}

	out := common.APIResponse{
		Error: errString,
	}
	outBytes, _ := json.Marshal(out)

	return C.CString(string(outBytes))
}

// Notify sends push notification by given token
//export Notify
func Notify(token *C.char) *C.char {
	err := statusAPI.Notify(C.GoString(token))

	res := true
	errString := ""
	if err != nil {
		res = false
		errString = err.Error()
	}

	out := common.NotifyResult{
		Status: res,
		Error:  errString,
	}

	outBytes, err := json.Marshal(out)
	if err != nil {
		log.Error("failed to marshal Notify output", "error", err.Error())
		return makeJSONResponse(err)
	}

	return C.CString(string(outBytes))
}
<|endoftext|>


<commit_before>package command

import (
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"os"
	"path"
	"sort"
	"strings"

	"github.com/funkygao/gafka/ctx"
	gzk "github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/samuel/go-zookeeper/zk"
)

type Dump struct {
	Ui  cli.Ui
	Cmd string

	zone       string
	path       string
	importFile string
	infile     string
	outfile    string
	outdir     string
	f          *os.File
}

func (this *Dump) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("dump", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&this.outfile, "o", "zk.dump", "")
	cmdFlags.StringVar(&this.infile, "in", "", "")
	cmdFlags.StringVar(&this.importFile, "import", "", "")
	cmdFlags.StringVar(&this.path, "p", "/", "")
	cmdFlags.StringVar(&this.outdir, "dir", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		requireAdminRights("-import").
		invalid(args) {
		return 2
	}

	if this.zone == "" {
		this.Ui.Error("unknown zone")
		return 2
	}

	if this.importFile != "" {
		if this.path == "" {
			this.Ui.Error("-p required")
			return 2
		}

		zkzone := gzk.NewZkZone(gzk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
		defer zkzone.Close()

		this.importFromFile(zkzone)
		return
	}

	if this.infile != "" {
		// display mode
		this.diplayDumppedFile()
		return
	}

	// dump mode
	this.outfile = this.zone + "." + this.outfile

	var err error
	if this.outdir != "" {
		err = os.MkdirAll(this.outdir, 0755)
		must(err)

		this.outfile = fmt.Sprintf("%s/%s", this.outdir, this.outfile)

		_, err = os.Lstat(this.outfile)
		if err == nil {
			// file exists, find next available number
			num := 1
			fname := ""
			for ; err == nil && num <= 9999; num++ { // 30 year is enough
				fname = this.outfile + fmt.Sprintf(".%04d", num)
				_, err = os.Lstat(fname)
			}
			if err == nil {
				panic("Not able to rotate, 30 years passed?")
			}

			err = os.Rename(this.outfile, fname)
			must(err)

			this.Ui.Info(fmt.Sprintf("rename %s -> %s", this.outfile, fname))
		}
	}

	this.f, err = os.OpenFile(this.outfile,
		os.O_WRONLY|os.O_CREATE|os.O_APPEND|os.O_TRUNC, 0666)
	must(err)

	zkzone := gzk.NewZkZone(gzk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	defer zkzone.Close()

	this.dump(zkzone.Conn(), this.path)
	this.f.Close()

	this.Ui.Info(fmt.Sprintf("dumpped to %s", this.outfile))

	return
}

func (this *Dump) importFromFile(zkzone *gzk.ZkZone) {
	f, err := os.Open(this.importFile)
	must(err)

	for {
		// read line, got the znode path
		var buf [1]byte
		zpath := make([]byte, 0, 8<<10)
		for {
			b := buf[:]
			_, err := f.Read(b)
			if err == io.EOF {
				return
			}
			must(err)

			if b[0] == '\n' {
				break
			}
			zpath = append(zpath, b[0])
		}

		newPath := path.Join(this.path, string(zpath))
		this.Ui.Infof("%50s -> [%s]%s", string(zpath), this.zone, newPath)

		// read the znode data
		// 1. data len
		// 2. data itself
		var dataLen int32
		err = binary.Read(f, binary.BigEndian, &dataLen)
		must(err)

		zdata := make([]byte, dataLen)
		_, err = io.ReadFull(f, zdata)
		must(err)

		must(zkzone.CreatePermenantZnode(newPath, zdata))
	}
}

func (this *Dump) diplayDumppedFile() {
	f, err := os.Open(this.infile)
	must(err)

	for {
		// read line, got the znode path
		var buf [1]byte
		zpath := make([]byte, 0, 8<<10)
		for {
			b := buf[:]
			_, err := f.Read(b)
			if err == io.EOF {
				return
			}
			must(err)

			if b[0] == '\n' {
				break
			}
			zpath = append(zpath, b[0])
		}

		this.Ui.Info(string(zpath))

		// read the znode data
		// 1. data len
		// 2. data itself
		var dataLen int32
		err = binary.Read(f, binary.BigEndian, &dataLen)
		must(err)

		zdata := make([]byte, dataLen)
		_, err = io.ReadFull(f, zdata)
		must(err)

		this.Ui.Output(string(zdata))
	}
}

// serialization format:
// path
// len(4) value
func (this *Dump) dump(conn *zk.Conn, path string) {
	children, _, err := conn.Children(path)
	if err != nil {
		must(err)
		return
	}

	sort.Strings(children)
	var buf [4]byte
	for _, child := range children {
		if path == "/" {
			path = ""
		}

		znode := fmt.Sprintf("%s/%s", path, child)

		// display znode content
		data, stat, err := conn.Get(znode)
		must(err)
		if stat.EphemeralOwner > 0 {
			// ignore ephemeral znodes
			continue
		}

		_, err = this.f.Write([]byte(znode))
		must(err)
		_, err = this.f.Write([]byte{'\n'})
		must(err)
		v := buf[0:4]
		binary.BigEndian.PutUint32(v, uint32(len(data)))
		_, err = this.f.Write(v)

		if len(data) > 0 {
			_, err = this.f.Write(data)
			must(err)
		}

		this.dump(conn, znode)
	}
}

func (*Dump) Synopsis() string {
	return "Dump/Import permanent znodes of Zookeeper"
}

func (this *Dump) Help() string {
	help := fmt.Sprintf(`
Usage: %s dump -z zone [options]

    %s

Options:

    -p path
      znode root path

    -o outfile
      Default zk.dump
      zone name will automatically prefix the final outfile.

    -import dumped file name

    -in dumpped input filename
      Display dumpped file contents in text format.

    -dir dir name
      Run daily dump to this directoy.
      zk will automatically rotate target dumps output.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<commit_msg>example given<commit_after>package command

import (
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"os"
	"path"
	"sort"
	"strings"

	"github.com/funkygao/gafka/ctx"
	gzk "github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/samuel/go-zookeeper/zk"
)

type Dump struct {
	Ui  cli.Ui
	Cmd string

	zone       string
	path       string
	importFile string
	infile     string
	outfile    string
	outdir     string
	f          *os.File
}

func (this *Dump) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("dump", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&this.outfile, "o", "zk.dump", "")
	cmdFlags.StringVar(&this.infile, "in", "", "")
	cmdFlags.StringVar(&this.importFile, "import", "", "")
	cmdFlags.StringVar(&this.path, "p", "/", "")
	cmdFlags.StringVar(&this.outdir, "dir", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		requireAdminRights("-import").
		invalid(args) {
		return 2
	}

	if this.zone == "" {
		this.Ui.Error("unknown zone")
		return 2
	}

	if this.importFile != "" {
		if this.path == "" {
			this.Ui.Error("-p required")
			return 2
		}

		zkzone := gzk.NewZkZone(gzk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
		defer zkzone.Close()

		this.importFromFile(zkzone)
		return
	}

	if this.infile != "" {
		// display mode
		this.diplayDumppedFile()
		return
	}

	// dump mode
	this.outfile = this.zone + "." + this.outfile

	var err error
	if this.outdir != "" {
		err = os.MkdirAll(this.outdir, 0755)
		must(err)

		this.outfile = fmt.Sprintf("%s/%s", this.outdir, this.outfile)

		_, err = os.Lstat(this.outfile)
		if err == nil {
			// file exists, find next available number
			num := 1
			fname := ""
			for ; err == nil && num <= 9999; num++ { // 30 year is enough
				fname = this.outfile + fmt.Sprintf(".%04d", num)
				_, err = os.Lstat(fname)
			}
			if err == nil {
				panic("Not able to rotate, 30 years passed?")
			}

			err = os.Rename(this.outfile, fname)
			must(err)

			this.Ui.Info(fmt.Sprintf("rename %s -> %s", this.outfile, fname))
		}
	}

	this.f, err = os.OpenFile(this.outfile,
		os.O_WRONLY|os.O_CREATE|os.O_APPEND|os.O_TRUNC, 0666)
	must(err)

	zkzone := gzk.NewZkZone(gzk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	defer zkzone.Close()

	this.dump(zkzone.Conn(), this.path)
	this.f.Close()

	this.Ui.Info(fmt.Sprintf("dumpped to %s", this.outfile))

	return
}

func (this *Dump) importFromFile(zkzone *gzk.ZkZone) {
	f, err := os.Open(this.importFile)
	must(err)

	for {
		// read line, got the znode path
		var buf [1]byte
		zpath := make([]byte, 0, 8<<10)
		for {
			b := buf[:]
			_, err := f.Read(b)
			if err == io.EOF {
				return
			}
			must(err)

			if b[0] == '\n' {
				break
			}
			zpath = append(zpath, b[0])
		}

		newPath := path.Join(this.path, string(zpath))
		this.Ui.Infof("%50s -> [%s]%s", string(zpath), this.zone, newPath)

		// read the znode data
		// 1. data len
		// 2. data itself
		var dataLen int32
		err = binary.Read(f, binary.BigEndian, &dataLen)
		must(err)

		zdata := make([]byte, dataLen)
		_, err = io.ReadFull(f, zdata)
		must(err)

		must(zkzone.CreatePermenantZnode(newPath, zdata))
	}
}

func (this *Dump) diplayDumppedFile() {
	f, err := os.Open(this.infile)
	must(err)

	for {
		// read line, got the znode path
		var buf [1]byte
		zpath := make([]byte, 0, 8<<10)
		for {
			b := buf[:]
			_, err := f.Read(b)
			if err == io.EOF {
				return
			}
			must(err)

			if b[0] == '\n' {
				break
			}
			zpath = append(zpath, b[0])
		}

		this.Ui.Info(string(zpath))

		// read the znode data
		// 1. data len
		// 2. data itself
		var dataLen int32
		err = binary.Read(f, binary.BigEndian, &dataLen)
		must(err)

		zdata := make([]byte, dataLen)
		_, err = io.ReadFull(f, zdata)
		must(err)

		this.Ui.Output(string(zdata))
	}
}

// serialization format:
// path
// len(4) value
func (this *Dump) dump(conn *zk.Conn, path string) {
	children, _, err := conn.Children(path)
	if err != nil {
		must(err)
		return
	}

	sort.Strings(children)
	var buf [4]byte
	for _, child := range children {
		if path == "/" {
			path = ""
		}

		znode := fmt.Sprintf("%s/%s", path, child)

		// display znode content
		data, stat, err := conn.Get(znode)
		must(err)
		if stat.EphemeralOwner > 0 {
			// ignore ephemeral znodes
			continue
		}

		_, err = this.f.Write([]byte(znode))
		must(err)
		_, err = this.f.Write([]byte{'\n'})
		must(err)
		v := buf[0:4]
		binary.BigEndian.PutUint32(v, uint32(len(data)))
		_, err = this.f.Write(v)

		if len(data) > 0 {
			_, err = this.f.Write(data)
			must(err)
		}

		this.dump(conn, znode)
	}
}

func (*Dump) Synopsis() string {
	return "Dump/Import permanent znodes of Zookeeper"
}

func (this *Dump) Help() string {
	help := fmt.Sprintf(`
Usage: %s dump -z zone [options]

    %s

Options:

    -p path
      znode root path

    -o outfile
      Default zk.dump
      zone name will automatically prefix the final outfile.

    -import dumped file name
      e,g.
      zk dump -z test -p /xxx/yyy
      zk dump -import test.zk.dump -z prod -p /
      this will clone /xxx/yyy of test zone into prod zone.

    -in dumpped input filename
      Display dumpped file contents in text format.

    -dir dir name
      Run daily dump to this directoy.
      zk will automatically rotate target dumps output.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package ilang

/*
	Scans a numeric statement such as:
		number = 3
		variable += 4
		letter = 'a'
*/
func (ic *Compiler) ScanNumericStatement() {
	var name = ic.Scan(0)
	var numeric = ic.GetVariable(name)

	if name == "error" {
		name = "ERROR"
	}

	var token = ic.Scan(0)

	switch token {
		case "=":
			value := ic.ScanExpression()
			if ic.ExpressionType.Push != "PUSH" {
				ic.RaiseError("Only numeric values can assigned to ",name,".")
			}
			if !ic.ExpressionType.Equals(numeric) {
				ic.RaiseError("Cannot add %s to %s", ic.ExpressionType, numeric)
			}

			ic.Assembly("ADD %v %v %v", name, 0, value)
		default:
			ic.ExpressionType = numeric
			ic.NextToken = token
			ic.Shunt(name)
			if ic.ExpressionType != Undefined {
				ic.RaiseError("blank expression! hmm?")
			}
			ic.ExpressionType = numeric
	}
}
<commit_msg>Mark numeric types when they have been renamed.<commit_after>package ilang

/*
	Scans a numeric statement such as:
		number = 3
		variable += 4
		letter = 'a'
*/
func (ic *Compiler) ScanNumericStatement() {
	var name = ic.Scan(0)
	var numeric = ic.GetVariable(name)

	if name == "error" {
		name = "ERROR"
	}

	var token = ic.Scan(0)

	switch token {
		case "=":
			value := ic.ScanExpression()
			if ic.ExpressionType.Push != "PUSH" {
				ic.RaiseError("Only numeric values can assigned to ",name,".")
			}
			if !ic.ExpressionType.Equals(numeric) {
				ic.RaiseError("Cannot add %s to %s", ic.ExpressionType, numeric)
			}

			ic.Assembly("ADD %v %v %v", name, 0, value)

			ic.MarkVariable(name, "renamed")

		default:
			ic.ExpressionType = numeric
			ic.NextToken = token
			ic.Shunt(name)
			if ic.ExpressionType != Undefined {
				ic.RaiseError("blank expression! hmm?")
			}
			ic.ExpressionType = numeric
	}
}
<|endoftext|>


<commit_before>package http_opentracing

import (
	"log"
	"strings"

	http_ctxtags "github.com/improbable-eng/go-httpwares/tags"
	opentracing "github.com/opentracing/opentracing-go"
)

const (
	TagTraceId           = "trace.traceid"
	TagSpanId            = "trace.spanid"
	TagSampled           = "trace.sampled"
	jaegerNotSampledFlag = "0"
)

// injectOpentracingIdsToTags writes the given context to the ctxtags.
// This is done in an incredibly hacky way, because the public-facing interface of opentracing doesn't give access to
// the TraceId and SpanId of the SpanContext. Only the Tracer's Inject/Extract methods know what these are.
// Most tracers have them encoded as keys with 'traceid' and 'spanid':
// https://github.com/openzipkin/zipkin-go-opentracing/blob/594640b9ef7e5c994e8d9499359d693c032d738c/propagation_ot.go#L29
// https://github.com/opentracing/basictracer-go/blob/1b32af207119a14b1b231d451df3ed04a72efebf/propagation_ot.go#L26
// Jaeger from Uber use one-key schema with next format '{trace-id}:{span-id}:{parent-span-id}:{flags}'
// https://www.jaegertracing.io/docs/client-libraries/#trace-span-identity
func injectOpentracingIdsToTags(span opentracing.Span, tags *http_ctxtags.Tags) {
	if err := span.Tracer().Inject(span.Context(), opentracing.HTTPHeaders, &tagsCarrier{tags}); err != nil {
		log.Printf("http_opentracing: failed extracting trace info into ctx %v", err)
	}
}

// tagsCarrier is a really hacky way of
type tagsCarrier struct {
	*http_ctxtags.Tags
}

func (t *tagsCarrier) Set(key, val string) {
	key = strings.ToLower(key)
	if strings.Contains(key, "traceid") {
		t.Tags.Set(TagTraceId, val) // this will most likely be base-16 (hex) encoded
	}

	if strings.Contains(key, "spanid") && !strings.Contains(strings.ToLower(key), "parent") {
		t.Tags.Set(TagSpanId, val) // this will most likely be base-16 (hex) encoded
	}

	if strings.Contains(key, "sampled") {
		t.Tags.Set(TagSampled, val)
	}

	if key == "uber-trace-id" {
		parts := strings.Split(val, ":")
		if len(parts) == 4 {
			t.Tags.Set(TagTraceId, parts[0])
			t.Tags.Set(TagSpanId, parts[1])

			if parts[3] != jaegerNotSampledFlag {
				t.Tags.Set(TagSampled, "true")
			} else {
				t.Tags.Set(TagSampled, "false")
			}
		}
	}
}
<commit_msg>update to fit with grpc (#53)<commit_after>package http_opentracing

import (
	"log"
	"strings"

	http_ctxtags "github.com/improbable-eng/go-httpwares/tags"
	opentracing "github.com/opentracing/opentracing-go"
)

const (
	TagTraceId           = "trace.traceid"
	TagSpanId            = "trace.spanid"
	TagSampled           = "trace.sampled"
	jaegerNotSampledFlag = "0"
)

// injectOpentracingIdsToTags writes the given context to the ctxtags.
// This is done in an incredibly hacky way, because the public-facing interface of opentracing doesn't give access to
// the TraceId and SpanId of the SpanContext. Only the Tracer's Inject/Extract methods know what these are.
// Most tracers have them encoded as keys with 'traceid' and 'spanid':
// https://github.com/openzipkin/zipkin-go-opentracing/blob/594640b9ef7e5c994e8d9499359d693c032d738c/propagation_ot.go#L29
// https://github.com/opentracing/basictracer-go/blob/1b32af207119a14b1b231d451df3ed04a72efebf/propagation_ot.go#L26
// Jaeger from Uber use one-key schema with next format '{trace-id}:{span-id}:{parent-span-id}:{flags}'
// https://www.jaegertracing.io/docs/client-libraries/#trace-span-identity
func injectOpentracingIdsToTags(span opentracing.Span, tags *http_ctxtags.Tags) {
	if err := span.Tracer().Inject(span.Context(), opentracing.HTTPHeaders, &tagsCarrier{tags}); err != nil {
		log.Printf("http_opentracing: failed extracting trace info into ctx %v", err)
	}
}

// tagsCarrier is a really hacky way of
type tagsCarrier struct {
	*http_ctxtags.Tags
}

func (t *tagsCarrier) Set(key, val string) {
	key = strings.ToLower(key)
	if strings.Contains(key, "traceid") {
		t.Tags.Set(TagTraceId, val) // this will most likely be base-16 (hex) encoded
	}

	if strings.Contains(key, "spanid") && !strings.Contains(strings.ToLower(key), "parent") {
		t.Tags.Set(TagSpanId, val) // this will most likely be base-16 (hex) encoded
	}

	if strings.Contains(key, "sampled") {
		switch val {
		case "true", "false":
			t.Tags.Set(TagSampled, val)
		}
	}

	if key == "uber-trace-id" {
		parts := strings.Split(val, ":")
		if len(parts) == 4 {
			t.Tags.Set(TagTraceId, parts[0])
			t.Tags.Set(TagSpanId, parts[1])

			if parts[3] != jaegerNotSampledFlag {
				t.Tags.Set(TagSampled, "true")
			} else {
				t.Tags.Set(TagSampled, "false")
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/hmac"
	"crypto/md4"
	"crypto/md5"
	"crypto/ripemd160"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"flag"
	"fmt"
	"hash"
	"hash/adler32"
	"hash/crc32"
	"hash/crc64"
	"hash/fnv"
	"io"
	"os"
	"sort"
	"strings"
)

const (
	version   = "1.0"
	chunkSize = 32768 // Read files in 32KB chunks
)

var (
	alg  *string = flag.String("a", "sha256", "algorithm")
	key  *string = flag.String("k", "", "key (for hashes that use a key, e.g. HMAC)")
	salt *string = flag.String("salt", "", "salt")
	hashStr  *bool = flag.Bool("s", false, "hash a string")
)

var (
	algDescs = map[string]string{
		"adler32":         "Adler-32 checksum (RFC 1950)",
		"crc32":           "32-bit cyclic redundancy check (CRC-32) checksum (defaults to IEEE polynomial)",
		"crc32ieee":       "CRC-32 using the IEEE polynomial (0xedb88320)",
		"crc32castagnoli": "CRC-32 using the Castagnoli polynomial (0x82f63b78)",
		"crc32koopman":    "CRC-32 using the Koopman polynomial (0xeb31d82e)",
		"crc64":           "64-bit cyclic redundancy check (CRC-64) checksum (defaults to ISO polynomial)",
		"crc64iso":        "CRC-64 using the ISO polynomial (0xD800000000000000)",
		"crc64ecma":       "CRC-64 using the ECMA polynomial (0xC96C5795D7870F42)",
		"fnv":             "FNV-1 non-cryptographic hash (defaults to fnv32)",
		"fnv32":           "32-bit FNV-1",
		"fnv32a":          "32-bit FNV-1a",
		"fnv64":           "64-bit FNV-1",
		"fnv64a":          "64-bit FNV-1a",
		"hmac":            "Keyed-Hash Message Authentication Code (HMAC) (requires -k <key>) (defaults to SHA-256)",
		"hmacmd5":         "HMAC using MD5 (requires -k <key>)",
		"hmacsha1":        "HMAC using SHA-1 (requires -k <key>)",
		"hmacsha256":      "HMAC using SHA-256 (requires -k <key>)",
		"hmacsha512":      "HMAC using SHA-512 (requires -k <key>)",
		"md4":             "MD4 hash (RFC 1320)",
		"md5":             "MD5 hash (RFC 1321)",
		"ripemd160":       "RIPEMD-160 hash",
		"sha1":            "SHA-1 hash (RFC 3174)",
		"sha224":          "SHA-224 hash (FIPS 180-2)",
		"sha256":          "SHA-256 hash (FIPS 180-2)",
		"sha384":          "SHA-384 hash (FIPS 180-2)",
		"sha512":          "SHA-512 hash (FIPS 180-2)",
	}
)

func GetGenerator(a string) (hash.Hash, os.Error) {
	var g hash.Hash
	switch a {
	default:
		return md5.New(), os.NewError("Invalid algorithm")
	case "adler32":
		g = adler32.New()
	case "crc32", "crc32ieee":
		g = crc32.New(crc32.MakeTable(crc32.IEEE))
	case "crc32castagnoli":
		g = crc32.New(crc32.MakeTable(crc32.Castagnoli))
	case "crc32koopman":
		g = crc32.New(crc32.MakeTable(crc32.Koopman))
	case "crc64", "crc64iso":
		g = crc64.New(crc64.MakeTable(crc64.ISO))
	case "crc64ecma":
		g = crc64.New(crc64.MakeTable(crc64.ECMA))
	case "fnv", "fnv32":
		g = fnv.New32()
	case "fnv32a":
		g = fnv.New32a()
	case "fnv64":
		g = fnv.New64()
	case "fnv64a":
		g = fnv.New64a()
	case "hmac", "hmacsha256":
		g = hmac.NewSHA256([]byte(*key))
	case "hmacmd5":
		g = hmac.NewMD5([]byte(*key))
	case "hmacsha1":
		g = hmac.NewSHA1([]byte(*key))
	case "hmacsha512":
		g = hmac.New(sha512.New, []byte(*key))
	case "md4":
		g = md4.New()
	case "md5":
		g = md5.New()
	case "ripemd160":
		g = ripemd160.New()
	case "sha1":
		g = sha1.New()
	case "sha224":
		g = sha256.New224()
	case "sha256":
		g = sha256.New()
	case "sha384":
		g = sha512.New384()
	case "sha512":
		g = sha512.New()
	}
	return g, nil
}

func HashString(gen hash.Hash, s string) string {
	gen.Write([]byte(s))
	return fmt.Sprintf("%x", gen.Sum())
}

func HashFile(gen hash.Hash, f io.Reader) (string, os.Error) {
	buf := make([]byte, chunkSize)
	gen.Write([]byte(*salt))
	for {
		bytesRead, err := f.Read(buf)
		if err != nil {
			if err == os.EOF && bytesRead == 0 { // Empty file
				gen.Write([]byte(""))
			} else {
				return "", err
			}
		} else {
			gen.Write(buf[:bytesRead])
		}
		if bytesRead < chunkSize { // EOF
			break
		}
	}
	return fmt.Sprintf("%x", gen.Sum()), nil
}

func Usage() {
	fmt.Println("Usage:", os.Args[0], "[-a <algorithm>] [-salt <salt>] [-s <string to hash>] / <file(s) to hash>\n")
	fmt.Println("Examples:")
	fmt.Println(" ", os.Args[0], `-a md5 document.txt               `, "Generate MD5 hash of a file")
	fmt.Println(" ", os.Args[0], `-a md5 *                          `, "Generate MD5 hash of all files in folder")
	fmt.Println(" ", os.Args[0], `-a sha1 -s hello world            `, "Generate SHA-1 hash of a string")
	fmt.Println(" ", os.Args[0], `-a sha1 -salt s4lt -s hello world `, "Generate salted SHA-1 hash of a string")
	fmt.Println(" ", os.Args[0], `-a hmacsha1 -k k3y -s untouched   `, "Generate HMAC of a string using SHA-1")
	fmt.Println("")
	fmt.Println("Available algorithms (default is SHA-256):")
	mk := make([]string, len(algDescs))
	i := 0
	for k, _ := range algDescs {
		mk[i] = k
		i++
	}
	sort.Strings(mk)
	for _, v := range mk {
		fmt.Printf("  %-15s %s\n", v, algDescs[v])
	}
	fmt.Println("")
	fmt.Println(`Note: For complex strings, put the string in a file, then run Picugen on
      the file. (Don't add newlines to the file as they will alter the output.)`)
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		Usage()
		return
	}
	*alg = strings.ToLower(*alg)
	gen, err := GetGenerator(*alg)
	if err != nil {
		fmt.Println("Error:", err)
		return
	}
	if *hashStr {
		var s string
		for i, word := range flag.Args() {
			if i > 0 {
				s += " "
			}
			s += word
		}
		fmt.Println(HashString(gen, *salt+s))
	} else {
		for _, path := range flag.Args() {
			var res string
			f, err := os.Open(path)
			if err != nil {
				res = err.String()
			} else {
				h, err := HashFile(gen, f)
				if err != nil {
					res = err.String()
				} else {
					res = h
				}
			}
			fmt.Println(res, "", path)
			f.Close()
			gen.Reset()
		}
	}
}
<commit_msg>Clarified terminology (hashes=digests)<commit_after>package main

import (
	"crypto/hmac"
	"crypto/md4"
	"crypto/md5"
	"crypto/ripemd160"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"flag"
	"fmt"
	"hash"
	"hash/adler32"
	"hash/crc32"
	"hash/crc64"
	"hash/fnv"
	"io"
	"os"
	"sort"
	"strings"
)

const (
	version   = "1.0"
	chunkSize = 32768 // Read files in 32KB chunks
)

var (
	alg  *string = flag.String("a", "sha256", "algorithm")
	key  *string = flag.String("k", "", "key (for hashes that use a key, e.g. HMAC)")
	salt *string = flag.String("salt", "", "salt")
	hashStr  *bool = flag.Bool("s", false, "hash a string")
)

var (
	algDescs = map[string]string{
		"adler32":         "Adler-32 checksum (RFC 1950)",
		"crc32":           "32-bit cyclic redundancy check (CRC-32) checksum (defaults to IEEE polynomial)",
		"crc32ieee":       "CRC-32 using the IEEE polynomial (0xedb88320)",
		"crc32castagnoli": "CRC-32 using the Castagnoli polynomial (0x82f63b78)",
		"crc32koopman":    "CRC-32 using the Koopman polynomial (0xeb31d82e)",
		"crc64":           "64-bit cyclic redundancy check (CRC-64) checksum (defaults to ISO polynomial)",
		"crc64iso":        "CRC-64 using the ISO polynomial (0xD800000000000000)",
		"crc64ecma":       "CRC-64 using the ECMA polynomial (0xC96C5795D7870F42)",
		"fnv":             "FNV-1 non-cryptographic hash (defaults to fnv32)",
		"fnv32":           "32-bit FNV-1",
		"fnv32a":          "32-bit FNV-1a",
		"fnv64":           "64-bit FNV-1",
		"fnv64a":          "64-bit FNV-1a",
		"hmac":            "Keyed-Hash Message Authentication Code (HMAC) (requires -k <key>) (defaults to SHA-256)",
		"hmacmd5":         "HMAC using MD5 (requires -k <key>)",
		"hmacsha1":        "HMAC using SHA-1 (requires -k <key>)",
		"hmacsha256":      "HMAC using SHA-256 (requires -k <key>)",
		"hmacsha512":      "HMAC using SHA-512 (requires -k <key>)",
		"md4":             "MD4 hash (RFC 1320)",
		"md5":             "MD5 hash (RFC 1321)",
		"ripemd160":       "RIPEMD-160 hash",
		"sha1":            "SHA-1 hash (RFC 3174)",
		"sha224":          "SHA-224 hash (FIPS 180-2)",
		"sha256":          "SHA-256 hash (FIPS 180-2)",
		"sha384":          "SHA-384 hash (FIPS 180-2)",
		"sha512":          "SHA-512 hash (FIPS 180-2)",
	}
)

func GetGenerator(a string) (hash.Hash, os.Error) {
	var g hash.Hash
	switch a {
	default:
		return md5.New(), os.NewError("Invalid algorithm")
	case "adler32":
		g = adler32.New()
	case "crc32", "crc32ieee":
		g = crc32.New(crc32.MakeTable(crc32.IEEE))
	case "crc32castagnoli":
		g = crc32.New(crc32.MakeTable(crc32.Castagnoli))
	case "crc32koopman":
		g = crc32.New(crc32.MakeTable(crc32.Koopman))
	case "crc64", "crc64iso":
		g = crc64.New(crc64.MakeTable(crc64.ISO))
	case "crc64ecma":
		g = crc64.New(crc64.MakeTable(crc64.ECMA))
	case "fnv", "fnv32":
		g = fnv.New32()
	case "fnv32a":
		g = fnv.New32a()
	case "fnv64":
		g = fnv.New64()
	case "fnv64a":
		g = fnv.New64a()
	case "hmac", "hmacsha256":
		g = hmac.NewSHA256([]byte(*key))
	case "hmacmd5":
		g = hmac.NewMD5([]byte(*key))
	case "hmacsha1":
		g = hmac.NewSHA1([]byte(*key))
	case "hmacsha512":
		g = hmac.New(sha512.New, []byte(*key))
	case "md4":
		g = md4.New()
	case "md5":
		g = md5.New()
	case "ripemd160":
		g = ripemd160.New()
	case "sha1":
		g = sha1.New()
	case "sha224":
		g = sha256.New224()
	case "sha256":
		g = sha256.New()
	case "sha384":
		g = sha512.New384()
	case "sha512":
		g = sha512.New()
	}
	return g, nil
}

func HashString(gen hash.Hash, s string) string {
	gen.Write([]byte(s))
	return fmt.Sprintf("%x", gen.Sum())
}

func HashFile(gen hash.Hash, f io.Reader) (string, os.Error) {
	buf := make([]byte, chunkSize)
	gen.Write([]byte(*salt))
	for {
		bytesRead, err := f.Read(buf)
		if err != nil {
			if err == os.EOF && bytesRead == 0 { // Empty file
				gen.Write([]byte(""))
			} else {
				return "", err
			}
		} else {
			gen.Write(buf[:bytesRead])
		}
		if bytesRead < chunkSize { // EOF
			break
		}
	}
	return fmt.Sprintf("%x", gen.Sum()), nil
}

func Usage() {
	fmt.Println("Usage:", os.Args[0], "[-a <algorithm>] [-salt <salt>] [-s <string to hash>] / <file(s) to hash>\n")
	fmt.Println("Examples:")
	fmt.Println(" ", os.Args[0], `-a md5 document.txt               `, "Generate MD5 digest of a file")
	fmt.Println(" ", os.Args[0], `-a md5 *                          `, "Generate MD5 digests of all files in folder")
	fmt.Println(" ", os.Args[0], `-a sha1 -s hello world            `, "Generate SHA-1 digest of a string")
	fmt.Println(" ", os.Args[0], `-a sha1 -salt s4lt -s hello world `, "Generate salted SHA-1 digest of a string")
	fmt.Println(" ", os.Args[0], `-a hmacsha1 -k k3y -s untouched   `, "Generate HMAC of a string using SHA-1")
	fmt.Println("")
	fmt.Println("Available algorithms (default is SHA-256):")
	mk := make([]string, len(algDescs))
	i := 0
	for k, _ := range algDescs {
		mk[i] = k
		i++
	}
	sort.Strings(mk)
	for _, v := range mk {
		fmt.Printf("  %-15s %s\n", v, algDescs[v])
	}
	fmt.Println("")
	fmt.Println(`Note: For complex strings, put the string in a file, then run Picugen on
      the file. (Don't add newlines to the file as they will alter the output.)`)
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		Usage()
		return
	}
	*alg = strings.ToLower(*alg)
	gen, err := GetGenerator(*alg)
	if err != nil {
		fmt.Println("Error:", err)
		return
	}
	if *hashStr {
		var s string
		for i, word := range flag.Args() {
			if i > 0 {
				s += " "
			}
			s += word
		}
		fmt.Println(HashString(gen, *salt+s))
	} else {
		for _, path := range flag.Args() {
			var res string
			f, err := os.Open(path)
			if err != nil {
				res = err.String()
			} else {
				h, err := HashFile(gen, f)
				if err != nil {
					res = err.String()
				} else {
					res = h
				}
			}
			fmt.Println(res, "", path)
			f.Close()
			gen.Reset()
		}
	}
}
<|endoftext|>


<commit_before>package functional

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"reflect"
	"testing"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/systemd"
	"github.com/coreos/fleet/unit"
)

func TestSystemdUnitFlow(t *testing.T) {
	uDir, err := ioutil.TempDir("", "fleet-")
	if err != nil {
		t.Fatalf("Failed creating tempdir: %v", err)
	}
	defer os.RemoveAll(uDir)

	mgr, err := systemd.NewSystemdManager(uDir)
	if err != nil {
		t.Fatalf("Failed initializing SystemdManager: %v", err)
	}

	units, err := mgr.Units()
	if err != nil {
		t.Fatalf("Failed calling Units(): %v", err)
	}

	if len(units) > 0 {
		t.Fatalf("Expected no units to be returned, got %v", units)
	}

	name := fmt.Sprintf("fleet-unit-%d.service", rand.Int63())
	uf := unit.NewUnit(`[Service]
ExecStart=/usr/bin/sleep 3000
`)
	j := job.NewJob(name, *uf)

	if err := mgr.Load(j.Name, j.Unit); err != nil {
		t.Fatalf("Failed loading job: %v", err)
	}

	units, err = mgr.Units()
	if err != nil {
		t.Fatalf("Failed calling Units(): %v", err)
	}

	if !reflect.DeepEqual([]string{name}, units){
		t.Fatalf("Expected [hello.service], got %v", units)
	}

	us, err := mgr.GetUnitState(name)
	if err == nil {
		expect := unit.UnitState{"loaded", "inactive", "dead", nil}
		if !reflect.DeepEqual(expect, *us) {
			t.Errorf("Expected UnitState %v, got %v", expect, *us)
		}
	} else {
		t.Errorf("Failed determining unit state: %v", err)
	}

	mgr.Unload(name)

	units, err = mgr.Units()
	if err != nil {
		t.Fatalf("Failed calling Units(): %v", err)
	}

	if len(units) > 0 {
		t.Fatalf("Expected no units to be returned, got %v", units)
	}
}
<commit_msg>test(systemd): Start and stop units, asserting state in between<commit_after>package functional

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"reflect"
	"testing"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/systemd"
	"github.com/coreos/fleet/unit"
)

func TestSystemdUnitFlow(t *testing.T) {
	uDir, err := ioutil.TempDir("", "fleet-")
	if err != nil {
		t.Fatalf("Failed creating tempdir: %v", err)
	}
	defer os.RemoveAll(uDir)

	mgr, err := systemd.NewSystemdManager(uDir)
	if err != nil {
		t.Fatalf("Failed initializing SystemdManager: %v", err)
	}

	units, err := mgr.Units()
	if err != nil {
		t.Fatalf("Failed calling Units(): %v", err)
	}

	if len(units) > 0 {
		t.Fatalf("Expected no units to be returned, got %v", units)
	}

	name := fmt.Sprintf("fleet-unit-%d.service", rand.Int63())
	uf := unit.NewUnit(`[Service]
ExecStart=/usr/bin/sleep 3000
`)
	j := job.NewJob(name, *uf)

	if err := mgr.Load(j.Name, j.Unit); err != nil {
		t.Fatalf("Failed loading job: %v", err)
	}

	units, err = mgr.Units()
	if err != nil {
		t.Fatalf("Failed calling Units(): %v", err)
	}

	if !reflect.DeepEqual([]string{name}, units){
		t.Fatalf("Expected [hello.service], got %v", units)
	}

	us, err := mgr.GetUnitState(name)
	if err == nil {
		expect := unit.UnitState{"loaded", "inactive", "dead", nil}
		if !reflect.DeepEqual(expect, *us) {
			t.Errorf("Expected UnitState %v, got %v", expect, *us)
		}
	} else {
		t.Errorf("Failed determining unit state: %v", err)
	}

	mgr.Start(name)

	us, err = mgr.GetUnitState(name)
	if err == nil {
		expect := unit.UnitState{"loaded", "active", "running", nil}
		if !reflect.DeepEqual(expect, *us) {
			t.Errorf("Expected UnitState %v, got %v", expect, *us)
		}
	} else {
		t.Errorf("Failed determining unit state: %v", err)
	}

	mgr.Stop(name)

	us, err = mgr.GetUnitState(name)
	if err == nil {
		expect := unit.UnitState{"loaded", "inactive", "dead", nil}
		if !reflect.DeepEqual(expect, *us) {
			t.Errorf("Expected UnitState %v, got %v", expect, *us)
		}
	} else {
		t.Errorf("Failed determining unit state: %v", err)
	}

	mgr.Unload(name)

	units, err = mgr.Units()
	if err != nil {
		t.Fatalf("Failed calling Units(): %v", err)
	}

	if len(units) > 0 {
		t.Fatalf("Expected no units to be returned, got %v", units)
	}
}
<|endoftext|>


<commit_before>package common

import (
	"sort"
	"strings"

	"go.skia.org/infra/go/sklog"
)

const (
	TEST_HARNESS_NAME = "fuzz"

	UNKNOWN_FUNCTION    = "UNKNOWN"
	UNKNOWN_FILE        = "UNKNOWN"
	ASSEMBLY_CODE_FILE  = "[assembly code]"
	UNSYMBOLIZED_RESULT = "UNSYMBOLIZED"
	UNKNOWN_LINE        = -1

	ASAN_OPTIONS        = "ASAN_OPTIONS=detect_leaks=0 symbolize=1"
	STABLE_FUZZER       = "stable"
	EXPERIMENTAL_FUZZER = "experimental"
	FUZZER_NOT_FOUND    = "FUZZER_NOT_FOUND"

	UNCLAIMED = "<unclaimed>"
)

// The list of architectures we fuzz on
var ARCHITECTURES = []string{"linux_x64"}

// By default, allow a generous amount of RAM and set a flexible approach to timeouts.
var defaultGenerationArgs = []string{"-m", "5000", "-t", "500+"}

var commonImpl CommonImpl

// FuzzerInfo contains all the configuration needed to display, execute and analyze a fuzzer.
type FuzzerInfo struct {
	// PrettyName is the human readable name for this fuzzer.
	PrettyName string
	// Status should be STABLE_FUZZER or EXPERIMENTAL_FUZZER
	Status string
	// The Groomer is responsible for stabilizing experimental fuzzers and will be the point
	// of contact for the sheriff for triaging any bad fuzzes found in stable fuzzers.
	Groomer string
	// ExtraBugLabels are any additional labels that should be included when making a bug.
	ExtraBugLabels []string
	// ArgsAfterExecutable is a slice of arguments that come after the executable
	// and before the path to the bytes file, that will be fuzzed.
	ArgsAfterExecutable []string
	// GenerationArgs is a slice of arguments that are used to adjust timeouts/memory usage
	// when generating files.
	GenerationArgs []string
}

// fuzzers is a map of fuzzer_name -> FuzzerInfo for all registered fuzzers.  This should be a
// centralized location to add a new fuzzer, i.e. adding an entry and information here, combined
// with modifying the fuzzer-be.service should be sufficient to add a new fuzzer into the system.
var fuzzers = map[string]FuzzerInfo{
	"api_draw_functions": {
		PrettyName:          "API - CanvasDrawFunctions",
		Status:              STABLE_FUZZER,
		Groomer:             "hcm",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "DrawFunctions", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_gradient": {
		PrettyName:          "API - Gradients",
		Status:              STABLE_FUZZER,
		Groomer:             "fmalita",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "Gradients", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_image_filter": {
		PrettyName:          "API - SerializedImageFilter",
		Status:              STABLE_FUZZER,
		Groomer:             "robertphillips",
		ExtraBugLabels:      []string{"Area-ImageFilter"},
		ArgsAfterExecutable: []string{"--type", "api", "--name", "SerializedImageFilter", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_parse_path": {
		PrettyName:          "API - ParsePath",
		Status:              STABLE_FUZZER,
		Groomer:             "caryclark",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "ParsePath", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_path_measure": {
		PrettyName:          "API - PathMeasure",
		Status:              STABLE_FUZZER,
		Groomer:             "caryclark",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "PathMeasure", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_pathop": {
		PrettyName:          "API - PathOp",
		Status:              STABLE_FUZZER,
		Groomer:             "caryclark",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "Pathop", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"color_deserialize": {
		PrettyName:          "SkColorSpace - Deserialize",
		Status:              STABLE_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      []string{"Area-ImageDecoder"},
		ArgsAfterExecutable: []string{"--type", "color_deserialize", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"image_filter_deserialize": {
		PrettyName:          "FilterFuzz Stub",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "reed",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "filter_fuzz", "--bytes"},
		GenerationArgs:      []string{"-m", "5000", "-t", "200+"},
	},
	"jpeg_encoder": {
		PrettyName:          "JPEG encoder",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "JPEGEncoder", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"mock_gpu_canvas": {
		PrettyName:          "Canvas to mock GL backend",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "robertphillips",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "MockGPUCanvas", "--bytes"},
		// For most of the canvases, some of the initial test cases take a long time. If they
		// time out (1000+ms ), afl-fuzz aborts during startup. To keep things moving, we
		// tell afl-fuzz to set the timout to 500ms and ignore any test cases that take longer
		// than that (the plus sign after 500).
		GenerationArgs: []string{"-m", "5000", "-t", "500+"},
	},
	"n32_canvas": {
		PrettyName:          "Canvas to raster n32 backend",
		Status:              STABLE_FUZZER,
		Groomer:             "halcanary",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "RasterN32Canvas", "--bytes"},
		GenerationArgs:      []string{"-m", "5000", "-t", "500+"},
	},
	"null_canvas": {
		PrettyName:          "Canvas to null canvas backend",
		Status:              STABLE_FUZZER,
		Groomer:             "halcanary",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "NullCanvas", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"null_gl_canvas": {
		PrettyName:          "Canvas to null gl canvas backend",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "robertphillips",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "NullGLCanvas", "--bytes"},
		GenerationArgs:      []string{"-m", "5000", "-t", "500+"},
	},
	"path_deserialize": {
		PrettyName:          "SkPath deserialize",
		Status:              STABLE_FUZZER,
		Groomer:             "reed",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "path_deserialize", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"pdf_canvas": {
		PrettyName:          "Canvas to PDF backend",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "halcanary",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "PDFCanvas", "--bytes"},
		GenerationArgs:      []string{"-m", "5000", "-t", "500+"},
	},
	"png_encoder": {
		PrettyName:          "PNG encoder",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "PNGEncoder", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"region_deserialize": {
		PrettyName:          "SkRegion deserialize",
		Status:              STABLE_FUZZER,
		Groomer:             "halcanary",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "region_deserialize", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"region_set_path": {
		PrettyName:          "SkRegion set_path",
		Status:              STABLE_FUZZER,
		Groomer:             "reed",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "region_set_path", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"skcodec_scale": {
		PrettyName:          "SkCodec (Scaling)",
		Status:              STABLE_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      []string{"Area-ImageDecoder"},
		ArgsAfterExecutable: []string{"--type", "image_scale", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"skcodec_mode": {
		PrettyName:          "SkCodec (Modes)",
		Status:              STABLE_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      []string{"Area-ImageDecoder"},
		ArgsAfterExecutable: []string{"--type", "image_mode", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"sksl2glsl": {
		PrettyName:          "SKSL Compiler (GLSL)",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "ethannicholas",
		ExtraBugLabels:      []string{},
		ArgsAfterExecutable: []string{"--type", "sksl2glsl", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"skottie_json": {
		PrettyName:          "Skottie from JSON",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "fmalita",
		ExtraBugLabels:      []string{},
		ArgsAfterExecutable: []string{"--type", "skottie_json", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"skp": {
		PrettyName:          "SKP from ReadBuffer",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "reed",
		ExtraBugLabels:      []string{},
		ArgsAfterExecutable: []string{"--type", "skp", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"textblob": {
		PrettyName:          "TextBlob deserialize",
		Status:              STABLE_FUZZER,
		Groomer:             "fmalita",
		ExtraBugLabels:      []string{},
		ArgsAfterExecutable: []string{"--type", "textblob", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"webp_encoder": {
		PrettyName:          "WEBP encoder",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "WEBPEncoder", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
}

// FUZZ_CATEGORIES is an alphabetized list of known fuzz categories.
var FUZZ_CATEGORIES = []string{}

var ANALYSIS_TYPES = []string{"ASAN_RELEASE", "ASAN_DEBUG", "CLANG_RELEASE", "CLANG_DEBUG"}

// TODO(kjlubick): Is it necessary to have two separate lists?
var STACKTRACE_ORDER = []string{"ASAN_RELEASE", "CLANG_RELEASE", "ASAN_DEBUG", "CLANG_DEBUG"}

func init() {
	commonImpl = &defaultImpl{}
	for k := range fuzzers {
		FUZZ_CATEGORIES = append(FUZZ_CATEGORIES, k)
	}
	sort.Strings(FUZZ_CATEGORIES)
}

func PrettifyCategory(category string) string {
	f, found := fuzzers[category]
	if !found {
		sklog.Errorf("Unknown category %s", category)
		return FUZZER_NOT_FOUND
	}
	return f.PrettyName
}

func ExtraBugLabels(category string) []string {
	f, found := fuzzers[category]
	if !found {
		sklog.Errorf("Unknown category %s", category)
		return nil
	}
	return f.ExtraBugLabels
}

// ReplicationArgs returns a space separated list of the args needed to replicate the crash
// of a fuzz of a given category.
func ReplicationArgs(category string) string {
	f, found := fuzzers[category]
	if !found {
		sklog.Errorf("Unknown category %s", category)
		return FUZZER_NOT_FOUND
	}
	return strings.Join(f.ArgsAfterExecutable, " ")
}

// CategoryReminder returns a short string to remind a human what fuzz category
// a fuzz belongs to.  For example "skcodec_scale" -> "image_scale" and
// "pdf_canvas" -> "api-PDFCanvas"
func CategoryReminder(category string) string {
	f, found := fuzzers[category]
	if !found {
		sklog.Errorf("Unknown category %s", category)
		return FUZZER_NOT_FOUND
	}
	// assume ['--type', type, {'--name', api_name}]
	cr := f.ArgsAfterExecutable[1]
	if cr == "api" {
		cr += "-" + f.ArgsAfterExecutable[3]
	}
	return cr
}

// HasCategory returns if a given string corresponds to a known fuzzer category.
func HasCategory(c string) bool {
	_, found := fuzzers[c]
	return found
}

// Status returns the status of a fuzz category (i.e. stable, experimental, etc)
func Status(c string) string {
	f, found := fuzzers[c]
	if !found {
		sklog.Errorf("Unknown category %s", c)
		return FUZZER_NOT_FOUND
	}
	return f.Status
}

// Groomer returns the groomer of a fuzz category
func Groomer(c string) string {
	f, found := fuzzers[c]
	if !found {
		sklog.Errorf("Unknown category %s", c)
		return FUZZER_NOT_FOUND
	}
	return f.Groomer
}

// Returns if fuzzer knows about a given architecture
func HasArchitecture(a string) bool {
	for _, ar := range ARCHITECTURES {
		if a == ar {
			return true
		}
	}
	return false
}

func SetMockCommon(c CommonImpl) {
	commonImpl = c
}

// Returns the Hostname of this machine. Clients should use this instead of
// os.Hostname because this function can be mocked out.
func Hostname() string {
	return commonImpl.Hostname()
}
<commit_msg>Add option for polyutils fuzzer<commit_after>package common

import (
	"sort"
	"strings"

	"go.skia.org/infra/go/sklog"
)

const (
	TEST_HARNESS_NAME = "fuzz"

	UNKNOWN_FUNCTION    = "UNKNOWN"
	UNKNOWN_FILE        = "UNKNOWN"
	ASSEMBLY_CODE_FILE  = "[assembly code]"
	UNSYMBOLIZED_RESULT = "UNSYMBOLIZED"
	UNKNOWN_LINE        = -1

	ASAN_OPTIONS        = "ASAN_OPTIONS=detect_leaks=0 symbolize=1"
	STABLE_FUZZER       = "stable"
	EXPERIMENTAL_FUZZER = "experimental"
	FUZZER_NOT_FOUND    = "FUZZER_NOT_FOUND"

	UNCLAIMED = "<unclaimed>"
)

// The list of architectures we fuzz on
var ARCHITECTURES = []string{"linux_x64"}

// By default, allow a generous amount of RAM and set a flexible approach to timeouts.
var defaultGenerationArgs = []string{"-m", "5000", "-t", "500+"}

var commonImpl CommonImpl

// FuzzerInfo contains all the configuration needed to display, execute and analyze a fuzzer.
type FuzzerInfo struct {
	// PrettyName is the human readable name for this fuzzer.
	PrettyName string
	// Status should be STABLE_FUZZER or EXPERIMENTAL_FUZZER
	Status string
	// The Groomer is responsible for stabilizing experimental fuzzers and will be the point
	// of contact for the sheriff for triaging any bad fuzzes found in stable fuzzers.
	Groomer string
	// ExtraBugLabels are any additional labels that should be included when making a bug.
	ExtraBugLabels []string
	// ArgsAfterExecutable is a slice of arguments that come after the executable
	// and before the path to the bytes file, that will be fuzzed.
	ArgsAfterExecutable []string
	// GenerationArgs is a slice of arguments that are used to adjust timeouts/memory usage
	// when generating files.
	GenerationArgs []string
}

// fuzzers is a map of fuzzer_name -> FuzzerInfo for all registered fuzzers.  This should be a
// centralized location to add a new fuzzer, i.e. adding an entry and information here, combined
// with modifying the fuzzer-be.service should be sufficient to add a new fuzzer into the system.
var fuzzers = map[string]FuzzerInfo{
	"api_draw_functions": {
		PrettyName:          "API - CanvasDrawFunctions",
		Status:              STABLE_FUZZER,
		Groomer:             "hcm",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "DrawFunctions", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_gradient": {
		PrettyName:          "API - Gradients",
		Status:              STABLE_FUZZER,
		Groomer:             "fmalita",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "Gradients", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_image_filter": {
		PrettyName:          "API - SerializedImageFilter",
		Status:              STABLE_FUZZER,
		Groomer:             "robertphillips",
		ExtraBugLabels:      []string{"Area-ImageFilter"},
		ArgsAfterExecutable: []string{"--type", "api", "--name", "SerializedImageFilter", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_parse_path": {
		PrettyName:          "API - ParsePath",
		Status:              STABLE_FUZZER,
		Groomer:             "caryclark",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "ParsePath", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_path_measure": {
		PrettyName:          "API - PathMeasure",
		Status:              STABLE_FUZZER,
		Groomer:             "caryclark",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "PathMeasure", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_pathop": {
		PrettyName:          "API - PathOp",
		Status:              STABLE_FUZZER,
		Groomer:             "caryclark",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "Pathop", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"api_pathop": {
		PrettyName:          "API - PolyUtils",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "jvanverth",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "PolyUtils", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"color_deserialize": {
		PrettyName:          "SkColorSpace - Deserialize",
		Status:              STABLE_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      []string{"Area-ImageDecoder"},
		ArgsAfterExecutable: []string{"--type", "color_deserialize", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"image_filter_deserialize": {
		PrettyName:          "FilterFuzz Stub",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "reed",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "filter_fuzz", "--bytes"},
		GenerationArgs:      []string{"-m", "5000", "-t", "200+"},
	},
	"jpeg_encoder": {
		PrettyName:          "JPEG encoder",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "JPEGEncoder", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"mock_gpu_canvas": {
		PrettyName:          "Canvas to mock GL backend",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "robertphillips",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "MockGPUCanvas", "--bytes"},
		// For most of the canvases, some of the initial test cases take a long time. If they
		// time out (1000+ms ), afl-fuzz aborts during startup. To keep things moving, we
		// tell afl-fuzz to set the timout to 500ms and ignore any test cases that take longer
		// than that (the plus sign after 500).
		GenerationArgs: []string{"-m", "5000", "-t", "500+"},
	},
	"n32_canvas": {
		PrettyName:          "Canvas to raster n32 backend",
		Status:              STABLE_FUZZER,
		Groomer:             "halcanary",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "RasterN32Canvas", "--bytes"},
		GenerationArgs:      []string{"-m", "5000", "-t", "500+"},
	},
	"null_canvas": {
		PrettyName:          "Canvas to null canvas backend",
		Status:              STABLE_FUZZER,
		Groomer:             "halcanary",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "NullCanvas", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"null_gl_canvas": {
		PrettyName:          "Canvas to null gl canvas backend",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "robertphillips",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "NullGLCanvas", "--bytes"},
		GenerationArgs:      []string{"-m", "5000", "-t", "500+"},
	},
	"path_deserialize": {
		PrettyName:          "SkPath deserialize",
		Status:              STABLE_FUZZER,
		Groomer:             "reed",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "path_deserialize", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"pdf_canvas": {
		PrettyName:          "Canvas to PDF backend",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "halcanary",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "PDFCanvas", "--bytes"},
		GenerationArgs:      []string{"-m", "5000", "-t", "500+"},
	},
	"png_encoder": {
		PrettyName:          "PNG encoder",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "PNGEncoder", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"region_deserialize": {
		PrettyName:          "SkRegion deserialize",
		Status:              STABLE_FUZZER,
		Groomer:             "halcanary",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "region_deserialize", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"region_set_path": {
		PrettyName:          "SkRegion set_path",
		Status:              STABLE_FUZZER,
		Groomer:             "reed",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "region_set_path", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"skcodec_scale": {
		PrettyName:          "SkCodec (Scaling)",
		Status:              STABLE_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      []string{"Area-ImageDecoder"},
		ArgsAfterExecutable: []string{"--type", "image_scale", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"skcodec_mode": {
		PrettyName:          "SkCodec (Modes)",
		Status:              STABLE_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      []string{"Area-ImageDecoder"},
		ArgsAfterExecutable: []string{"--type", "image_mode", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"sksl2glsl": {
		PrettyName:          "SKSL Compiler (GLSL)",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "ethannicholas",
		ExtraBugLabels:      []string{},
		ArgsAfterExecutable: []string{"--type", "sksl2glsl", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"skottie_json": {
		PrettyName:          "Skottie from JSON",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "fmalita",
		ExtraBugLabels:      []string{},
		ArgsAfterExecutable: []string{"--type", "skottie_json", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"skp": {
		PrettyName:          "SKP from ReadBuffer",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "reed",
		ExtraBugLabels:      []string{},
		ArgsAfterExecutable: []string{"--type", "skp", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"textblob": {
		PrettyName:          "TextBlob deserialize",
		Status:              STABLE_FUZZER,
		Groomer:             "fmalita",
		ExtraBugLabels:      []string{},
		ArgsAfterExecutable: []string{"--type", "textblob", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
	"webp_encoder": {
		PrettyName:          "WEBP encoder",
		Status:              EXPERIMENTAL_FUZZER,
		Groomer:             "scroggo",
		ExtraBugLabels:      nil,
		ArgsAfterExecutable: []string{"--type", "api", "--name", "WEBPEncoder", "--bytes"},
		GenerationArgs:      defaultGenerationArgs,
	},
}

// FUZZ_CATEGORIES is an alphabetized list of known fuzz categories.
var FUZZ_CATEGORIES = []string{}

var ANALYSIS_TYPES = []string{"ASAN_RELEASE", "ASAN_DEBUG", "CLANG_RELEASE", "CLANG_DEBUG"}

// TODO(kjlubick): Is it necessary to have two separate lists?
var STACKTRACE_ORDER = []string{"ASAN_RELEASE", "CLANG_RELEASE", "ASAN_DEBUG", "CLANG_DEBUG"}

func init() {
	commonImpl = &defaultImpl{}
	for k := range fuzzers {
		FUZZ_CATEGORIES = append(FUZZ_CATEGORIES, k)
	}
	sort.Strings(FUZZ_CATEGORIES)
}

func PrettifyCategory(category string) string {
	f, found := fuzzers[category]
	if !found {
		sklog.Errorf("Unknown category %s", category)
		return FUZZER_NOT_FOUND
	}
	return f.PrettyName
}

func ExtraBugLabels(category string) []string {
	f, found := fuzzers[category]
	if !found {
		sklog.Errorf("Unknown category %s", category)
		return nil
	}
	return f.ExtraBugLabels
}

// ReplicationArgs returns a space separated list of the args needed to replicate the crash
// of a fuzz of a given category.
func ReplicationArgs(category string) string {
	f, found := fuzzers[category]
	if !found {
		sklog.Errorf("Unknown category %s", category)
		return FUZZER_NOT_FOUND
	}
	return strings.Join(f.ArgsAfterExecutable, " ")
}

// CategoryReminder returns a short string to remind a human what fuzz category
// a fuzz belongs to.  For example "skcodec_scale" -> "image_scale" and
// "pdf_canvas" -> "api-PDFCanvas"
func CategoryReminder(category string) string {
	f, found := fuzzers[category]
	if !found {
		sklog.Errorf("Unknown category %s", category)
		return FUZZER_NOT_FOUND
	}
	// assume ['--type', type, {'--name', api_name}]
	cr := f.ArgsAfterExecutable[1]
	if cr == "api" {
		cr += "-" + f.ArgsAfterExecutable[3]
	}
	return cr
}

// HasCategory returns if a given string corresponds to a known fuzzer category.
func HasCategory(c string) bool {
	_, found := fuzzers[c]
	return found
}

// Status returns the status of a fuzz category (i.e. stable, experimental, etc)
func Status(c string) string {
	f, found := fuzzers[c]
	if !found {
		sklog.Errorf("Unknown category %s", c)
		return FUZZER_NOT_FOUND
	}
	return f.Status
}

// Groomer returns the groomer of a fuzz category
func Groomer(c string) string {
	f, found := fuzzers[c]
	if !found {
		sklog.Errorf("Unknown category %s", c)
		return FUZZER_NOT_FOUND
	}
	return f.Groomer
}

// Returns if fuzzer knows about a given architecture
func HasArchitecture(a string) bool {
	for _, ar := range ARCHITECTURES {
		if a == ar {
			return true
		}
	}
	return false
}

func SetMockCommon(c CommonImpl) {
	commonImpl = c
}

// Returns the Hostname of this machine. Clients should use this instead of
// os.Hostname because this function can be mocked out.
func Hostname() string {
	return commonImpl.Hostname()
}
<|endoftext|>


<commit_before>package gbench

import (
	"context"
	"fmt"
	"math/rand"
	"strings"
	"time"

	rpc "github.com/relab/gorums/dev"
	"github.com/tylertreat/bench"

	"google.golang.org/grpc"
)

// GorumsRequesterFactory implements RequesterFactory by creating a Requester which
// issues requests to a register using the Gorums framework.
type GorumsRequesterFactory struct {
	Addrs             []string
	ReadQuorum        int
	WriteQuorum       int
	PayloadSize       int
	QCTimeout         time.Duration
	WriteRatioPercent int
}

// GetRequester returns a new Requester, called for each Benchmark connection.
func (r *GorumsRequesterFactory) GetRequester(uint64) bench.Requester {
	return &gorumsRequester{
		addrs:       r.Addrs,
		readq:       r.ReadQuorum,
		writeq:      r.WriteQuorum,
		payloadSize: r.PayloadSize,
		timeout:     r.QCTimeout,
		writeRatio:  r.WriteRatioPercent,
		dialOpts: []grpc.DialOption{
			grpc.WithInsecure(),
			grpc.WithBlock(),
			grpc.WithTimeout(time.Second),
		},
	}
}

type gorumsRequester struct {
	addrs       []string
	readq       int
	writeq      int
	payloadSize int
	timeout     time.Duration
	writeRatio  int

	dialOpts []grpc.DialOption

	mgr    *rpc.Manager
	config *rpc.Configuration
	state  *rpc.State
}

func (gr *gorumsRequester) Setup() error {
	var err error
	gr.mgr, err = rpc.NewManager(
		gr.addrs,
		rpc.WithGrpcDialOptions(gr.dialOpts...),
	)
	if err != nil {
		return err
	}

	ids := gr.mgr.NodeIDs()
	qspec := newRegisterQSpec(gr.readq, gr.writeq)
	gr.config, err = gr.mgr.NewConfiguration(ids, qspec)
	if err != nil {
		return err
	}

	// Set initial state.
	gr.state = &rpc.State{
		Value:     strings.Repeat("x", gr.payloadSize),
		Timestamp: time.Now().UnixNano(),
	}
	ctx, cancel := context.WithTimeout(context.Background(), gr.timeout)
	defer cancel()
	wreply, err := gr.config.Write(ctx, gr.state)
	if err != nil {
		return fmt.Errorf("write rpc error: %v", err)
	}
	if !wreply.New {
		return fmt.Errorf("intital write reply was not marked as new")
	}

	return nil
}

func (gr *gorumsRequester) Request() error {
	var err error
	ctx, cancel := context.WithTimeout(context.Background(), gr.timeout)
	defer cancel()
	switch gr.writeRatio {
	case 0:
		_, err = gr.config.Read(ctx, &rpc.ReadRequest{})
	case 100:
		gr.state.Timestamp = time.Now().UnixNano()
		_, err = gr.config.Write(ctx, gr.state)
	default:
		x := rand.Intn(100)
		if x < gr.writeRatio {
			gr.state.Timestamp = time.Now().UnixNano()
			_, err = gr.config.Write(ctx, gr.state)
		} else {
			_, err = gr.config.Read(ctx, &rpc.ReadRequest{})
		}
	}

	return err
}

func (gr *gorumsRequester) Teardown() error {
	gr.mgr.Close()
	gr.mgr = nil
	gr.config = nil
	return nil
}

type registerQSpec struct {
	rq, wq int
}

func newRegisterQSpec(rq, wq int) rpc.QuorumSpec {
	return &registerQSpec{
		rq: rq,
		wq: wq,
	}
}

func (rqs *registerQSpec) ReadQF(replies []*rpc.State) (*rpc.State, bool) {
	if len(replies) < rqs.rq {
		return nil, false
	}
	return replies[0], true
}

func (rqs *registerQSpec) ReadFutureQF(replies []*rpc.State) (*rpc.State, bool) {
	return rqs.ReadQF(replies)
}

func (rqs *registerQSpec) ReadCorrectableQF(replies []*rpc.State) (*rpc.State, int, bool) {
	panic("not implemented")
}

func (rqs *registerQSpec) ReadPrelimQF(replies []*rpc.State) (*rpc.State, int, bool) {
	panic("not implemented")
}

func (rqs *registerQSpec) WriteQF(req *rpc.State, replies []*rpc.WriteResponse) (*rpc.WriteResponse, bool) {
	if len(replies) < rqs.wq {
		return nil, false
	}
	return replies[0], true
}

func (rqs *registerQSpec) WriteFutureQF(req *rpc.State, replies []*rpc.WriteResponse) (*rpc.WriteResponse, bool) {
	return rqs.WriteQF(req, replies)
}
<commit_msg>fixed a few more missing QFs<commit_after>package gbench

import (
	"context"
	"fmt"
	"math/rand"
	"strings"
	"time"

	rpc "github.com/relab/gorums/dev"
	"github.com/tylertreat/bench"

	"google.golang.org/grpc"
)

// GorumsRequesterFactory implements RequesterFactory by creating a Requester which
// issues requests to a register using the Gorums framework.
type GorumsRequesterFactory struct {
	Addrs             []string
	ReadQuorum        int
	WriteQuorum       int
	PayloadSize       int
	QCTimeout         time.Duration
	WriteRatioPercent int
}

// GetRequester returns a new Requester, called for each Benchmark connection.
func (r *GorumsRequesterFactory) GetRequester(uint64) bench.Requester {
	return &gorumsRequester{
		addrs:       r.Addrs,
		readq:       r.ReadQuorum,
		writeq:      r.WriteQuorum,
		payloadSize: r.PayloadSize,
		timeout:     r.QCTimeout,
		writeRatio:  r.WriteRatioPercent,
		dialOpts: []grpc.DialOption{
			grpc.WithInsecure(),
			grpc.WithBlock(),
			grpc.WithTimeout(time.Second),
		},
	}
}

type gorumsRequester struct {
	addrs       []string
	readq       int
	writeq      int
	payloadSize int
	timeout     time.Duration
	writeRatio  int

	dialOpts []grpc.DialOption

	mgr    *rpc.Manager
	config *rpc.Configuration
	state  *rpc.State
}

func (gr *gorumsRequester) Setup() error {
	var err error
	gr.mgr, err = rpc.NewManager(
		gr.addrs,
		rpc.WithGrpcDialOptions(gr.dialOpts...),
	)
	if err != nil {
		return err
	}

	ids := gr.mgr.NodeIDs()
	qspec := newRegisterQSpec(gr.readq, gr.writeq)
	gr.config, err = gr.mgr.NewConfiguration(ids, qspec)
	if err != nil {
		return err
	}

	// Set initial state.
	gr.state = &rpc.State{
		Value:     strings.Repeat("x", gr.payloadSize),
		Timestamp: time.Now().UnixNano(),
	}
	ctx, cancel := context.WithTimeout(context.Background(), gr.timeout)
	defer cancel()
	wreply, err := gr.config.Write(ctx, gr.state)
	if err != nil {
		return fmt.Errorf("write rpc error: %v", err)
	}
	if !wreply.New {
		return fmt.Errorf("intital write reply was not marked as new")
	}

	return nil
}

func (gr *gorumsRequester) Request() error {
	var err error
	ctx, cancel := context.WithTimeout(context.Background(), gr.timeout)
	defer cancel()
	switch gr.writeRatio {
	case 0:
		_, err = gr.config.Read(ctx, &rpc.ReadRequest{})
	case 100:
		gr.state.Timestamp = time.Now().UnixNano()
		_, err = gr.config.Write(ctx, gr.state)
	default:
		x := rand.Intn(100)
		if x < gr.writeRatio {
			gr.state.Timestamp = time.Now().UnixNano()
			_, err = gr.config.Write(ctx, gr.state)
		} else {
			_, err = gr.config.Read(ctx, &rpc.ReadRequest{})
		}
	}

	return err
}

func (gr *gorumsRequester) Teardown() error {
	gr.mgr.Close()
	gr.mgr = nil
	gr.config = nil
	return nil
}

type registerQSpec struct {
	rq, wq int
}

func newRegisterQSpec(rq, wq int) rpc.QuorumSpec {
	return &registerQSpec{
		rq: rq,
		wq: wq,
	}
}

func (rqs *registerQSpec) ReadQF(replies []*rpc.State) (*rpc.State, bool) {
	if len(replies) < rqs.rq {
		return nil, false
	}
	return replies[0], true
}

func (rqs *registerQSpec) ReadFutureQF(replies []*rpc.State) (*rpc.State, bool) {
	return rqs.ReadQF(replies)
}

func (rqs *registerQSpec) ReadCustomReturnQF(replies []*rpc.State) (*rpc.State, bool) {
	return rqs.ReadQF(replies)
}

func (rqs *registerQSpec) ReadCorrectableQF(replies []*rpc.State) (*rpc.State, int, bool) {
	panic("not implemented")
}

func (rqs *registerQSpec) ReadPrelimQF(replies []*rpc.State) (*rpc.State, int, bool) {
	panic("not implemented")
}

func (rqs *registerQSpec) WriteQF(req *rpc.State, replies []*rpc.WriteResponse) (*rpc.WriteResponse, bool) {
	if len(replies) < rqs.wq {
		return nil, false
	}
	return replies[0], true
}

func (rqs *registerQSpec) WriteFutureQF(req *rpc.State, replies []*rpc.WriteResponse) (*rpc.WriteResponse, bool) {
	return rqs.WriteQF(req, replies)
}

func (rqs *registerQSpec) WritePerNodeQF(replies []*rpc.WriteResponse) (*rpc.WriteResponse, bool) {
	if len(replies) < rqs.wq {
		return nil, false
	}
	return replies[0], true
}
<|endoftext|>