github.com/jxskiss/gopkg/v2@v2.14.9-0.20240514120614-899f3e7952b4/utils/strutil/random.go (about)

     1  package strutil
     2  
     3  import (
     4  	cryptorand "crypto/rand"
     5  	"encoding/hex"
     6  	"math/big"
     7  	"math/bits"
     8  	"unsafe"
     9  
    10  	"github.com/jxskiss/gopkg/v2/perf/fastrand"
    11  )
    12  
    13  func random(table string, length int) []byte {
    14  	buf := make([]byte, length)
    15  	tabLen := len(table)
    16  	for i := range buf {
    17  		buf[i] = table[fastrand.N(tabLen)]
    18  	}
    19  	return buf
    20  }
    21  
    22  // Random returns a random string of length consisting of characters
    23  // from table.
    24  // It panics if length <= 0 or len(table) <= 1.
    25  func Random(table string, length int) string {
    26  	if length <= 0 || len(table) <= 1 {
    27  		panic("strutil: invalid argument to Random")
    28  	}
    29  	buf := random(table, length)
    30  	return b2s(buf)
    31  }
    32  
    33  // See [crypto/rand.Int] about the implementation details.
    34  func cryptoRandom(table string, length int) []byte {
    35  	ret := make([]byte, 0, length)
    36  	_max := big.NewInt(int64(len(table)))
    37  
    38  	// bitLen is the maximum bit length needed to encode a value < max.
    39  	// k is the maximum byte length needed to encode a value < max.
    40  	// b is the number of bits in the most significant byte of max-1.
    41  	bitLen := bits.Len(uint(len(table) - 1))
    42  	k := (bitLen + 7) / 8
    43  	b := uint(bitLen % 8)
    44  	if b == 0 {
    45  		b = 8
    46  	}
    47  
    48  	buf := make([]byte, k*(length+10))
    49  	n := new(big.Int)
    50  
    51  	for {
    52  		_, err := cryptorand.Read(buf)
    53  		if err != nil {
    54  			panic(err)
    55  		}
    56  		for i := 0; i+k <= len(buf); i += k {
    57  			x := buf[i : i+k]
    58  
    59  			// Clear bits in the first byte to increase the probability
    60  			// that the candidate is < max.
    61  			x[0] &= uint8(int(1<<b) - 1)
    62  
    63  			n.SetBytes(x)
    64  			if n.Cmp(_max) < 0 {
    65  				ret = append(ret, table[n.Int64()])
    66  				if len(ret) == length {
    67  					return ret
    68  				}
    69  			}
    70  		}
    71  	}
    72  }
    73  
    74  // RandomCrypto returns a random string of length consisting of
    75  // characters from table.
    76  // It panics if length <= 0 or len(table) <= 1.
    77  func RandomCrypto(table string, length int) string {
    78  	if length <= 0 || len(table) <= 1 {
    79  		panic("strutil: invalid argument to RandomCrypto")
    80  	}
    81  	buf := cryptoRandom(table, length)
    82  	return b2s(buf)
    83  }
    84  
    85  func b2s(b []byte) string {
    86  	return *(*string)(unsafe.Pointer(&b))
    87  }
    88  
    89  // RandomHex returns a random hex string of length consisting of
    90  // cryptographic-safe random bytes.
    91  func RandomHex(length int) string {
    92  	if length <= 0 {
    93  		panic("strutil: invalid argument to RandomHex")
    94  	}
    95  	n := length/2 + 1
    96  	buf := make([]byte, n)
    97  	_, err := cryptorand.Read(buf)
    98  	if err != nil {
    99  		panic(err)
   100  	}
   101  	return hex.EncodeToString(buf)[:length]
   102  }