rpc: add marshal/unmarshal functions for ProofType RPC type

lightninglabs · Oct 3, 2023 · 7b7b1e5 · 7b7b1e5
1 parent 953d165
commit 7b7b1e5
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 11 deletions.
diff --git a/rpcserver.go b/rpcserver.go
@@ -2447,7 +2447,27 @@ func (r *rpcServer) FetchAssetMeta(ctx context.Context,
 	}, nil
 }
 
-func marshalUniID(id universe.Identifier) *unirpc.ID {
+// MarshalUniProofType marshals the universe proof type into the RPC
+// counterpart.
+func MarshalUniProofType(
+	proofType universe.ProofType) (unirpc.ProofType, error) {
+
+	switch proofType {
+	case universe.ProofTypeNone:
+		return unirpc.ProofType_NONE, nil
+	case universe.ProofTypeIssuance:
+		return unirpc.ProofType_ISSUANCE, nil
+	case universe.ProofTypeTransfer:
+		return unirpc.ProofType_TRANSFER, nil
+
+	default:
+		return unirpc.ProofType_NONE, fmt.Errorf("unknown universe "+
+			"proof type: %v", proofType)
+	}
+}
+
+// MarshalUniID marshals the universe ID into the RPC counterpart.
+func MarshalUniID(id universe.Identifier) (*unirpc.ID, error) {
 	var uniID unirpc.ID
 
 	if id.GroupKey != nil {
@@ -2460,7 +2480,13 @@ func marshalUniID(id universe.Identifier) *unirpc.ID {
 		}
 	}
 
-	return &uniID
+	proofTypeRpc, err := MarshalUniProofType(id.ProofType)
+	if err != nil {
+		return nil, fmt.Errorf("unable to marshal proof type: %w", err)
+	}
+	uniID.ProofType = proofTypeRpc
+
+	return &uniID, nil
 }
 
 // marshalMssmtNode marshals a MS-SMT node into the RPC counterpart.
@@ -2486,8 +2512,13 @@ func marshalUniverseRoot(node universe.BaseRoot) (*unirpc.UniverseRoot, error) {
 		rpcGroupedAssets[assetID.String()] = amount
 	}
 
+	uniID, err := MarshalUniID(node.ID)
+	if err != nil {
+		return nil, err
+	}
+
 	return &unirpc.UniverseRoot{
-		Id:               marshalUniID(node.ID),
+		Id:               uniID,
 		MssmtRoot:        mssmtRoot,
 		AssetName:        node.AssetName,
 		AmountsByAssetId: rpcGroupedAssets,
@@ -2523,15 +2554,41 @@ func (r *rpcServer) AssetRoots(ctx context.Context,
 	return resp, nil
 }
 
+func UnmarshalUniProofType(rpcType unirpc.ProofType) (universe.ProofType,
+	error) {
+
+	switch rpcType {
+	case unirpc.ProofType_NONE:
+		return universe.ProofTypeNone, nil
+
+	case unirpc.ProofType_ISSUANCE:
+		return universe.ProofTypeIssuance, nil
+
+	case unirpc.ProofType_TRANSFER:
+		return universe.ProofTypeTransfer, nil
+
+	default:
+		return 0, fmt.Errorf("unknown universe proof type: %v", rpcType)
+	}
+}
+
 // unmarshalUniID parses the RPC universe ID into the native counterpart.
 func unmarshalUniID(rpcID *unirpc.ID) (universe.Identifier, error) {
+	// Unmarshal the proof type.
+	proofType, err := UnmarshalUniProofType(rpcID.ProofType)
+	if err != nil {
+		return universe.Identifier{}, fmt.Errorf("unable to unmarshal "+
+			"proof type: %w", err)
+	}
+
 	switch {
 	case rpcID.GetAssetId() != nil:
 		var assetID asset.ID
 		copy(assetID[:], rpcID.GetAssetId())
 
 		return universe.Identifier{
-			AssetID: assetID,
+			AssetID:   assetID,
+			ProofType: proofType,
 		}, nil
 
 	case rpcID.GetAssetIdStr() != "":
@@ -2546,7 +2603,8 @@ func unmarshalUniID(rpcID *unirpc.ID) (universe.Identifier, error) {
 		copy(assetID[:], assetIDBytes)
 
 		return universe.Identifier{
-			AssetID: assetID,
+			AssetID:   assetID,
+			ProofType: proofType,
 		}, nil
 
 	case rpcID.GetGroupKey() != nil:
@@ -2556,7 +2614,8 @@ func unmarshalUniID(rpcID *unirpc.ID) (universe.Identifier, error) {
 		}
 
 		return universe.Identifier{
-			GroupKey: groupKey,
+			GroupKey:  groupKey,
+			ProofType: proofType,
 		}, nil
 
 	case rpcID.GetGroupKeyStr() != "":
@@ -2573,7 +2632,8 @@ func unmarshalUniID(rpcID *unirpc.ID) (universe.Identifier, error) {
 		}
 
 		return universe.Identifier{
-			GroupKey: groupKey,
+			GroupKey:  groupKey,
+			ProofType: proofType,
 		}, nil
 
 	default:

diff --git a/universe_rpc_diff.go b/universe_rpc_diff.go
@@ -95,8 +95,12 @@ func (r *RpcUniverseDiff) RootNodes(
 func (r *RpcUniverseDiff) RootNode(ctx context.Context,
 	id universe.Identifier) (universe.BaseRoot, error) {
 
+	uniID, err := MarshalUniID(id)
+	if err != nil {
+		return universe.BaseRoot{}, err
+	}
 	rootReq := &universerpc.AssetRootQuery{
-		Id: marshalUniID(id),
+		Id: uniID,
 	}
 
 	universeRoot, err := r.conn.QueryAssetRoots(ctx, rootReq)
@@ -111,7 +115,11 @@ func (r *RpcUniverseDiff) RootNode(ctx context.Context,
 func (r *RpcUniverseDiff) UniverseLeafKeys(ctx context.Context,
 	id universe.Identifier) ([]universe.LeafKey, error) {
 
-	assetKeys, err := r.conn.AssetLeafKeys(ctx, marshalUniID(id))
+	uniID, err := MarshalUniID(id)
+	if err != nil {
+		return nil, err
+	}
+	assetKeys, err := r.conn.AssetLeafKeys(ctx, uniID)
 	if err != nil {
 		return nil, err
 	}
@@ -140,8 +148,13 @@ func (r *RpcUniverseDiff) FetchIssuanceProof(ctx context.Context,
 	id universe.Identifier,
 	key universe.LeafKey) ([]*universe.Proof, error) {
 
+	uniID, err := MarshalUniID(id)
+	if err != nil {
+		return nil, err
+	}
+
 	uProofs, err := r.conn.QueryProof(ctx, &universerpc.UniverseKey{
-		Id:      marshalUniID(id),
+		Id:      uniID,
 		LeafKey: marshalLeafKey(key),
 	})
 	if err != nil {

diff --git a/universe_rpc_registrar.go b/universe_rpc_registrar.go
@@ -84,9 +84,14 @@ func (r *RpcUniverseRegistrar) RegisterIssuance(ctx context.Context,
 	id universe.Identifier, key universe.LeafKey,
 	leaf *universe.Leaf) (*universe.Proof, error) {
 
+	uniID, err := MarshalUniID(id)
+	if err != nil {
+		return nil, err
+	}
+
 	// First, we'll parse the proofs and key into their RPC counterparts.
 	uniKey := &unirpc.UniverseKey{
-		Id:      marshalUniID(id),
+		Id:      uniID,
 		LeafKey: marshalLeafKey(key),
 	}