<commit_before>package controllers

import (
	"bytes"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/eirka/eirka-libs/config"
	"github.com/eirka/eirka-libs/db"
	"github.com/eirka/eirka-libs/user"

	local "github.com/eirka/eirka-post/config"
)

func init() {

	// Database connection settings
	dbase := db.Database{

		User:           local.Settings.Database.User,
		Password:       local.Settings.Database.Password,
		Proto:          local.Settings.Database.Proto,
		Host:           local.Settings.Database.Host,
		Database:       local.Settings.Database.Database,
		MaxIdle:        local.Settings.Database.MaxIdle,
		MaxConnections: local.Settings.Database.MaxConnections,
	}

	// Set up DB connection
	dbase.NewDb()

	// Get limits and stuff from database
	config.GetDatabaseSettings()

	user.Secret = "secret"
}

func performRequest(r http.Handler, method, path string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func performJwtFormRequest(r http.Handler, method, path, token string, body *io.Reader) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, &body)
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func TestEmailController(t *testing.T) {

	var err error

	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	router.Use(user.Auth(true))

	router.POST("/email", EmailController)

	first := performRequest(router, "POST", "/email")

	assert.Equal(t, first.Code, 401, "HTTP request code should match")

	u := user.DefaultUser()
	u.SetId(2)
	u.SetAuthenticated()

	_, err = user.HashPassword("testpassword")
	assert.NoError(t, err, "An error was not expected")

	token, err := u.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, token, "token should be returned")
	}

	var b bytes.Buffer

	mw := multipart.NewWriter(&b)
	mw.WriteField("ib", "1")
	mw.WriteField("email", "test@cool.com")
	mw.Close()

	second := performJwtFormRequest(router, "POST", "/email", token, b)

}
<commit_msg>add email test<commit_after>package controllers

import (
	"bytes"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/eirka/eirka-libs/config"
	"github.com/eirka/eirka-libs/db"
	"github.com/eirka/eirka-libs/user"

	local "github.com/eirka/eirka-post/config"
)

func init() {

	// Database connection settings
	dbase := db.Database{

		User:           local.Settings.Database.User,
		Password:       local.Settings.Database.Password,
		Proto:          local.Settings.Database.Proto,
		Host:           local.Settings.Database.Host,
		Database:       local.Settings.Database.Database,
		MaxIdle:        local.Settings.Database.MaxIdle,
		MaxConnections: local.Settings.Database.MaxConnections,
	}

	// Set up DB connection
	dbase.NewDb()

	// Get limits and stuff from database
	config.GetDatabaseSettings()

	user.Secret = "secret"
}

func performRequest(r http.Handler, method, path string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func performJwtFormRequest(r http.Handler, method, path, token string, body io.Reader) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, body)
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func TestEmailController(t *testing.T) {

	var err error

	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	router.Use(user.Auth(true))

	router.POST("/email", EmailController)

	first := performRequest(router, "POST", "/email")

	assert.Equal(t, first.Code, 401, "HTTP request code should match")

	u := user.DefaultUser()
	u.SetId(2)
	u.SetAuthenticated()

	_, err = user.HashPassword("testpassword")
	assert.NoError(t, err, "An error was not expected")

	token, err := u.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, token, "token should be returned")
	}

	var b bytes.Buffer

	mw := multipart.NewWriter(&b)
	mw.WriteField("ib", "1")
	mw.WriteField("email", "test@cool.com")
	mw.Close()

	second := performJwtFormRequest(router, "POST", "/email", token, b)

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
)

// property represents property for each field.
type property struct {
	Type       string
	Length     int
	Prefix     string
	Value      interface{}
	Multi      bool
	Max        float64
	Properties map[string]*property
	Num        int
	List       []interface{}
}

// gen generates and returns value of the property.
func (p *property) gen(seq int) interface{} {
	if p.Properties != nil {
		s := make([]interface{}, p.Num)

		for i := range s {
			s[i] = genProps(p.Properties, seq)
		}

		return s
	}

	switch p.Value {
	case "$seq":
		if p.Length == 0 {
			if p.Prefix == "" {
				return p.Value
			}

			return p.Prefix + p.Value.(string)
		}

		var s string

		if p.Max == 0 {
			s = strconv.Itoa(seq)
		} else {
			s = strconv.Itoa(seq%int(p.Max) + 1)
		}

		return p.Prefix + strings.Repeat("0", p.Length-len(p.Prefix)-len(s)) + s
	case "$seq_int":
		return seq
	case "$rand_num":
		if p.Multi {
			s := make([]string, rand.Intn(5)+1)

			for i := range s {
				s[i] = randNum(p.Length)
			}

			return s
		}

		return randNum(p.Length)
	case "$rand_int":
		return randInt(int(p.Max))
	case "$rand_double":
		return randDouble(p.Max)
	case "$rand_kana":
		if p.Multi {
			s := make([]string, rand.Intn(5)+1)

			for i := range s {
				s[i] = randKana(p.Length / 2)
			}

			return s
		}

		return randKana(p.Length / 2)
	case "$rand_bool":
		return randBool()
	case "$rand_date":
		return randDate()
	case "$rand_list":
		if p.Multi {
			s := make([]interface{}, rand.Intn(5)+1)

			for i := range s {
				s[i] = p.List[rand.Intn(len(p.List))]
			}

			return s
		}

		return p.List[rand.Intn(len(p.List))]
	default:
		return p.Value
	}
}

// config represents configuration for the processing.
type config struct {
	Action        string
	Index         string
	Type          string
	Num           int
	Props         map[string]*property
	MaxNumPerFile int `json:"max_num_per_file"`
}

// Flags
var (
	inPath  = flag.String("i", "", "input file path")
	outPath = flag.String("o", "", "output file path")
)

// Kana
var (
	kanas = []string{
		"あ", "い", "う", "え", "お", "か", "き", "く", "け", "こ",
		"さ", "し", "す", "せ", "そ", "た", "ち", "つ", "て", "と",
		"な", "に", "ぬ", "ね", "の", "は", "ひ", "ふ", "へ", "ほ",
		"ま", "み", "む", "め", "も", "や", "ゆ", "よ",
		"ら", "り", "る", "れ", "ろ", "わ", "を", "ん",
		"が", "ぎ", "ぐ", "げ", "ご", "ざ", "じ", "ず", "ぜ", "ぞ",
		"だ", "ぢ", "づ", "で", "ど", "ば", "び", "ぶ", "べ", "ぼ",
		"ぱ", "ぴ", "ぷ", "ぺ", "ぽ",
	}

	kanaLen = len(kanas)
)

// The maximum number of CPUs
var maxprocs int

// Wait group
var wg sync.WaitGroup

// Mutex
var mu sync.Mutex

// LF
var lf = []byte("\n")

// Files
var files = make([]*os.File, 0)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	maxprocs = runtime.GOMAXPROCS(0)

	flag.Parse()
}

type context struct {
	w  io.Writer
	bs [][]byte
}

func write(ctxc <-chan *context) {
	for ctx := range ctxc {
		mu.Lock()

		for _, b := range ctx.bs {
			if _, err := ctx.w.Write(b); err != nil {
				mu.Unlock()
				panic(err)
			}

			ctx.w.Write(lf)
		}

		mu.Unlock()
	}

	wg.Done()
}

func main() {
	in, err := ioutil.ReadFile(*inPath)
	if err != nil {
		panic(err)
	}

	var conf config
	if err := json.Unmarshal(in, &conf); err != nil {
		panic(err)
	}

	ctxc := make(chan *context)

	for i := 0; i < maxprocs; i++ {
		wg.Add(1)
		go write(ctxc)
	}

	defer func() {
		close(ctxc)

		wg.Wait()

		for _, f := range files {
			f.Close()
		}
	}()

	var f *os.File

	var numPerFile int
	var seqFile int

	for seq := 1; seq <= conf.Num; seq++ {
		numPerFile++

		if seq == 1 || numPerFile > conf.MaxNumPerFile {
			seqFile++

			strSeq := strconv.Itoa(seqFile)

			f, err = os.Create(*outPath + "." + strings.Repeat("0", 4-len(strSeq)) + strSeq)
			if err != nil {
				panic(err)
			}

			files = append(files, f)

			numPerFile = 1
		}

		meta := make(map[string]string)

		meta["_index"] = conf.Index
		meta["_type"] = conf.Type
		meta["_id"] = conf.Props["_id"].gen(seq).(string)

		if p, exist := conf.Props["_parent"]; exist {
			meta["_parent"] = p.gen(seq).(string)
		}

		action := map[string]map[string]string{
			conf.Action: meta,
		}

		actionOut, err := json.Marshal(action)
		if err != nil {
			panic(err)
		}

		srcOut, err := json.Marshal(genProps(conf.Props, seq))
		if err != nil {
			panic(err)
		}

		ctxc <- &context{
			w:  f,
			bs: [][]byte{actionOut, srcOut},
		}
	}
}

func genProps(props map[string]*property, seq int) map[string]interface{} {
	src := make(map[string]interface{})

	for k, p := range props {
		if k == "_id" || k == "_parent" {
			continue
		}

		src[k] = p.gen(seq)
	}

	return src
}

// randNum generates and returns a random number.
func randNum(l int) string {
	var s string

	for i := 0; i < l; i++ {
		s += strconv.Itoa(rand.Intn(10))
	}

	return s
}

// randInt generates and returns a random integer value.
func randInt(n int) int {
	return rand.Intn(n)
}

// randDouble generates and returns a random double value.
func randDouble(n float64) float64 {
	return rand.Float64() * n
}

// randKana generates and returns a random kana.
func randKana(l int) string {
	var s string

	for i := 0; i < l; i++ {
		s += kanas[rand.Intn(kanaLen)]
	}

	return s
}

// randBool generates and returns a random boolean value.
func randBool() bool {
	return rand.Intn(2) == 1
}

// randDate generates and returns a random date value.
func randDate() string {
	m := strconv.Itoa(rand.Intn(12) + 1)

	if len(m) < 2 {
		m = "0" + m
	}

	var maxD int

	switch m {
	case "01", "03", "05", "07", "08", "10", "12":
		maxD = 31
	case "02":
		maxD = 28
	default:
		maxD = 30
	}

	d := strconv.Itoa(rand.Intn(maxD) + 1)

	if len(d) < 2 {
		d = "0" + d
	}

	return "2015" + m + d
}
<commit_msg>Update cmd/esgen/main.go<commit_after>package main

import (
	"encoding/json"
	"flag"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
)

// property represents property for each field.
type property struct {
	Type       string
	Length     int
	Prefix     string
	Value      interface{}
	Multi      bool
	Max        float64
	Properties map[string]*property
	Num        int
	List       []interface{}
}

// gen generates and returns value of the property.
func (p *property) gen(seq int) interface{} {
	if p.Properties != nil {
		s := make([]interface{}, p.Num)

		for i := range s {
			s[i] = genProps(p.Properties, seq)
		}

		return s
	}

	switch p.Value {
	case "$seq":
		if p.Length == 0 {
			if p.Prefix == "" {
				return p.Value
			}

			return p.Prefix + p.Value.(string)
		}

		var s string

		if p.Max == 0 {
			s = strconv.Itoa(seq)
		} else {
			s = strconv.Itoa(seq%int(p.Max) + 1)
		}

		return p.Prefix + strings.Repeat("0", p.Length-len(p.Prefix)-len(s)) + s
	case "$seq_int":
		return seq
	case "$rand_num":
		if p.Multi {
			s := make([]string, rand.Intn(5)+1)

			for i := range s {
				s[i] = randNum(p.Length)
			}

			return s
		}

		return randNum(p.Length)
	case "$rand_int":
		return randInt(int(p.Max))
	case "$rand_double":
		return randDouble(p.Max)
	case "$rand_kana":
		if p.Multi {
			s := make([]string, rand.Intn(5)+1)

			for i := range s {
				s[i] = randKana(p.Length / 2)
			}

			return s
		}

		return randKana(p.Length / 2)
	case "$rand_bool":
		return randBool()
	case "$rand_date":
		return randDate()
	case "$rand_list":
		if p.Multi {
			s := make([]interface{}, rand.Intn(5)+1)

			for i := range s {
				s[i] = p.List[rand.Intn(len(p.List))]
			}

			return s
		}

		return p.List[rand.Intn(len(p.List))]
	default:
		return p.Value
	}
}

// config represents configuration for the processing.
type config struct {
	Action        string
	Index         string
	Type          string
	Num           int
	Props         map[string]*property
	MaxNumPerFile int `json:"max_num_per_file"`
}

// Flags
var (
	inPath  = flag.String("i", "", "input file path")
	outPath = flag.String("o", "", "output file path")
)

