From e4b1c3c2b5bf349dc6cbb4a22d77444b76276f45 Mon Sep 17 00:00:00 2001
From: Denis Galeev <dengaleev@gmail.com>
Date: Thu, 6 Feb 2025 17:34:13 +0300
Subject: [PATCH] feat recorder: add rotate options

---
 config/config.go                 |  5 +++--
 config/parsing/recorder/parse.go | 35 +++++++++++++++++++++++++++++---
 recorder/file.go                 | 24 +++++++++-------------
 3 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/config/config.go b/config/config.go
index 42b1eb36..ffda6f92 100644
--- a/config/config.go
+++ b/config/config.go
@@ -278,8 +278,9 @@ type RecorderConfig struct {
 }
 
 type FileRecorder struct {
-	Path string `json:"path"`
-	Sep  string `yaml:",omitempty" json:"sep,omitempty"`
+	Path     string             `json:"path"`
+	Sep      string             `yaml:",omitempty" json:"sep,omitempty"`
+	Rotation *LogRotationConfig `yaml:",omitempty" json:"rotation,omitempty"`
 }
 
 type TCPRecorder struct {
diff --git a/config/parsing/recorder/parse.go b/config/parsing/recorder/parse.go
index c8e7a421..41944e0b 100644
--- a/config/parsing/recorder/parse.go
+++ b/config/parsing/recorder/parse.go
@@ -2,16 +2,26 @@ package recorder
 
 import (
 	"crypto/tls"
+	"io"
 	"net/http"
+	"os"
+	"path/filepath"
 	"strings"
 
+	"github.com/go-gost/core/logger"
 	"github.com/go-gost/core/recorder"
 	"github.com/go-gost/x/config"
 	"github.com/go-gost/x/internal/plugin"
 	xrecorder "github.com/go-gost/x/recorder"
 	recorder_plugin "github.com/go-gost/x/recorder/plugin"
+	"gopkg.in/natefinch/lumberjack.v2"
 )
 
+type discardCloser struct{}
+
+func (discardCloser) Write(p []byte) (n int, err error) { return len(p), nil }
+func (discardCloser) Close() error                      { return nil }
+
 func ParseRecorder(cfg *config.RecorderConfig) (r recorder.Recorder) {
 	if cfg == nil {
 		return nil
@@ -42,9 +52,28 @@ func ParseRecorder(cfg *config.RecorderConfig) (r recorder.Recorder) {
 	}
 
 	if cfg.File != nil && cfg.File.Path != "" {
-		return xrecorder.FileRecorder(cfg.File.Path,
-			xrecorder.SepRecorderOption(cfg.File.Sep),
-		)
+		var out io.WriteCloser = discardCloser{}
+
+		if cfg.File.Rotation != nil {
+			out = &lumberjack.Logger{
+				Filename:   cfg.File.Path,
+				MaxSize:    cfg.File.Rotation.MaxSize,
+				MaxAge:     cfg.File.Rotation.MaxAge,
+				MaxBackups: cfg.File.Rotation.MaxBackups,
+				LocalTime:  cfg.File.Rotation.LocalTime,
+				Compress:   cfg.File.Rotation.Compress,
+			}
+		} else {
+			os.MkdirAll(filepath.Dir(cfg.File.Path), 0755)
+			f, err := os.OpenFile(cfg.File.Path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
+			if err != nil {
+				logger.Default().Warn(err)
+			} else {
+				out = f
+			}
+		}
+
+		return xrecorder.FileRecorder(out, xrecorder.SepRecorderOption(cfg.File.Sep))
 	}
 
 	if cfg.TCP != nil && cfg.TCP.Addr != "" {
diff --git a/recorder/file.go b/recorder/file.go
index cbd4c901..5d24ae57 100644
--- a/recorder/file.go
+++ b/recorder/file.go
@@ -2,7 +2,7 @@ package recorder
 
 import (
 	"context"
-	"os"
+	"io"
 
 	"github.com/go-gost/core/recorder"
 )
@@ -20,40 +20,36 @@ func SepRecorderOption(sep string) FileRecorderOption {
 }
 
 type fileRecorder struct {
-	filename string
-	sep      string
+	out io.WriteCloser
+	sep string
 }
 
 // FileRecorder records data to file.
-func FileRecorder(filename string, opts ...FileRecorderOption) recorder.Recorder {
+func FileRecorder(out io.WriteCloser, opts ...FileRecorderOption) recorder.Recorder {
 	var options fileRecorderOptions
 	for _, opt := range opts {
 		opt(&options)
 	}
 
 	return &fileRecorder{
-		filename: filename,
-		sep:      options.sep,
+		out: out,
+		sep: options.sep,
 	}
 }
 
 func (r *fileRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error {
-	f, err := os.OpenFile(r.filename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
-	if err != nil {
+	if _, err := r.out.Write(b); err != nil {
 		return err
 	}
-	defer f.Close()
 
-	if _, err = f.Write(b); err != nil {
-		return err
-	}
 	if r.sep != "" {
-		_, err := f.WriteString(r.sep)
+		_, err := io.WriteString(r.out, r.sep)
 		return err
 	}
+
 	return nil
 }
 
 func (r *fileRecorder) Close() error {
-	return nil
+	return r.out.Close()
 }