v2ray-core/common/session/session.go

package session

import (
	"context"
	"math/rand"
)

type ID uint32

func NewID() ID {
	for {
		id := ID(rand.Uint32())
		if id != 0 {
			return id
		}
	}
}

type sessionKey int

const (
	idSessionKey sessionKey = iota
)

func ContextWithID(ctx context.Context, id ID) context.Context {
	return context.WithValue(ctx, idSessionKey, id)
}

func IDFromContext(ctx context.Context) ID {
	if id, ok := ctx.Value(idSessionKey).(ID); ok {
		return id
	}
	return 0
}