// Kana
var (
	kanas = []string{
		"あ", "い", "う", "え", "お", "か", "き", "く", "け", "こ",
		"さ", "し", "す", "せ", "そ", "た", "ち", "つ", "て", "と",
		"な", "に", "ぬ", "ね", "の", "は", "ひ", "ふ", "へ", "ほ",
		"ま", "み", "む", "め", "も", "や", "ゆ", "よ",
		"ら", "り", "る", "れ", "ろ", "わ", "を", "ん",
		"が", "ぎ", "ぐ", "げ", "ご", "ざ", "じ", "ず", "ぜ", "ぞ",
		"だ", "ぢ", "づ", "で", "ど", "ば", "び", "ぶ", "べ", "ぼ",
		"ぱ", "ぴ", "ぷ", "ぺ", "ぽ",
	}

	kanaLen = len(kanas)
)

// The maximum number of CPUs
var maxprocs int

// Wait group
var wg sync.WaitGroup

// Mutex
var mu sync.Mutex

// LF
var lf = []byte("\n")

// Files
var files = make([]*os.File, 0)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	maxprocs = runtime.GOMAXPROCS(0)

	flag.Parse()
}

type context struct {
	w  io.Writer
	bs [][]byte
}

func write(ctxc <-chan *context) {
	for ctx := range ctxc {
		mu.Lock()

		for _, b := range ctx.bs {
			if _, err := ctx.w.Write(b); err != nil {
				mu.Unlock()
				panic(err)
			}

			ctx.w.Write(lf)
		}

		mu.Unlock()
	}

	wg.Done()
}

func main() {
	in, err := ioutil.ReadFile(*inPath)
	if err != nil {
		panic(err)
	}

	var conf config
	if err := json.Unmarshal(in, &conf); err != nil {
		panic(err)
	}

	ctxc := make(chan *context)

	for i := 0; i < maxprocs; i++ {
		wg.Add(1)
		go write(ctxc)
	}

	defer func() {
		close(ctxc)

		wg.Wait()

		for _, f := range files {
			f.Close()
		}
	}()

	var f *os.File

	var numPerFile int
	var seqFile int

	for seq := 1; seq <= conf.Num; seq++ {
		numPerFile++

		if seq == 1 || numPerFile > conf.MaxNumPerFile {
			seqFile++

			strSeq := strconv.Itoa(seqFile)

			f, err = os.Create(*outPath + "." + strings.Repeat("0", 4-len(strSeq)) + strSeq)
			if err != nil {
				panic(err)
			}

			files = append(files, f)

			numPerFile = 1
		}

		meta := make(map[string]string)

		meta["_index"] = conf.Index
		meta["_type"] = conf.Type
		meta["_id"] = conf.Props["_id"].gen(seq).(string)

		if p, exist := conf.Props["_parent"]; exist {
			meta["_parent"] = p.gen(seq).(string)
		}

		action := map[string]map[string]string{
			conf.Action: meta,
		}

		actionOut, err := json.Marshal(action)
		if err != nil {
			panic(err)
		}

		srcOut, err := json.Marshal(genProps(conf.Props, seq))
		if err != nil {
			panic(err)
		}

		ctxc <- &context{
			w:  f,
			bs: [][]byte{actionOut, srcOut},
		}
	}
}

func genProps(props map[string]*property, seq int) map[string]interface{} {
	src := make(map[string]interface{})

	for k, p := range props {
		if k == "_id" || k == "_parent" {
			continue
		}

		src[k] = p.gen(seq)
	}

	return src
}

// randNum generates and returns a random number.
func randNum(l int) string {
	var s string

	for i := 0; i < l; i++ {
		s += strconv.Itoa(rand.Intn(10))
	}

	return s
}

// randInt generates and returns a random integer value.
func randInt(n int) int {
	return rand.Intn(n)
}

// randDouble generates and returns a random double value.
func randDouble(n float64) float64 {
	return rand.Float64() * n
}

// randKana generates and returns a random kana.
func randKana(l int) string {
	var s string

	for i := 0; i < l; i++ {
		s += kanas[rand.Intn(kanaLen)]
	}

	return s
}

// randBool generates and returns a random boolean value.
func randBool() bool {
	return rand.Intn(2) == 1
}

// randDate generates and returns a random date value.
func randDate() string {
	m := strconv.Itoa(rand.Intn(3) + 1)

	if len(m) < 2 {
		m = "0" + m
	}

	var maxD int

	switch m {
	case "01", "03", "05", "07", "08", "10", "12":
		maxD = 31
	case "02":
		maxD = 28
	default:
		maxD = 30
	}

	d := strconv.Itoa(rand.Intn(maxD) + 1)

	if len(d) < 2 {
		d = "0" + d
	}

	return "2015" + m + d
}
<|endoftext|>


<commit_before>// Copyright 2017-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package bootconfig

import (
	"archive/zip"
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha512"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	urootcrypto "github.com/u-root/u-root/pkg/crypto"
)

// memoryZipReader is used to unpack a zip file from a byte sequence in memory.
type memoryZipReader struct {
	Content []byte
}

func (r *memoryZipReader) ReadAt(p []byte, offset int64) (n int, err error) {
	cLen := int64(len(r.Content))
	if offset > cLen {
		return 0, io.EOF
	}
	if cLen-offset >= int64(len(p)) {
		n = len(p)
		err = nil
	} else {
		err = io.EOF
		n = int(int64(cLen) - offset)
	}
	copy(p, r.Content[offset:int(offset)+n])
	return n, err
}

// FIXME:
// FromZip tries to extract a boot configuration from a ZIP file after verifying
// its signature with the provided public key file. The signature is expected to
// be appended to the ZIP file and have fixed length `ed25519.SignatureSize` .
// The returned string argument is the temporary directory where the files were
// extracted, if successful.
// No decoder (e.g. JSON, ZIP) or other function parsing the input file is called
// before verifying the signature.
func FromZip(filename string) (*Manifest, string, error) {
	// load the whole zip file in memory - we need it anyway for the signature
	// matching.
	// TODO refuse to read if too big?
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, "", err
	}
	urootcrypto.TryMeasureData(urootcrypto.BlobPCR, data, filename)
	zipbytes := data

	r, err := zip.NewReader(&memoryZipReader{Content: zipbytes}, int64(len(zipbytes)))
	if err != nil {
		return nil, "", err
	}
	tempDir, err := ioutil.TempDir(os.TempDir(), "bootconfig")
	if err != nil {
		return nil, "", err
	}
	log.Printf("Created temporary directory %s", tempDir)
	var manifest *Manifest
	for _, f := range r.File {
		if f.FileInfo().IsDir() {
			// Dont care - will be handled later
			continue
		}

		destination := path.Join(tempDir, f.Name)
		if len(f.Name) == 0 {
			log.Printf("Warning: skipping zero-length file name (flags: %d, mode: %s)", f.Flags, f.Mode())
			continue
		}
		// Check if folder exists
		if _, err := os.Stat(destination); os.IsNotExist(err) {
			if err := os.MkdirAll(path.Dir(destination), os.ModeDir|os.FileMode(0700)); err != nil {
				return nil, "", err
			}
		}
		fd, err := f.Open()
		if err != nil {
			return nil, "", err
		}
		buf, err := ioutil.ReadAll(fd)
		if err != nil {
			return nil, "", err
		}
		if f.Name == "manifest.json" {
			// make sure it's not a duplicate manifest within the ZIP file
			// and inform the user otherwise
			if manifest != nil {
				log.Printf("Warning: duplicate manifest.json found, the last found wins")
			}
			// parse the Manifest containing the boot configurations
			manifest, err = ManifestFromBytes(buf)
			if err != nil {
				return nil, "", err
			}
		}
		if err := ioutil.WriteFile(destination, buf, f.Mode()); err != nil {
			return nil, "", err
		}
		log.Printf("Extracted file %s (flags: %d, mode: %s)", f.Name, f.Flags, f.Mode())
	}
	if manifest == nil {
		return nil, "", errors.New("no manifest found")
	}
	return manifest, tempDir, nil
}

// FIXME:
// ToZip tries to pack all files specified in the the provided manifest.json
// into a zip archive. An error is returned, if the files (kernel, initrd, etc)
// doesn't exist at the paths written inside the manifest.json relative to its
// location. Optionally , if privkeyfile is not nil, after creating the archive an ed25519 signature is added to the
// archive. A copy of manifest.json is included in the final with adopted path of the
// bootfiles to mach teir location relative to the archive root.
func ToZip(output string, manifest string) error {
	// Get manifest from file. Make sure the file is named accordingliy, since
	// FromZip will search 'manifest.json' while extraction.

	if base := path.Base(manifest); base != "manifest.json" {
		return fmt.Errorf("Invalid manifest name. Want 'manifest.json', got: %s", base)
	}
	manifestBody, err := ioutil.ReadFile(manifest)
	if err != nil {
		return err
	}
	mf, err := ManifestFromBytes(manifestBody)
	if err != nil {
		return err
	} else if !mf.IsValid() {
		return errors.New("Manifest is not valid")
	}

	// Create a buffer to write the archive to.
	buf := new(bytes.Buffer)
	// Create a new zip archive.
	z := zip.NewWriter(buf)

	var dest, origin string
	//Archive boot files
	for i, cfg := range mf.Configs {
		dir := fmt.Sprintf("bootconfig_%d/", i)
		z.Create(dir)
		if cfg.Kernel != "" {
			dest = path.Join(dir, path.Base(cfg.Kernel))
			origin = path.Join(path.Dir(manifest), cfg.Kernel)
			toZip(z, dest, origin)
			cfg.Kernel = dest
		}
		if cfg.Initramfs != "" {
			dest = path.Join(dir, path.Base(cfg.Initramfs))
			origin = path.Join(path.Dir(manifest), cfg.Initramfs)
			toZip(z, dest, origin)
			cfg.Initramfs = dest
		}
		if cfg.DeviceTree != "" {
			dest = path.Join(dir, path.Base(cfg.DeviceTree))
			origin = path.Join(path.Dir(manifest), cfg.DeviceTree)
			toZip(z, dest, origin)
			cfg.DeviceTree = dest
		}
		mf.Configs[i] = cfg
	}

	// Archive root certificate
	z.Create("certs/")
	dest = "certs/root.cert"
	origin = path.Join(path.Dir(manifest), mf.RootCertPath)
	err = toZip(z, dest, origin)
	if err != nil {
		log.Fatal("DEBUG Error:", err)
	}
	mf.RootCertPath = dest

	// Archive manifest
	newManifest, err := ManifestToBytes(mf)
	if err != nil {
		return err
	}
	dst, err := z.Create(path.Base(manifest))
	if err != nil {
		return err
	}
	_, err = io.Copy(dst, bytes.NewReader(newManifest))
	if err != nil {
		return err
	}

	// Write central directory of archive
	err = z.Close()
	if err != nil {
		return err
	}

	// Write buf to disk
	if path.Ext(output) != ".zip" {
		output = output + ".zip"
	}
	err = ioutil.WriteFile(output, buf.Bytes(), 0777)
	if err != nil {
		return err
	}
	return nil
}

func toZip(w *zip.Writer, newPath, originPath string) error {
	dst, err := w.Create(newPath)
	if err != nil {
		return err
	}
	// Copy content from inputpath to new file
	src, err := os.Open(originPath)
	if err != nil {
		return fmt.Errorf("Cannot find %s specified in manifest", originPath)
	}
	_, err = io.Copy(dst, src)
	if err != nil {
		return err
	}
	return src.Close()
}

