github.com/livekit/protocol@v1.16.1-0.20240517185851-47e4c6bba773/utils/id.go (about)

     1  // Copyright 2023 LiveKit, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package utils
    16  
    17  import (
    18  	"crypto/rand"
    19  	"crypto/sha1"
    20  	"fmt"
    21  	mrand "math/rand/v2"
    22  	"os"
    23  	"sync"
    24  
    25  	"github.com/jxskiss/base62"
    26  	"github.com/lithammer/shortuuid/v4"
    27  
    28  	"github.com/livekit/protocol/livekit"
    29  	"github.com/livekit/protocol/utils/must"
    30  )
    31  
    32  const (
    33  	GuidSize        = 12
    34  	guidScratchSize = GuidSize + 10
    35  )
    36  
    37  const (
    38  	RoomPrefix            = "RM_"
    39  	NodePrefix            = "ND_"
    40  	ParticipantPrefix     = "PA_"
    41  	TrackPrefix           = "TR_"
    42  	APIKeyPrefix          = "API"
    43  	EgressPrefix          = "EG_"
    44  	IngressPrefix         = "IN_"
    45  	SIPTrunkPrefix        = "ST_"
    46  	SIPDispatchRulePrefix = "SDR_"
    47  	SIPCallPrefix         = "SCL_"
    48  	RPCPrefix             = "RPC_"
    49  	WHIPResourcePrefix    = "WH_"
    50  	RTMPResourcePrefix    = "RT_"
    51  	URLResourcePrefix     = "UR_"
    52  	AgentWorkerPrefix     = "AW_"
    53  	AgentJobPrefix        = "AJ_"
    54  )
    55  
    56  var guidGeneratorPool = sync.Pool{
    57  	New: func() any {
    58  		return must.Get(newGuidGenerator(guidScratchSize))
    59  	},
    60  }
    61  
    62  func NewGuid(prefix string) string {
    63  	g := guidGeneratorPool.Get().(*guidGenerator)
    64  	defer guidGeneratorPool.Put(g)
    65  	return g.NewGuid(prefix)
    66  }
    67  
    68  // HashedID creates a hashed ID from a unique string
    69  func HashedID(id string) string {
    70  	h := sha1.New()
    71  	h.Write([]byte(id))
    72  	val := h.Sum(nil)
    73  
    74  	return base62.EncodeToString(val)
    75  }
    76  
    77  func LocalNodeID() (string, error) {
    78  	hostname, err := os.Hostname()
    79  	if err != nil {
    80  		return "", err
    81  	}
    82  	return fmt.Sprintf("%s%s", NodePrefix, HashedID(hostname)[:8]), nil
    83  }
    84  
    85  var b57Index = newB57Index()
    86  var b57Chars = []byte(shortuuid.DefaultAlphabet)
    87  
    88  func newB57Index() [256]byte {
    89  	var index [256]byte
    90  	for i := 0; i < len(b57Chars); i++ {
    91  		index[b57Chars[i]] = byte(i)
    92  	}
    93  	return index
    94  }
    95  
    96  type guidGenerator struct {
    97  	scratch []byte
    98  	rng     *mrand.ChaCha8
    99  }
   100  
   101  func newGuidGenerator(scratchSize int) (*guidGenerator, error) {
   102  	var seed [32]byte
   103  	if _, err := rand.Read(seed[:]); err != nil {
   104  		return nil, err
   105  	}
   106  
   107  	return &guidGenerator{
   108  		scratch: make([]byte, scratchSize),
   109  		rng:     mrand.NewChaCha8(seed),
   110  	}, nil
   111  }
   112  
   113  func (g *guidGenerator) readIDChars(b []byte) {
   114  	var n int
   115  	for {
   116  		r := g.rng.Uint64()
   117  		for i := 0; i < 10; i++ {
   118  			if int(r&0x3f) < len(b57Chars) {
   119  				b[n] = b57Chars[r&0x3f]
   120  				n++
   121  				if n == len(b) {
   122  					return
   123  				}
   124  			}
   125  			r >>= 6
   126  		}
   127  	}
   128  }
   129  
   130  func (g *guidGenerator) NewGuid(prefix string) string {
   131  	s := append(g.scratch[:0], make([]byte, len(prefix)+GuidSize)...)
   132  	copy(s, prefix)
   133  	g.readIDChars(s[len(prefix):])
   134  	return string(s)
   135  }
   136  
   137  func guidPrefix[T livekit.Guid]() string {
   138  	var id T
   139  	switch any(id).(type) {
   140  	case livekit.TrackID:
   141  		return TrackPrefix
   142  	case livekit.ParticipantID:
   143  		return ParticipantPrefix
   144  	case livekit.RoomID:
   145  		return RoomPrefix
   146  	default:
   147  		panic("unreachable")
   148  	}
   149  }
   150  
   151  func MarshalGuid[T livekit.Guid](id T) livekit.GuidBlock {
   152  	var b livekit.GuidBlock
   153  	idb := []byte(id)[len(id)-GuidSize:]
   154  	for i := 0; i < 3; i++ {
   155  		j := i * 3
   156  		k := i * 4
   157  		b[j] = b57Index[idb[k]]<<2 | b57Index[idb[k+1]]>>4
   158  		b[j+1] = b57Index[idb[k+1]]<<4 | b57Index[idb[k+2]]>>2
   159  		b[j+2] = b57Index[idb[k+2]]<<6 | b57Index[idb[k+3]]
   160  	}
   161  	return b
   162  }
   163  
   164  func UnmarshalGuid[T livekit.Guid](b livekit.GuidBlock) T {
   165  	prefix := guidPrefix[T]()
   166  	id := make([]byte, len(prefix)+GuidSize)
   167  	copy(id, []byte(prefix))
   168  	idb := id[len(prefix):]
   169  	for i := 0; i < 3; i++ {
   170  		j := i * 3
   171  		k := i * 4
   172  		idb[k] = b57Chars[b[j]>>2]
   173  		idb[k+1] = b57Chars[(b[j]&3)<<4|b[j+1]>>4]
   174  		idb[k+2] = b57Chars[(b[j+1]&15)<<2|b[j+2]>>6]
   175  		idb[k+3] = b57Chars[b[j+2]&63]
   176  	}
   177  	return T(id)
   178  }