TUN-7700: Implement feature selector to determine if connections will prefer post quantum cryptography

features/features.go

@@ -28,3 +28,22 @@ func Contains(feature string) bool {
 	}
 	return false
 }
+
+// Remove any duplicates from the slice
+func Dedup(slice []string) []string {
+
+	// Convert the slice into a set
+	set := make(map[string]bool, 0)
+	for _, str := range slice {
+		set[str] = true
+	}
+
+	// Convert the set back into a slice
+	keys := make([]string, len(set))
+	i := 0
+	for str := range set {
+		keys[i] = str
+		i++
+	}
+	return keys
+}

features/selector.go

@@ -0,0 +1,164 @@
+package features
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"hash/fnv"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/rs/zerolog"
+)
+
+const (
+	featureSelectorHostname = "cfd-features.argotunnel.com"
+	defaultRefreshFreq      = time.Hour * 6
+	lookupTimeout           = time.Second * 10
+)
+
+type PostQuantumMode uint8
+
+const (
+	PostQuantumDisabled PostQuantumMode = iota
+	// Prefer post quantum, but fallback if connection cannot be established
+	PostQuantumPrefer
+	// If the user passes the --post-quantum flag, we override
+	// CurvePreferences to only support hybrid post-quantum key agreements.
+	PostQuantumStrict
+)
+
+// If the TXT record adds other fields, the umarshal logic will ignore those keys
+// If the TXT record is missing a key, the field will unmarshal to the default Go value
+type featuresRecord struct {
+	PostQuantumPercentage int32 `json:"pq"`
+}
+
+func NewFeatureSelector(ctx context.Context, accountTag string, staticFeatures StaticFeatures, logger *zerolog.Logger) (*FeatureSelector, error) {
+	return newFeatureSelector(ctx, accountTag, logger, newDNSResolver(), staticFeatures, defaultRefreshFreq)
+}
+
+// FeatureSelector determines if this account will try new features. It preiodically queries a DNS TXT record
+// to see which features are turned on
+type FeatureSelector struct {
+	accountHash int32
+	logger      *zerolog.Logger
+	resolver    resolver
+
+	staticFeatures StaticFeatures
+
+	// lock protects concurrent access to dynamic features
+	lock     sync.RWMutex
+	features featuresRecord
+}
+
+// Features set by user provided flags
+type StaticFeatures struct {
+	PostQuantumMode *PostQuantumMode
+}
+
+func newFeatureSelector(ctx context.Context, accountTag string, logger *zerolog.Logger, resolver resolver, staticFeatures StaticFeatures, refreshFreq time.Duration) (*FeatureSelector, error) {
+	selector := &FeatureSelector{
+		accountHash:    switchThreshold(accountTag),
+		logger:         logger,
+		resolver:       resolver,
+		staticFeatures: staticFeatures,
+	}
+
+	if err := selector.refresh(ctx); err != nil {
+		logger.Err(err).Msg("Failed to fetch features, default to disable")
+	}
+
+	go selector.refreshLoop(ctx, refreshFreq)
+
+	return selector, nil
+}
+
+func (fs *FeatureSelector) PostQuantumMode() PostQuantumMode {
+	if fs.staticFeatures.PostQuantumMode != nil {
+		return *fs.staticFeatures.PostQuantumMode
+	}
+
+	fs.lock.RLock()
+	defer fs.lock.RUnlock()
+
+	if fs.features.PostQuantumPercentage > fs.accountHash {
+		return PostQuantumPrefer
+	}
+	return PostQuantumDisabled
+}
+
+func (fs *FeatureSelector) refreshLoop(ctx context.Context, refreshFreq time.Duration) {
+	ticker := time.NewTicker(refreshFreq)
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			err := fs.refresh(ctx)
+			if err != nil {
+				fs.logger.Err(err).Msg("Failed to refresh feature selector")
+			}
+		}
+	}
+}
+
+func (fs *FeatureSelector) refresh(ctx context.Context) error {
+	record, err := fs.resolver.lookupRecord(ctx)
+	if err != nil {
+		return err
+	}
+
+	var features featuresRecord
+	if err := json.Unmarshal(record, &features); err != nil {
+		return err
+	}
+
+	pq_enabled := features.PostQuantumPercentage > fs.accountHash
+	fs.logger.Debug().Int32("account_hash", fs.accountHash).Int32("pq_perct", features.PostQuantumPercentage).Bool("pq_enabled", pq_enabled).Msg("Refreshed feature")
+
+	fs.lock.Lock()
+	defer fs.lock.Unlock()
+
+	fs.features = features
+
+	return nil
+}
+
+// resolver represents an object that can look up featuresRecord
+type resolver interface {
+	lookupRecord(ctx context.Context) ([]byte, error)
+}
+
+type dnsResolver struct {
+	resolver *net.Resolver
+}
+
+func newDNSResolver() *dnsResolver {
+	return &dnsResolver{
+		resolver: net.DefaultResolver,
+	}
+}
+
+func (dr *dnsResolver) lookupRecord(ctx context.Context) ([]byte, error) {
+	ctx, cancel := context.WithTimeout(ctx, lookupTimeout)
+	defer cancel()
+
+	records, err := dr.resolver.LookupTXT(ctx, featureSelectorHostname)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(records) == 0 {
+		return nil, fmt.Errorf("No TXT record found for %s to determine which features to opt-in", featureSelectorHostname)
+	}
+
+	return []byte(records[0]), nil
+}
+
+func switchThreshold(accountTag string) int32 {
+	h := fnv.New32a()
+	_, _ = h.Write([]byte(accountTag))
+	return int32(h.Sum32() % 100)
+}

