optionize shard storage; enable custom shard generation

storage/shard/shard.go

@@ -16,27 +16,52 @@ const (
 	DeclarationIdentifier = "com.github.jessepeterson.kmfddm.storage.shard.v1"
 )
 
+// ShardFunc computes a shard value. A string of a decimal value
+// between 0 and 100 inclusive should be returned. I.e. "42".
+type ShardFunc func(string) string
+
+// FNV1Shard hashes input using FNV1 modulo 101.
+func FNV1Shard(input string) string {
+	hash := fnv.New32()
+	hash.Write([]byte(input))
+	result := hash.Sum32()
+	return strconv.Itoa(int(result % 101))
+}
+
 // ShardStorage is a dynamic storage backend that synthesizes a shard declaration.
 // The declaration is a management properties declaration that sets the "shard"
 // property to the shard number for the enrollment. This can then be used in
 // activation predicates.
-// The shard number is computed from the enrollment ID, hashed using FNV1 modulo 101.
-type ShardStorage struct{}
+// The shard number is computed from the enrollment ID.
+type ShardStorage struct {
+	shardFunc ShardFunc
+}
 
-// NewShardStorage creates a new shard storage.
-func NewShardStorage() *ShardStorage {
-	return new(ShardStorage)
+type Option func(*ShardStorage)
+
+// WithShardFunc sets the shard function to f.
+func WithShardFunc(f ShardFunc) Option {
+	if f == nil {
+		panic("nil shard func")
+	}
+	return func(s *ShardStorage) {
+		s.shardFunc = f
+	}
 }
 
-func computeShard(input string) string {
-	hash := fnv.New32()
-	hash.Write([]byte(input))
-	result := hash.Sum32()
-	return strconv.Itoa(int(result % 101))
+// NewShardStorage creates a new shard storage.
+// By default the shard function is hashed using [FNV1Shard].
+// Use [WithShardFunc] to change it.
+func NewShardStorage(opts ...Option) *ShardStorage {
+	s := &ShardStorage{shardFunc: FNV1Shard}
+	for _, opt := range opts {
+		opt(s)
+	}
+	return s
 }
 
-func serverToken(enrollmentID string) string {
-	return "shard=" + computeShard(enrollmentID) + ";version=1"
+func (s *ShardStorage) serverToken(enrollmentID string) string {
+	return "shard=" + s.shardFunc(enrollmentID) + ";version=1"
 }
 
 // RetrieveDeclarationItems synthesizes a dynamic shard declaration.
@@ -45,7 +70,7 @@ func (s *ShardStorage) RetrieveDeclarationItems(_ context.Context, enrollmentID
 	return []*ddm.Declaration{{
 		Type:        DeclarationType,
 		Identifier:  DeclarationIdentifier,
-		ServerToken: serverToken(enrollmentID),
+		ServerToken: s.serverToken(enrollmentID),
 	}}, nil
 }
 
@@ -61,9 +86,9 @@ func (s *ShardStorage) RetrieveEnrollmentDeclarationJSON(_ context.Context, decl
 	"Type": "` + DeclarationType + `",
 	"Identifier": "` + DeclarationIdentifier + `",
 	"Payload": {
-		"shard": ` + computeShard(enrollmentID) + `
+		"shard": ` + s.shardFunc(enrollmentID) + `
 	},
-	"ServerToken": "` + serverToken(enrollmentID) + `"
+	"ServerToken": "` + s.serverToken(enrollmentID) + `"
 }`
 	return []byte(json), nil
 }