github.com/benz9527/xboot@v0.0.0-20240504061247-c23f15593274/lib/id/nano_id.go (about)

     1  package id
     2  
     3  import (
     4  	crand "crypto/rand"
     5  	"encoding/binary"
     6  	"fmt"
     7  	"strconv"
     8  	"sync"
     9  
    10  	"github.com/benz9527/xboot/lib/infra"
    11  )
    12  
    13  var classicNanoIDAlphabet = [64]byte{
    14  	'A', 'B', 'C', 'D', 'E',
    15  	'F', 'G', 'H', 'I', 'J',
    16  	'K', 'L', 'M', 'N', 'O',
    17  	'P', 'Q', 'R', 'S', 'T',
    18  	'U', 'V', 'W', 'X', 'Y',
    19  	'Z', 'a', 'b', 'c', 'd',
    20  	'e', 'f', 'g', 'h', 'i',
    21  	'j', 'k', 'l', 'm', 'n',
    22  	'o', 'p', 'q', 'r', 's',
    23  	't', 'u', 'v', 'w', 'x',
    24  	'y', 'z', '0', '1', '2',
    25  	'3', '4', '5', '6', '7',
    26  	'8', '9', '-', '_',
    27  }
    28  
    29  func rngUint32() uint32 {
    30  	randUint32 := [4]byte{}
    31  	if _, err := crand.Read(randUint32[:]); err != nil {
    32  		panic(err)
    33  	}
    34  	if randUint32[3]&0x8 == 0x0 {
    35  		return binary.LittleEndian.Uint32(randUint32[:])
    36  	}
    37  	return binary.BigEndian.Uint32(randUint32[:])
    38  }
    39  
    40  func shuffle(arr []byte) {
    41  	size := len(arr)
    42  	count := uint32(size >> 1)
    43  	for i := uint32(0); i < count; i++ {
    44  		j := rngUint32() % uint32(size)
    45  		arr[i], arr[j] = arr[j], arr[i]
    46  	}
    47  }
    48  
    49  func init() {
    50  	shuffle(classicNanoIDAlphabet[:])
    51  }
    52  
    53  func ClassicNanoID(length int) (NanoIDGen, error) {
    54  	if length < 1 || length > 255 {
    55  		return nil, infra.NewErrorStack("invalid nano-id length: " + strconv.Itoa(length))
    56  	}
    57  
    58  	preAllocSize := length * length * 8
    59  	bytes := make([]byte, preAllocSize)
    60  	if _, err := crand.Read(bytes); err != nil {
    61  		return nil, infra.WrapErrorStackWithMessage(err, "[nano-id] pre-allocate bytes failed")
    62  	}
    63  	nanoID := make([]byte, length)
    64  	offset := 0
    65  	mask := byte(len(classicNanoIDAlphabet) - 1)
    66  
    67  	var mu sync.Mutex
    68  	return func() string {
    69  		mu.Lock()
    70  		defer mu.Unlock()
    71  
    72  		if offset == preAllocSize {
    73  			if _, err := crand.Read(bytes); /* impossible */ err != nil {
    74  				panic(fmt.Errorf("[nano-id] pre-allocate bytes failed (run out of data), %w", err))
    75  			}
    76  			offset = 0
    77  		}
    78  
    79  		for i := 0; i < length; i++ {
    80  			nanoID[i] = classicNanoIDAlphabet[bytes[i+offset]&mask]
    81  		}
    82  		offset += length
    83  		return string(nanoID)
    84  	}, nil
    85  }