// AddSignature signes the bootfiles inside an stboot.zip and inserts the
// signatures into the archive along with the respective certificate
func AddSignature(archive, privKey, certificate string) error {

	mf, dir, err := FromZip(archive)
	if err != nil {
		return err
	}

	// collect boot binaries
	// XXX Refactor if we remove bootconfig from manifest
	// Maybe just walk through certs/folders and match do root/bootconfig
	for i := range mf.Configs {

		bootconfigDir := path.Join(dir, fmt.Sprintf("bootconfig_%d", i))

		bcHash, err := HashBootconfigDir(bootconfigDir)
		if err != nil {
			log.Println(fmt.Sprintf("Failed to hash bootconfig - Err %s", err))
			return err
		}

		// Sign hash with Key
		buff, err := ioutil.ReadFile(privKey)
		privPem, _ := pem.Decode(buff)
		rsaPrivKey, err := x509.ParsePKCS1PrivateKey(privPem.Bytes)

		if rsaPrivKey == nil {
			panic("RSA Key is nil")
		}

		opts := &rsa.PSSOptions{SaltLength: rsa.PSSSaltLengthEqualsHash}

		log.Printf("bootconfig hash is: %x", bcHash)
		signature, err := rsa.SignPSS(rand.Reader, rsaPrivKey, crypto.SHA512, bcHash, opts)
		if signature == nil {
			panic("Signing failed.")
		}

		log.Println("Signing..")
		log.Println(fmt.Sprintf("%x", signature))

		// Create dir for signature
		err = os.MkdirAll(path.Join(dir, fmt.Sprintf("certs/bootconfig_%d/", i)), os.ModeDir|os.FileMode(0700))
		if err != nil {
			log.Println(fmt.Sprintf("Creating directories in %s for signatures failed - Check permissions.", dir))
			return err
		}

		// Extract part of Public Key for identification
		certificateString, err := ioutil.ReadFile(certificate)
		if err != nil {
			log.Println(fmt.Sprintf("Failed to read certificate - Err %s", err))
			return err
		}

		cert, err := parseCertificate(certificateString)
		if err != nil {
			log.Println(fmt.Sprintf("Failed to parse certificate %s", certificateString))
		}

		// Write signature to folder
		err = ioutil.WriteFile(path.Join(dir, fmt.Sprintf("certs/bootconfig_%d/%s.signature", i, fmt.Sprintf("%x", cert.PublicKey)[2:18])), signature, 0644)
		if err != nil {
			log.Println(fmt.Sprintf("Writing into %s failed - Check permissions.", dir))
			return err
		}

		// cp cert to folder
		err = ioutil.WriteFile(path.Join(dir, fmt.Sprintf("certs/bootconfig_%d/%s.cert", i, fmt.Sprintf("%x", cert.PublicKey)[2:18])), certificateString, 0644)
		if err != nil {
			log.Println(fmt.Sprintf("Copying certificate %s to .zip failed - Check permissions.", certificate))
			return err
		}
	}

	// Pack it again
	// Create a buffer to write the archive to.
	buf := new(bytes.Buffer)
	// Create a new zip archive.
	z := zip.NewWriter(buf)

	// Walk the directory and pack it.
	filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			err := toZip(z, strings.Replace(path, dir, "", -1)[1:], path)
			if err != nil {
				log.Println(fmt.Sprintf("Error adding file %s to .zip archive again", strings.Replace(path, dir, "", -1)))
			}
		}

		return nil
	})

	z.Close()

	pathToZip := fmt.Sprintf("./.original/%d", time.Now().Unix())
	os.MkdirAll(pathToZip, os.ModePerm)
	os.Rename(archive, pathToZip+"/stboot.zip")
	log.Println("Backed up old stboot.zip to " + pathToZip)

	err = ioutil.WriteFile(archive, buf.Bytes(), 0777)
	if err != nil {
		log.Println(fmt.Sprintf("Unable to write new stboot.zip file - recover old from %s", pathToZip))
		return err
	}
	log.Println("Stboot file has been written to " + archive)

	return nil

}

// parseCertificate parses certificate from raw certificate
func parseCertificate(rawCertificate []byte) (x509.Certificate, error) {

	block, _ := pem.Decode(rawCertificate)
	if block == nil {
		panic("failed to parse PEM block containing the public key")
	}

	pub, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		fmt.Println("failed to parse DER encoded public key: " + err.Error())
		return *pub, err
	}

	return *pub, nil
}

// HashBootconfigDir hashes every file inside bootconigDir and returns a
// SHA512 hash
func HashBootconfigDir(bootconfigDir string) ([]byte, error) {

	hash := sha512.New()
	hash.Reset()

	files, err := ioutil.ReadDir(bootconfigDir)
	if err != nil {
		return nil, err
	}
	for _, file := range files {
		if !file.IsDir() {
			p := path.Join(bootconfigDir, file.Name())
			buf, err := ioutil.ReadFile(p)
			if err != nil {
				return nil, err
			}
			hash.Write(buf)

		}
	}
	return hash.Sum(nil), nil
}
<commit_msg>stconfig tool: remove backups after signing<commit_after>// Copyright 2017-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package bootconfig

import (
	"archive/zip"
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha512"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	urootcrypto "github.com/u-root/u-root/pkg/crypto"
)

// memoryZipReader is used to unpack a zip file from a byte sequence in memory.
type memoryZipReader struct {
	Content []byte
}

func (r *memoryZipReader) ReadAt(p []byte, offset int64) (n int, err error) {
	cLen := int64(len(r.Content))
	if offset > cLen {
		return 0, io.EOF
	}
	if cLen-offset >= int64(len(p)) {
		n = len(p)
		err = nil
	} else {
		err = io.EOF
		n = int(int64(cLen) - offset)
	}
	copy(p, r.Content[offset:int(offset)+n])
	return n, err
}

// FIXME:
// FromZip tries to extract a boot configuration from a ZIP file after verifying
// its signature with the provided public key file. The signature is expected to
// be appended to the ZIP file and have fixed length `ed25519.SignatureSize` .
// The returned string argument is the temporary directory where the files were
// extracted, if successful.
// No decoder (e.g. JSON, ZIP) or other function parsing the input file is called
// before verifying the signature.
func FromZip(filename string) (*Manifest, string, error) {
	// load the whole zip file in memory - we need it anyway for the signature
	// matching.
	// TODO refuse to read if too big?
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, "", err
	}
	urootcrypto.TryMeasureData(urootcrypto.BlobPCR, data, filename)
	zipbytes := data

	r, err := zip.NewReader(&memoryZipReader{Content: zipbytes}, int64(len(zipbytes)))
	if err != nil {
		return nil, "", err
	}
	tempDir, err := ioutil.TempDir(os.TempDir(), "bootconfig")
	if err != nil {
		return nil, "", err
	}
	log.Printf("Created temporary directory %s", tempDir)
	var manifest *Manifest
	for _, f := range r.File {
		if f.FileInfo().IsDir() {
			// Dont care - will be handled later
			continue
		}

		destination := path.Join(tempDir, f.Name)
		if len(f.Name) == 0 {
			log.Printf("Warning: skipping zero-length file name (flags: %d, mode: %s)", f.Flags, f.Mode())
			continue
		}
		// Check if folder exists
		if _, err := os.Stat(destination); os.IsNotExist(err) {
			if err := os.MkdirAll(path.Dir(destination), os.ModeDir|os.FileMode(0700)); err != nil {
				return nil, "", err
			}
		}
		fd, err := f.Open()
		if err != nil {
			return nil, "", err
		}
		buf, err := ioutil.ReadAll(fd)
		if err != nil {
			return nil, "", err
		}
		if f.Name == "manifest.json" {
			// make sure it's not a duplicate manifest within the ZIP file
			// and inform the user otherwise
			if manifest != nil {
				log.Printf("Warning: duplicate manifest.json found, the last found wins")
			}
			// parse the Manifest containing the boot configurations
			manifest, err = ManifestFromBytes(buf)
			if err != nil {
				return nil, "", err
			}
		}
		if err := ioutil.WriteFile(destination, buf, f.Mode()); err != nil {
			return nil, "", err
		}
		log.Printf("Extracted file %s (flags: %d, mode: %s)", f.Name, f.Flags, f.Mode())
	}
	if manifest == nil {
		return nil, "", errors.New("no manifest found")
	}
	return manifest, tempDir, nil
}

// FIXME:
// ToZip tries to pack all files specified in the the provided manifest.json
// into a zip archive. An error is returned, if the files (kernel, initrd, etc)
// doesn't exist at the paths written inside the manifest.json relative to its
// location. Optionally , if privkeyfile is not nil, after creating the archive an ed25519 signature is added to the
// archive. A copy of manifest.json is included in the final with adopted path of the
// bootfiles to mach teir location relative to the archive root.
func ToZip(output string, manifest string) error {
	// Get manifest from file. Make sure the file is named accordingliy, since
	// FromZip will search 'manifest.json' while extraction.

	if base := path.Base(manifest); base != "manifest.json" {
		return fmt.Errorf("Invalid manifest name. Want 'manifest.json', got: %s", base)
	}
	manifestBody, err := ioutil.ReadFile(manifest)
	if err != nil {
		return err
	}
	mf, err := ManifestFromBytes(manifestBody)
	if err != nil {
		return err
	} else if !mf.IsValid() {
		return errors.New("Manifest is not valid")
	}

	// Create a buffer to write the archive to.
	buf := new(bytes.Buffer)
	// Create a new zip archive.
	z := zip.NewWriter(buf)

	var dest, origin string
	//Archive boot files
	for i, cfg := range mf.Configs {
		dir := fmt.Sprintf("bootconfig_%d/", i)
		z.Create(dir)
		if cfg.Kernel != "" {
			dest = path.Join(dir, path.Base(cfg.Kernel))
			origin = path.Join(path.Dir(manifest), cfg.Kernel)
			toZip(z, dest, origin)
			cfg.Kernel = dest
		}
		if cfg.Initramfs != "" {
			dest = path.Join(dir, path.Base(cfg.Initramfs))
			origin = path.Join(path.Dir(manifest), cfg.Initramfs)
			toZip(z, dest, origin)
			cfg.Initramfs = dest
		}
		if cfg.DeviceTree != "" {
			dest = path.Join(dir, path.Base(cfg.DeviceTree))
			origin = path.Join(path.Dir(manifest), cfg.DeviceTree)
			toZip(z, dest, origin)
			cfg.DeviceTree = dest
		}
		mf.Configs[i] = cfg
	}

	// Archive root certificate
	z.Create("certs/")
	dest = "certs/root.cert"
	origin = path.Join(path.Dir(manifest), mf.RootCertPath)
	err = toZip(z, dest, origin)
	if err != nil {
		log.Fatal("DEBUG Error:", err)
	}
	mf.RootCertPath = dest

	// Archive manifest
	newManifest, err := ManifestToBytes(mf)
	if err != nil {
		return err
	}
	dst, err := z.Create(path.Base(manifest))
	if err != nil {
		return err
	}
	_, err = io.Copy(dst, bytes.NewReader(newManifest))
	if err != nil {
		return err
	}

	// Write central directory of archive
	err = z.Close()
	if err != nil {
		return err
	}

	// Write buf to disk
	if path.Ext(output) != ".zip" {
		output = output + ".zip"
	}
	err = ioutil.WriteFile(output, buf.Bytes(), 0777)
	if err != nil {
		return err
	}
	return nil
}

func toZip(w *zip.Writer, newPath, originPath string) error {
	dst, err := w.Create(newPath)
	if err != nil {
		return err
	}
	// Copy content from inputpath to new file
	src, err := os.Open(originPath)
	if err != nil {
		return fmt.Errorf("Cannot find %s specified in manifest", originPath)
	}
	_, err = io.Copy(dst, src)
	if err != nil {
		return err
	}
	return src.Close()
}

// AddSignature signes the bootfiles inside an stboot.zip and inserts the
// signatures into the archive along with the respective certificate
func AddSignature(archive, privKey, certificate string) error {

	mf, dir, err := FromZip(archive)
	if err != nil {
		return err
	}

	// collect boot binaries
	// XXX Refactor if we remove bootconfig from manifest
	// Maybe just walk through certs/folders and match do root/bootconfig
	for i := range mf.Configs {

		bootconfigDir := path.Join(dir, fmt.Sprintf("bootconfig_%d", i))

		bcHash, err := HashBootconfigDir(bootconfigDir)
		if err != nil {
			log.Println(fmt.Sprintf("Failed to hash bootconfig - Err %s", err))
			return err
		}

		// Sign hash with Key
		buff, err := ioutil.ReadFile(privKey)
		privPem, _ := pem.Decode(buff)
		rsaPrivKey, err := x509.ParsePKCS1PrivateKey(privPem.Bytes)

		if rsaPrivKey == nil {
			panic("RSA Key is nil")
		}

		opts := &rsa.PSSOptions{SaltLength: rsa.PSSSaltLengthEqualsHash}

		log.Printf("bootconfig hash is: %x", bcHash)
		signature, err := rsa.SignPSS(rand.Reader, rsaPrivKey, crypto.SHA512, bcHash, opts)
		if signature == nil {
			panic("Signing failed.")
		}

		log.Println("Signing..")
		log.Println(fmt.Sprintf("%x", signature))

		// Create dir for signature
		err = os.MkdirAll(path.Join(dir, fmt.Sprintf("certs/bootconfig_%d/", i)), os.ModeDir|os.FileMode(0700))
		if err != nil {
			log.Println(fmt.Sprintf("Creating directories in %s for signatures failed - Check permissions.", dir))
			return err
		}

		// Extract part of Public Key for identification
		certificateString, err := ioutil.ReadFile(certificate)
		if err != nil {
			log.Println(fmt.Sprintf("Failed to read certificate - Err %s", err))
			return err
		}

		cert, err := parseCertificate(certificateString)
		if err != nil {
			log.Println(fmt.Sprintf("Failed to parse certificate %s", certificateString))
		}

		// Write signature to folder
		err = ioutil.WriteFile(path.Join(dir, fmt.Sprintf("certs/bootconfig_%d/%s.signature", i, fmt.Sprintf("%x", cert.PublicKey)[2:18])), signature, 0644)
		if err != nil {
			log.Println(fmt.Sprintf("Writing into %s failed - Check permissions.", dir))
			return err
		}

		// cp cert to folder
		err = ioutil.WriteFile(path.Join(dir, fmt.Sprintf("certs/bootconfig_%d/%s.cert", i, fmt.Sprintf("%x", cert.PublicKey)[2:18])), certificateString, 0644)
		if err != nil {
			log.Println(fmt.Sprintf("Copying certificate %s to .zip failed - Check permissions.", certificate))
			return err
		}
	}

	// Pack it again
	// Create a buffer to write the archive to.
	buf := new(bytes.Buffer)
	// Create a new zip archive.
	z := zip.NewWriter(buf)

	// Walk the directory and pack it.
	filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			err := toZip(z, strings.Replace(path, dir, "", -1)[1:], path)
			if err != nil {
				log.Println(fmt.Sprintf("Error adding file %s to .zip archive again", strings.Replace(path, dir, "", -1)))
			}
		}

		return nil
	})

	z.Close()

	pathToZip := fmt.Sprintf("./.original/%d", time.Now().Unix())
	os.MkdirAll(pathToZip, os.ModePerm)
	os.Rename(archive, pathToZip+"/stboot.zip")

	err = ioutil.WriteFile(archive, buf.Bytes(), 0777)
	if err != nil {
		log.Println(fmt.Sprintf("Unable to write new stboot.zip file - recover old from %s", pathToZip))
		return err
	}
	log.Println("Updated Stboot file has been written to " + archive)

	os.RemoveAll(pathToZip)

	return nil

}

