lc_smart_admin/server/service/tcp/tcp.go

package tcp

import (
	"errors"
	"net"
	"server/utils/logger"
	"strings"
	"sync"
	"time"
)

func ListenTcp() {
	var listen net.Listener
	var err error
	// 监听当前的tcp连接
	for {
		listen, err = net.Listen("tcp", "0.0.0.0:9200")
		if err != nil {
			logger.Logger.Errorf("Listen failed, err: %v. Retrying in 5 seconds...", err)
			time.Sleep(5 * time.Second) // 休眠一段时间后重试
			continue
		}
		break // 成功监听后退出循环
	}

	tracker := NewConnectionTracker() //创建连接检测器

	for {
		conn, err := listen.Accept()
		if err != nil {
			logger.Logger.Errorf("Accept failed, err:%v", err)
			continue
		}
		err = CheckConn(conn, tracker)
		if err != nil {
			conn.Close() // 如果是恶意连接，则关闭连接
			continue
		}
	}
}

func CheckConn(conn net.Conn, tracker *ConnectionTracker) error {
	logger.Logger.Debugf("StartDevice addr:%s", conn.RemoteAddr().String())
	arr := strings.Split(conn.RemoteAddr().String(), ":")
	ip := arr[0]
	// 记录连接
	tracker.recordConnection(ip)

	// 检查是否为恶意连接
	if tracker.isMalicious(ip) {
		logger.Logger.Debugf("恶意连接检测到 ip: %s\n", ip)
		return errors.New("connection is Malicious")
	}
	device := Device{}
	device.Start(conn)
	return nil
}

type ConnectionTracker struct {
	mu          sync.Mutex
	connections map[string][]time.Time // 存储每个 IP 的连接时间戳
}

func NewConnectionTracker() *ConnectionTracker {
	return &ConnectionTracker{
		connections: make(map[string][]time.Time),
	}
}

func (ct *ConnectionTracker) recordConnection(ip string) {
	ct.mu.Lock()
	defer ct.mu.Unlock()

	now := time.Now()
	ct.connections[ip] = append(ct.connections[ip], now)

	// 清理过期的连接记录
	ct.cleanUpExpired(ip, now)
}

func (ct *ConnectionTracker) cleanUpExpired(ip string, now time.Time) {
	threshold := now.Add(-3 * time.Minute)
	if timestamps, exists := ct.connections[ip]; exists {
		var filtered []time.Time
		for _, t := range timestamps {
			if t.After(threshold) { // 检查时间戳是否在三分钟内
				filtered = append(filtered, t) // 如果在范围内，保存到 filtered 列表
			}
		}
		ct.connections[ip] = filtered
	}
}

// 判断是否是恶意连接
func (ct *ConnectionTracker) isMalicious(ip string) bool {
	ct.mu.Lock()
	defer ct.mu.Unlock()

	if timestamps, exists := ct.connections[ip]; exists {
		return len(timestamps) >= 10 // 定义恶意连接的阈值
	}
	return false
}