features/selector_test.go

@@ -0,0 +1,128 @@
+package features
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/require"
+)
+
+func TestUnmarshalFeaturesRecord(t *testing.T) {
+	tests := []struct {
+		record             []byte
+		expectedPercentage int32
+		expectedErr        bool
+	}{
+		{
+			record:             []byte(`{"pq":0}`),
+			expectedPercentage: 0,
+		},
+		{
+			record:             []byte(`{"pq":39}`),
+			expectedPercentage: 39,
+		},
+		{
+			record:             []byte(`{"pq":100}`),
+			expectedPercentage: 100,
+		},
+		{
+			record:             []byte(`{}`),
+			expectedPercentage: 0, // Unmarshal to default struct if key is not present
+		},
+		{
+			record:             []byte(`{"kyber":768}`),
+			expectedPercentage: 0, // Unmarshal to default struct if key is not present
+		},
+		{
+			record:      []byte(`{"pq":"kyber768"}`),
+			expectedErr: true,
+		},
+	}
+
+	for _, test := range tests {
+		var features featuresRecord
+		err := json.Unmarshal(test.record, &features)
+		if test.expectedErr {
+			require.Error(t, err, test)
+		} else {
+			require.NoError(t, err)
+			require.Equal(t, test.expectedPercentage, features.PostQuantumPercentage, test)
+		}
+	}
+}
+
+func TestRefreshFeaturesRecord(t *testing.T) {
+	// The hash of the accountTag is 82
+	accountTag := t.Name()
+	threshold := switchThreshold(accountTag)
+
+	percentages := []int32{0, 10, 80, 83, 100}
+	refreshFreq := time.Millisecond * 10
+	selector := newTestSelector(t, percentages, nil, refreshFreq)
+
+	for _, percentage := range percentages {
+		if percentage > threshold {
+			require.Equal(t, PostQuantumPrefer, selector.PostQuantumMode())
+		} else {
+			require.Equal(t, PostQuantumDisabled, selector.PostQuantumMode())
+		}
+
+		time.Sleep(refreshFreq + time.Millisecond)
+	}
+
+	// Make sure error doesn't override the last fetched features
+	require.Equal(t, PostQuantumPrefer, selector.PostQuantumMode())
+}
+
+func TestStaticFeatures(t *testing.T) {
+	percentages := []int32{0}
+	pqMode := PostQuantumStrict
+	selector := newTestSelector(t, percentages, &pqMode, time.Millisecond*10)
+	require.Equal(t, PostQuantumStrict, selector.PostQuantumMode())
+}
+
+// Verify that if the first refresh fails, the selector will use default features
+func TestFailedRefreshInitToDefault(t *testing.T) {
+	selector := newTestSelector(t, []int32{}, nil, time.Second)
+	require.Equal(t, featuresRecord{}, selector.features)
+	require.Equal(t, PostQuantumDisabled, selector.PostQuantumMode())
+}
+
+func newTestSelector(t *testing.T, percentages []int32, pqMode *PostQuantumMode, refreshFreq time.Duration) *FeatureSelector {
+	accountTag := t.Name()
+	logger := zerolog.Nop()
+
+	resolver := &mockResolver{
+		percentages: percentages,
+	}
+
+	staticFeatures := StaticFeatures{
+		PostQuantumMode: pqMode,
+	}
+	selector, err := newFeatureSelector(context.Background(), accountTag, &logger, resolver, staticFeatures, refreshFreq)
+	require.NoError(t, err)
+
+	return selector
+}
+
+type mockResolver struct {
+	nextIndex   int
+	percentages []int32
+}
+
+func (mr *mockResolver) lookupRecord(ctx context.Context) ([]byte, error) {
+	if mr.nextIndex >= len(mr.percentages) {
+		return nil, fmt.Errorf("no more record to lookup")
+	}
+
+	record, err := json.Marshal(featuresRecord{
+		PostQuantumPercentage: mr.percentages[mr.nextIndex],
+	})
+	mr.nextIndex++
+
+	return record, err
+}