// parseCertificate parses certificate from raw certificate
func parseCertificate(rawCertificate []byte) (x509.Certificate, error) {

	block, _ := pem.Decode(rawCertificate)
	if block == nil {
		panic("failed to parse PEM block containing the public key")
	}

	pub, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		fmt.Println("failed to parse DER encoded public key: " + err.Error())
		return *pub, err
	}

	return *pub, nil
}

// HashBootconfigDir hashes every file inside bootconigDir and returns a
// SHA512 hash
func HashBootconfigDir(bootconfigDir string) ([]byte, error) {

	hash := sha512.New()
	hash.Reset()

	files, err := ioutil.ReadDir(bootconfigDir)
	if err != nil {
		return nil, err
	}
	for _, file := range files {
		if !file.IsDir() {
			p := path.Join(bootconfigDir, file.Name())
			buf, err := ioutil.ReadFile(p)
			if err != nil {
				return nil, err
			}
			hash.Write(buf)

		}
	}
	return hash.Sum(nil), nil
}
<|endoftext|>


<commit_before>package bevel

import (
	"errors"
	"fmt"
	"log"
	"time"
)

// Closer is an interface that defines the operations
// that Closer implementors must adhere to.
type Closer interface {
	Close() error
}

// Poster is an interface that defines the operations
// that Poster implementors must adhere to.
type Poster interface {
	Post(Message)
}

// AddWriterer is an interface that defines the operation
// AddWriter.
type AddWriterer interface {
	AddWriter(Writer)
}

// EventBusManager is an interface that combines the operations
// of an Event Bus Manager.
type EventBusManager interface {
	AddWriterer
	Poster
	Closer
	fmt.Stringer
}

// Manager holds the properties needed by the business event logger.
// It is the receiver of utility methods for consumers.
type Manager struct {
	done        chan bool
	bus         chan Message
	writersPool WriterPool
	msgCounter  Counter
}

// StartNewListener creates a new business event bus and adds
// a Writer to the WriterPool.
func StartNewListener(w WriteCloser) EventBusManager {
	wp := *NewWriterPool()
	wp.AddWriter(w)

	bem := Manager{
		done:        make(chan bool),
		bus:         make(chan Message),
		writersPool: wp,
		msgCounter:  Counter{0},
	}

	go bem.listen()

	return &bem
}

// Post sends a Message to the business event message bus
// for ingestion by all Writer's in the WriterPool.
func (bem *Manager) Post(m Message) {
	time.Sleep(time.Second)

	// post the envelop on the bus (i.e. "post the letter")
	if bem.bus == nil || bem.done == nil {
		log.Printf("the event bus is closed - lost message: %#v", m)
		return
	}

	bem.bus <- NewMesageEnvelop(m)
}

// AddWriter adds a Writer to the WriterPool.
func (bem *Manager) AddWriter(w Writer) {
	bem.writersPool.AddWriter(w)
}

func (bem *Manager) writeMessage(m Message) {
	bem.writersPool.Write(m)
}

// String implements Stringer.
func (bem *Manager) String() string {
	s := bem.writersPool.String()
	s += fmt.Sprintf(" - total number of messages posted: %d", bem.msgCounter.Get())

	return fmt.Sprintf("%s", s)
}

// listen is the main loop of the business event loop.
func (bem *Manager) listen() {
	defer func() { bem.done <- true }()

ListenerLoop:
	for {
		select {
		case m, ok := <-bem.bus:
			if ok {
				// Received a Message wrapped in a MessageEnvelop.
				// Call the writer to write it to destination - in this case to Kafka.
				bem.msgCounter.Inc()
				bem.writeMessage(m)
			}
		case <-bem.done:
			// Received "Termination Ping" request.
			break ListenerLoop
		}
	}
}

// Close closes the channels in the Manager.
// The recommended approach is for a channel to be used unidirectionally and
// be closed by the sender rather than the receivers.
// This means that it is the responsibility of the Posters to close the
// event bus when no more messages are being posted.
// See an example implementation in main_test using a sync.WaitGroup.
func (bem *Manager) Close() error {
	if bem.done == nil {
		return errors.New("this event bus manager is already closed")
	}

	close(bem.bus)
	bem.done <- true
	<-bem.done
	bem.bus = nil

	close(bem.done)
	bem.done = nil

	return nil
}
<commit_msg>[BusDrainAtShutdown] Add draining of bus at close time (#13)<commit_after>package bevel

import (
	"errors"
	"fmt"
	"log"
)

// Closer is an interface that defines the operations
// that Closer implementors must adhere to.
type Closer interface {
	Close() error
}

// Poster is an interface that defines the operations
// that Poster implementors must adhere to.
type Poster interface {
	Post(Message)
}

// AddWriterer is an interface that defines the operation
// AddWriter.
type AddWriterer interface {
	AddWriter(Writer)
}

// EventBusManager is an interface that combines the operations
// of an Event Bus Manager.
type EventBusManager interface {
	AddWriterer
	Poster
	Closer
	fmt.Stringer
}

// Manager holds the properties needed by the business event logger.
// It is the receiver of utility methods for consumers.
type Manager struct {
	done        chan bool
	bus         chan Message
	writersPool WriterPool
	msgCounter  Counter
}

// StartNewListener creates a new business event bus and adds
// a Writer to the WriterPool.
func StartNewListener(w WriteCloser) EventBusManager {
	wp := *NewWriterPool()
	wp.AddWriter(w)

	bem := Manager{
		done:        make(chan bool),
		bus:         make(chan Message),
		writersPool: wp,
		msgCounter:  Counter{0},
	}

	go bem.listen()

	return &bem
}

// Post sends a Message to the business event message bus
// for ingestion by all Writer's in the WriterPool.
func (bem *Manager) Post(m Message) {
	// ensure the bus is open for messages (i.e. "post office is open")
	if bem.bus == nil || bem.done == nil {
		log.Printf("the event bus is closed - lost message: %#v", m)
		return
	}

	bem.bus <- NewMesageEnvelop(m)
}

// AddWriter adds a Writer to the WriterPool.
func (bem *Manager) AddWriter(w Writer) {
	bem.writersPool.AddWriter(w)
}

func (bem *Manager) writeMessage(m Message) {
	bem.writersPool.Write(m)
}

// String implements Stringer.
func (bem *Manager) String() string {
	s := bem.writersPool.String()
	s += fmt.Sprintf(" - total number of messages posted: %d", bem.msgCounter.Get())

	return fmt.Sprintf("%s", s)
}

// listen is the main loop of the business event loop.
func (bem *Manager) listen() {
	defer func() {
		bem.done <- true // Sending "Termination Pong" response
	}()

ListenerLoop:
	for {
		select {
		case m, ok := <-bem.bus:
			if ok {
				// Received a Message wrapped in a MessageEnvelop.
				// Call the writer to write it to destination - in this case to Kafka.
				bem.msgCounter.Inc()
				bem.writeMessage(m)
			}
		case <-bem.done:
			// Received "Termination Ping" request.
			// Drain the remaining messages on the bus and break out.
			for m := range bem.bus {
				bem.msgCounter.Inc()
				bem.writeMessage(m)
			}
			break ListenerLoop
		}
	}
}

// Close closes the channels in the Manager.
// The recommended approach is for a channel to be used unidirectionally and
// be closed by the sender rather than the receivers.
// This means that it is the responsibility of the Posters to close the
// event bus when no more messages are being posted.
// See an example implementation in main_test using a sync.WaitGroup.
func (bem *Manager) Close() error {
	if bem.done == nil {
		return errors.New("this event bus manager is already closed")
	}

	close(bem.bus)
	bem.done <- true
	<-bem.done
	bem.bus = nil

	close(bem.done)
	bem.done = nil

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"os"
	"time"

	"github.com/ethereum/go-ethereum/cmd/utils"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/state"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/rpc"
	"github.com/ethereum/go-ethereum/xeth"
	"github.com/robertkrimen/otto"
)

/*
node admin bindings
*/

func (js *jsre) adminBindings() {
	js.re.Set("admin", struct{}{})
	t, _ := js.re.Get("admin")
	admin := t.Object()
	admin.Set("suggestPeer", js.suggestPeer)
	admin.Set("startRPC", js.startRPC)
	admin.Set("nodeInfo", js.nodeInfo)
	admin.Set("peers", js.peers)
	admin.Set("newAccount", js.newAccount)
	admin.Set("unlock", js.unlock)
	admin.Set("import", js.importChain)
	admin.Set("export", js.exportChain)
	admin.Set("verbosity", js.verbosity)
	admin.Set("backtrace", js.backtrace)
	admin.Set("progress", js.downloadProgress)

	admin.Set("miner", struct{}{})
	t, _ = admin.Get("miner")
	miner := t.Object()
	miner.Set("start", js.startMining)
	miner.Set("stop", js.stopMining)
	miner.Set("hashrate", js.hashrate)
	miner.Set("setExtra", js.setExtra)

	admin.Set("debug", struct{}{})
	t, _ = admin.Get("debug")
	debug := t.Object()
	debug.Set("printBlock", js.printBlock)
	debug.Set("dumpBlock", js.dumpBlock)
	debug.Set("getBlockRlp", js.getBlockRlp)
	debug.Set("setHead", js.setHead)
}

func (js *jsre) getBlock(call otto.FunctionCall) (*types.Block, error) {
	var block *types.Block
	if len(call.ArgumentList) > 0 {
		if call.Argument(0).IsNumber() {
			num, _ := call.Argument(0).ToInteger()
			block = js.ethereum.ChainManager().GetBlockByNumber(uint64(num))
		} else if call.Argument(0).IsString() {
			hash, _ := call.Argument(0).ToString()
			block = js.ethereum.ChainManager().GetBlock(common.HexToHash(hash))
		} else {
			return nil, errors.New("invalid argument for dump. Either hex string or number")
		}
		return block, nil
	}

	return nil, errors.New("requires block number or block hash as argument")
}

func (js *jsre) setHead(call otto.FunctionCall) otto.Value {
	block, err := js.getBlock(call)
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}

	if block == nil {
		fmt.Println("block not found")
		return otto.UndefinedValue()
	}

	js.ethereum.ChainManager().SetHead(block)
	return otto.UndefinedValue()
}

func (js *jsre) downloadProgress(call otto.FunctionCall) otto.Value {
	current, max := js.ethereum.Downloader().Stats()

	return js.re.ToVal(fmt.Sprintf("%d/%d", current, max))
}

func (js *jsre) getBlockRlp(call otto.FunctionCall) otto.Value {
	block, err := js.getBlock(call)
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}

	if block == nil {
		fmt.Println("block not found")
		return otto.UndefinedValue()
	}

	encoded, _ := rlp.EncodeToBytes(block)
	return js.re.ToVal(fmt.Sprintf("%x", encoded))
}

