better batch-generation (#1219)

2019-02-06 09:04:12 +01:00 · 2019-02-06 09:04:12 +01:00 · 7b7e6c43f8
commit 7b7e6c43f8
parent 218e95728c
7 changed files with 469 additions and 67 deletions
--- a/cmd/identity/batch.go
+++ b/cmd/identity/batch.go
@ -0,0 +1,202 @@
+// Copyright (C) 2019 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package main
+
+import (
+	"archive/tar"
+	"bytes"
+	"context"
+	"crypto/ecdsa"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"sync/atomic"
+	"text/tabwriter"
+
+	"github.com/spf13/cobra"
+	"github.com/zeebo/errs"
+	"golang.org/x/sync/errgroup"
+
+	"storj.io/storj/internal/cui"
+	"storj.io/storj/pkg/cfgstruct"
+	"storj.io/storj/pkg/identity"
+	"storj.io/storj/pkg/peertls"
+	"storj.io/storj/pkg/process"
+	"storj.io/storj/pkg/storj"
+)
+
+var (
+	certMode       int64 = 0644
+	keyMode        int64 = 0600
+	keyGenerateCmd       = &cobra.Command{
+		Use:         "batch-generate",
+		Short:       "generate lots of keys",
+		RunE:        cmdKeyGenerate,
+		Annotations: map[string]string{"type": "setup"},
+	}
+
+	keyCfg struct {
+		MinDifficulty int    `help:"minimum difficulty to output" default:"30"`
+		Concurrency   int    `help:"worker concurrency" default:"4"`
+		OutputDir     string `help:"output directory to place keys" default:"."`
+	}
+)
+
+func init() {
+	rootCmd.AddCommand(keyGenerateCmd)
+	cfgstruct.Bind(keyGenerateCmd.Flags(), &keyCfg)
+}
+
+func cmdKeyGenerate(cmd *cobra.Command, args []string) (err error) {
+	ctx, cancel := context.WithCancel(process.Ctx(cmd))
+	defer cancel()
+
+	err = os.MkdirAll(keyCfg.OutputDir, 0700)
+	if err != nil {
+		return err
+	}
+
+	var group errgroup.Group
+	defer func() {
+		err = errs.Combine(err, group.Wait())
+	}()
+
+	screen, err := cui.NewScreen()
+	if err != nil {
+		return err
+	}
+	group.Go(func() error {
+		defer cancel()
+		err := screen.Run()
+		return errs.Combine(err, screen.Close())
+	})
+
+	counter := new(uint32)
+	diffCounts := [256]uint32{}
+	return identity.GenerateKeys(ctx, uint16(keyCfg.MinDifficulty), keyCfg.Concurrency,
+		func(k *ecdsa.PrivateKey, id storj.NodeID) (done bool, err error) {
+			difficulty, err := id.Difficulty()
+			if err != nil {
+				return false, err
+			}
+
+			if int(difficulty) > len(diffCounts) {
+				atomic.AddUint32(&diffCounts[len(diffCounts)-1], 1)
+			} else {
+				atomic.AddUint32(&diffCounts[difficulty], 1)
+			}
+
+			if err := renderStats(screen, diffCounts[:]); err != nil {
+				return false, err
+			}
+
+			genName := fmt.Sprintf("gen-%02d-%d", difficulty, atomic.AddUint32(counter, 1))
+			err = saveIdentityTar(filepath.Join(keyCfg.OutputDir, genName), k, id)
+			return false, err
+		})
+}
+
+func saveIdentityTar(path string, key *ecdsa.PrivateKey, id storj.NodeID) error {
+	ct, err := peertls.CATemplate()
+	if err != nil {
+		return err
+	}
+
+	caCert, err := peertls.NewCert(key, nil, ct, nil)
+	if err != nil {
+		return err
+	}
+
+	ca := &identity.FullCertificateAuthority{
+		Cert: caCert,
+		ID:   id,
+		Key:  key,
+	}
+
+	ident, err := ca.NewIdentity()
+	if err != nil {
+		return err
+	}
+
+	tarData := new(bytes.Buffer)
+	tw := tar.NewWriter(tarData)
+
+	caCertBytes, caCertErr := peertls.ChainBytes(ca.Cert)
+	caKeyBytes, caKeyErr := peertls.KeyBytes(ca.Key)
+	identCertBytes, identCertErr := peertls.ChainBytes(ident.Leaf, ident.CA)
+	identKeyBytes, identKeyErr := peertls.KeyBytes(ident.Key)
+	if err := errs.Combine(caCertErr, caKeyErr, identCertErr, identKeyErr); err != nil {
+		return err
+	}
+
+	if err := errs.Combine(
+		writeToTar(tw, "ca.cert", certMode, caCertBytes),
+		writeToTar(tw, "ca.key", keyMode, caKeyBytes),
+		writeToTar(tw, "identity.cert", certMode, identCertBytes),
+		writeToTar(tw, "identity.key", keyMode, identKeyBytes),
+	); err != nil {
+		return err
+	}
+
+	if err := tw.Close(); err != nil {
+		return errs.Wrap(err)
+	}
+
+	if err = ioutil.WriteFile(path+".tar", tarData.Bytes(), 0600); err != nil {
+		return errs.Wrap(err)
+	}
+	return nil
+}
+
+func writeToTar(tw *tar.Writer, name string, mode int64, data []byte) error {
+	err := tw.WriteHeader(&tar.Header{
+		Name: name,
+		Mode: mode,
+		Size: int64(len(data)),
+	})
+	if err != nil {
+		return err
+	}
+
+	_, err = tw.Write(data)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func renderStats(screen *cui.Screen, stats []uint32) error {
+	screen.Lock()
+	defer screen.Unlock()
+
+	var err error
+	printf := func(w io.Writer, format string, args ...interface{}) {
+		if err == nil {
+			_, err = fmt.Fprintf(w, format, args...)
+		}
+	}
+
+	printf(screen, "Batch Identity Creation\n\n\n")
+
+	w := tabwriter.NewWriter(screen, 0, 2, 2, ' ', 0)
+	printf(w, "Difficulty\tCount\n")
+
+	total := uint32(0)
+
+	for difficulty := len(stats); difficulty > 0; difficulty-- {
+		count := atomic.LoadUint32(&stats[difficulty-1])
+		total += count
+		if count == 0 {
+			continue
+		}
+		printf(w, "%d\t%d\n", difficulty, count)
+	}
+	printf(w, "Total\t%d\n", total)
+
+	err = errs.Combine(err, w.Flush())
+
+	return screen.Flush()
+}
--- a/cmd/identity/keys.go
+++ b/cmd/identity/keys.go
@ -1,65 +0,0 @@
-// Copyright (C) 2019 Storj Labs, Inc.
-// See LICENSE for copying information.
-
-package main
-
-import (
-	"crypto/ecdsa"
-	"crypto/x509"
-	"fmt"
-	"io/ioutil"
-	"os"
-	"path/filepath"
-	"sync/atomic"
-
-	"github.com/spf13/cobra"
-
-	"storj.io/storj/pkg/cfgstruct"
-	"storj.io/storj/pkg/identity"
-	"storj.io/storj/pkg/process"
-	"storj.io/storj/pkg/storj"
-)
-
-var (
-	keyGenerateCmd = &cobra.Command{
-		Use:         "batch-generate",
-		Short:       "generate lots of keys",
-		RunE:        cmdKeyGenerate,
-		Annotations: map[string]string{"type": "setup"},
-	}
-
-	keyCfg struct {
-		MinDifficulty int    `help:"minimum difficulty to output" default:"30"`
-		Concurrency   int    `help:"worker concurrency" default:"4"`
-		OutputDir     string `help:"output directory to place keys" default:"."`
-	}
-)
-
-func init() {
-	rootCmd.AddCommand(keyGenerateCmd)
-	cfgstruct.Bind(keyGenerateCmd.Flags(), &keyCfg)
-}
-
-func cmdKeyGenerate(cmd *cobra.Command, args []string) (err error) {
-	ctx := process.Ctx(cmd)
-	err = os.MkdirAll(keyCfg.OutputDir, 0700)
-	if err != nil {
-		return err
-	}
-	counter := new(uint32)
-	return identity.GenerateKeys(ctx, uint16(keyCfg.MinDifficulty), keyCfg.Concurrency,
-		func(k *ecdsa.PrivateKey, id storj.NodeID) (done bool, err error) {
-			data, err := x509.MarshalECPrivateKey(k)
-			if err != nil {
-				return false, err
-			}
-			difficulty, err := id.Difficulty()
-			if err != nil {
-				return false, err
-			}
-			filename := fmt.Sprintf("gen-%02d-%d.key", difficulty, atomic.AddUint32(counter, 1))
-			fmt.Println("writing", filename)
-			err = ioutil.WriteFile(filepath.Join(keyCfg.OutputDir, filename), data, 0600)
-			return false, err
-		})
-}
--- a/go.mod
+++ b/go.mod
@ -82,6 +82,7 @@ require (
 	github.com/nats-io/nats-streaming-server v0.11.0 // indirect
 	github.com/nats-io/nuid v1.0.0 // indirect
 	github.com/nsf/jsondiff v0.0.0-20160203110537-7de28ed2b6e3
+	github.com/nsf/termbox-go v0.0.0-20190121233118-02980233997d
 	github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c // indirect
 	github.com/pierrec/lz4 v2.0.5+incompatible // indirect
 	github.com/pkg/errors v0.8.1 // indirect
--- a/go.sum
+++ b/go.sum
@ -241,6 +241,8 @@ github.com/nats-io/nuid v1.0.0 h1:44QGdhbiANq8ZCbUkdn6W5bqtg+mHuDE4wOUuxxndFs=
 github.com/nats-io/nuid v1.0.0/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c=
 github.com/nsf/jsondiff v0.0.0-20160203110537-7de28ed2b6e3 h1:OqFSgO6CJ8heZRAbXLpT+ojX+jnnGij4qZwUz/SJJ9I=
 github.com/nsf/jsondiff v0.0.0-20160203110537-7de28ed2b6e3/go.mod h1:uFMI8w+ref4v2r9jz+c9i1IfIttS/OkmLfrk1jne5hs=
+github.com/nsf/termbox-go v0.0.0-20190121233118-02980233997d h1:x3S6kxmy49zXVVyhcnrFqxvNVCBPb2KZ9hV2RBdS840=
+github.com/nsf/termbox-go v0.0.0-20190121233118-02980233997d/go.mod h1:IuKpRQcYE1Tfu+oAQqaLisqDeXgjyyltCfsaoYN18NQ=
 github.com/onsi/ginkgo v1.6.0 h1:Ix8l273rp3QzYgXSR+c8d1fTG7UPgYkOSELPhiY/YGw=
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/gomega v1.4.2 h1:3mYCb7aPxS/RU7TI1y4rkEn1oKmPRjNJLNEXgw7MH2I=
@ -318,8 +320,6 @@ github.com/tidwall/gjson v1.1.3 h1:u4mspaByxY+Qk4U1QYYVzGFI8qxN/3jtEV0ZDb2vRic=
 github.com/tidwall/gjson v1.1.3/go.mod h1:c/nTNbUr0E0OrXEhq1pwa8iEgc2DOt4ZZqAt1HtCkPA=
 github.com/tidwall/match v0.0.0-20171002075945-1731857f09b1 h1:pWIN9LOlFRCJFqWIOEbHLvY0WWJddsjH2FQ6N0HKZdU=
 github.com/tidwall/match v0.0.0-20171002075945-1731857f09b1/go.mod h1:LujAq0jyVjBy028G1WhWfIzbpQfMO8bBZ6Tyb0+pL9E=
-github.com/vivint/infectious v0.0.0-20180906161625-e155e6eb3575 h1:t0v1w3EiqMhDYBFbzwSfUHivx3yOMJG7R7YUm1Amlh8=
-github.com/vivint/infectious v0.0.0-20180906161625-e155e6eb3575/go.mod h1:5oyMAv4hrBEKqBwORFsiqIrCNCmL2qcZLQTdJLYeYIc=
 github.com/vivint/infectious v0.0.0-20190108171102-2455b059135b h1:dLkqBELopfQNhe8S9ucnSf+HhiUCgK/hPIjVG0f9GlY=
 github.com/vivint/infectious v0.0.0-20190108171102-2455b059135b/go.mod h1:5oyMAv4hrBEKqBwORFsiqIrCNCmL2qcZLQTdJLYeYIc=
 github.com/yuin/gopher-lua v0.0.0-20180918061612-799fa34954fb h1:Jmfk7z2f/+gxVFAgPsJMuczO1uEIxZy6wytTdeZ49lg=
--- a/internal/cui/example.go
+++ b/internal/cui/example.go
@ -0,0 +1,51 @@
+// Copyright (C) 2019 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+// +build ignore
+
+package main
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"golang.org/x/sync/errgroup"
+
+	"storj.io/storj/internal/cui"
+)
+
+func main() {
+	screen, err := cui.NewScreen()
+	if err != nil {
+		fmt.Println(err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	var group errgroup.Group
+	group.Go(func() error {
+		defer cancel()
+		return screen.Run()
+	})
+
+	counter := 0
+	for ctx.Err() == nil {
+		width, _ := screen.Size()
+		fmt.Fprintf(screen, "%2d\n", counter)
+		fmt.Fprintf(screen, "%s\n", strings.Repeat("=", width))
+		screen.Flush()
+
+		counter++
+		time.Sleep(time.Second)
+	}
+	if err := screen.Close(); err != nil {
+		fmt.Println(err)
+	}
+
+	if err := group.Wait(); err != nil {
+		fmt.Println(err)
+	}
+}
--- a/internal/cui/screen.go
+++ b/internal/cui/screen.go
@ -0,0 +1,212 @@
+// Copyright (C) 2019 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package cui
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"errors"
+	"sync"
+
+	termbox "github.com/nsf/termbox-go"
+)
+
+var initialized = false
+
+const padding = 2
+
+// Point is a 2D coordinate in console
+//   X is the column
+//   Y is the row
+type Point struct{ X, Y int }
+
+// Rect is a 2D rectangle in console, excluding Max edge
+type Rect struct{ Min, Max Point }
+
+// Screen is a writable area on screen
+type Screen struct {
+	rendering sync.Mutex
+
+	blitting sync.Mutex
+	closed   bool
+	flushed  frame
+	pending  frame
+}
+
+type frame struct {
+	size    Point
+	content []byte
+}
+
+// NewScreen returns a new screen, only one screen can be use at a time.
+func NewScreen() (*Screen, error) {
+	if initialized {
+		return nil, errors.New("only one screen allowed at a time")
+	}
+	initialized = true
+	if err := termbox.Init(); err != nil {
+		initialized = false
+		return nil, err
+	}
+
+	termbox.SetInputMode(termbox.InputEsc)
+	termbox.HideCursor()
+	screen := &Screen{}
+	screen.flushed.size.X, screen.flushed.size.Y = termbox.Size()
+	screen.pending.size = screen.flushed.size
+	return screen, nil
+}
+
+func (screen *Screen) markClosed() {
+	screen.blitting.Lock()
+	screen.closed = true
+	screen.blitting.Unlock()
+}
+
+func (screen *Screen) isClosed() bool {
+	screen.blitting.Lock()
+	defer screen.blitting.Unlock()
+	return screen.closed
+}
+
+// Close closes the screen.
+func (screen *Screen) Close() error {
+	screen.markClosed()
+
+	// shutdown termbox
+	termbox.Close()
+	initialized = false
+	return nil
+}
+
+// Run runs the event loop
+func (screen *Screen) Run() error {
+	defer screen.markClosed()
+
+	for !screen.isClosed() {
+		switch ev := termbox.PollEvent(); ev.Type {
+		case termbox.EventInterrupt:
+			// either screen refresh or close
+		case termbox.EventKey:
+			switch ev.Key {
+			case termbox.KeyCtrlC, termbox.KeyEsc:
+				return nil
+			default:
+				// ignore key presses
+			}
+		case termbox.EventError:
+			return ev.Err
+		case termbox.EventResize:
+			screen.blitting.Lock()
+			screen.flushed.size.X, screen.flushed.size.Y = ev.Width, ev.Height
+			err := screen.blit(&screen.flushed)
+			screen.blitting.Unlock()
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+// Size returns the current size of the screen.
+func (screen *Screen) Size() (width, height int) {
+	width, height = screen.pending.size.X-2*padding, screen.pending.size.Y-2*padding
+	if width < 0 {
+		width = 0
+	}
+	if height < 0 {
+		height = 0
+	}
+	return width, height
+}
+
+// Lock screen for exclusive rendering
+func (screen *Screen) Lock() { screen.rendering.Lock() }
+
+// Unlock screen
+func (screen *Screen) Unlock() { screen.rendering.Unlock() }
+
+// Write writes to the screen.
+func (screen *Screen) Write(data []byte) (int, error) {
+	screen.pending.content = append(screen.pending.content, data...)
+	return len(data), nil
+}
+
+// Flush flushes pending content to the console and clears for new frame.
+func (screen *Screen) Flush() error {
+	screen.blitting.Lock()
+	var err error
+	if !screen.closed {
+		err = screen.blit(&screen.pending)
+	} else {
+		err = context.Canceled
+	}
+	screen.pending.content = nil
+	screen.pending.size = screen.flushed.size
+	screen.blitting.Unlock()
+
+	return err
+}
+
+// blit writes content to the console
+func (screen *Screen) blit(frame *frame) error {
+	screen.flushed.content = frame.content
+	size := screen.flushed.size
+
+	if err := termbox.Clear(termbox.ColorDefault, termbox.ColorDefault); err != nil {
+		return err
+	}
+
+	drawRect(Rect{
+		Min: Point{0, 0},
+		Max: size,
+	}, lightStyle)
+
+	scanner := bufio.NewScanner(bytes.NewReader(frame.content))
+	y := padding
+	for scanner.Scan() && y <= size.Y-2*padding {
+		x := padding
+		for _, r := range scanner.Text() {
+			if x > size.X-2*padding {
+				break
+			}
+			termbox.SetCell(x, y, r, termbox.ColorDefault, termbox.ColorDefault)
+			x++
+		}
+		y++
+	}
+
+	return termbox.Flush()
+}
+
+type rectStyle [3][3]rune
+
+var lightStyle = rectStyle{
+	{'┌', '─', '┐'},
+	{'│', ' ', '│'},
+	{'└', '─', '┘'},
+}
+
+// drawRect draws a rectangle using termbox
+func drawRect(r Rect, style rectStyle) {
+	attr := termbox.ColorDefault
+
+	termbox.SetCell(r.Min.X, r.Min.Y, style[0][0], attr, attr)
+	termbox.SetCell(r.Max.X-1, r.Min.Y, style[0][2], attr, attr)
+	termbox.SetCell(r.Max.X-1, r.Max.Y-1, style[2][2], attr, attr)
+	termbox.SetCell(r.Min.X, r.Max.Y-1, style[2][0], attr, attr)
+
+	for x := r.Min.X + 1; x < r.Max.X-1; x++ {
+		termbox.SetCell(x, r.Min.Y, style[0][1], attr, attr)
+		termbox.SetCell(x, r.Max.Y-1, style[2][1], attr, attr)
+	}
+
+	for y := r.Min.Y + 1; y < r.Max.Y-1; y++ {
+		termbox.SetCell(r.Min.X, y, style[1][0], attr, attr)
+		termbox.SetCell(r.Max.X-1, y, style[1][2], attr, attr)
+	}
+}
--- a/pkg/identity/generate.go
+++ b/pkg/identity/generate.go
@ -59,6 +59,7 @@ func GenerateKeys(ctx context.Context, minDifficulty uint16, concurrency int, fo
 					errchan <- err
 					return
 				}
+
 				done, err := found(k, id)
 				if err != nil {
 					errchan <- err