Adding read command

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

README.md

@@ -11,6 +11,6 @@ This is still fairly immature. Here's what I want to be able to do with it:
 - [x] print the version (`cmdv`)
 - [x] print the ID (`cmdI`)
 - [x] verify the image (`cmdX` & verify PGP signature)
-- [ ] generate some amount of entropy
+- [x] generate some amount of entropy (`onerng read` command)
 - [ ] add extra AES128-whitening
 - [ ] run as a daemon and integrate with `rngd`

cmd/onerng/main.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"os"
 	"os/signal"
-	"time"
 
 	"github.com/hairyhenderson/go-onerng/cmd"
 )
@@ -12,7 +11,7 @@ import (
 func main() {
 	ctx := context.Background()
 
-	ctx, cancel := context.WithDeadline(ctx, time.Now().Add(3*time.Second))
+	ctx, cancel := context.WithCancel(ctx)
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt)
 	defer func() {

cmd/read.go

@@ -0,0 +1,64 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/pkg/errors"
+
+	"github.com/dustin/go-humanize"
+	"github.com/hairyhenderson/go-onerng"
+	"github.com/spf13/cobra"
+)
+
+func readCmd(ctx context.Context) *cobra.Command {
+	readOut := ""
+	disableAvalanche := false
+	enableRF := false
+	disableWhitener := false
+	count := int64(-1)
+	cmd := &cobra.Command{
+		Use:   "read",
+		Short: "read some random data from the OneRNG",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			o := onerng.OneRNG{Path: opts.Device}
+			err := o.Init(ctx)
+			if err != nil {
+				return errors.Wrapf(err, "init failed")
+			}
+			var out *os.File
+			if readOut == "-" {
+				out = os.Stdout
+			} else {
+				out, err = os.OpenFile(readOut, os.O_RDWR|os.O_CREATE, 0644)
+				if err != nil {
+					return err
+				}
+			}
+			flags := onerng.ReadMode(onerng.Default)
+			if disableAvalanche {
+				flags |= onerng.DisableAvalanche
+			}
+			if enableRF {
+				flags |= onerng.EnableRF
+			}
+			if disableWhitener {
+				flags |= onerng.DisableWhitener
+			}
+			start := time.Now()
+			written, err := o.Read(ctx, out, count, flags)
+			delta := time.Since(start)
+			rate := float64(written) / delta.Seconds()
+			fmt.Fprintf(os.Stderr, "%s written in %s (%s/s)\n", humanize.Bytes(uint64(written)), delta, humanize.Bytes(uint64(rate)))
+			return err
+		},
+	}
+	cmd.Flags().StringVarP(&readOut, "out", "o", "-", "output file for data (use - for stdout)")
+	cmd.Flags().BoolVar(&disableAvalanche, "disable-avalanche", false, "Disable noise generation from the Avalanche Diode")
+	cmd.Flags().BoolVar(&enableRF, "enable-rf", false, "Enable noise generation from RF")
+	cmd.Flags().BoolVar(&disableWhitener, "disable-whitener", false, "Disable the on-board CRC16 generator")
+	cmd.Flags().Int64VarP(&count, "count", "n", -1, "Read only N bytes (use -1 for unlimited)")
+	return cmd
+}

cmd/root.go

