Merge pull request #1658 from lightninglabs/more-group-key-refactors

[group key addrs 6/7]: bug fixes and refactors

Author: guggero

address/book.go

@@ -138,6 +138,11 @@ type Storage interface {
 	SetAddrManaged(ctx context.Context, addr *AddrWithKeyInfo,
 		managedFrom time.Time) error
 
+	// LastEventHeightByVersion returns the last event height for a given
+	// address version.
+	LastEventHeightByVersion(ctx context.Context,
+		version Version) (uint32, error)
+
 	// InsertInternalKey inserts an internal key into the database to make
 	// sure it is identified as a local key later on when importing proofs.
 	// The key can be an internal key for an asset script key or the
@@ -701,6 +706,14 @@ func (b *Book) SetAddrManaged(ctx context.Context, addr *AddrWithKeyInfo,
 	return b.cfg.Store.SetAddrManaged(ctx, addr, managedFrom)
 }
 
+// LastEventHeightByVersion returns the last event height for a given address
+// version.
+func (b *Book) LastEventHeightByVersion(ctx context.Context,
+	version Version) (uint32, error) {
+
+	return b.cfg.Store.LastEventHeightByVersion(ctx, version)
+}
+
 // GetOrCreateEvent creates a new address event for the given status, address
 // and transaction. If an event for that address and transaction already exists,
 // then the status and transaction information is updated instead.

address/book_test.go

@@ -146,8 +146,8 @@ func TestBook_QueryAssetInfo(t *testing.T) {
 		expectsError: false,
 	}}
 
