ceremonyclient/node/store/data_proof.go
2023-09-27 04:05:39 -05:00

275 lines
6.5 KiB
Go

package store
import (
"encoding/binary"
"github.com/cockroachdb/pebble"
"github.com/pkg/errors"
"go.uber.org/zap"
"golang.org/x/crypto/sha3"
"source.quilibrium.com/quilibrium/monorepo/node/protobufs"
)
type DataProofStore interface {
NewTransaction() (Transaction, error)
GetAggregateProof(
filter []byte,
commitment []byte,
frameNumber uint64,
inclusionReassembler func(typeUrl string, data [][]byte) ([]byte, error),
) (
*protobufs.InclusionAggregateProof,
error,
)
PutAggregateProof(
txn Transaction,
aggregateProof *protobufs.InclusionAggregateProof,
commitment []byte,
inclusionSplitter func(typeUrl string, data []byte) ([][]byte, error),
) error
}
type PebbleDataProofStore struct {
db *pebble.DB
logger *zap.Logger
}
func NewPebbleDataProofStore(
db *pebble.DB,
logger *zap.Logger,
) *PebbleDataProofStore {
return &PebbleDataProofStore{
db,
logger,
}
}
const (
DATA_PROOF = 0x04
DATA_PROOF_METADATA = 0x00
DATA_PROOF_INCLUSION = 0x01
DATA_PROOF_SEGMENT = 0x02
)
func dataProofMetadataKey(filter []byte, commitment []byte) []byte {
key := []byte{DATA_PROOF, DATA_PROOF_METADATA}
key = append(key, commitment...)
key = append(key, filter...)
return key
}
func dataProofInclusionKey(
filter []byte,
commitment []byte,
seqNo uint64,
) []byte {
key := []byte{DATA_PROOF, DATA_PROOF_INCLUSION}
key = append(key, commitment...)
key = binary.BigEndian.AppendUint64(key, seqNo)
key = append(key, filter...)
return key
}
func dataProofSegmentKey(
filter []byte,
hash []byte,
) []byte {
key := []byte{DATA_PROOF, DATA_PROOF_SEGMENT}
key = append(key, hash...)
key = append(key, filter...)
return key
}
func (p *PebbleDataProofStore) NewTransaction() (Transaction, error) {
return &PebbleTransaction{
b: p.db.NewBatch(),
}, nil
}
func internalGetAggregateProof(
db *pebble.DB,
filter []byte,
commitment []byte,
frameNumber uint64,
inclusionReassembler func(typeUrl string, data [][]byte) ([]byte, error),
) (*protobufs.InclusionAggregateProof, error) {
value, closer, err := db.Get(dataProofMetadataKey(filter, commitment))
if err != nil {
if errors.Is(err, pebble.ErrNotFound) {
return nil, ErrNotFound
}
return nil, errors.Wrap(err, "get aggregate proof")
}
defer closer.Close()
copied := make([]byte, len(value[8:]))
limit := binary.BigEndian.Uint64(value[0:8])
copy(copied, value[8:])
aggregate := &protobufs.InclusionAggregateProof{
Filter: filter,
FrameNumber: frameNumber,
InclusionCommitments: []*protobufs.InclusionCommitment{},
Proof: copied,
}
iter := db.NewIter(&pebble.IterOptions{
LowerBound: dataProofInclusionKey(filter, commitment, 0),
UpperBound: dataProofInclusionKey(filter, commitment, limit+1),
})
i := uint32(0)
for iter.First(); iter.Valid(); iter.Next() {
incCommit := iter.Value()
urlLength := binary.BigEndian.Uint16(incCommit[:2])
commitLength := binary.BigEndian.Uint16(incCommit[2:4])
url := make([]byte, urlLength)
copy(url, incCommit[4:urlLength+4])
commit := make([]byte, commitLength)
copy(commit, incCommit[urlLength+4:urlLength+4+commitLength])
remainder := int(urlLength + 4 + commitLength)
inclusionCommitment := &protobufs.InclusionCommitment{
Filter: filter,
FrameNumber: frameNumber,
Position: i,
TypeUrl: string(url),
Commitment: commit,
}
chunks := [][]byte{}
for j := 0; j < (len(incCommit)-remainder)/32; j++ {
start := remainder + (j * 32)
end := remainder + ((j + 1) * 32)
segValue, dataCloser, err := db.Get(
dataProofSegmentKey(filter, incCommit[start:end]),
)
if err != nil {
if errors.Is(err, pebble.ErrNotFound) {
// If we've lost this key it means we're in a corrupted state
return nil, ErrInvalidData
}
return nil, errors.Wrap(err, "get aggregate proof")
}
segCopy := make([]byte, len(segValue))
copy(segCopy, segValue)
chunks = append(chunks, segCopy)
if err = dataCloser.Close(); err != nil {
return nil, errors.Wrap(err, "get aggregate proof")
}
}
inclusionCommitment.Data, err = inclusionReassembler(string(url), chunks)
if err != nil {
return nil, errors.Wrap(err, "get aggregate proof")
}
aggregate.InclusionCommitments = append(
aggregate.InclusionCommitments,
inclusionCommitment,
)
i++
}
if err = iter.Close(); err != nil {
return nil, errors.Wrap(err, "get aggregate proof")
}
return aggregate, nil
}
func (p *PebbleDataProofStore) GetAggregateProof(
filter []byte,
commitment []byte,
frameNumber uint64,
inclusionReassembler func(typeUrl string, data [][]byte) ([]byte, error),
) (*protobufs.InclusionAggregateProof, error) {
return internalGetAggregateProof(
p.db,
filter,
commitment,
frameNumber,
inclusionReassembler,
)
}
func internalPutAggregateProof(
db *pebble.DB,
txn Transaction,
aggregateProof *protobufs.InclusionAggregateProof,
commitment []byte,
inclusionSplitter func(typeUrl string, data []byte) ([][]byte, error),
) error {
buf := binary.BigEndian.AppendUint64(
nil,
uint64(len(aggregateProof.InclusionCommitments)),
)
buf = append(buf, aggregateProof.Proof...)
for i, inc := range aggregateProof.InclusionCommitments {
segments, err := inclusionSplitter(inc.TypeUrl, inc.Data)
if err != nil {
return errors.Wrap(err, "get aggregate proof")
}
urlLength := len(inc.TypeUrl)
commitLength := len(inc.Commitment)
encoded := binary.BigEndian.AppendUint16(nil, uint16(urlLength))
encoded = binary.BigEndian.AppendUint16(encoded, uint16(commitLength))
encoded = append(encoded, []byte(inc.TypeUrl)...)
encoded = append(encoded, inc.Commitment...)
for _, segment := range segments {
hash := sha3.Sum256(segment)
if err = txn.Set(
dataProofSegmentKey(aggregateProof.Filter, hash[:]),
segment,
); err != nil {
return errors.Wrap(err, "put aggregate proof")
}
encoded = append(encoded, hash[:]...)
}
if err = txn.Set(
dataProofInclusionKey(aggregateProof.Filter, commitment, uint64(i)),
encoded,
); err != nil {
return errors.Wrap(err, "put aggregate proof")
}
}
if err := txn.Set(
dataProofMetadataKey(aggregateProof.Filter, commitment),
buf,
); err != nil {
return errors.Wrap(err, "put aggregate proof")
}
return nil
}
func (p *PebbleDataProofStore) PutAggregateProof(
txn Transaction,
aggregateProof *protobufs.InclusionAggregateProof,
commitment []byte,
inclusionSplitter func(typeUrl string, data []byte) ([][]byte, error),
) error {
return internalPutAggregateProof(
p.db,
txn,
aggregateProof,
commitment,
inclusionSplitter,
)
}