func (js *jsre) setExtra(call otto.FunctionCall) otto.Value {
	extra, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}

	if len(extra) > 1024 {
		fmt.Println("error: cannot exceed 1024 bytes")
		return otto.UndefinedValue()
	}

	js.ethereum.Miner().SetExtra([]byte(extra))
	return otto.UndefinedValue()
}

func (js *jsre) hashrate(otto.FunctionCall) otto.Value {
	return js.re.ToVal(js.ethereum.Miner().HashRate())
}

func (js *jsre) backtrace(call otto.FunctionCall) otto.Value {
	tracestr, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}
	glog.GetTraceLocation().Set(tracestr)

	return otto.UndefinedValue()
}

func (js *jsre) verbosity(call otto.FunctionCall) otto.Value {
	v, err := call.Argument(0).ToInteger()
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}

	glog.SetV(int(v))
	return otto.UndefinedValue()
}

func (js *jsre) startMining(call otto.FunctionCall) otto.Value {
	_, err := call.Argument(0).ToInteger()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	// threads now ignored
	err = js.ethereum.StartMining()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	return otto.TrueValue()
}

func (js *jsre) stopMining(call otto.FunctionCall) otto.Value {
	js.ethereum.StopMining()
	return otto.TrueValue()
}

func (js *jsre) startRPC(call otto.FunctionCall) otto.Value {
	addr, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	port, err := call.Argument(1).ToInteger()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}

	config := rpc.RpcConfig{
		ListenAddress: addr,
		ListenPort:    uint(port),
		// CorsDomain:    ctx.GlobalString(RPCCORSDomainFlag.Name),
	}

	xeth := xeth.New(js.ethereum, nil)
	err = rpc.Start(xeth, config)

	if err != nil {
		fmt.Printf(err.Error())
		return otto.FalseValue()
	}

	return otto.TrueValue()
}

func (js *jsre) suggestPeer(call otto.FunctionCall) otto.Value {
	nodeURL, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	err = js.ethereum.SuggestPeer(nodeURL)
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	return otto.TrueValue()
}

func (js *jsre) unlock(call otto.FunctionCall) otto.Value {
	addr, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	seconds, err := call.Argument(2).ToInteger()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	arg := call.Argument(1)
	var passphrase string
	if arg.IsUndefined() {
		fmt.Println("Please enter a passphrase now.")
		passphrase, err = readPassword("Passphrase: ", true)
		if err != nil {
			utils.Fatalf("%v", err)
		}
	} else {
		passphrase, err = arg.ToString()
		if err != nil {
			fmt.Println(err)
			return otto.FalseValue()
		}
	}
	am := js.ethereum.AccountManager()
	// err := am.Unlock(common.FromHex(split[0]), split[1])
	// if err != nil {
	// 	utils.Fatalf("Unlock account failed '%v'", err)
	// }
	err = am.TimedUnlock(common.FromHex(addr), passphrase, time.Duration(seconds)*time.Second)
	if err != nil {
		fmt.Printf("Unlock account failed '%v'\n", err)
		return otto.FalseValue()
	}
	return otto.TrueValue()
}

func (js *jsre) newAccount(call otto.FunctionCall) otto.Value {
	arg := call.Argument(0)
	var passphrase string
	if arg.IsUndefined() {
		fmt.Println("The new account will be encrypted with a passphrase.")
		fmt.Println("Please enter a passphrase now.")
		auth, err := readPassword("Passphrase: ", true)
		if err != nil {
			utils.Fatalf("%v", err)
		}
		confirm, err := readPassword("Repeat Passphrase: ", false)
		if err != nil {
			utils.Fatalf("%v", err)
		}
		if auth != confirm {
			utils.Fatalf("Passphrases did not match.")
		}
		passphrase = auth
	} else {
		var err error
		passphrase, err = arg.ToString()
		if err != nil {
			fmt.Println(err)
			return otto.FalseValue()
		}
	}
	acct, err := js.ethereum.AccountManager().NewAccount(passphrase)
	if err != nil {
		fmt.Printf("Could not create the account: %v", err)
		return otto.UndefinedValue()
	}
	return js.re.ToVal(common.Bytes2Hex(acct.Address))
}

func (js *jsre) nodeInfo(call otto.FunctionCall) otto.Value {
	return js.re.ToVal(js.ethereum.NodeInfo())
}

func (js *jsre) peers(call otto.FunctionCall) otto.Value {
	return js.re.ToVal(js.ethereum.PeersInfo())
}

func (js *jsre) importChain(call otto.FunctionCall) otto.Value {
	if len(call.ArgumentList) == 0 {
		fmt.Println("err: require file name")
		return otto.FalseValue()
	}

	fn, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}

	var fh *os.File
	fh, err = os.OpenFile(fn, os.O_RDONLY, os.ModePerm)
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	defer fh.Close()

	var blocks types.Blocks
	if err = rlp.Decode(fh, &blocks); err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}

	js.ethereum.ChainManager().Reset()
	if err = js.ethereum.ChainManager().InsertChain(blocks); err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}

	return otto.TrueValue()
}

func (js *jsre) exportChain(call otto.FunctionCall) otto.Value {
	if len(call.ArgumentList) == 0 {
		fmt.Println("err: require file name")
		return otto.FalseValue()
	}

	fn, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	if err := utils.ExportChain(js.ethereum.ChainManager(), fn); err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	return otto.TrueValue()
}

func (js *jsre) printBlock(call otto.FunctionCall) otto.Value {
	var block *types.Block
	if len(call.ArgumentList) > 0 {
		if call.Argument(0).IsNumber() {
			num, _ := call.Argument(0).ToInteger()
			block = js.ethereum.ChainManager().GetBlockByNumber(uint64(num))
		} else if call.Argument(0).IsString() {
			hash, _ := call.Argument(0).ToString()
			block = js.ethereum.ChainManager().GetBlock(common.HexToHash(hash))
		} else {
			fmt.Println("invalid argument for dump. Either hex string or number")
		}

	} else {
		block = js.ethereum.ChainManager().CurrentBlock()
	}
	if block == nil {
		fmt.Println("block not found")
		return otto.UndefinedValue()
	}

	fmt.Println(block)

	return otto.UndefinedValue()
}

func (js *jsre) dumpBlock(call otto.FunctionCall) otto.Value {
	var block *types.Block
	if len(call.ArgumentList) > 0 {
		if call.Argument(0).IsNumber() {
			num, _ := call.Argument(0).ToInteger()
			block = js.ethereum.ChainManager().GetBlockByNumber(uint64(num))
		} else if call.Argument(0).IsString() {
			hash, _ := call.Argument(0).ToString()
			block = js.ethereum.ChainManager().GetBlock(common.HexToHash(hash))
		} else {
			fmt.Println("invalid argument for dump. Either hex string or number")
		}

	} else {
		block = js.ethereum.ChainManager().CurrentBlock()
	}
	if block == nil {
		fmt.Println("block not found")
		return otto.UndefinedValue()
	}

	statedb := state.New(block.Root(), js.ethereum.StateDb())
	dump := statedb.RawDump()
	return js.re.ToVal(dump)

}
<commit_msg>geth: added a `admin.debug.block` command which reprocess the block<commit_after>package main

import (
	"errors"
	"fmt"
	"os"
	"time"

	"github.com/ethereum/go-ethereum/cmd/utils"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/state"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/core/vm"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/ethereum/go-ethereum/rpc"
	"github.com/ethereum/go-ethereum/xeth"
	"github.com/robertkrimen/otto"
)

/*
node admin bindings
*/

func (js *jsre) adminBindings() {
	js.re.Set("admin", struct{}{})
	t, _ := js.re.Get("admin")
	admin := t.Object()
	admin.Set("suggestPeer", js.suggestPeer)
	admin.Set("startRPC", js.startRPC)
	admin.Set("nodeInfo", js.nodeInfo)
	admin.Set("peers", js.peers)
	admin.Set("newAccount", js.newAccount)
	admin.Set("unlock", js.unlock)
	admin.Set("import", js.importChain)
	admin.Set("export", js.exportChain)
	admin.Set("verbosity", js.verbosity)
	admin.Set("backtrace", js.backtrace)
	admin.Set("progress", js.downloadProgress)

	admin.Set("miner", struct{}{})
	t, _ = admin.Get("miner")
	miner := t.Object()
	miner.Set("start", js.startMining)
	miner.Set("stop", js.stopMining)
	miner.Set("hashrate", js.hashrate)
	miner.Set("setExtra", js.setExtra)

	admin.Set("debug", struct{}{})
	t, _ = admin.Get("debug")
	debug := t.Object()
	debug.Set("printBlock", js.printBlock)
	debug.Set("dumpBlock", js.dumpBlock)
	debug.Set("getBlockRlp", js.getBlockRlp)
	debug.Set("setHead", js.setHead)
	debug.Set("block", js.debugBlock)
}

func (js *jsre) getBlock(call otto.FunctionCall) (*types.Block, error) {
	var block *types.Block
	if len(call.ArgumentList) > 0 {
		if call.Argument(0).IsNumber() {
			num, _ := call.Argument(0).ToInteger()
			block = js.ethereum.ChainManager().GetBlockByNumber(uint64(num))
		} else if call.Argument(0).IsString() {
			hash, _ := call.Argument(0).ToString()
			block = js.ethereum.ChainManager().GetBlock(common.HexToHash(hash))
		} else {
			return nil, errors.New("invalid argument for dump. Either hex string or number")
		}
		return block, nil
	}

	return nil, errors.New("requires block number or block hash as argument")
}

func (js *jsre) debugBlock(call otto.FunctionCall) otto.Value {
	block, err := js.getBlock(call)
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}

	if block == nil {
		fmt.Println("block not found")
		return otto.UndefinedValue()
	}

	old := vm.Debug
	vm.Debug = true
	_, err = js.ethereum.BlockProcessor().RetryProcess(block)
	if err != nil {
		glog.Infoln(err)
	}
	vm.Debug = old

	return otto.UndefinedValue()
}

func (js *jsre) setHead(call otto.FunctionCall) otto.Value {
	block, err := js.getBlock(call)
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}

	if block == nil {
		fmt.Println("block not found")
		return otto.UndefinedValue()
	}

	js.ethereum.ChainManager().SetHead(block)
	return otto.UndefinedValue()
}

func (js *jsre) downloadProgress(call otto.FunctionCall) otto.Value {
	current, max := js.ethereum.Downloader().Stats()

	return js.re.ToVal(fmt.Sprintf("%d/%d", current, max))
}

func (js *jsre) getBlockRlp(call otto.FunctionCall) otto.Value {
	block, err := js.getBlock(call)
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}

	if block == nil {
		fmt.Println("block not found")
		return otto.UndefinedValue()
	}

	encoded, _ := rlp.EncodeToBytes(block)
	return js.re.ToVal(fmt.Sprintf("%x", encoded))
}

func (js *jsre) setExtra(call otto.FunctionCall) otto.Value {
	extra, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}

	if len(extra) > 1024 {
		fmt.Println("error: cannot exceed 1024 bytes")
		return otto.UndefinedValue()
	}

	js.ethereum.Miner().SetExtra([]byte(extra))
	return otto.UndefinedValue()
}

func (js *jsre) hashrate(otto.FunctionCall) otto.Value {
	return js.re.ToVal(js.ethereum.Miner().HashRate())
}

func (js *jsre) backtrace(call otto.FunctionCall) otto.Value {
	tracestr, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}
	glog.GetTraceLocation().Set(tracestr)

	return otto.UndefinedValue()
}

func (js *jsre) verbosity(call otto.FunctionCall) otto.Value {
	v, err := call.Argument(0).ToInteger()
	if err != nil {
		fmt.Println(err)
		return otto.UndefinedValue()
	}

	glog.SetV(int(v))
	return otto.UndefinedValue()
}

func (js *jsre) startMining(call otto.FunctionCall) otto.Value {
	_, err := call.Argument(0).ToInteger()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	// threads now ignored
	err = js.ethereum.StartMining()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	return otto.TrueValue()
}

func (js *jsre) stopMining(call otto.FunctionCall) otto.Value {
	js.ethereum.StopMining()
	return otto.TrueValue()
}

func (js *jsre) startRPC(call otto.FunctionCall) otto.Value {
	addr, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	port, err := call.Argument(1).ToInteger()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}

	config := rpc.RpcConfig{
		ListenAddress: addr,
		ListenPort:    uint(port),
		// CorsDomain:    ctx.GlobalString(RPCCORSDomainFlag.Name),
	}

	xeth := xeth.New(js.ethereum, nil)
	err = rpc.Start(xeth, config)

	if err != nil {
		fmt.Printf(err.Error())
		return otto.FalseValue()
	}

	return otto.TrueValue()
}

