ai-security-xdr/haobang-security-dm/syslog-serve/src/main/java/com/haobang/dispatch/DeviceContextResolver.java

package com.haobang.dispatch;

import com.common.entity.DeviceDevice;
import com.common.entity.DeviceUnknown;
import com.common.service.DeviceDeviceService;
import com.common.service.DeviceUnknownService;
import com.haobang.config.AppProperties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.util.List;
import java.time.LocalDateTime;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.LongAdder;

/**
 * 来源 IP 到设备接入上下文的进程内缓存。
 *
 * <p>原实现每条日志至少查询两次设备表；该缓存使正常热路径不访问达梦。缓存失效时同一
 * source_ip 使用条带锁合并并发查询，避免恢复洪峰击穿数据库。</p>
 */
@Component
public class DeviceContextResolver {
    private static final Logger log = LoggerFactory.getLogger(DeviceContextResolver.class);
    private final DeviceDeviceService deviceService;
    private final DeviceUnknownService unknownService;
    private final AppProperties properties;
    private final ConcurrentHashMap<String, CacheEntry> cache = new ConcurrentHashMap<>();
    private final Object[] locks = new Object[64];
    private final LongAdder hits = new LongAdder();
    private final LongAdder misses = new LongAdder();

    public DeviceContextResolver(DeviceDeviceService deviceService, DeviceUnknownService unknownService,
                                 AppProperties properties) {
        this.deviceService = deviceService;
        this.unknownService = unknownService;
        this.properties = properties;
        for (int i = 0; i < locks.length; i++) locks[i] = new Object();
    }

    public Resolution resolve(String sourceIp, String protocol) {
        long now = System.currentTimeMillis();
        CacheEntry cached = cache.get(sourceIp);
        if (cached != null && cached.expiresAt > now) {
            hits.increment();
            return cached.resolution;
        }
        Object lock = locks[(sourceIp.hashCode() & Integer.MAX_VALUE) % locks.length];
        synchronized (lock) {
            cached = cache.get(sourceIp);
            if (cached != null && cached.expiresAt > now) {
                hits.increment();
                return cached.resolution;
            }
            misses.increment();
            List<DeviceDevice> devices = deviceService.getByIpSafely(sourceIp);
            Resolution resolution;
            if (devices == null || devices.isEmpty()) {
                recordUnknown(sourceIp, protocol);
                resolution = Resolution.terminal(Type.UNKNOWN_DEVICE, null,
                        "来源 IP 未注册: " + sourceIp);
            } else if (devices.size() != 1) {
                resolution = Resolution.terminal(Type.DUPLICATE_DEVICE, null,
                        "来源 IP 存在多条设备记录: " + sourceIp);
            } else {
                DeviceDevice device = devices.get(0);
                Integer assignedCollector = device.getDeviceCollectId();
                if (assignedCollector == null
                        || assignedCollector != properties.getAppServiceDeviceCollectId()) {
                    resolution = Resolution.terminal(Type.NOT_ASSIGNED, device,
                            "设备不属于当前采集探针: deviceId=" + device.getId());
                } else {
                    resolution = Resolution.registered(device);
                }
            }
            cache.put(sourceIp, new CacheEntry(resolution,
                    now + properties.getSyslogDeviceContextCacheTtlMs()));
            if (resolution.type != Type.REGISTERED) {
                log.warn("Syslog 设备接入终态: type={}, sourceIp={}, reason={}",
                        resolution.type, sourceIp, resolution.reason);
            }
            return resolution;
        }
    }

    public void invalidate(String sourceIp) {
        if (sourceIp != null) cache.remove(sourceIp);
    }

    public void invalidateAll() { cache.clear(); }
    public long getHitCount() { return hits.sum(); }
    public long getMissCount() { return misses.sum(); }
    public int getCacheSize() { return cache.size(); }

    /** 未知设备仍保留原业务登记能力，但最多每个本地缓存 TTL 执行一次，不占用网络线程。 */
    private void recordUnknown(String sourceIp, String protocol) {
        List<DeviceUnknown> existing = unknownService.getDevicesByIp(sourceIp);
        LocalDateTime now = LocalDateTime.now();
        if (existing != null && !existing.isEmpty()) {
            unknownService.updateLastTime(existing.get(0).getId(), now);
            return;
        }
        DeviceUnknown unknown = new DeviceUnknown();
        unknown.setDeviceCollectId(properties.getAppServiceDeviceCollectId());
        unknown.setDeviceCollectName(properties.getAppServiceDeviceCollectName());
        unknown.setDeviceIp(sourceIp);
        unknown.setFirstTime(now);
        unknown.setLastTime(now);
        unknown.setOrganizationId(0);
        unknown.setNetworkProtocol(protocol);
        unknown.setSourceMethod("SYSTEM");
        unknownService.createDevice(unknown);
    }

    private static final class CacheEntry {
        private final Resolution resolution;
        private final long expiresAt;
        private CacheEntry(Resolution resolution, long expiresAt) {
            this.resolution = resolution;
            this.expiresAt = expiresAt;
        }
    }

    public enum Type { REGISTERED, UNKNOWN_DEVICE, DUPLICATE_DEVICE, NOT_ASSIGNED }

    public static final class Resolution {
        private final Type type;
        private final DeviceDevice device;
        private final String reason;

        private Resolution(Type type, DeviceDevice device, String reason) {
            this.type = type;
            this.device = device;
            this.reason = reason;
        }
        static Resolution registered(DeviceDevice device) {
            return new Resolution(Type.REGISTERED, device, null);
        }
        static Resolution terminal(Type type, DeviceDevice device, String reason) {
            return new Resolution(type, device, reason);
        }
        public boolean isRegistered() { return type == Type.REGISTERED; }
        public Type getType() { return type; }
        public DeviceDevice getDevice() { return device; }
        public String getReason() { return reason; }
    }
}