sets: convert set types to generics (#3519)

* sets: convert set types to generics

* sets: use internal sets package

authorize/access_tracker.go

@@ -58,8 +58,8 @@ func (tracker *AccessTracker) Run(ctx context.Context) {
 	ticker := time.NewTicker(tracker.debouncePeriod)
 	defer ticker.Stop()
 
-	sessionAccesses := sets.NewSizeLimitedStringSet(tracker.maxSize)
-	serviceAccountAccesses := sets.NewSizeLimitedStringSet(tracker.maxSize)
+	sessionAccesses := sets.NewSizeLimited[string](tracker.maxSize)
+	serviceAccountAccesses := sets.NewSizeLimited[string](tracker.maxSize)
 	runTrackSessionAccess := func(sessionID string) {
 		sessionAccesses.Add(sessionID)
 	}
@@ -95,8 +95,8 @@ func (tracker *AccessTracker) Run(ctx context.Context) {
 			return
 		}
 
-		sessionAccesses = sets.NewSizeLimitedStringSet(tracker.maxSize)
-		serviceAccountAccesses = sets.NewSizeLimitedStringSet(tracker.maxSize)
+		sessionAccesses = sets.NewSizeLimited[string](tracker.maxSize)
+		serviceAccountAccesses = sets.NewSizeLimited[string](tracker.maxSize)
 	}
 
 	for {

config/envoyconfig/listeners.go

@@ -719,7 +719,7 @@ func (b *Builder) buildDownstreamValidationContext(ctx context.Context,
 }
 
 func getRouteableDomainsForTLSServerName(options *config.Options, addr string, tlsServerName string) ([]string, error) {
-	allDomains := sets.NewSortedString()
+	allDomains := sets.NewSorted[string]()
 
 	if addr == options.Addr {
 		domains, err := options.GetAllRouteableHTTPDomainsForTLSServerName(tlsServerName)
@@ -741,7 +741,7 @@ func getRouteableDomainsForTLSServerName(options *config.Options, addr string, t
 }
 
 func getAllRouteableDomains(options *config.Options, addr string) ([]string, error) {
-	allDomains := sets.NewSortedString()
+	allDomains := sets.NewSorted[string]()
 
 	if addr == options.Addr {
 		domains, err := options.GetAllRouteableHTTPDomains()
@@ -768,7 +768,7 @@ func getAllTLSDomains(options *config.Options, addr string) ([]string, error) {
 		return nil, err
 	}
 
-	domains := sets.NewSortedString()
+	domains := sets.NewSorted[string]()
 	for _, hp := range allDomains {
 		if d, _, err := net.SplitHostPort(hp); err == nil {
 			domains.Add(d)

config/options.go

@@ -1062,7 +1062,7 @@ func (o *Options) GetAllRouteableGRPCDomains() ([]string, error) {
 // GetAllRouteableGRPCDomainsForTLSServerName  returns all the possible gRPC domains handled by the Pomerium options
 // for the given TLS server name.
 func (o *Options) GetAllRouteableGRPCDomainsForTLSServerName(tlsServerName string) ([]string, error) {
-	domains := sets.NewSortedString()
+	domains := sets.NewSorted[string]()
 
 	// authorize urls
 	if IsAll(o.Services) {
@@ -1134,7 +1134,7 @@ func (o *Options) GetAllRouteableHTTPDomainsForTLSServerName(tlsServerName strin
 		return nil, err
 	}
 
-	domains := sets.NewSortedString()
+	domains := sets.NewSorted[string]()
 	if IsAuthenticate(o.Services) {
 		authenticateURL, err := o.GetInternalAuthenticateURL()
 		if err != nil {

go.mod

@@ -17,7 +17,6 @@ require (
 	github.com/docker/docker v20.10.17+incompatible
 	github.com/envoyproxy/go-control-plane v0.10.2-0.20220325020618-49ff273808a1
 	github.com/envoyproxy/protoc-gen-validate v0.6.7
-	github.com/fsnotify/fsnotify v1.5.4 // indirect
 	github.com/go-chi/chi/v5 v5.0.7
 	github.com/go-jose/go-jose/v3 v3.0.0
 	github.com/go-redis/redis/v8 v8.11.5
@@ -54,9 +53,7 @@ require (
 	github.com/rjeczalik/notify v0.9.3-0.20201210012515-e2a77dcc14cf
 	github.com/rs/cors v1.8.2
 	github.com/rs/zerolog v1.27.0
-	github.com/scylladb/go-set v1.0.2
 	github.com/shirou/gopsutil/v3 v3.22.6
-	github.com/spf13/cobra v1.5.0 // indirect
 	github.com/spf13/viper v1.12.0
 	github.com/stretchr/testify v1.8.0
 	github.com/tniswong/go.rfcx v0.0.0-20181019234604-07783c52761f
@@ -67,6 +64,7 @@ require (
 	go.opencensus.io v0.23.0
 	go.uber.org/zap v1.21.0
 	golang.org/x/crypto v0.0.0-20220525230936-793ad666bf5e
+	golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e
 	golang.org/x/net v0.0.0-20220630215102-69896b714898
 	golang.org/x/oauth2 v0.0.0-20220622183110-fd043fe589d2
 	golang.org/x/sync v0.0.0-20220601150217-0de741cfad7f
@@ -127,6 +125,7 @@ require (
 	github.com/fatih/structtag v1.2.0 // indirect
 	github.com/felixge/httpsnoop v1.0.2 // indirect
 	github.com/firefart/nonamedreturns v1.0.4 // indirect
+	github.com/fsnotify/fsnotify v1.5.4 // indirect
 	github.com/fxamacker/cbor/v2 v2.3.0 // indirect
 	github.com/fzipp/gocyclo v0.6.0 // indirect
 	github.com/ghodss/yaml v1.0.0 // indirect
@@ -246,6 +245,7 @@ require (
 	github.com/sourcegraph/go-diff v0.6.1 // indirect
 	github.com/spf13/afero v1.8.2 // indirect
 	github.com/spf13/cast v1.5.0 // indirect
+	github.com/spf13/cobra v1.5.0 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/ssgreg/nlreturn/v2 v2.2.1 // indirect

go.sum

@@ -518,8 +518,6 @@ github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5Kwzbycv
 github.com/fatih/color v1.10.0/go.mod h1:ELkj/draVOlAH/xkhN6mQ50Qd0MPOk5AAr3maGEBuJM=
 github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
-github.com/fatih/set v0.2.1 h1:nn2CaJyknWE/6txyUDGwysr3G5QC6xWB/PtVjPBbeaA=
-github.com/fatih/set v0.2.1/go.mod h1:+RKtMCH+favT2+3YecHGxcc0b4KyVWA1QWWJUs4E0CI=
 github.com/fatih/structtag v1.2.0 h1:/OdNE99OxoI/PqaW/SuSK9uxxT3f/tcSZgon/ssNSx4=
 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
 github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
@@ -1414,8 +1412,6 @@ github.com/sanposhiho/wastedassign/v2 v2.0.6/go.mod h1:KyZ0MWTwxxBmfwn33zh3k1dms
 github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0=
 github.com/sclevine/agouti v3.0.0+incompatible/go.mod h1:b4WX9W9L1sfQKXeJf1mUTLZKJ48R1S7H23Ji7oFO5Bw=
 github.com/sclevine/spec v1.2.0/go.mod h1:W4J29eT/Kzv7/b9IWLB055Z+qvVC9vt0Arko24q7p+U=
-github.com/scylladb/go-set v1.0.2 h1:SkvlMCKhP0wyyct6j+0IHJkBkSZL+TDzZ4E7f7BCcRE=
-github.com/scylladb/go-set v1.0.2/go.mod h1:DkpGd78rljTxKAnTDPFqXSGxvETQnJyuSOQwsHycqfs=
 github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc=
 github.com/seccomp/libseccomp-golang v0.9.1/go.mod h1:GbW5+tmTXfcxTToHLXlScSlAvWlF4P2Ca7zGrPiEpWo=
 github.com/seccomp/libseccomp-golang v0.9.2-0.20210429002308-3879420cc921/go.mod h1:JA8cRccbGaA1s33RQf7Y1+q9gHmZX1yB/z9WDN1C6fg=
@@ -1755,6 +1751,8 @@ golang.org/x/exp v0.0.0-20200119233911-0405dc783f0a/go.mod h1:2RIsYlXP63K8oxa1u0
 golang.org/x/exp v0.0.0-20200207192155-f17229e696bd/go.mod h1:J/WKrq2StrnmMY6+EHIKF9dgMWnmCNThgcyBT1FY9mM=
 golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6/go.mod h1:3jZMyOhIsHpP37uCMkUooju7aAi5cS1Q23tOzKc+0MU=
 golang.org/x/exp v0.0.0-20200331195152-e8c3332aa8e5/go.mod h1:4M0jN8W1tt0AVLNr8HDosyJCDCDuyL9N9+3m7wDWgKw=
+golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e h1:+WEEuIdZHnUeJJmEUjyYC2gfUMj69yZXw17EnHg/otA=
+golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e/go.mod h1:Kr81I6Kryrl9sr8s2FK3vxD90NdsKWRuOIl2O4CvYbA=
 golang.org/x/exp/typeparams v0.0.0-20220613132600-b0d781184e0d h1:+W8Qf4iJtMGKkyAygcKohjxTk4JPsL9DpzApJ22m5Ic=
 golang.org/x/exp/typeparams v0.0.0-20220613132600-b0d781184e0d/go.mod h1:AbB0pIl9nAr9wVwH+Z2ZpaocVmF5I4GyWCDIsVjR0bk=
 golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js=

internal/redisutil/client.go

@@ -10,28 +10,29 @@ import (
 	"time"
 
 	"github.com/go-redis/redis/v8"
-	"github.com/scylladb/go-set"
+
+	"github.com/pomerium/pomerium/internal/sets"
 )
 
 var (
-	standardSchemes = set.NewStringSet("redis", "rediss", "unix")
-	clusterSchemes  = set.NewStringSet(
+	standardSchemes = sets.NewHash("redis", "rediss", "unix")
+	clusterSchemes  = sets.NewHash(
 		"redis+cluster", "redis-cluster",
 		"rediss+cluster", "rediss-cluster",
 		"redis+clusters", "redis-clusters",
 	)
-	sentinelSchemes = set.NewStringSet(
+	sentinelSchemes = sets.NewHash(
 		"redis+sentinel", "redis-sentinel",
 		"rediss+sentinel", "rediss-sentinel",
 		"redis+sentinels", "redis-sentinels",
 	)
-	sentinelClusterSchemes = set.NewStringSet(
+	sentinelClusterSchemes = sets.NewHash(
 		"redis+sentinel+cluster", "redis-sentinel-cluster",
 		"rediss+sentinel+cluster", "rediss-sentinel-cluster",
 		"redis+sentinels+cluster", "redis-sentinels-cluster",
 		"redis+sentinel+clusters", "redis-sentinel-clusters",
 	)
-	tlsSchemes = set.NewStringSet(
+	tlsSchemes = sets.NewHash(
 		"rediss",
 		"rediss+cluster", "rediss-cluster",
 		"redis+clusters", "redis-clusters",

internal/sets/hash.go

@@ -6,10 +6,12 @@ type Hash[T comparable] struct {
 }
 
 // NewHash creates a new Hash set.
-func NewHash[T comparable]() *Hash[T] {
-	return &Hash[T]{
+func NewHash[T comparable](initialValues ...T) *Hash[T] {
+	s := &Hash[T]{
 		m: make(map[T]struct{}),
 	}
+	s.Add(initialValues...)
+	return s
 }
 
 // Add adds a value to the set.

internal/sets/size_limited.go

@@ -1,22 +1,22 @@
 package sets
 
-// A SizeLimitedStringSet is a StringSet which is limited to a given size. Once
+// A SizeLimited is a Set which is limited to a given size. Once
 // the capacity is reached an element will be removed at random.
-type SizeLimitedStringSet struct {
-	m        map[string]struct{}
+type SizeLimited[T comparable] struct {
+	m        map[T]struct{}
 	capacity int
 }
 
-// NewSizeLimitedStringSet create a new SizeLimitedStringSet.
-func NewSizeLimitedStringSet(capacity int) *SizeLimitedStringSet {
-	return &SizeLimitedStringSet{
-		m:        make(map[string]struct{}),
+// NewSizeLimited create a new SizeLimited.
+func NewSizeLimited[T comparable](capacity int) *SizeLimited[T] {
+	return &SizeLimited[T]{
+		m:        make(map[T]struct{}),
 		capacity: capacity,
 	}
 }
 
 // Add adds an element to the set.
-func (s *SizeLimitedStringSet) Add(element string) {
+func (s *SizeLimited[T]) Add(element T) {
 	s.m[element] = struct{}{}
 	for len(s.m) > s.capacity {
 		for k := range s.m {
@@ -27,7 +27,7 @@ func (s *SizeLimitedStringSet) Add(element string) {
 }
 
 // ForEach iterates over all the elements in the set.
-func (s *SizeLimitedStringSet) ForEach(callback func(element string) bool) {
+func (s *SizeLimited[T]) ForEach(callback func(element T) bool) {
 	for k := range s.m {
 		if !callback(k) {
 			return

internal/sets/sorted.go

@@ -1,59 +1,65 @@
 package sets
 
-import "github.com/google/btree"
+import (
+	"github.com/google/btree"
+	"golang.org/x/exp/constraints"
+)
 
-func lessFn(a, b string) bool { return a < b }
-
-// A SortedString is a set of strings with sorted iteration.
-type SortedString struct {
-	b *btree.BTreeG[string]
+// A Sorted is a set with sorted iteration.
+type Sorted[T any] struct {
+	b    *btree.BTreeG[T]
+	less func(a, b T) bool
 }
 
-// NewSortedString creates a new sorted string set.
-func NewSortedString() *SortedString {
-	return &SortedString{
-		b: btree.NewG[string](8, lessFn),
+// NewSorted creates a new sorted string set.
+func NewSorted[T constraints.Ordered]() *Sorted[T] {
+	less := func(a, b T) bool {
+		return a < b
+	}
+	return &Sorted[T]{
+		b:    btree.NewG(8, less),
+		less: less,
 	}
 }
 
 // Add adds a string to the set.
-func (s *SortedString) Add(elements ...string) {
+func (s *Sorted[T]) Add(elements ...T) {
 	for _, element := range elements {
 		s.b.ReplaceOrInsert(element)
 	}
 }
 
 // Clear clears the set.
-func (s *SortedString) Clear() {
-	s.b = btree.NewG[string](8, lessFn)
+func (s *Sorted[T]) Clear() {
+	s.b = btree.NewG(8, s.less)
 }
 
 // Delete deletes an element from the set.
-func (s *SortedString) Delete(element string) {
+func (s *Sorted[T]) Delete(element T) {
 	s.b.Delete(element)
 }
 
 // ForEach iterates over the set in ascending order.
-func (s *SortedString) ForEach(callback func(element string) bool) {
-	s.b.Ascend(func(item string) bool {
+func (s *Sorted[T]) ForEach(callback func(element T) bool) {
+	s.b.Ascend(func(item T) bool {
 		return callback(item)
 	})
 }
 
 // Has returns true if the elment is in the set.
-func (s *SortedString) Has(element string) bool {
+func (s *Sorted[T]) Has(element T) bool {
 	return s.b.Has(element)
 }
 
 // Size returns the size of the set.
-func (s *SortedString) Size() int {
+func (s *Sorted[T]) Size() int {
 	return s.b.Len()
 }
 
 // ToSlice returns a slice of all the elements in the set.
-func (s *SortedString) ToSlice() []string {
-	arr := make([]string, 0, s.Size())
-	s.b.Ascend(func(item string) bool {
+func (s *Sorted[T]) ToSlice() []T {
+	arr := make([]T, 0, s.Size())
+	s.b.Ascend(func(item T) bool {
 		arr = append(arr, item)
 		return true
 	})