func (js *jsre) suggestPeer(call otto.FunctionCall) otto.Value {
	nodeURL, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	err = js.ethereum.SuggestPeer(nodeURL)
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	return otto.TrueValue()
}

func (js *jsre) unlock(call otto.FunctionCall) otto.Value {
	addr, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	seconds, err := call.Argument(2).ToInteger()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	arg := call.Argument(1)
	var passphrase string
	if arg.IsUndefined() {
		fmt.Println("Please enter a passphrase now.")
		passphrase, err = readPassword("Passphrase: ", true)
		if err != nil {
			utils.Fatalf("%v", err)
		}
	} else {
		passphrase, err = arg.ToString()
		if err != nil {
			fmt.Println(err)
			return otto.FalseValue()
		}
	}
	am := js.ethereum.AccountManager()
	// err := am.Unlock(common.FromHex(split[0]), split[1])
	// if err != nil {
	// 	utils.Fatalf("Unlock account failed '%v'", err)
	// }
	err = am.TimedUnlock(common.FromHex(addr), passphrase, time.Duration(seconds)*time.Second)
	if err != nil {
		fmt.Printf("Unlock account failed '%v'\n", err)
		return otto.FalseValue()
	}
	return otto.TrueValue()
}

func (js *jsre) newAccount(call otto.FunctionCall) otto.Value {
	arg := call.Argument(0)
	var passphrase string
	if arg.IsUndefined() {
		fmt.Println("The new account will be encrypted with a passphrase.")
		fmt.Println("Please enter a passphrase now.")
		auth, err := readPassword("Passphrase: ", true)
		if err != nil {
			utils.Fatalf("%v", err)
		}
		confirm, err := readPassword("Repeat Passphrase: ", false)
		if err != nil {
			utils.Fatalf("%v", err)
		}
		if auth != confirm {
			utils.Fatalf("Passphrases did not match.")
		}
		passphrase = auth
	} else {
		var err error
		passphrase, err = arg.ToString()
		if err != nil {
			fmt.Println(err)
			return otto.FalseValue()
		}
	}
	acct, err := js.ethereum.AccountManager().NewAccount(passphrase)
	if err != nil {
		fmt.Printf("Could not create the account: %v", err)
		return otto.UndefinedValue()
	}
	return js.re.ToVal(common.Bytes2Hex(acct.Address))
}

func (js *jsre) nodeInfo(call otto.FunctionCall) otto.Value {
	return js.re.ToVal(js.ethereum.NodeInfo())
}

func (js *jsre) peers(call otto.FunctionCall) otto.Value {
	return js.re.ToVal(js.ethereum.PeersInfo())
}

func (js *jsre) importChain(call otto.FunctionCall) otto.Value {
	if len(call.ArgumentList) == 0 {
		fmt.Println("err: require file name")
		return otto.FalseValue()
	}

	fn, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}

	var fh *os.File
	fh, err = os.OpenFile(fn, os.O_RDONLY, os.ModePerm)
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	defer fh.Close()

	var blocks types.Blocks
	if err = rlp.Decode(fh, &blocks); err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}

	js.ethereum.ChainManager().Reset()
	if err = js.ethereum.ChainManager().InsertChain(blocks); err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}

	return otto.TrueValue()
}

func (js *jsre) exportChain(call otto.FunctionCall) otto.Value {
	if len(call.ArgumentList) == 0 {
		fmt.Println("err: require file name")
		return otto.FalseValue()
	}

	fn, err := call.Argument(0).ToString()
	if err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	if err := utils.ExportChain(js.ethereum.ChainManager(), fn); err != nil {
		fmt.Println(err)
		return otto.FalseValue()
	}
	return otto.TrueValue()
}

func (js *jsre) printBlock(call otto.FunctionCall) otto.Value {
	var block *types.Block
	if len(call.ArgumentList) > 0 {
		if call.Argument(0).IsNumber() {
			num, _ := call.Argument(0).ToInteger()
			block = js.ethereum.ChainManager().GetBlockByNumber(uint64(num))
		} else if call.Argument(0).IsString() {
			hash, _ := call.Argument(0).ToString()
			block = js.ethereum.ChainManager().GetBlock(common.HexToHash(hash))
		} else {
			fmt.Println("invalid argument for dump. Either hex string or number")
		}

	} else {
		block = js.ethereum.ChainManager().CurrentBlock()
	}
	if block == nil {
		fmt.Println("block not found")
		return otto.UndefinedValue()
	}

	fmt.Println(block)

	return otto.UndefinedValue()
}

func (js *jsre) dumpBlock(call otto.FunctionCall) otto.Value {
	var block *types.Block
	if len(call.ArgumentList) > 0 {
		if call.Argument(0).IsNumber() {
			num, _ := call.Argument(0).ToInteger()
			block = js.ethereum.ChainManager().GetBlockByNumber(uint64(num))
		} else if call.Argument(0).IsString() {
			hash, _ := call.Argument(0).ToString()
			block = js.ethereum.ChainManager().GetBlock(common.HexToHash(hash))
		} else {
			fmt.Println("invalid argument for dump. Either hex string or number")
		}

	} else {
		block = js.ethereum.ChainManager().CurrentBlock()
	}
	if block == nil {
		fmt.Println("block not found")
		return otto.UndefinedValue()
	}

	statedb := state.New(block.Root(), js.ethereum.StateDb())
	dump := statedb.RawDump()
	return js.re.ToVal(dump)

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"runtime"

	"github.com/codegangsta/cli"
	"github.com/grafana/grafana/pkg/cmd/grafana-cli/commands"
	"github.com/grafana/grafana/pkg/cmd/grafana-cli/log"
)

var version = "master"

func getGrafanaPluginDir() string {
	os := runtime.GOOS
	if os == "windows" {
		return "C:\\opt\\grafana\\plugins"
	} else {
		return "/var/lib/grafana/plugins"
	}
}

func main() {
	SetupLogging()

	app := cli.NewApp()
	app.Name = "Grafana cli"
	app.Usage = ""
	app.Author = "Grafana Project"
	app.Email = "https://github.com/grafana/grafana"
	app.Version = version
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "pluginsDir",
			Usage:  "path to the grafana plugin directory",
			Value:  getGrafanaPluginDir(),
			EnvVar: "GF_PLUGIN_DIR",
		},
		cli.StringFlag{
			Name:   "repo",
			Usage:  "url to the plugin repository",
			Value:  "https://grafana.net/api/plugins",
			EnvVar: "GF_PLUGIN_REPO",
		},
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "enable debug logging",
		},
	}

	app.Commands = commands.Commands
	app.CommandNotFound = cmdNotFound

	if err := app.Run(os.Args); err != nil {
		log.Errorf("%v", err)
	}
}

func SetupLogging() {
	for _, f := range os.Args {
		if f == "-D" || f == "--debug" || f == "-debug" {
			log.SetDebug(true)
		}
	}
}

func cmdNotFound(c *cli.Context, command string) {
	fmt.Printf(
		"%s: '%s' is not a %s command. See '%s --help'.\n",
		c.App.Name,
		command,
		c.App.Name,
		os.Args[0],
	)
	os.Exit(1)
}
<commit_msg>feat(cli): adds better support for plugin folder in dev<commit_after>package main

import (
	"fmt"
	"os"
	"runtime"

	"github.com/codegangsta/cli"
	"github.com/grafana/grafana/pkg/cmd/grafana-cli/commands"
	"github.com/grafana/grafana/pkg/cmd/grafana-cli/log"
	"strings"
)

var version = "master"

func getGrafanaPluginDir() string {
	currentOS := runtime.GOOS
	defaultNix := "/var/lib/grafana/plugins"

	if currentOS == "windows" {
		return "C:\\opt\\grafana\\plugins"
	}

	pwd, err := os.Getwd()

	if err != nil {
		log.Error("Could not get current path. using default")
		return defaultNix
	}

	if isDevenvironment(pwd) {
		return "../../../data/plugins"
	}

	return defaultNix
}

func isDevenvironment(pwd string) bool {
	// if grafana-cli is executed from the cmd folder we can assume
	// that its in development environment.
	return strings.HasSuffix(pwd, "/pkg/cmd/grafana-cli")
}

func main() {
	SetupLogging()

	app := cli.NewApp()
	app.Name = "Grafana cli"
	app.Usage = ""
	app.Author = "Grafana Project"
	app.Email = "https://github.com/grafana/grafana"
	app.Version = version
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "pluginsDir",
			Usage:  "path to the grafana plugin directory",
			Value:  getGrafanaPluginDir(),
			EnvVar: "GF_PLUGIN_DIR",
		},
		cli.StringFlag{
			Name:   "repo",
			Usage:  "url to the plugin repository",
			Value:  "https://grafana.net/api/plugins",
			EnvVar: "GF_PLUGIN_REPO",
		},
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "enable debug logging",
		},
	}

	app.Commands = commands.Commands
	app.CommandNotFound = cmdNotFound

	if err := app.Run(os.Args); err != nil {
		log.Errorf("%v", err)
	}
}

func SetupLogging() {
	for _, f := range os.Args {
		if f == "-D" || f == "--debug" || f == "-debug" {
			log.SetDebug(true)
		}
	}
}

func cmdNotFound(c *cli.Context, command string) {
	fmt.Printf(
		"%s: '%s' is not a %s command. See '%s --help'.\n",
		c.App.Name,
		command,
		c.App.Name,
		os.Args[0],
	)
	os.Exit(1)
}
<|endoftext|>


<commit_before>// Command gobin  Go commands by reading their import paths, go-getting
// their repositories and building new executables.
//
// Command gobin can search for and list Go executables, create a $GOPATH workspace
// out of a binary. It can also update an executable which does not have its
// sources fetched in current $GOPATH.
//
// Source
//
// Command gobin can guess an origin of the sources used to build Go executables
// it finds on a system. It can be used to create precise mirror of the sources
// for system Go executables, without any unneeded packages.
//
//   ~ $ gobin -s /tmp/gopath godoc golint goimports gotree gowhich
//   code.google.com/p/go.tools (download)
//   github.com/rjeczalik/tools (download)
//   github.com/rjeczalik/which (download)
//   (...)
//   github.com/rjeczalik/which/cmd/gowhich
//   code.google.com/p/go.tools/cmd/godoc
//   ~ $ tree -L 3 /tmp/gopath/src/
//   /tmp/gopath/src/
//   âââ code.google.com
//   âÂ Â  âââ p
//   âÂ Â      âââ go.tools
//   âââ github.com
//       âââ rjeczalik
//           âââ tools
//           âââ which
//
//   7 directories, 0 files
//
// Update
//
// Command gobin can update single executable or automagically discover all of
// them in directories specified in $GOPATH, $GOBIN and $GOPATH environment
// variables.
//
// Executing gobin without any arguments makes it list all Go executables found
// in $GOPATH, $GOBIN and $GOPATH.
//
// Updating multiple executables is performed on multiple goroutines, bumping
// up $GOMAXPROCS environment variable may speed up the overall run-time
// significantly.
//
// Example
//
//   ~ $ GOMAXPROCS=2 gobin -u
//   ok 	/home/rjeczalik/bin/goimports	(code.google.com/p/go.tools/cmd/goimports)	13.966s
//   ok 	/home/rjeczalik/bin/godoc	(code.google.com/p/go.tools/cmd/godoc)	17.960s
//   ok 	/home/rjeczalik/bin/pulsecli	(github.com/x-formation/pulsekit/cmd/pulsecli)	13.052s
//   ok 	/home/rjeczalik/workspace/bin/pulsecli	(github.com/x-formation/pulsekit/cmd/pulsecli)	13.052s
//
// Usage
//
//   NAME:
//       gobin - looks for Go executables system-wide ($PATH/$GOBIN/$GOPATH),
//               lists them, fetches their sources and updates them
//
//   USAGE:
//       gobin [-u] [-s=.|gopath] [path|package...]
//
//   FLAGS:
//       -u                Updates Go binaries
//       -s <dir>          Go-gets sources for Go specified binaries into <dir> $GOPATH
//                         (use '.' for current $GOPATH)
//       -ldflags=<flags>  passes "-ldflags=flags" to "go install"
//
//   EXAMPLES:
//       gobin                    Lists all Go binaries (looks up $PATH/$GOBIN/$GOPATH)
//       gobin -s=. ~/bin         Go-gets sources used to build all Go binaries in ~/bin
//                                into current $GOPATH
//       gobin -s=/var/mirror     Go-gets all sources used to build all Go binaries found
//                                on system into new /var/mirror $GOPATH
//       gobin -u                 Updates all Go binaries
//       gobin -u github.com      Updates all Go binaries installed from github.com
//       gobin ~/bin              Lists all Go binaries from the ~/bin directory
package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/rjeczalik/bin"
)

