Unify changes loading in change loader

data/aclstatebuilder.go

@@ -12,6 +12,7 @@ type ACLStateBuilder struct {
 	aclState *ACLState
 	identity string
 	key      threadmodels.EncryptionPrivKey
+	decoder  threadmodels.SigningPubKeyDecoder
 }
 
 type decreasedPermissionsParameters struct {
@@ -19,28 +20,10 @@ type decreasedPermissionsParameters struct {
 	startChange string
 }
 
-func NewACLStateBuilder(
+func NewACLStateBuilder(decoder threadmodels.SigningPubKeyDecoder) *ACLStateBuilder {
-	tree *Tree,
-	identity string,
-	key threadmodels.EncryptionPrivKey,
-	decoder threadmodels.SigningPubKeyDecoder) (*ACLStateBuilder, error) {
-	root := tree.Root()
-	if !root.IsSnapshot {
-		return nil, fmt.Errorf("root should always be a snapshot")
-	}
-
-	snapshot := root.Content.GetAclData().GetAclSnapshot()
-	state, err := NewACLStateFromSnapshot(snapshot, identity, key, decoder)
-	if err != nil {
-		return nil, fmt.Errorf("could not build aclState from snapshot: %w", err)
-	}
-
 	return &ACLStateBuilder{
-		tree:     tree,
+		decoder: decoder,
-		aclState: state,
+	}
-		identity: identity,
-		key:      key,
-	}, nil
 }
 
 func (sb *ACLStateBuilder) Build() (*ACLState, error) {
@@ -48,6 +31,31 @@ func (sb *ACLStateBuilder) Build() (*ACLState, error) {
 	return state, err
 }
 
+func (sb *ACLStateBuilder) Init(
+	tree *Tree,
+	accountData *AccountData) error {
+	root := tree.Root()
+	if !root.IsSnapshot {
+		return fmt.Errorf("root should always be a snapshot")
+	}
+
+	snapshot := root.Content.GetAclData().GetAclSnapshot()
+	state, err := NewACLStateFromSnapshot(
+		snapshot,
+		accountData.Identity,
+		accountData.EncKey,
+		sb.decoder)
+	if err != nil {
+		return fmt.Errorf("could not build aclState from snapshot: %w", err)
+	}
+	sb.tree = tree
+	sb.identity = accountData.Identity
+	sb.key = accountData.EncKey
+	sb.aclState = state
+
+	return nil
+}
+
 // TODO: we can probably have only one state builder, because we can build both at the same time
 func (sb *ACLStateBuilder) BuildBefore(beforeId string) (*ACLState, bool, error) {
 	var (

data/acltreebuilder.go

@@ -1,14 +1,10 @@
 package data
 
 import (
-	"context"
 	"fmt"
-	"github.com/anytypeio/go-anytype-infrastructure-experiments/data/pb"
 	"github.com/anytypeio/go-anytype-infrastructure-experiments/data/threadmodels"
 	"github.com/anytypeio/go-anytype-infrastructure-experiments/util/slice"
-	"github.com/gogo/protobuf/proto"
 	"github.com/textileio/go-threads/core/thread"
-	"time"
 )
 
 type ACLTreeBuilder struct {
@@ -17,64 +13,23 @@ type ACLTreeBuilder struct {
 	signingPubKeyDecoder threadmodels.SigningPubKeyDecoder
 	tree                 *Tree
 	thread               threadmodels.Thread
+
+	*changeLoader
 }
 
 func NewACLTreeBuilder(t threadmodels.Thread, decoder threadmodels.SigningPubKeyDecoder) *ACLTreeBuilder {
 	return &ACLTreeBuilder{
-		cache:                make(map[string]*Change),
-		identityKeys:         make(map[string]threadmodels.SigningPubKey),
 		signingPubKeyDecoder: decoder,
-		tree:                 &Tree{}, // TODO: add NewTree method
 		thread:               t,
+		changeLoader:         newChangeLoader(t, decoder),
 	}
 }
 
-func (tb *ACLTreeBuilder) loadChange(id string) (ch *Change, err error) {
+func (tb *ACLTreeBuilder) Init() {
-	if ch, ok := tb.cache[id]; ok {
+	tb.cache = make(map[string]*Change)
-		return ch, nil
+	tb.identityKeys = make(map[string]threadmodels.SigningPubKey)
-	}
+	tb.tree = &Tree{}
-
+	tb.changeLoader.init(tb.cache, tb.identityKeys)
-	// TODO: Add virtual changes logic
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
-	defer cancel()
-
-	change, err := tb.thread.GetChange(ctx, id)
-	if err != nil {
-		return nil, err
-	}
-
-	aclChange := new(pb.ACLChange)
-
-	// TODO: think what should we do with such cases, because this can be used by attacker to break our tree
-	if err = proto.Unmarshal(change.Payload, aclChange); err != nil {
-		return
-	}
-	var verified bool
-	verified, err = tb.verify(aclChange.Identity, change.Payload, change.Signature)
-	if err != nil {
-		return
-	}
-	if !verified {
-		err = fmt.Errorf("the signature of the payload cannot be verified")
-		return
-	}
-
-	ch, err = NewACLChange(id, aclChange)
-	tb.cache[id] = ch
-
-	return ch, nil
-}
-
-func (tb *ACLTreeBuilder) verify(identity string, payload, signature []byte) (isVerified bool, err error) {
-	identityKey, exists := tb.identityKeys[identity]
-	if !exists {
-		identityKey, err = tb.signingPubKeyDecoder.DecodeFromString(identity)
-		if err != nil {
-			return
-		}
-		tb.identityKeys[identity] = identityKey
-	}
-	return identityKey.Verify(payload, signature)
 }
 
 func (tb *ACLTreeBuilder) Build() (*Tree, error) {

data/changeloader.go

@@ -0,0 +1,80 @@
+package data
+
+import (
+	"context"
+	"fmt"
+	"github.com/anytypeio/go-anytype-infrastructure-experiments/data/pb"
+	"github.com/anytypeio/go-anytype-infrastructure-experiments/data/threadmodels"
+	"github.com/gogo/protobuf/proto"
+	"time"
+)
+
+type changeLoader struct {
+	cache                map[string]*Change
+	identityKeys         map[string]threadmodels.SigningPubKey
+	signingPubKeyDecoder threadmodels.SigningPubKeyDecoder
+	thread               threadmodels.Thread
+}
+
+func newChangeLoader(
+	thread threadmodels.Thread,
+	signingPubKeyDecoder threadmodels.SigningPubKeyDecoder) *changeLoader {
+	return &changeLoader{
+		signingPubKeyDecoder: signingPubKeyDecoder,
+		thread:               thread,
+	}
+}
+
+func (c *changeLoader) init(cache map[string]*Change,
+	identityKeys map[string]threadmodels.SigningPubKey) {
+	c.cache = cache
+	c.identityKeys = identityKeys
+}
+
+func (c *changeLoader) loadChange(id string) (ch *Change, err error) {
+	if ch, ok := c.cache[id]; ok {
+		return ch, nil
+	}
+
+	// TODO: Add virtual changes logic
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+	defer cancel()
+
+	change, err := c.thread.GetChange(ctx, id)
+	if err != nil {
+		return nil, err
+	}
+
+	aclChange := new(pb.ACLChange)
+
+	// TODO: think what should we do with such cases, because this can be used by attacker to break our tree
+	if err = proto.Unmarshal(change.Payload, aclChange); err != nil {
+		return
+	}
+	var verified bool
+	verified, err = c.verify(aclChange.Identity, change.Payload, change.Signature)
+	if err != nil {
+		return
+	}
+	if !verified {
+		err = fmt.Errorf("the signature of the payload cannot be verified")
+		return
+	}
+
+	ch, err = NewACLChange(id, aclChange)
+	c.cache[id] = ch
+
+	return ch, nil
+}
+
+func (c *changeLoader) verify(identity string, payload, signature []byte) (isVerified bool, err error) {
+	identityKey, exists := c.identityKeys[identity]
+	if !exists {
+		identityKey, err = c.signingPubKeyDecoder.DecodeFromString(identity)
+		if err != nil {
+			return
+		}
+		c.identityKeys[identity] = identityKey
+	}
+	return identityKey.Verify(payload, signature)
+}

data/treebuilder.go

@@ -1,17 +1,13 @@
 package data
 
 import (
-	"context"
 	"errors"
 	"fmt"
-	"github.com/anytypeio/go-anytype-infrastructure-experiments/data/pb"
 	"github.com/anytypeio/go-anytype-infrastructure-experiments/data/threadmodels"
 	//"github.com/anytypeio/go-anytype-infrastructure-experiments/pkg/lib/logging"
 	"github.com/anytypeio/go-anytype-infrastructure-experiments/util/slice"
-	"github.com/gogo/protobuf/proto"
 	"github.com/prometheus/common/log"
 	"github.com/textileio/go-threads/core/thread"
-	"time"
 )
 
 var (
@@ -25,64 +21,23 @@ type TreeBuilder struct {
 	signingPubKeyDecoder threadmodels.SigningPubKeyDecoder
 	tree                 *Tree
 	thread               threadmodels.Thread
+
+	*changeLoader
 }
 
 func NewTreeBuilder(t threadmodels.Thread, decoder threadmodels.SigningPubKeyDecoder) *TreeBuilder {
 	return &TreeBuilder{
-		cache:                make(map[string]*Change),
-		identityKeys:         make(map[string]threadmodels.SigningPubKey),
 		signingPubKeyDecoder: decoder,
-		tree:                 &Tree{}, // TODO: add NewTree method
 		thread:               t,
+		changeLoader:         newChangeLoader(t, decoder),
 	}
 }
 
-func (tb *TreeBuilder) loadChange(id string) (ch *Change, err error) {
+func (tb *TreeBuilder) Init() {
-	if ch, ok := tb.cache[id]; ok {
+	tb.cache = make(map[string]*Change)
-		return ch, nil
+	tb.identityKeys = make(map[string]threadmodels.SigningPubKey)
-	}
+	tb.tree = &Tree{}
-
+	tb.changeLoader.init(tb.cache, tb.identityKeys)
-	// TODO: Add virtual changes logic
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
-	defer cancel()
-
-	change, err := tb.thread.GetChange(ctx, id)
-	if err != nil {
-		return nil, err
-	}
-
-	aclChange := new(pb.ACLChange)
-
-	// TODO: think what should we do with such cases, because this can be used by attacker to break our tree
-	if err = proto.Unmarshal(change.Payload, aclChange); err != nil {
-		return
-	}
-	var verified bool
-	verified, err = tb.verify(aclChange.Identity, change.Payload, change.Signature)
-	if err != nil {
-		return
-	}
-	if !verified {
-		err = fmt.Errorf("the signature of the payload cannot be verified")
-		return
-	}
-
-	ch, err = NewChange(id, aclChange)
-	tb.cache[id] = ch
-
-	return ch, nil
-}
-
-func (tb *TreeBuilder) verify(identity string, payload, signature []byte) (isVerified bool, err error) {
-	identityKey, exists := tb.identityKeys[identity]
-	if !exists {
-		identityKey, err = tb.signingPubKeyDecoder.DecodeFromString(identity)
-		if err != nil {
-			return
-		}
-		tb.identityKeys[identity] = identityKey
-	}
-	return identityKey.Verify(payload, signature)
 }
 
 func (tb *TreeBuilder) Build(fromStart bool) (*Tree, error) {