github.com/u-root/u-root@v7.0.1-0.20200915234505-ad7babab0a8e+incompatible/pkg/rand/random_urandom.go (about)

     1  // Copyright 2019 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // +build aix darwin dragonfly freebsd nacl netbsd openbsd plan9 solaris linux
     6  
     7  // Package rand implements cancelable reads from a cryptographically safe
     8  // random number source.
     9  package rand
    10  
    11  import (
    12  	"context"
    13  	"fmt"
    14  	"sync"
    15  	"syscall"
    16  
    17  	"golang.org/x/sys/unix"
    18  )
    19  
    20  // Reader is a cryptographically safe random number source.
    21  var Reader = DefaultReaderWithContext(context.Background())
    22  
    23  // Read blockingly reads from a random number source.
    24  func Read(b []byte) (int, error) {
    25  	return Reader.Read(b)
    26  }
    27  
    28  // ReadContext is a context-aware reader for random numbers.
    29  func ReadContext(ctx context.Context, b []byte) (int, error) {
    30  	return Reader.ReadContext(ctx, b)
    31  }
    32  
    33  // ContextReader is a cancelable io.Reader.
    34  type ContextReader interface {
    35  	// Read behaves like a blocking io.Reader.Read.
    36  	//
    37  	// Read wraps ReadContext with a background context.
    38  	Read(b []byte) (n int, err error)
    39  
    40  	// ReadContext is an io.Reader that blocks until data is available or
    41  	// until ctx is done.
    42  	ReadContext(ctx context.Context, b []byte) (n int, err error)
    43  }
    44  
    45  // contextReader is a cancelable io.Reader.
    46  type contextReader interface {
    47  	ReadContext(context.Context, []byte) (int, error)
    48  }
    49  
    50  // ctxReader takes a contextReader and turns it into a ContextReader.
    51  type ctxReader struct {
    52  	contextReader
    53  	ctx context.Context
    54  }
    55  
    56  func (cr ctxReader) Read(b []byte) (int, error) {
    57  	return cr.contextReader.ReadContext(cr.ctx, b)
    58  }
    59  
    60  // DefaultReaderWithContext returns a context-aware io.Reader.
    61  //
    62  // Because this stores the context, only use this in situations where an
    63  // io.Reader is unavoidable.
    64  func DefaultReaderWithContext(ctx context.Context) ContextReader {
    65  	return ctxReader{
    66  		ctx:           ctx,
    67  		contextReader: defaultContextReader,
    68  	}
    69  }
    70  
    71  // urandomReader is a contextReader.
    72  type urandomReader struct {
    73  	once sync.Once
    74  
    75  	// fd is expected to be non-blocking.
    76  	fd int
    77  }
    78  
    79  func (r *urandomReader) init() error {
    80  	var realErr error
    81  	r.once.Do(func() {
    82  		fd, err := unix.Open("/dev/urandom", unix.O_RDONLY, 0)
    83  		if err != nil {
    84  			realErr = fmt.Errorf("open(/dev/urandom): %v", err)
    85  			return
    86  		}
    87  		r.fd = fd
    88  	})
    89  	return realErr
    90  }
    91  
    92  // ReadContext implements a cancelable read from /dev/urandom.
    93  func (r *urandomReader) ReadContext(ctx context.Context, b []byte) (int, error) {
    94  	if err := r.init(); err != nil {
    95  		return 0, err
    96  	}
    97  	for {
    98  		n, err := unix.Read(r.fd, b)
    99  		if err == nil {
   100  			return n, err
   101  		}
   102  		select {
   103  		case <-ctx.Done():
   104  			return 0, ctx.Err()
   105  
   106  		default:
   107  			if err != nil && err != syscall.EAGAIN && err != syscall.EINTR {
   108  				return n, err
   109  			}
   110  		}
   111  	}
   112  }