waf-platform/EdgeDNS/internal/accesslogs/dns_file_writer.go

package accesslogs

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"github.com/TeaOSLab/EdgeDNS/internal/remotelogs"
)

const (
	defaultDNSLogDir = "/var/log/edge/edge-dns"
	envDNSLogDir     = "EDGE_DNS_LOG_DIR"
)

var (
	sharedDNSFileWriter     *DNSFileWriter
	sharedDNSFileWriterOnce sync.Once
)

// SharedDNSFileWriter 返回 DNS 本地日志写入器（单例）.
func SharedDNSFileWriter() *DNSFileWriter {
	sharedDNSFileWriterOnce.Do(func() {
		sharedDNSFileWriter = NewDNSFileWriter()
	})
	return sharedDNSFileWriter
}

// DNSFileWriter 将 DNS 访问日志以 JSON Lines 写入本地文件，供 Fluent Bit 采集.
type DNSFileWriter struct {
	dir    string
	mu     sync.Mutex
	file   *os.File
	inited bool
}

// NewDNSFileWriter 创建 DNS 本地日志写入器.
func NewDNSFileWriter() *DNSFileWriter {
	logDir := resolveDefaultDNSLogDir()
	return &DNSFileWriter{dir: logDir}
}

func resolveDefaultDNSLogDir() string {
	logDir := strings.TrimSpace(os.Getenv(envDNSLogDir))
	if logDir == "" {
		return defaultDNSLogDir
	}
	return logDir
}

func resolveDNSDirFromPolicyPath(policyPath string) string {
	policyPath = strings.TrimSpace(policyPath)
	if policyPath == "" {
		return ""
	}

	if strings.HasSuffix(policyPath, "/") || strings.HasSuffix(policyPath, "\\") {
		return filepath.Clean(policyPath)
	}

	baseName := filepath.Base(policyPath)
	if strings.Contains(baseName, ".") || strings.Contains(baseName, "${") {
		return filepath.Clean(filepath.Dir(policyPath))
	}

	return filepath.Clean(policyPath)
}

// Dir 返回当前日志目录.
func (w *DNSFileWriter) Dir() string {
	return w.dir
}

// SetDirByPolicyPath 使用公用日志策略 path 更新目录，空值时回退到 EDGE_DNS_LOG_DIR/default。
func (w *DNSFileWriter) SetDirByPolicyPath(policyPath string) {
	dir := resolveDNSDirFromPolicyPath(policyPath)
	w.SetDir(dir)
}

// SetDir 更新目录并重置文件句柄。
func (w *DNSFileWriter) SetDir(dir string) {
	if strings.TrimSpace(dir) == "" {
		dir = resolveDefaultDNSLogDir()
	}

	w.mu.Lock()
	defer w.mu.Unlock()

	if dir == w.dir {
		return
	}

	if w.file != nil {
		_ = w.file.Close()
		w.file = nil
	}
	w.inited = false
	w.dir = dir
}

// EnsureInit 在启动时预创建目录与 access.log.
func (w *DNSFileWriter) EnsureInit() error {
	if w.dir == "" {
		return nil
	}
	return w.init()
}

func (w *DNSFileWriter) init() error {
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.inited && w.file != nil {
		return nil
	}

	if w.dir == "" {
		return nil
	}

	if err := os.MkdirAll(w.dir, 0755); err != nil {
		remotelogs.Error("DNS_ACCESS_LOG_FILE", "mkdir log dir failed: "+err.Error())
		return err
	}

	fp, err := os.OpenFile(filepath.Join(w.dir, "access.log"), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		remotelogs.Error("DNS_ACCESS_LOG_FILE", "open access.log failed: "+err.Error())
		return err
	}

	w.file = fp
	w.inited = true
	return nil
}

// WriteBatch 批量写入 DNS 访问日志.
func (w *DNSFileWriter) WriteBatch(logs []*pb.NSAccessLog, clusterId int64) {
	if len(logs) == 0 || w.dir == "" {
		return
	}
	if err := w.init(); err != nil {
		return
	}

	w.mu.Lock()
	fp := w.file
	w.mu.Unlock()
	if fp == nil {
		return
	}

	for _, log := range logs {
		ingestLog := FromNSAccessLog(log, clusterId)
		if ingestLog == nil {
			continue
		}
		line, err := json.Marshal(ingestLog)
		if err != nil {
			continue
		}
		_, _ = fp.Write(append(line, '\n'))
	}
}

// Reopen 关闭并重新打开日志文件（配合 logrotate）.
func (w *DNSFileWriter) Reopen() error {
	w.mu.Lock()
	if w.file != nil {
		_ = w.file.Close()
		w.file = nil
	}
	w.inited = false
	w.mu.Unlock()
	return w.init()
}

// Close 关闭日志文件.
func (w *DNSFileWriter) Close() error {
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.file == nil {
		return nil
	}

	err := w.file.Close()
	w.file = nil
	w.inited = false
	if err != nil {
		remotelogs.Error("DNS_ACCESS_LOG_FILE", fmt.Sprintf("close access.log failed: %v", err))
		return err
	}
	return nil
}