-	for _, test := range tests {
-		t.Run(test.name, func(t *testing.T) {
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
 			mockStorage := &MockStorage{}
 			mockSyncer := &MockAssetSyncer{}
 			book := NewBook(BookConfig{
@@ -157,9 +157,9 @@ func TestBook_QueryAssetInfo(t *testing.T) {
 				Chain:        ChainParams{},
 				StoreTimeout: time.Second,
 			})
-			test.setupMocks(mockStorage, mockSyncer)
-			_, err := book.QueryAssetInfo(ctx, test.specifier)
-			if test.expectsError {
+			tc.setupMocks(mockStorage, mockSyncer)
+			_, err := book.QueryAssetInfo(ctx, tc.specifier)
+			if tc.expectsError {
 				require.Error(t, err)
 			} else {
 				require.NoError(t, err)
@@ -182,6 +182,13 @@ func (m *MockStorage) FetchAllAssetMeta(
 	return args.Get(0).(map[asset.ID]*proof.MetaReveal), args.Error(1)
 }
 
+func (m *MockStorage) LastEventHeightByVersion(ctx context.Context,
+	version Version) (uint32, error) {
+
+	args := m.Called(ctx, version)
+	return args.Get(0).(uint32), args.Error(1)
+}
+
 func (m *MockStorage) AddrByScriptKeyAndVersion(ctx context.Context,
 	key *btcec.PublicKey, version Version) (*AddrWithKeyInfo, error) {
 

authmailbox/client_test.go

@@ -33,7 +33,6 @@ type serverHarness struct {
 	clientCfg    *ClientConfig
 	mockSigner   *test.MockSigner
 	mockMsgStore *MockMsgStore
-	mockTxStore  proof.TxProofStore
 	srv          *Server
 	grpcServer   *grpc.Server
 	cleanup      func()
@@ -45,7 +44,6 @@ func newServerHarness(t *testing.T) *serverHarness {
 	signer.Signature = test.RandBytes(64)
 
 	inMemMsgStore := NewMockStore()
-	inMemTxStore := proof.NewMockTxProofStore()
 
 	nextPort := port.NextAvailablePort()
 	listenAddr := fmt.Sprintf(test.ListenAddrTemplate, nextPort)
@@ -56,7 +54,6 @@ func newServerHarness(t *testing.T) *serverHarness {
 		HeaderVerifier: proof.MockHeaderVerifier,
 		MerkleVerifier: proof.DefaultMerkleVerifier,
 		MsgStore:       inMemMsgStore,
-		TxProofStore:   inMemTxStore,
 	}
 	h := &serverHarness{
 		listenAddr: listenAddr,
@@ -70,7 +67,6 @@ func newServerHarness(t *testing.T) *serverHarness {
 		},
 		mockSigner:   signer,
 		mockMsgStore: inMemMsgStore,
-		mockTxStore:  inMemTxStore,
 	}
 	h.start(t)
 
@@ -217,6 +213,15 @@ func (h *clientHarness) stop(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func randProof(t *testing.T) proof.TxProof {
+	t.Helper()
+
+	randOp := test.RandOp(t)
+	return proof.TxProof{
+		ClaimedOutPoint: randOp,
+	}
+}
+
 // TestServerClientAuthAndRestart tests the server and client authentication
 // process, and that the client can re-connect to the server after it has
 // restarted. It also tests that the client can receive messages from the
@@ -226,7 +231,6 @@ func TestServerClientAuthAndRestart(t *testing.T) {
 	harness := newServerHarness(t)
 	clientCfg := harness.clientCfg
 
-	randOp := test.RandOp(t)
 	clientKey1, _ := test.RandKeyDesc(t)
 	clientKey2, _ := test.RandKeyDesc(t)
 	filter := MessageFilter{}
@@ -279,7 +283,7 @@ func TestServerClientAuthAndRestart(t *testing.T) {
 	}
 
 	// We also store the message in the store, so we can retrieve it later.
-	_, err = harness.mockMsgStore.StoreMessage(ctx, randOp, msg1)
+	_, err = harness.mockMsgStore.StoreMessage(ctx, randProof(t), msg1)
 	require.NoError(t, err)
 
 	harness.srv.publishMessage(msg1)
@@ -310,7 +314,7 @@ func TestServerClientAuthAndRestart(t *testing.T) {
 		ReceiverKey:      *clientKey1.PubKey,
 		ArrivalTimestamp: time.Now(),
 	}
-	_, err = harness.mockMsgStore.StoreMessage(ctx, randOp, msg2)
+	_, err = harness.mockMsgStore.StoreMessage(ctx, randProof(t), msg2)
 	require.NoError(t, err)
 
 	harness.srv.publishMessage(msg2)
@@ -379,6 +383,7 @@ func TestSendMessage(t *testing.T) {
 
 	clientKey1, _ := test.RandKeyDesc(t)
 	clientKey2, _ := test.RandKeyDesc(t)
+	clientKey3, _ := test.RandKeyDesc(t)
 
 	proofWithHeight := func(p proof.TxProof, h uint32) proof.TxProof {
 		p.BlockHeight = h
@@ -391,7 +396,7 @@ func TestSendMessage(t *testing.T) {
 	testCases := []struct {
 		name         string
 		txProofs     []proof.TxProof
-		recvKey      keychain.KeyDescriptor
+		recvKeys     []keychain.KeyDescriptor
 		sendKey      keychain.KeyDescriptor
 		msgs         [][]byte
 		expiryHeight uint32
@@ -400,15 +405,15 @@ func TestSendMessage(t *testing.T) {
 		{
 			name:         "empty payload",
 			txProofs:     []proof.TxProof{*txProof1},
-			recvKey:      clientKey2,
+			recvKeys:     []keychain.KeyDescriptor{clientKey2},
 			sendKey:      clientKey1,
 			msgs:         [][]byte{nil},
 			expectedErrs: []string{"empty payload"},
 		},
 		{
 			name:     "long payload",
 			txProofs: []proof.TxProof{*txProof1},
-			recvKey:  clientKey2,
+			recvKeys: []keychain.KeyDescriptor{clientKey2},
 			sendKey:  clientKey1,
 			msgs: [][]byte{
 				bytes.Repeat([]byte("foo"), MsgMaxSize),
@@ -418,7 +423,7 @@ func TestSendMessage(t *testing.T) {
 		{
 			name:         "missing expiry height",
 			txProofs:     []proof.TxProof{*txProof1},
-			recvKey:      clientKey2,
+			recvKeys:     []keychain.KeyDescriptor{clientKey2},
 			sendKey:      clientKey1,
 			msgs:         [][]byte{[]byte("yoooo")},
 			expectedErrs: []string{"missing expiry block height"},
@@ -428,7 +433,7 @@ func TestSendMessage(t *testing.T) {
 			txProofs: []proof.TxProof{
 				proofWithHeight(*txProof1, 100002),
 			},
-			recvKey:      clientKey2,
+			recvKeys:     []keychain.KeyDescriptor{clientKey2},
 			sendKey:      clientKey1,
 			msgs:         [][]byte{[]byte("yoooo")},
 			expiryHeight: 123,
@@ -442,7 +447,7 @@ func TestSendMessage(t *testing.T) {
 			txProofs: []proof.TxProof{
 				proofWithHeight(*txProof1, 100002),
 			},
-			recvKey:      clientKey2,
+			recvKeys:     []keychain.KeyDescriptor{clientKey2},
 			sendKey:      clientKey1,
 			msgs:         [][]byte{[]byte("yoooo")},
 			expiryHeight: 100002 + 123,
@@ -453,7 +458,10 @@ func TestSendMessage(t *testing.T) {
 				proofWithHeight(*txProof1, 100002),
 				proofWithHeight(*txProof1, 100002),
 			},
-			recvKey: clientKey2,
+			recvKeys: []keychain.KeyDescriptor{
+				clientKey2,
+				clientKey3,
+			},
 			sendKey: clientKey1,
 			msgs: [][]byte{
 				[]byte("yoooo"),
@@ -472,7 +480,11 @@ func TestSendMessage(t *testing.T) {
 				proofWithHeight(*txProof2, 100002),
 				proofWithHeight(*txProof3, 100002),
 			},
-			recvKey: clientKey2,
+			recvKeys: []keychain.KeyDescriptor{
+				clientKey2,
+				clientKey2,
+				clientKey2,
+			},
 			sendKey: clientKey1,
 			msgs: [][]byte{
 				[]byte("yoooo"),
@@ -507,9 +519,10 @@ func TestSendMessage(t *testing.T) {
 			for idx := range tc.msgs {
 				msg := tc.msgs[idx]
 				txProof := tc.txProofs[idx]
+				recvKey := tc.recvKeys[idx]
 
 				msgID, err := client1.client.SendMessage(
-					ctx, *tc.recvKey.PubKey, msg, txProof,
+					ctx, *recvKey.PubKey, msg, txProof,
 					tc.expiryHeight,
 				)
 
@@ -528,6 +541,18 @@ func TestSendMessage(t *testing.T) {
 				// We should be able to read the message if
 				// there was no error sending it.
 				client2.readMessages(t, msgID)
+
+				// Sending the same message again should result
+				// in the same message ID, but should not cause
+				// another message to be sent to any recipients.
+				msgIDReSend, err := client1.client.SendMessage(
+					ctx, *recvKey.PubKey, msg, txProof,
+					tc.expiryHeight,
+				)
+				require.NoError(t, err)
+
+				require.Equal(t, msgID, msgIDReSend)
+				client2.expectNoMessage(t)
 			}
 		})
 	}
@@ -569,7 +594,6 @@ func TestReceiveBacklog(t *testing.T) {
 	harness := newServerHarness(t)
 
 	ctx := context.Background()
-	randOp := test.RandOp(t)
 	receiver1, _ := test.RandKeyDesc(t)
 	receiver2, _ := test.RandKeyDesc(t)
 	receiver3, _ := test.RandKeyDesc(t)
@@ -611,7 +635,9 @@ func TestReceiveBacklog(t *testing.T) {
 	}
 
 	for _, msg := range messages {
-		_, err := harness.mockMsgStore.StoreMessage(ctx, randOp, msg)
+		_, err := harness.mockMsgStore.StoreMessage(
+			ctx, randProof(t), msg,
+		)
 		require.NoError(t, err)
 	}
 

authmailbox/message.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/lightninglabs/taproot-assets/proof"
 )
 
 const (
@@ -19,6 +20,10 @@ var (
 	// allowed length.
 	ErrMessageTooLong = fmt.Errorf("message too long, max %d bytes",
 		MsgMaxSize)
+
+	// ErrMessageNotFound is returned when a message with the given ID or
+	// outpoint cannot be found in the mailbox.
+	ErrMessageNotFound = fmt.Errorf("message not found")
 )
 
 // Message represents a message in the mailbox.
@@ -90,12 +95,19 @@ type MsgStore interface {
 	// outpoint of the transaction that was used to prove the message's
 	// authenticity. If a message with the same outpoint already exists,
 	// it returns proof.ErrTxMerkleProofExists.
-	StoreMessage(ctx context.Context, claimedOp wire.OutPoint,
+	StoreMessage(ctx context.Context, proof proof.TxProof,
 		msg *Message) (uint64, error)
 
 	// FetchMessage retrieves a message from the mailbox by its ID.
 	FetchMessage(ctx context.Context, id uint64) (*Message, error)
 
+	// FetchMessageByOutPoint retrieves a message from the mailbox by its
+	// claimed outpoint of the TX proof that was used to send it. If no
+	// message with the given outpoint exists, it returns
+	// ErrMessageNotFound.
+	FetchMessageByOutPoint(ctx context.Context,
+		claimedOp wire.OutPoint) (*Message, error)
+
 	// QueryMessages retrieves messages based on a query.
 	QueryMessages(ctx context.Context, filter MessageFilter) ([]*Message,
 		error)

authmailbox/mock.go

@@ -7,30 +7,42 @@ import (
 	"sync/atomic"
 
 	"github.com/btcsuite/btcd/wire"
+	"github.com/lightninglabs/taproot-assets/proof"
 )
 
 type MockMsgStore struct {
-	messages      map[uint64]*Message
-	nextMessageID atomic.Uint64
-	mu            sync.RWMutex
+	messages          map[uint64]*Message
+	outpointToMessage map[wire.OutPoint]uint64
+	nextMessageID     atomic.Uint64
+	proofs            map[wire.OutPoint]struct{}
+	mu                sync.RWMutex
 }
 
 var _ MsgStore = (*MockMsgStore)(nil)
 
 func NewMockStore() *MockMsgStore {
 	return &MockMsgStore{
-		messages: make(map[uint64]*Message),
+		messages:          make(map[uint64]*Message),
+		outpointToMessage: make(map[wire.OutPoint]uint64),
+		proofs:            make(map[wire.OutPoint]struct{}),
 	}
 }
 
-func (s *MockMsgStore) StoreMessage(_ context.Context, _ wire.OutPoint,
+func (s *MockMsgStore) StoreMessage(_ context.Context, txProof proof.TxProof,
 	msg *Message) (uint64, error) {
 
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	if _, exists := s.proofs[txProof.ClaimedOutPoint]; exists {
+		return 0, proof.ErrTxMerkleProofExists
+	}
+
+	s.proofs[txProof.ClaimedOutPoint] = struct{}{}
+
 	id := s.nextMessageID.Add(1)
 	s.messages[id] = msg
+	s.outpointToMessage[txProof.ClaimedOutPoint] = id
 
 	return id, nil
 }
@@ -49,6 +61,20 @@ func (s *MockMsgStore) FetchMessage(_ context.Context,
 	return msg, nil
 }
 
+func (s *MockMsgStore) FetchMessageByOutPoint(ctx context.Context,
+	claimedOp wire.OutPoint) (*Message, error) {
+
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	msgID, exists := s.outpointToMessage[claimedOp]
+	if !exists {
+		return nil, ErrMessageNotFound
+	}
+
+	return s.FetchMessage(ctx, msgID)
+}
+
 func (s *MockMsgStore) QueryMessages(_ context.Context,
 	filter MessageFilter) ([]*Message, error) {