@@ -52,6 +52,7 @@ func initConfig(ctx context.Context, cmd *cobra.Command) {
 		flushCmd(ctx),
 		imageCmd(ctx),
 		initCmd(ctx),
+		readCmd(ctx),
 	)
 
 	if cfgFile != "" { // enable ability to specify config file via flag

commands.go

@@ -31,3 +31,17 @@ const (
 	// CmdRF - Raw RF noise
 	CmdRF = "cmd7\n"
 )
+
+const (
+	// DisableWhitener - Disable the on-board CRC16 generator - no effect if both noise generators are disabled
+	DisableWhitener ReadMode = 1 << iota
+	// EnableRF - Enable noise generation from RF
+	EnableRF
+	// DisableAvalanche - Disable noise generation from the Avalanche Diode
+	DisableAvalanche
+
+	// Default - Avalanche enabled, RF disabled, Whitener enabled.
+	Default ReadMode = 0
+	// Silent - a convenience - everything disabled
+	Silent ReadMode = 4
+)

copy.go

@@ -0,0 +1,46 @@
+package onerng
+
+import (
+	"context"
+	"io"
+	"os"
+	"strings"
+	"time"
+)
+
+type readerFunc func(p []byte) (n int, err error)
+
+func (rf readerFunc) Read(p []byte) (n int, err error) { return rf(p) }
+
+// io.CopyN/io.Copy with context support
+func copyWithContext(ctx context.Context, dst io.Writer, src *os.File, n int64) (int64, error) {
+	// allow 10 500ms timeouts, for a total of 5s. After this, it's probably worth just giving up
+	allowedTimeouts := 10
+
+	rf := func(p []byte) (int, error) {
+		// we don't want reads to block forever, but we also don't want to time out immediately
+		err := src.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
+		if err != nil {
+			return 0, err
+		}
+
+		select {
+		case <-ctx.Done():
+			return 0, ctx.Err()
+		default:
+			n, err := src.Read(p)
+			if allowedTimeouts > 0 {
+				if err != nil && strings.HasSuffix(err.Error(), "i/o timeout") {
+					allowedTimeouts--
+					return n, nil
+				}
+			}
+			return n, err
+		}
+	}
+
+	if n < 0 {
+		return io.Copy(dst, readerFunc(rf))
+	}
+	return io.CopyN(dst, readerFunc(rf), n)
+}

onerng.go

@@ -18,6 +18,9 @@ type OneRNG struct {
 	Path string
 }
 
+// ReadMode -
+type ReadMode uint32
+
 func (o *OneRNG) cmd(ctx context.Context, d *os.File, c ...string) (err error) {
 	for _, v := range c {
 		_, err = d.WriteString(v)
@@ -219,6 +222,25 @@ func (o *OneRNG) Init(ctx context.Context) error {
 	return nil
 }
 
+// Read -
+func (o *OneRNG) Read(ctx context.Context, out io.WriteCloser, n int64, flags ReadMode) (written int64, err error) {
+	d, err := os.OpenFile(o.Path, os.O_RDWR, 0600)
+	if err != nil {
+		return 0, err
+	}
+	defer d.Close()
+
+	err = o.cmd(ctx, d, NoiseCommand(flags), CmdRun)
+	if err != nil {
+		return 0, err
+	}
+
+	defer o.cmd(ctx, d, CmdPause)
+
+	written, err = copyWithContext(ctx, out, d, n)
+	return written, err
+}
+
 // readData - try to read some data from the RNG
 func (o *OneRNG) readData(ctx context.Context) (int, error) {
 	d, err := os.OpenFile(o.Path, os.O_RDWR, 0600)
@@ -290,3 +312,9 @@ func scan(ctx context.Context, d *os.File, buf chan string, errc chan error) {
 		}
 	}
 }
+
+// NoiseCommand - returns the appropriate noise-generation command for the given flags
+func NoiseCommand(flags ReadMode) string {
+	num := strconv.Itoa(int(flags))
+	return "cmd" + num + "\n"
+}

onerng_test.go

@@ -0,0 +1,34 @@
+package onerng
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNoiseCommand(t *testing.T) {
+	fmt.Println(DisableWhitener)
+	fmt.Println(EnableRF)
+	fmt.Println(DisableAvalanche)
+
+	fmt.Println(Default)
+	fmt.Println(Silent)
+
+	testdata := []struct {
+		flags ReadMode
+		cmd   string
+	}{
+		{Default, "cmd0\n"},
+		{DisableWhitener, "cmd1\n"},
+		{EnableRF, "cmd2\n"},
+		{EnableRF | DisableWhitener, "cmd3\n"},
+		{DisableAvalanche, "cmd4\n"},
+		{DisableAvalanche | DisableWhitener, "cmd5\n"},
+		{DisableAvalanche | EnableRF, "cmd6\n"},
+		{DisableAvalanche | EnableRF | DisableWhitener, "cmd7\n"},
+	}
+	for _, d := range testdata {
+		assert.Equal(t, d.cmd, NoiseCommand(d.flags), d.cmd, d.flags)
+	}
+}