NDSS检测重入攻击的插件代码

TheDao_plugin.go

package main

import (
	"../../pluginlog"
	"fmt"
	"github.com/ethereum/collector"
	"github.com/json-iterator/go"
	"math/big"
	"strings"
	// 调试用到的库
	"strconv"
	"os"
)

var logger pluginlog.ErrTxLog
var json = jsoniter.ConfigCompatibleWithStandardLibrary
// 待调试的所有交易，会把一下交易的所有信息写到日志，调试用
var debugtxs = []string {
	"0x1928404bf10ec33cf5ca887f1fbf6e4c611183ba9292ce54a9406bc8f5838128",
	"0x1722c4afe7cdafc834177369789da9fab983374cf64911d3176cbb775213c460",
	"0x88bac3361684ef55be82e7e4ab34ff9396af762d59ed7441d1c14ac039e2379c",
}

type RunOpCode struct {
	MethodName string   `json:"methodname"`
	OpCode     []string `json:"option"`
}

type DaoInfo struct {
	BlockNumber      string              `json:"blocknumber"`
	TxHash           string              `json:"txhash"`
	GasUsed          uint64              `json:"gasused"`
	Cycle            string              `json:"cycle"`
	Victim           string              `json:"victim"`
	TotalCycleCount  uint64              `json:"totalcyclecount"`
	TotalValueCount  string              `json:"totalvaluecount"`
	InternalLog      string              `json:"internallog"`
}

type Node struct {
	address     string  `json:"address"`
	invalue     big.Int `json:"invalue"`
	outvalue    big.Int `json:"outvalue"`
	parent      *Node   `json:"parent`
	children    []*Node `json:"children"`
	ancestors   []*Node `json:"ancestors"`
}

// 全局变量
var (
	root Node // 交易调用树的根节点
	cur_p *Node // 指向当前所在的合约节点指针
	txhash string // 交易hash
	gasused uint64 // 交易的gas
	blocknumber string // 交易的块高

	collectors []*collector.CollectorDataT // 当前交易产生的所有collector数据
)

// 写日志文件函数封装，可能有多个日志要记
func Log(logger *pluginlog.ErrTxLog, message string) {
	// 准备写日志
	logger.CheckIfCreateNewFile()
	logger.OpenFile()
	// 写日志
	logger.WriteLog(message + "\n")
	logger.CloseFile()
}

// 插件入口函数：梦开始的地方
func Run() []byte {
	var data = RunOpCode{
		MethodName: "TheDao",
		OpCode:     []string {
			"EXTERNALINFOSTART", "EXTERNALINFOEND", "CALLSTART",
			"CALLEND", "CALLCODESTART", "CALLCODEEND",
		},
	}

	logger.InitialFileLog("./plugin_log/thedaolog/thedaolog")
	b, err := json.Marshal(&data)
	if err != nil {
		return []byte{}
	}

	return b
}

// 数据处理函数：做梦的地方
func Recv(m *collector.CollectorDataT) (byte, string) {
	switch m.Option {
	// 该OpCode 是外部调用开始	
	case "EXTERNALINFOSTART":
		txhash = m.ExternalInfo.TxHash
		blocknumber = m.ExternalInfo.BlockNumber
		collectors = []*collector.CollectorDataT{m}
		// StopSync(3694500) // 3694500处停止同步
		StopSync(4000000) // 4000000处停止同步
		var val big.Int
		val.SetString(m.ExternalInfo.Value, 10)
		// 保持代码有效性，删除head节点
		root = Node {
			// 统一小写
			address:     strings.ToLower(m.ExternalInfo.To),
			invalue:     val,
			outvalue:    *big.NewInt(int64(0)),
			parent:      nil,
			children:    nil,
			ancestors:   nil,
		}
		cur_p = &root

	// 该OpCode 是内部调用开始
	case "CALLSTART", "CALLCODESTART":
		collectors = append(collectors, m)
		var val big.Int
		val.SetString(m.Command.Value, 10)
		node := Node {
			// 统一小写
			address:     strings.ToLower(m.Command.To),
			invalue:     val,
			outvalue:    *big.NewInt(int64(0)),
			parent:      cur_p,
			children:    nil,
			ancestors:   nil,
		}
		node.ancestors = append([]*Node{cur_p}, cur_p.ancestors...)
		cur_p.children = append(cur_p.children, &node)
		cur_p.outvalue.Add(&cur_p.outvalue, &val)
		cur_p = &node

	// 该OpCode 是内部调用结束
	case "CALLEND", "CALLCODEEND":
		collectors = append(collectors, m)
		cur_p = cur_p.parent
		// 如果当前调用失败，在树中删除相应节点
		if !m.Command.IsInternalSucceeded || !m.Command.IsCallValid {
			// 记录失败的调用，调试用
			Log(&logger, fmt.Sprintf(
				"[FAILINFO]: @%s\t%s--%s failed!", txhash, cur_p.address,
				cur_p.children[len(cur_p.children)-1].address,
			))
			cur_p.children = cur_p.children[:len(cur_p.children)-1]
		}

	// 该OpCode 是外部调用结束
	case "EXTERNALINFOEND":
		collectors = append(collectors, m)
		// 使用的gas
		gasused = m.ExternalInfo.GasUsed
		// 如果该交易失败
		if !m.ExternalInfo.IsSuccess {
			// 记录失败的交易，调试用
			// Log(&logger, "[FAILINFO]: @" + txhash + " failed!")
			return 0x00, ""
		}
		// 处理环 && 调用关系
		ProcCycleInfo()
		// 处理本次交易的全部信息，调试用
		ProcTxInfo()
	}
	return 0x00, ""
}