func die(v interface{}) {
	fmt.Fprintln(os.Stderr, v)
	os.Exit(1)
}

const usage = `NAME:
	gobin - performs discovery of Go executables ($PATH/$GOBIN/$GOPATH),
	        lists them, fetches their sources and updates them

USAGE:
	gobin [-u] [-s=.|gopath] [path|package...]

FLAGS:
	-u                Updates Go binaries
	-s <dir>          Go-gets sources for Go specified binaries into <dir> $GOPATH
	                  (use '.' for current $GOPATH)
	-ldflags=<flags>  passes "-ldflags=flags" to "go install"

EXAMPLES:
	gobin                    Lists all Go binaries (looks up $PATH/$GOBIN/$GOPATH)
	gobin -s=. ~/bin         Go-gets sources used to build all Go binaries in ~/bin
	                         into current $GOPATH
	gobin -s=/var/mirror     Go-gets all sources used to build all Go binaries found
	                         on system into new /var/mirror $GOPATH
	gobin -u                 Updates all Go binaries in-place
	gobin -u github.com      Updates all Go binaries installed from github.com
	gobin ~/bin              Lists all Go binaries from the ~/bin directory
	gobin -u -ldflags='-w -s'	Updates all go binaries in-place, using "go install -ldflags='-w -s'"`

var (
	source, ldflags string
	update          bool
)

func ishelp(s string) bool {
	return s == "-h" || s == "-help" || s == "help" || s == "--help" || s == "/?"
}

func parse() []string {
	flag.Usage = func() { die(usage) }
	flag.StringVar(&source, "s", "", "")
	flag.BoolVar(&update, "u", false, "")
	flag.StringVar(&ldflags, "ldflags", "", "")
	flag.Parse()
	return flag.Args()
}

func self() string {
	if strings.Contains(os.Args[0], string(os.PathSeparator)) {
		if self, err := filepath.Abs(os.Args[0]); err == nil {
			if fiself, err := os.Stat(self); err == nil {
				if fiargs, err := os.Stat(os.Args[0]); err == nil && os.SameFile(fiself, fiargs) {
					return self
				}
			}
		}
	}
	if self, err := exec.LookPath(filepath.Base(os.Args[0])); err == nil {
		return self
	}
	return ""
}

func log(b *bin.Bin, d time.Duration, err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "fail\t%s\t(%s)\n", b.Path, b.Package)
		fmt.Fprintf(os.Stderr, "\terror: %v\n", err)
	} else {
		fmt.Printf("ok\t%s\t(%s)\t%.3fs\n", b.Path, b.Package, d.Seconds())
	}
}

// TODO(rjeczalik): Bin.CanWrite needs a Flock here
func main() {
	if len(os.Args) == 2 && ishelp(os.Args[1]) {
		fmt.Println(usage)
		return
	}
	var b, e = bin.Search(parse())
	if e != nil {
		die(e)
	}
	var installFlags []string
	switch {
	case update:
		if self := self(); self != "" {
			for i := range b {
				if b[i].Path == self {
					b[i], b = b[len(b)-1], b[:len(b)-1]
					break
				}
			}
		}
		if ldflags != "" {
			installFlags = append(installFlags, "-ldflags="+ldflags)
		}
		bin.Update(b, log, installFlags...)
	case source != "":
		if source == "." {
			source = os.Getenv("GOPATH")
			if source == "" {
				die("bin: unable to read current $GOPATH or $GOPATH is empty")
			}
			if i := strings.Index(source, string(os.PathListSeparator)); i != -1 {
				source = source[:i]
			}
		}
		if bin.Source(b, source) != nil {
			os.Exit(1)
		}
	default:
		for i := range b {
			fmt.Printf("%s\t(%s)\n", b[i].Path, b[i].Package)
		}
	}
}
<commit_msg>revert tree output borking<commit_after>// Command gobin  Go commands by reading their import paths, go-getting
// their repositories and building new executables.
//
// Command gobin can search for and list Go executables, create a $GOPATH workspace
// out of a binary. It can also update an executable which does not have its
// sources fetched in current $GOPATH.
//
// Source
//
// Command gobin can guess an origin of the sources used to build Go executables
// it finds on a system. It can be used to create precise mirror of the sources
// for system Go executables, without any unneeded packages.
//
//   ~ $ gobin -s /tmp/gopath godoc golint goimports gotree gowhich
//   code.google.com/p/go.tools (download)
//   github.com/rjeczalik/tools (download)
//   github.com/rjeczalik/which (download)
//   (...)
//   github.com/rjeczalik/which/cmd/gowhich
//   code.google.com/p/go.tools/cmd/godoc
//   ~ $ tree -L 3 /tmp/gopath/src/
//   /tmp/gopath/src/
//   ├── code.google.com
//   │   └── p
//   │       └── go.tools
//   └── github.com
//       └── rjeczalik
//           ├── tools
//           └── which
//
//   7 directories, 0 files
//
// Update
//
// Command gobin can update single executable or automagically discover all of
// them in directories specified in $GOPATH, $GOBIN and $GOPATH environment
// variables.
//
// Executing gobin without any arguments makes it list all Go executables found
// in $GOPATH, $GOBIN and $GOPATH.
//
// Updating multiple executables is performed on multiple goroutines, bumping
// up $GOMAXPROCS environment variable may speed up the overall run-time
// significantly.
//
// Example
//
//   ~ $ GOMAXPROCS=2 gobin -u
//   ok 	/home/rjeczalik/bin/goimports	(code.google.com/p/go.tools/cmd/goimports)	13.966s
//   ok 	/home/rjeczalik/bin/godoc	(code.google.com/p/go.tools/cmd/godoc)	17.960s
//   ok 	/home/rjeczalik/bin/pulsecli	(github.com/x-formation/pulsekit/cmd/pulsecli)	13.052s
//   ok 	/home/rjeczalik/workspace/bin/pulsecli	(github.com/x-formation/pulsekit/cmd/pulsecli)	13.052s
//
// Usage
//
//   NAME:
//       gobin - looks for Go executables system-wide ($PATH/$GOBIN/$GOPATH),
//               lists them, fetches their sources and updates them
//
//   USAGE:
//       gobin [-u] [-s=.|gopath] [path|package...]
//
//   FLAGS:
//       -u                Updates Go binaries
//       -s <dir>          Go-gets sources for Go specified binaries into <dir> $GOPATH
//                         (use '.' for current $GOPATH)
//       -ldflags=<flags>  passes "-ldflags=flags" to "go install"
//
//   EXAMPLES:
//       gobin                    Lists all Go binaries (looks up $PATH/$GOBIN/$GOPATH)
//       gobin -s=. ~/bin         Go-gets sources used to build all Go binaries in ~/bin
//                                into current $GOPATH
//       gobin -s=/var/mirror     Go-gets all sources used to build all Go binaries found
//                                on system into new /var/mirror $GOPATH
//       gobin -u                 Updates all Go binaries
//       gobin -u github.com      Updates all Go binaries installed from github.com
//       gobin ~/bin              Lists all Go binaries from the ~/bin directory
package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/rjeczalik/bin"
)

func die(v interface{}) {
	fmt.Fprintln(os.Stderr, v)
	os.Exit(1)
}

const usage = `NAME:
	gobin - performs discovery of Go executables ($PATH/$GOBIN/$GOPATH),
	        lists them, fetches their sources and updates them

USAGE:
	gobin [-u] [-s=.|gopath] [path|package...]

FLAGS:
	-u                Updates Go binaries
	-s <dir>          Go-gets sources for Go specified binaries into <dir> $GOPATH
	                  (use '.' for current $GOPATH)
	-ldflags=<flags>  passes "-ldflags=flags" to "go install"

EXAMPLES:
	gobin                    Lists all Go binaries (looks up $PATH/$GOBIN/$GOPATH)
	gobin -s=. ~/bin         Go-gets sources used to build all Go binaries in ~/bin
	                         into current $GOPATH
	gobin -s=/var/mirror     Go-gets all sources used to build all Go binaries found
	                         on system into new /var/mirror $GOPATH
	gobin -u                 Updates all Go binaries in-place
	gobin -u github.com      Updates all Go binaries installed from github.com
	gobin ~/bin              Lists all Go binaries from the ~/bin directory
	gobin -u -ldflags='-w -s'	Updates all go binaries in-place, using "go install -ldflags='-w -s'"`

var (
	source, ldflags string
	update          bool
)

func ishelp(s string) bool {
	return s == "-h" || s == "-help" || s == "help" || s == "--help" || s == "/?"
}

func parse() []string {
	flag.Usage = func() { die(usage) }
	flag.StringVar(&source, "s", "", "")
	flag.BoolVar(&update, "u", false, "")
	flag.StringVar(&ldflags, "ldflags", "", "")
	flag.Parse()
	return flag.Args()
}

func self() string {
	if strings.Contains(os.Args[0], string(os.PathSeparator)) {
		if self, err := filepath.Abs(os.Args[0]); err == nil {
			if fiself, err := os.Stat(self); err == nil {
				if fiargs, err := os.Stat(os.Args[0]); err == nil && os.SameFile(fiself, fiargs) {
					return self
				}
			}
		}
	}
	if self, err := exec.LookPath(filepath.Base(os.Args[0])); err == nil {
		return self
	}
	return ""
}

func log(b *bin.Bin, d time.Duration, err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "fail\t%s\t(%s)\n", b.Path, b.Package)
		fmt.Fprintf(os.Stderr, "\terror: %v\n", err)
	} else {
		fmt.Printf("ok\t%s\t(%s)\t%.3fs\n", b.Path, b.Package, d.Seconds())
	}
}

// TODO(rjeczalik): Bin.CanWrite needs a Flock here
func main() {
	if len(os.Args) == 2 && ishelp(os.Args[1]) {
		fmt.Println(usage)
		return
	}
	var b, e = bin.Search(parse())
	if e != nil {
		die(e)
	}
	var installFlags []string
	switch {
	case update:
		if self := self(); self != "" {
			for i := range b {
				if b[i].Path == self {
					b[i], b = b[len(b)-1], b[:len(b)-1]
					break
				}
			}
		}
		if ldflags != "" {
			installFlags = append(installFlags, "-ldflags="+ldflags)
		}
		bin.Update(b, log, installFlags...)
	case source != "":
		if source == "." {
			source = os.Getenv("GOPATH")
			if source == "" {
				die("bin: unable to read current $GOPATH or $GOPATH is empty")
			}
			if i := strings.Index(source, string(os.PathListSeparator)); i != -1 {
				source = source[:i]
			}
		}
		if bin.Source(b, source) != nil {
			os.Exit(1)
		}
	default:
		for i := range b {
			fmt.Printf("%s\t(%s)\n", b[i].Path, b[i].Package)
		}
	}
}
<|endoftext|>


<commit_before>// Package httpproxy is a cache implementation that can proxy artifacts
// from/to another HTTP-based remote cache.
package httpproxy

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/buchgr/bazel-remote/cache"
	"github.com/buchgr/bazel-remote/cache/disk/casblob"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

type uploadReq struct {
	hash string
	size int64
	kind cache.EntryKind
	rc   io.ReadCloser
}

type remoteHTTPProxyCache struct {
	remote       *http.Client
	baseURL      string
	uploadQueue  chan<- uploadReq
	accessLogger cache.Logger
	errorLogger  cache.Logger
	requestURL   func(hash string, kind cache.EntryKind) string
	v2mode       bool
}

var (
	cacheHits = promauto.NewCounter(prometheus.CounterOpts{
		Name: "bazel_remote_http_cache_hits",
		Help: "The total number of HTTP backend cache hits",
	})
	cacheMisses = promauto.NewCounter(prometheus.CounterOpts{
		Name: "bazel_remote_http_cache_misses",
		Help: "The total number of HTTP backend cache misses",
	})
)

