Merge pull request #566 from tonkeeper/extract_risk

rewrite extract risk from messages
tonkeeper · Dec 11, 2024 · aa7c028 · aa7c028
2 parents 8d750de + bfd41b3
commit aa7c028
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 61 deletions.
diff --git a/pkg/api/converters.go b/pkg/api/converters.go
@@ -268,15 +268,21 @@ func (h *Handler) convertMultisig(ctx context.Context, item core.Multisig) (*oas
 		for _, account := range order.Signers {
 			signers = append(signers, account.ToRaw())
 		}
-		messages, err := convertMultisigActionsToRawMessages(order.Actions)
-		if err != nil {
-			return nil, err
+		risk := walletPkg.Risk{
+			TransferAllRemainingBalance: false,
+			Jettons:                     map[tongo.AccountID]big.Int{},
 		}
-		risk, err := walletPkg.ExtractRiskFromRawMessages(messages)
-		if err != nil {
-			return nil, err
+		for _, action := range order.Actions {
+			switch action.SumType {
+			case "SendMessage":
+				var err error
+				risk, err = walletPkg.ExtractRiskFromMessage(action.SendMessage.Field0.Message, risk, action.SendMessage.Field0.Mode)
+				if err != nil {
+					return nil, err
+				}
+			}
 		}
-		oasRisk, err := h.convertRisk(ctx, *risk, item.AccountID)
+		oasRisk, err := h.convertRisk(ctx, risk, item.AccountID)
 		if err != nil {
 			return nil, err
 		}

diff --git a/pkg/api/multisig_handlers.go b/pkg/api/multisig_handlers.go
@@ -7,31 +7,9 @@ import (
 
 	"github.com/tonkeeper/opentonapi/pkg/core"
 	"github.com/tonkeeper/opentonapi/pkg/oas"
-	"github.com/tonkeeper/tongo/abi"
-	"github.com/tonkeeper/tongo/boc"
-	"github.com/tonkeeper/tongo/tlb"
 	"github.com/tonkeeper/tongo/ton"
-	tongoWallet "github.com/tonkeeper/tongo/wallet"
 )
 
-func convertMultisigActionsToRawMessages(actions []abi.MultisigSendMessageAction) ([]tongoWallet.RawMessage, error) {
-	var messages []tongoWallet.RawMessage
-	for _, action := range actions {
-		switch action.SumType {
-		case "SendMessage":
-			msg := boc.NewCell()
-			if err := tlb.Marshal(msg, action.SendMessage.Field0.Message); err != nil {
-				return nil, err
-			}
-			messages = append(messages, tongoWallet.RawMessage{
-				Message: msg,
-				Mode:    action.SendMessage.Field0.Mode,
-			})
-		}
-	}
-	return messages, nil
-}
-
 func (h *Handler) GetMultisigAccount(ctx context.Context, params oas.GetMultisigAccountParams) (*oas.Multisig, error) {
 	accountID, err := ton.ParseAccountID(params.AccountID)
 	if err != nil {

diff --git a/pkg/wallet/risk.go b/pkg/wallet/risk.go
@@ -45,42 +45,51 @@ func ExtractRiskFromRawMessages(rawMessages []tongoWallet.RawMessage) (*Risk, er
 		if err := tlb.Unmarshal(rawMsg.Message, &m); err != nil {
 			return nil, err
 		}
-		tonValue := uint64(m.MessageInternal.Value.Grams)
-		destination, err := ton.AccountIDFromTlb(m.MessageInternal.Dest)
+		var err error
+		risk, err = ExtractRiskFromMessage(m, risk, rawMsg.Mode)
 		if err != nil {
 			return nil, err
 		}
-		risk.Ton += tonValue
-		msgBody := m.MessageInternal.Body.Value.Value
-		if err != nil {
-			continue
+	}
+	return &risk, nil
+}
+
+func ExtractRiskFromMessage(m abi.MessageRelaxed, risk Risk, mode byte) (Risk, error) {
+	if tongoWallet.IsMessageModeSet(int(mode), tongoWallet.AttachAllRemainingBalance) {
+		risk.TransferAllRemainingBalance = true
+	}
+	tonValue := uint64(m.MessageInternal.Value.Grams)
+	destination, err := ton.AccountIDFromTlb(m.MessageInternal.Dest)
+	if err != nil {
+		return Risk{}, err
+	}
+	risk.Ton += tonValue
+	msgBody := m.MessageInternal.Body.Value.Value
+	switch x := msgBody.(type) {
+	case abi.NftTransferMsgBody:
+		if destination == nil {
+			return risk, err
 		}
-		switch x := msgBody.(type) {
-		case abi.NftTransferMsgBody:
-			if destination == nil {
-				continue
-			}
-			// here, destination is an NFT
-			risk.Nfts = append(risk.Nfts, *destination)
-		case abi.JettonBurnMsgBody:
-			if destination == nil {
-				continue
-			}
-			// here, destination is a jetton wallet
-			amount := big.Int(x.Amount)
-			currentJettons := risk.Jettons[*destination]
-			var total big.Int
-			risk.Jettons[*destination] = *total.Add(&currentJettons, &amount)
-		case abi.JettonTransferMsgBody:
-			if destination == nil {
-				continue
-			}
-			// here, destination is a jetton wallet
-			amount := big.Int(x.Amount)
-			currentJettons := risk.Jettons[*destination]
-			var total big.Int
-			risk.Jettons[*destination] = *total.Add(&currentJettons, &amount)
+		// here, destination is an NFT
+		risk.Nfts = append(risk.Nfts, *destination)
+	case abi.JettonBurnMsgBody:
+		if destination == nil {
+			return risk, err
 		}
+		// here, destination is a jetton wallet
+		amount := big.Int(x.Amount)
+		currentJettons := risk.Jettons[*destination]
+		var total big.Int
+		risk.Jettons[*destination] = *total.Add(&currentJettons, &amount)
+	case abi.JettonTransferMsgBody:
+		if destination == nil {
+			return risk, err
+		}
+		// here, destination is a jetton wallet
+		amount := big.Int(x.Amount)
+		currentJettons := risk.Jettons[*destination]
+		var total big.Int
+		risk.Jettons[*destination] = *total.Add(&currentJettons, &amount)
 	}
-	return &risk, nil
+	return risk, err
 }