// 判断有没有环，顺便把所有的调用记录下来
func ProcCycleInfo() bool {
	cur_p = &root // 换成EOA调用的第一个地址
	nodes := []*Node{cur_p}
	// 该交易下总的情况（环，总的交易额，总的环个数，总的环
	cycles := make(map[string]bool)
	totalValueCount := *big.NewInt(int64(0))
	var totalCycleCount uint64
	totalCycleCount = 0
	totalCycleStr := "["
	totalCallStr := "["
	victim := ""
	// 通过迭代的方式后序遍历树结构
	for len(nodes) > 0 {
		cur_p = nodes[len(nodes)-1]
		nodes = nodes[:len(nodes)-1]
		// 遍历的时候记录一下调用关系（没有前后关系）
		if cur_p.parent != nil {
			totalCallStr += "{\"addr\":\"" + cur_p.parent.address + "--" + cur_p.address +
				"\",\"value\":" + fmt.Sprintf("%v", cur_p.invalue.String()) + "},"
		}
		// 寻找环啦
		if cur_p.ancestors != nil && cur_p.children != nil {
			for _, child := range cur_p.children {
				// 自己调用自己不算，新增一个，要是没向外转钱（自己和儿子都没转钱）也不算
				if child.address == cur_p.address || (child.outvalue.Int64() == 0 && !IsChildrenOutValue(child)) {
					continue
				}
				// 看看之前有没有调用过当前节点调用过的节点，
				// 若有，则形成了一个环A--...--B(cur_p)--A
				for _, ancestor := range cur_p.ancestors {
					// 找到一个环了（可能有多个），需要继续在这下面找一个一样的
					if ancestor.address == child.address {
						cycleArr := []*Node{child}
						tem_p := cur_p
						// 找到调用的路线
						for tem_p != ancestor {
							cycleArr = append(cycleArr, tem_p)
							tem_p = tem_p.parent
						}
						cycleArr = append(cycleArr, ancestor)
						// 将环的调用字符串存储到map 中
						cycleStr := ""
						for _, call := range cycleArr {
							if cycleStr == "" {
								cycleStr = call.address
							} else {
								cycleStr = call.address + "--" + cycleStr
							}
						}
						// 在这下面找还有没有一样的环，直接搜索吧…… 短路逻辑先判断是否含有第二个环，有就不用继续找了
						if IsPathExist(cycleArr, child) || IsContinuedCall(cycleArr[len(cycleArr)-2].address, child) {
							var value *big.Int
							victim, value = ProcCycleVictims(cycleArr)
							totalValueCount.Add(&totalValueCount, value)
							cycles[cycleStr] = true
							break // 找到直接退出
						}
					}
				}
			}
		}
		// 孩子节点进入栈中
		for _, child := range cur_p.children {
			nodes = append(nodes, child)
		}
	}
	// 判断有没有环数大于等于2的环，事实上我每次环数大于等于2时才入字典
	// 因此替换成是否存在cycles[string]bool
	for cycle, _ := range cycles {
		// 此时认为是一个重入攻击
		totalCycleCount += 1
		totalCycleStr += "{\"detail\":\"" + cycle + "\"},"
	}
	totalCycleStr += "]"
	totalCallStr += "]"

	// 没有环
	if totalCycleCount == 0 {
		return false
	}

	var daoInfo = DaoInfo {
		BlockNumber:      blocknumber,
		TxHash:           txhash,
		GasUsed:          gasused,
		Cycle:            totalCycleStr,
		Victim:           victim,
		TotalCycleCount:  totalCycleCount,
		TotalValueCount:  totalValueCount.String(),
		InternalLog:      "", // 先不写，太多了 占空间
		// InternalLog:      totalCallStr,
	}
	daoJsonData, err := json.Marshal(daoInfo)
	if err != nil {
		panic(err)
	}
	// 写日志
	Log(&logger, string(daoJsonData))
	return true
}

