github.com/safing/portbase@v0.19.5/rng/test/main.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"crypto/aes"
     6  	"crypto/cipher"
     7  	"crypto/rand"
     8  	"encoding/binary"
     9  	"encoding/hex"
    10  	"fmt"
    11  	"io"
    12  	"os"
    13  	"runtime"
    14  	"strconv"
    15  	"time"
    16  
    17  	"github.com/safing/portbase/log"
    18  	"github.com/safing/portbase/modules"
    19  	"github.com/safing/portbase/rng"
    20  	"github.com/safing/portbase/run"
    21  )
    22  
    23  var (
    24  	module *modules.Module
    25  
    26  	outputFile *os.File
    27  	outputSize uint64 = 1000000
    28  )
    29  
    30  func init() {
    31  	module = modules.Register("main", prep, start, nil, "rng")
    32  }
    33  
    34  func main() {
    35  	runtime.GOMAXPROCS(1)
    36  	os.Exit(run.Run())
    37  }
    38  
    39  func prep() error {
    40  	if len(os.Args) < 3 {
    41  		fmt.Printf("usage: ./%s {fortuna|tickfeeder} <file> [output size in MB]", os.Args[0])
    42  		return modules.ErrCleanExit
    43  	}
    44  
    45  	switch os.Args[1] {
    46  	case "fortuna":
    47  	case "tickfeeder":
    48  	default:
    49  		return fmt.Errorf("usage: %s {fortuna|tickfeeder}", os.Args[0])
    50  	}
    51  
    52  	if len(os.Args) > 3 {
    53  		n, err := strconv.ParseUint(os.Args[3], 10, 64)
    54  		if err != nil {
    55  			return fmt.Errorf("failed to parse output size: %w", err)
    56  		}
    57  		outputSize = n * 1000000
    58  	}
    59  
    60  	var err error
    61  	outputFile, err = os.OpenFile(os.Args[2], os.O_CREATE|os.O_WRONLY, 0o0644) //nolint:gosec
    62  	if err != nil {
    63  		return fmt.Errorf("failed to open output file: %w", err)
    64  	}
    65  
    66  	return nil
    67  }
    68  
    69  //nolint:gocognit
    70  func start() error {
    71  	// generates 1MB and writes to stdout
    72  
    73  	log.Infof("writing %dMB to stdout, a \".\" will be printed at every 1024 bytes.", outputSize/1000000)
    74  
    75  	switch os.Args[1] {
    76  	case "fortuna":
    77  		module.StartWorker("fortuna", fortuna)
    78  
    79  	case "tickfeeder":
    80  		module.StartWorker("noise", noise)
    81  		module.StartWorker("tickfeeder", tickfeeder)
    82  
    83  	default:
    84  		return fmt.Errorf("usage: ./%s {fortuna|tickfeeder}", os.Args[0])
    85  	}
    86  
    87  	return nil
    88  }
    89  
    90  func fortuna(_ context.Context) error {
    91  	var bytesWritten uint64
    92  
    93  	for {
    94  		if module.IsStopping() {
    95  			return nil
    96  		}
    97  
    98  		b, err := rng.Bytes(64)
    99  		if err != nil {
   100  			return err
   101  		}
   102  		_, err = outputFile.Write(b)
   103  		if err != nil {
   104  			return err
   105  		}
   106  
   107  		bytesWritten += 64
   108  		if bytesWritten%1024 == 0 {
   109  			_, _ = os.Stderr.WriteString(".")
   110  		}
   111  		if bytesWritten%65536 == 0 {
   112  			fmt.Fprintf(os.Stderr, "\n%d bytes written\n", bytesWritten)
   113  		}
   114  		if bytesWritten >= outputSize {
   115  			_, _ = os.Stderr.WriteString("\n")
   116  			break
   117  		}
   118  	}
   119  
   120  	go modules.Shutdown() //nolint:errcheck
   121  	return nil
   122  }
   123  
   124  func tickfeeder(ctx context.Context) error {
   125  	var bytesWritten uint64
   126  	var value int64
   127  	var pushes int
   128  
   129  	for {
   130  		if module.IsStopping() {
   131  			return nil
   132  		}
   133  
   134  		time.Sleep(10 * time.Nanosecond)
   135  
   136  		value = (value << 1) | (time.Now().UnixNano() % 2)
   137  		pushes++
   138  
   139  		if pushes >= 64 {
   140  			b := make([]byte, 8)
   141  			binary.LittleEndian.PutUint64(b, uint64(value))
   142  			_, err := outputFile.Write(b)
   143  			if err != nil {
   144  				return err
   145  			}
   146  			bytesWritten += 8
   147  			if bytesWritten%1024 == 0 {
   148  				_, _ = os.Stderr.WriteString(".")
   149  			}
   150  			if bytesWritten%65536 == 0 {
   151  				fmt.Fprintf(os.Stderr, "\n%d bytes written\n", bytesWritten)
   152  			}
   153  			pushes = 0
   154  		}
   155  
   156  		if bytesWritten >= outputSize {
   157  			_, _ = os.Stderr.WriteString("\n")
   158  			break
   159  		}
   160  	}
   161  
   162  	go modules.Shutdown() //nolint:errcheck
   163  	return nil
   164  }
   165  
   166  func noise(ctx context.Context) error {
   167  	// do some aes ctr for noise
   168  
   169  	key, _ := hex.DecodeString("6368616e676520746869732070617373")
   170  	data := []byte("some plaintext x")
   171  
   172  	block, err := aes.NewCipher(key)
   173  	if err != nil {
   174  		panic(err)
   175  	}
   176  
   177  	iv := make([]byte, aes.BlockSize)
   178  	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
   179  		panic(err)
   180  	}
   181  
   182  	stream := cipher.NewCTR(block, iv)
   183  	for {
   184  		select {
   185  		case <-ctx.Done():
   186  			return nil
   187  		default:
   188  			stream.XORKeyStream(data, data)
   189  		}
   190  	}
   191  }