func (r *remoteHTTPProxyCache) uploadFile(item uploadReq) {

	if item.size == 0 {
		item.rc.Close()
		// See https://github.com/golang/go/issues/20257#issuecomment-299509391
		item.rc = http.NoBody
	}

	url := r.requestURL(item.hash, item.kind)

	req, err := http.NewRequestWithContext(context.Background(), http.MethodHead, url, nil)
	if err != nil {
		r.errorLogger.Printf("INTERNAL ERROR, FAILED TO SETUP HTTP PROXY UPLOAD %s: %s", url, err)
		item.rc.Close()
		return
	}

	rsp, err := r.remote.Do(req)
	if err == nil && rsp.StatusCode == http.StatusOK {
		r.accessLogger.Printf("SKIP UPLOAD %s", item.hash)
		item.rc.Close()
		return
	}

	req, err = http.NewRequestWithContext(context.Background(), http.MethodPut, url, item.rc)
	if err != nil {
		r.errorLogger.Printf("INTERNAL ERROR, FAILED TO SETUP HTTP PROXY UPLOAD %s: %s", url, err)

		// item.rc will be closed if we call req.Do(), but not if we
		// return earlier.
		item.rc.Close()

		return
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	req.ContentLength = item.size

	rsp, err = r.remote.Do(req)
	if err != nil {
		r.errorLogger.Printf("HTTP %s UPLOAD: %s", url, err.Error())
		return
	}
	_, err = io.Copy(ioutil.Discard, rsp.Body)
	if err != nil {
		r.errorLogger.Printf("HTTP %s UPLOAD: %s", url, err.Error())
		return
	}
	rsp.Body.Close()

	logResponse(r.accessLogger, "UPLOAD", rsp.StatusCode, url)
}

// New creates a cache that proxies requests to a HTTP remote cache.
// `storageMode` must be one of "uncompressed" (which expects legacy
// CAS blobs) or "zstd" (which expects cas.v2 blobs).
func New(baseURL *url.URL, storageMode string, remote *http.Client,
	accessLogger cache.Logger, errorLogger cache.Logger,
	numUploaders, maxQueuedUploads int) (cache.Proxy, error) {

	proxy := &remoteHTTPProxyCache{
		remote:       remote,
		baseURL:      strings.TrimRight(baseURL.String(), "/"),
		accessLogger: accessLogger,
		errorLogger:  errorLogger,
		v2mode:       storageMode == "zstd",
	}

	if storageMode == "zstd" {
		proxy.requestURL = func(hash string, kind cache.EntryKind) string {
			if kind == cache.CAS {
				return fmt.Sprintf("%s/cas.v2/%s", proxy.baseURL, hash)
			}

			return fmt.Sprintf("%s/%s/%s", proxy.baseURL, kind, hash)
		}
	} else if storageMode == "uncompressed" {
		proxy.requestURL = func(hash string, kind cache.EntryKind) string {
			return fmt.Sprintf("%s/%s/%s", proxy.baseURL, kind, hash)
		}
	} else {
		return nil, fmt.Errorf("Invalid http_proxy.mode specified: %q",
			storageMode)
	}

	if maxQueuedUploads > 0 && numUploaders > 0 {
		uploadQueue := make(chan uploadReq, maxQueuedUploads)

		for i := 0; i < numUploaders; i++ {
			go func() {
				for item := range uploadQueue {
					proxy.uploadFile(item)
				}
			}()
		}

		proxy.uploadQueue = uploadQueue
	}

	return proxy, nil
}

// Helper function for logging responses
func logResponse(logger cache.Logger, method string, code int, url string) {
	logger.Printf("HTTP %s %d %s", method, code, url)
}

func (r *remoteHTTPProxyCache) Put(ctx context.Context, kind cache.EntryKind, hash string, size int64, rc io.ReadCloser) {
	if r.uploadQueue == nil {
		rc.Close()
		return
	}

	item := uploadReq{
		hash: hash,
		size: size,
		kind: kind,
		rc:   rc,
	}

	select {
	case r.uploadQueue <- item:
	default:
		r.errorLogger.Printf("too many uploads queued")
		rc.Close()
	}
}

func (r *remoteHTTPProxyCache) Get(ctx context.Context, kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
	url := r.requestURL(hash, kind)

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		cacheMisses.Inc()
		return nil, -1, err
	}

	rsp, err := r.remote.Do(req)
	if err != nil {
		cacheMisses.Inc()
		return nil, -1, err
	}

	logResponse(r.accessLogger, "DOWNLOAD", rsp.StatusCode, url)

	if rsp.StatusCode == http.StatusNotFound {
		cacheMisses.Inc()
		return nil, -1, nil
	}

	if rsp.StatusCode != http.StatusOK {
		// If the failed http response contains some data then
		// forward up to 1 KiB.
		var errorBytes []byte
		errorBytes, err = ioutil.ReadAll(io.LimitReader(rsp.Body, 1024))
		var errorText string
		if err == nil {
			errorText = string(errorBytes)
		}

		cacheMisses.Inc()
		return nil, -1, &cache.Error{
			Code: rsp.StatusCode,
			Text: errorText,
		}
	}

	if kind == cache.CAS && r.v2mode {
		cacheHits.Inc()
		return casblob.ExtractLogicalSize(rsp.Body)
	}

	sizeBytesStr := rsp.Header.Get("Content-Length")
	if sizeBytesStr == "" {
		err = errors.New("Missing Content-Length header")
		cacheMisses.Inc()
		return nil, -1, err
	}

	sizeBytesInt, err := strconv.Atoi(sizeBytesStr)
	if err != nil {
		cacheMisses.Inc()
		return nil, -1, err
	}
	sizeBytes := int64(sizeBytesInt)

	cacheHits.Inc()

	return rsp.Body, sizeBytes, nil
}

func (r *remoteHTTPProxyCache) Contains(ctx context.Context, kind cache.EntryKind, hash string) (bool, int64) {

	url := r.requestURL(hash, kind)

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return false, -1
	}

	rsp, err := r.remote.Do(req)
	if err == nil && rsp.StatusCode == http.StatusOK {
		if kind != cache.CAS {
			return true, rsp.ContentLength
		}

		// We don't know the content size without reading the file header
		// and that could be very costly for the backend server. So return
		// "unknown size".
		return true, -1
	}

	return false, -1
}
<commit_msg>Make remoteHTTPProxyCache.Contains do HEAD request<commit_after>// Package httpproxy is a cache implementation that can proxy artifacts
// from/to another HTTP-based remote cache.
package httpproxy

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/buchgr/bazel-remote/cache"
	"github.com/buchgr/bazel-remote/cache/disk/casblob"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

type uploadReq struct {
	hash string
	size int64
	kind cache.EntryKind
	rc   io.ReadCloser
}

type remoteHTTPProxyCache struct {
	remote       *http.Client
	baseURL      string
	uploadQueue  chan<- uploadReq
	accessLogger cache.Logger
	errorLogger  cache.Logger
	requestURL   func(hash string, kind cache.EntryKind) string
	v2mode       bool
}

var (
	cacheHits = promauto.NewCounter(prometheus.CounterOpts{
		Name: "bazel_remote_http_cache_hits",
		Help: "The total number of HTTP backend cache hits",
	})
	cacheMisses = promauto.NewCounter(prometheus.CounterOpts{
		Name: "bazel_remote_http_cache_misses",
		Help: "The total number of HTTP backend cache misses",
	})
)

func (r *remoteHTTPProxyCache) uploadFile(item uploadReq) {

	if item.size == 0 {
		item.rc.Close()
		// See https://github.com/golang/go/issues/20257#issuecomment-299509391
		item.rc = http.NoBody
	}

	url := r.requestURL(item.hash, item.kind)

	req, err := http.NewRequestWithContext(context.Background(), http.MethodHead, url, nil)
	if err != nil {
		r.errorLogger.Printf("INTERNAL ERROR, FAILED TO SETUP HTTP PROXY UPLOAD %s: %s", url, err)
		item.rc.Close()
		return
	}

	rsp, err := r.remote.Do(req)
	if err == nil && rsp.StatusCode == http.StatusOK {
		r.accessLogger.Printf("SKIP UPLOAD %s", item.hash)
		item.rc.Close()
		return
	}

	req, err = http.NewRequestWithContext(context.Background(), http.MethodPut, url, item.rc)
	if err != nil {
		r.errorLogger.Printf("INTERNAL ERROR, FAILED TO SETUP HTTP PROXY UPLOAD %s: %s", url, err)

		// item.rc will be closed if we call req.Do(), but not if we
		// return earlier.
		item.rc.Close()

		return
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	req.ContentLength = item.size

	rsp, err = r.remote.Do(req)
	if err != nil {
		r.errorLogger.Printf("HTTP %s UPLOAD: %s", url, err.Error())
		return
	}
	_, err = io.Copy(ioutil.Discard, rsp.Body)
	if err != nil {
		r.errorLogger.Printf("HTTP %s UPLOAD: %s", url, err.Error())
		return
	}
	rsp.Body.Close()

	logResponse(r.accessLogger, "UPLOAD", rsp.StatusCode, url)
}

// New creates a cache that proxies requests to a HTTP remote cache.
// `storageMode` must be one of "uncompressed" (which expects legacy
// CAS blobs) or "zstd" (which expects cas.v2 blobs).
func New(baseURL *url.URL, storageMode string, remote *http.Client,
	accessLogger cache.Logger, errorLogger cache.Logger,
	numUploaders, maxQueuedUploads int) (cache.Proxy, error) {

	proxy := &remoteHTTPProxyCache{
		remote:       remote,
		baseURL:      strings.TrimRight(baseURL.String(), "/"),
		accessLogger: accessLogger,
		errorLogger:  errorLogger,
		v2mode:       storageMode == "zstd",
	}

	if storageMode == "zstd" {
		proxy.requestURL = func(hash string, kind cache.EntryKind) string {
			if kind == cache.CAS {
				return fmt.Sprintf("%s/cas.v2/%s", proxy.baseURL, hash)
			}

			return fmt.Sprintf("%s/%s/%s", proxy.baseURL, kind, hash)
		}
	} else if storageMode == "uncompressed" {
		proxy.requestURL = func(hash string, kind cache.EntryKind) string {
			return fmt.Sprintf("%s/%s/%s", proxy.baseURL, kind, hash)
		}
	} else {
		return nil, fmt.Errorf("Invalid http_proxy.mode specified: %q",
			storageMode)
	}

	if maxQueuedUploads > 0 && numUploaders > 0 {
		uploadQueue := make(chan uploadReq, maxQueuedUploads)

		for i := 0; i < numUploaders; i++ {
			go func() {
				for item := range uploadQueue {
					proxy.uploadFile(item)
				}
			}()
		}

		proxy.uploadQueue = uploadQueue
	}

	return proxy, nil
}

// Helper function for logging responses
func logResponse(logger cache.Logger, method string, code int, url string) {
	logger.Printf("HTTP %s %d %s", method, code, url)
}

func (r *remoteHTTPProxyCache) Put(ctx context.Context, kind cache.EntryKind, hash string, size int64, rc io.ReadCloser) {
	if r.uploadQueue == nil {
		rc.Close()
		return
	}

	item := uploadReq{
		hash: hash,
		size: size,
		kind: kind,
		rc:   rc,
	}

	select {
	case r.uploadQueue <- item:
	default:
		r.errorLogger.Printf("too many uploads queued")
		rc.Close()
	}
}

func (r *remoteHTTPProxyCache) Get(ctx context.Context, kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
	url := r.requestURL(hash, kind)

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		cacheMisses.Inc()
		return nil, -1, err
	}

	rsp, err := r.remote.Do(req)
	if err != nil {
		cacheMisses.Inc()
		return nil, -1, err
	}

	logResponse(r.accessLogger, "DOWNLOAD", rsp.StatusCode, url)

	if rsp.StatusCode == http.StatusNotFound {
		cacheMisses.Inc()
		return nil, -1, nil
	}

	if rsp.StatusCode != http.StatusOK {
		// If the failed http response contains some data then
		// forward up to 1 KiB.
		var errorBytes []byte
		errorBytes, err = ioutil.ReadAll(io.LimitReader(rsp.Body, 1024))
		var errorText string
		if err == nil {
			errorText = string(errorBytes)
		}

		cacheMisses.Inc()
		return nil, -1, &cache.Error{
			Code: rsp.StatusCode,
			Text: errorText,
		}
	}

	if kind == cache.CAS && r.v2mode {
		cacheHits.Inc()
		return casblob.ExtractLogicalSize(rsp.Body)
	}

	sizeBytesStr := rsp.Header.Get("Content-Length")
	if sizeBytesStr == "" {
		err = errors.New("Missing Content-Length header")
		cacheMisses.Inc()
		return nil, -1, err
	}

	sizeBytesInt, err := strconv.Atoi(sizeBytesStr)
	if err != nil {
		cacheMisses.Inc()
		return nil, -1, err
	}
	sizeBytes := int64(sizeBytesInt)

	cacheHits.Inc()

	return rsp.Body, sizeBytes, nil
}

func (r *remoteHTTPProxyCache) Contains(ctx context.Context, kind cache.EntryKind, hash string) (bool, int64) {

	url := r.requestURL(hash, kind)

	req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
	if err != nil {
		return false, -1
	}

	rsp, err := r.remote.Do(req)
	if err == nil && rsp.StatusCode == http.StatusOK {
		if kind != cache.CAS {
			return true, rsp.ContentLength
		}

		// We don't know the content size without reading the file header
		// and that could be very costly for the backend server. So return
		// "unknown size".
		return true, -1
	}

	return false, -1
}
<|endoftext|>