// 判断路径是否存在（找有没有第二个相同的环），从后往前搜索，因为我的cycleArr是倒着的
func IsPathExist(cycleArr []*Node, node *Node) bool {
	nodes := []*Node{node}
	res := false
	for len(nodes) > 0 {
		tem_p := nodes[len(nodes)-1]
		nodes = nodes[:len(nodes)-1]
		res = res || IsPathExist_(cycleArr, tem_p, len(cycleArr)-1)
		for _, child := range tem_p.children {
			nodes = append(nodes, child)
		}
	}
	return res
}

// 递归子函数，新增一个，第二个环也要转了钱钱才算
func IsPathExist_(cycleArr []*Node, node *Node, index int) bool {
	if node.address != cycleArr[index].address {
		return false
	} else if index == 0 {
		// 向外转了钱钱才算，或者它调用的合约转了钱钱
		return node.outvalue.Int64() != 0 || IsChildrenOutValue(node)
	}
	res := false
	for _, child := range node.children {
		res = res || IsPathExist_(cycleArr, child, index-1)
	}
	return res
}

// 孩子有没有向外转过钱钱
func IsChildrenOutValue(node *Node) bool {
	for _, child := range node.children {
		if child.outvalue.Int64() != 0 {
			return true
		}
	}
	return false
}

// 有没有继续调用某个地址的合约
func IsContinuedCall(address string, node *Node) bool {
	for _, child := range node.children {
		if child.address == address {
			return true
		}
	}
	return false
}

// 找出一个环中的受害者
func ProcCycleVictims(cycleArr []*Node) (string, *big.Int) {
	victims := make(map[string]*big.Int)
	victim := ""
	value := big.NewInt(int64(0))
	for _, node := range cycleArr {
		// 转出的钱钱比收入的多，它就是受害者了
		if node.outvalue.Cmp(&(node.invalue)) == 1 {
			if value, ok := victims[node.address]; ok {
				value.Add(value, &(node.outvalue))
			} else {
				victims[node.address] = big.NewInt(int64(0))
				victims[node.address].Add(victims[node.address], &(node.outvalue))
			}
		}
	}
	if len(victims) == 0 {
		victims[cycleArr[0].address] = big.NewInt(int64(0))
		victims[cycleArr[0].address].Add(victims[cycleArr[0].address], &(cycleArr[0].outvalue))
	}
	// 只要被偷钱最多的
	for k, v := range victims {
		if value.Cmp(v) != 1 {
			victim, value = k, v
		}
	}
	// 虽然这样很难看，但事实上0xd2这个合约和0xbb是一样的
	if victim == "0xd2e16a20dd7b1ae54fb0312209784478d069c7b0" {
		victim = "0xbb9bc244d798123fde783fcc1c72d3bb8c189413"
	}
	return victim, value
}

// 处理交易的全部信息，调试用
func ProcTxInfo() {
	for _, tx := range debugtxs {
		if txhash == tx {
			Log(&logger, "[DEBUG] @" + txhash + ":\n")
			LogCollectorsJson() // 先把这条交易的所有collectors的json打印出来
			Log(&logger, GetTree(&root)) // 打印出树来
		}
	}
}

// 打印出树，调试用
func GetTree(node *Node) string {
	if node.children == nil {
		var ancestorsAddr []string
		for _, ancestor := range node.ancestors {
			ancestorsAddr = append(ancestorsAddr, ancestor.address)
		}
		return fmt.Sprintf("{\"address\": \"%s\", \"ancestors\": \"%s\"}",
			node.address, strings.Join(ancestorsAddr, ", "),
		)
	} else {
		var ancestorsAddr []string
		for _, ancestor := range node.ancestors {
			ancestorsAddr = append(ancestorsAddr, ancestor.address)
		}
		res := fmt.Sprintf("{\"address\": \"%s\", \"ancestors\": \"%s\", \"children\": [",
			node.address, strings.Join(ancestorsAddr, ", "),
		)
		var childJsonArr []string
		for _, child := range node.children {
			childJsonArr = append(childJsonArr, GetTree(child))
		}
		res += strings.Join(childJsonArr, ", ")
		res += "]}"
		return res
	}
}

// 记录一条交易中的所有collector的日志，调试用
func LogCollectorsJson() {
	for _, m := range collectors {
		if mJson, err := json.Marshal(m); err == nil {
			Log(&logger, string(mJson))
		}
	}
}

// 到某个区块高度停止同步，调试用
func StopSync(stopNumber int) {
	if curNumber, err := strconv.Atoi(blocknumber); err == nil {
		if curNumber > stopNumber {
			fmt.Println("stop sync...")
			os.Exit(0)
		}
	}
}