From f768ac4bfd7c11482da9e1c0e45a533a6e848ae8 Mon Sep 17 00:00:00 2001
From: Justin Cormack <justin@specialbusservice.com>
Date: Wed, 2 Aug 2017 20:55:36 +0100
Subject: [PATCH] Add flags to disable the different random sources

This makes testing the different options much easier.

Signed-off-by: Justin Cormack <justin@specialbusservice.com>
---
 pkg/rngd/cmd/rngd/main.go            |  8 +++++---
 pkg/rngd/cmd/rngd/rng_linux_amd64.go | 12 ++++++++----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/pkg/rngd/cmd/rngd/main.go b/pkg/rngd/cmd/rngd/main.go
index 6ef1cae1f..ecea2c2d1 100644
--- a/pkg/rngd/cmd/rngd/main.go
+++ b/pkg/rngd/cmd/rngd/main.go
@@ -3,6 +3,7 @@
 package main
 
 import (
+	"flag"
 	"log"
 	"os"
 
@@ -10,10 +11,11 @@ import (
 )
 
 func main() {
-	oneshot := len(os.Args) > 1 && os.Args[1] == "-1"
+	oneshot := flag.Bool("1", false, "Enable oneshot mode")
+	flag.Parse()
 
 	timeout := -1
-	if oneshot {
+	if *oneshot {
 		timeout = 0
 	}
 
@@ -61,7 +63,7 @@ func main() {
 		if nevents == 1 && events[0].Events&unix.EPOLLOUT == unix.EPOLLOUT {
 			continue
 		}
-		if oneshot {
+		if *oneshot {
 			log.Printf("Wrote %d bytes of entropy, exiting as oneshot\n", count)
 			break
 		}
diff --git a/pkg/rngd/cmd/rngd/rng_linux_amd64.go b/pkg/rngd/cmd/rngd/rng_linux_amd64.go
index 91f2c5135..ad3a5ffed 100644
--- a/pkg/rngd/cmd/rngd/rng_linux_amd64.go
+++ b/pkg/rngd/cmd/rngd/rng_linux_amd64.go
@@ -36,12 +36,16 @@ import "C"
 
 import (
 	"errors"
+	"flag"
 	"os"
 	"unsafe"
 
 	"golang.org/x/sys/unix"
 )
 
+var disableRdrand = flag.Bool("disable-rdrand", false, "Disable use of RDRAND")
+var disableRdseed = flag.Bool("disable-rdseed", false, "Disable use of RDSEED")
+
 var hasRdrand, hasRdseed bool
 
 type randInfo struct {
@@ -51,19 +55,19 @@ type randInfo struct {
 }
 
 func initRand() bool {
-	hasRdrand = C.hasrdrand() == 1
-	hasRdseed = C.hasrdseed() == 1
+	hasRdrand = C.hasrdrand() == 1 && !*disableRdrand
+	hasRdseed = C.hasrdseed() == 1 && !*disableRdseed
 	return hasRdrand || hasRdseed
 }
 
 func rand() (uint64, error) {
 	var x C.uint64_t
 	// prefer rdseed as that is correct seed
-	if hasRdseed && C.rdseed(&x) == 1 {
+	if hasRdseed && C.rdseed(&x) == 1 && !*disableRdseed {
 		return uint64(x), nil
 	}
 	// failed rdseed, rdrand better than nothing
-	if hasRdrand && C.rdrand(&x) == 1 {
+	if hasRdrand && C.rdrand(&x) == 1 && !*disableRdrand {
 		return uint64(x), nil
 	}
 	return 0, errors.New("No randomness available")