[优化] Phase3

openflare_relay/internal/config/config.go

@@ -0,0 +1,173 @@
+package config
+
+import (
+	"encoding/json"
+	"errors"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+)
+
+type MillisecondDuration time.Duration
+
+func (d *MillisecondDuration) UnmarshalJSON(b []byte) error {
+	var v interface{}
+	if err := json.Unmarshal(b, &v); err != nil {
+		return err
+	}
+	switch value := v.(type) {
+	case float64:
+		*d = MillisecondDuration(time.Duration(value) * time.Millisecond)
+		return nil
+	case string:
+		duration, err := time.ParseDuration(value)
+		if err != nil {
+			return err
+		}
+		*d = MillisecondDuration(duration)
+		return nil
+	default:
+		return errors.New("invalid duration format")
+	}
+}
+
+func (d MillisecondDuration) Duration() time.Duration {
+	return time.Duration(d)
+}
+
+func (d MillisecondDuration) String() string {
+	return time.Duration(d).String()
+}
+
+type Config struct {
+	ServerURL         string              `json:"server_url"`
+	AgentToken        string              `json:"agent_token"`
+	DiscoveryToken    string              `json:"discovery_token"`
+	NodeName          string              `json:"node_name"`
+	NodeIP            string              `json:"node_ip"`
+	FrpsPath          string              `json:"frps_path"`
+	DataDir           string              `json:"data_dir"`
+	StatePath         string              `json:"state_path"`
+	HeartbeatInterval MillisecondDuration `json:"heartbeat_interval"`
+	RequestTimeout    MillisecondDuration `json:"request_timeout"`
+	configPath        string
+}
+
+func Load(path string) (*Config, error) {
+	data, err := os.ReadFile(path)
+	if err != nil && !os.IsNotExist(err) {
+		return nil, err
+	}
+	cfg := &Config{}
+	if err == nil {
+		if err = json.Unmarshal(data, cfg); err != nil {
+			return nil, err
+		}
+	}
+	if err != nil && !hasEnvConfig() {
+		return nil, err
+	}
+	cfg.configPath = path
+	applyEnvOverrides(cfg)
+	applyDefaults(cfg, filepath.Dir(path))
+	if err = validate(cfg); err != nil {
+		return nil, err
+	}
+	return cfg, nil
+}
+
+func hasEnvConfig() bool {
+	for _, key := range []string{
+		"OPENFLARE_SERVER_URL",
+		"OPENFLARE_AGENT_TOKEN",
+		"OPENFLARE_DISCOVERY_TOKEN",
+		"OPENFLARE_NODE_NAME",
+		"OPENFLARE_NODE_IP",
+		"OPENFLARE_DATA_DIR",
+		"OPENFLARE_FRPS_PATH",
+	} {
+		if strings.TrimSpace(os.Getenv(key)) != "" {
+			return true
+		}
+	}
+	return false
+}
+
+func applyEnvOverrides(cfg *Config) {
+	if cfg == nil {
+		return
+	}
+	overrideString := func(key string, target *string) {
+		if value := strings.TrimSpace(os.Getenv(key)); value != "" {
+			*target = value
+		}
+	}
+	overrideString("OPENFLARE_SERVER_URL", &cfg.ServerURL)
+	overrideString("OPENFLARE_AGENT_TOKEN", &cfg.AgentToken)
+	overrideString("OPENFLARE_DISCOVERY_TOKEN", &cfg.DiscoveryToken)
+	overrideString("OPENFLARE_NODE_NAME", &cfg.NodeName)
+	overrideString("OPENFLARE_NODE_IP", &cfg.NodeIP)
+	overrideString("OPENFLARE_DATA_DIR", &cfg.DataDir)
+	overrideString("OPENFLARE_FRPS_PATH", &cfg.FrpsPath)
+}
+
+func applyDefaults(cfg *Config, baseDir string) {
+	baseDir = filepath.Clean(baseDir)
+	if cfg.FrpsPath == "" {
+		cfg.FrpsPath = "frps" // rely on PATH
+	}
+	if cfg.DataDir == "" {
+		cfg.DataDir = filepath.Join(baseDir, "data")
+	}
+	if cfg.NodeName == "" {
+		host, _ := os.Hostname()
+		cfg.NodeName = strings.TrimSpace(host)
+	}
+	if cfg.StatePath == "" {
+		cfg.StatePath = filepath.Join(cfg.DataDir, "relay-state.json")
+	}
+	if cfg.HeartbeatInterval <= 0 {
+		cfg.HeartbeatInterval = MillisecondDuration(10 * time.Second)
+	}
+	if cfg.RequestTimeout <= 0 {
+		cfg.RequestTimeout = MillisecondDuration(10 * time.Second)
+	}
+}
+
+func validate(cfg *Config) error {
+	if cfg.ServerURL == "" {
+		return errors.New("server_url 不能为空")
+	}
+	if strings.TrimSpace(cfg.AgentToken) == "" && strings.TrimSpace(cfg.DiscoveryToken) == "" {
+		return errors.New("agent_token 和 discovery_token 不能同时为空")
+	}
+	if cfg.NodeName == "" {
+		return errors.New("node_name 不能为空")
+	}
+	return nil
+}
+
+func (cfg *Config) InitialAuthToken() string {
+	if cfg == nil {
+		return ""
+	}
+	if token := strings.TrimSpace(cfg.AgentToken); token != "" {
+		return token
+	}
+	return strings.TrimSpace(cfg.DiscoveryToken)
+}
+
+func (cfg *Config) Save() error {
+	if cfg == nil {
+		return errors.New("config 不能为空")
+	}
+	if cfg.configPath == "" {
+		return errors.New("config path 未初始化")
+	}
+	data, err := json.MarshalIndent(cfg, "", "  ")
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(cfg.configPath, data, 0o644)
+}

openflare_relay/internal/frps/manager.go

@@ -0,0 +1,148 @@
+package frps
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"log/slog"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"sync"
+	"time"
+
+	"openflare/service"
+)
+
+type Manager struct {
+	frpsPath   string
+	dataDir    string
+	configPath string
+
+	mu           sync.RWMutex
+	activeConfig *service.RelayConfig
+	cmd          *exec.Cmd
+	status       string
+}
+
+func NewManager(frpsPath string, dataDir string) *Manager {
+	return &Manager{
+		frpsPath:   frpsPath,
+		dataDir:    dataDir,
+		configPath: filepath.Join(dataDir, "frps.toml"),
+		status:     "unhealthy",
+	}
+}
+
+func (m *Manager) GetVersion() string {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	cmd := exec.CommandContext(ctx, m.frpsPath, "-v")
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		slog.Error("failed to get frps version", "error", err)
+		return ""
+	}
+	return strings.TrimSpace(string(out))
+}
+
+func (m *Manager) GetStatus() string {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.status
+}
+
+func (m *Manager) UpdateConfig(cfg *service.RelayConfig) {
+	if cfg == nil {
+		return
+	}
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	// Check if config changed
+	if m.activeConfig != nil &&
+		m.activeConfig.BindPort == cfg.BindPort &&
+		m.activeConfig.VhostHTTPPort == cfg.VhostHTTPPort &&
+		m.activeConfig.AuthToken == cfg.AuthToken {
+		return // No change
+	}
+
+	m.activeConfig = cfg
+	slog.Info("relay config updated, reloading frps")
+
+	if err := m.renderConfig(cfg); err != nil {
+		slog.Error("failed to render frps config", "error", err)
+		m.status = "unhealthy"
+		return
+	}
+
+	if err := m.restartProcess(); err != nil {
+		slog.Error("failed to restart frps", "error", err)
+		m.status = "unhealthy"
+	} else {
+		m.status = "healthy"
+	}
+}
+
+func (m *Manager) renderConfig(cfg *service.RelayConfig) error {
+	if err := os.MkdirAll(m.dataDir, 0755); err != nil {
+		return err
+	}
+	var buf bytes.Buffer
+	buf.WriteString(fmt.Sprintf("bindPort = %d\n", cfg.BindPort))
+	if cfg.VhostHTTPPort > 0 {
+		buf.WriteString(fmt.Sprintf("vhostHTTPPort = %d\n", cfg.VhostHTTPPort))
+	}
+	if cfg.AuthToken != "" {
+		buf.WriteString("[auth]\n")
+		buf.WriteString("method = \"token\"\n")
+		buf.WriteString(fmt.Sprintf("token = \"%s\"\n", cfg.AuthToken))
+	}
+
+	return os.WriteFile(m.configPath, buf.Bytes(), 0644)
+}
+
+func (m *Manager) restartProcess() error {
+	if m.cmd != nil && m.cmd.Process != nil {
+		slog.Debug("stopping existing frps process")
+		_ = m.cmd.Process.Kill()
+		_ = m.cmd.Wait()
+		m.cmd = nil
+	}
+
+	cmd := exec.Command(m.frpsPath, "-c", m.configPath)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	if err := cmd.Start(); err != nil {
+		return err
+	}
+
+	m.cmd = cmd
+
+	// Start a goroutine to monitor process exit
+	go func(c *exec.Cmd) {
+		err := c.Wait()
+		slog.Warn("frps process exited", "error", err)
+		m.mu.Lock()
+		if m.cmd == c {
+			m.cmd = nil
+			m.status = "unhealthy"
+		}
+		m.mu.Unlock()
+	}(cmd)
+
+	return nil
+}
+
+func (m *Manager) Stop() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.cmd != nil && m.cmd.Process != nil {
+		_ = m.cmd.Process.Kill()
+		_ = m.cmd.Wait()
+		m.cmd = nil
+	}
+}

openflare_relay/internal/heartbeat/service.go

@@ -0,0 +1,65 @@
+package heartbeat
+
+import (
+	"context"
+	"log/slog"
+	"time"
+
+	"openflare-relay/internal/config"
+	"openflare-relay/internal/frps"
+	"openflare-relay/internal/httpclient"
+	"openflare/service"
+)
+
+type Service struct {
+	client      *httpclient.Client
+	frpsManager *frps.Manager
+	config      *config.Config
+}
+
+func New(client *httpclient.Client, manager *frps.Manager, cfg *config.Config) *Service {
+	return &Service{
+		client:      client,
+		frpsManager: manager,
+		config:      cfg,
+	}
+}
+
+func (s *Service) Run(ctx context.Context) {
+	ticker := time.NewTicker(s.config.HeartbeatInterval.Duration())
+	defer ticker.Stop()
+
+	// initial heartbeat
+	s.doHeartbeat(ctx)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			s.doHeartbeat(ctx)
+		}
+	}
+}
+
+func (s *Service) doHeartbeat(ctx context.Context) {
+	slog.Debug("sending heartbeat")
+
+	payload := service.RelayHeartbeatPayload{
+		RelayVersion:   "0.1.0", // TODO dynamically inject build version
+		FrpVersion:     s.frpsManager.GetVersion(),
+		RelayStatus:    s.frpsManager.GetStatus(),
+		FrpsConnCount:  0,
+		FrpsProxyCount: 0,
+	}
+
+	resp, err := s.client.Heartbeat(ctx, payload)
+	if err != nil {
+		slog.Error("heartbeat failed", "error", err)
+		return
+	}
+	slog.Debug("heartbeat succeeded")
+
+	// Update configs if changed
+	s.frpsManager.UpdateConfig(resp.RelayConfig)
+}

openflare_relay/internal/httpclient/client.go

@@ -0,0 +1,111 @@
+package httpclient
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"log/slog"
+	"net/http"
+	"strings"
+	"time"
+
+	"openflare/service"
+)
+
+type APIResponse[T any] struct {
+	Success bool   `json:"success"`
+	Message string `json:"message"`
+	Data    T      `json:"data"`
+}
+
+type Client struct {
+	baseURL    string
+	token      string
+	httpClient *http.Client
+}
+
+func New(baseURL string, token string, timeout time.Duration) *Client {
+	return &Client{
+		baseURL: strings.TrimRight(baseURL, "/"),
+		token:   token,
+		httpClient: &http.Client{
+			Timeout: timeout,
+		},
+	}
+}
+
+func (c *Client) Heartbeat(ctx context.Context, payload service.RelayHeartbeatPayload) (*service.RelayHeartbeatResponse, error) {
+	resp := APIResponse[service.RelayHeartbeatResponse]{}
+	if err := c.postJSON(ctx, "/api/relay/heartbeat", payload, &resp); err != nil {
+		return nil, err
+	}
+	if !resp.Success {
+		return nil, errors.New(resp.Message)
+	}
+	return &resp.Data, nil
+}
+
+func (c *Client) SetToken(token string) {
+	c.token = strings.TrimSpace(token)
+	slog.Debug("http client token updated")
+}
+
+func (c *Client) getJSON(ctx context.Context, path string, target any) error {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+path, nil)
+	if err != nil {
+		return err
+	}
+	req.Header.Set("X-Agent-Token", c.token)
+	return c.do(req, target)
+}
+
+func (c *Client) postJSON(ctx context.Context, path string, body any, target any) error {
+	data, err := json.Marshal(body)
+	if err != nil {
+		return err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+path, bytes.NewReader(data))
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("X-Agent-Token", c.token)
+	return c.do(req, target)
+}
+
+func (c *Client) do(req *http.Request, target any) error {
+	res, err := c.httpClient.Do(req)
+	if err != nil {
+		slog.Error("http request failed", "method", req.Method, "path", req.URL.Path, "error", err)
+		return err
+	}
+	defer func(Body io.ReadCloser) {
+		err := Body.Close()
+		if err != nil {
+			slog.Error("failed to close response body", "error", err)
+		}
+	}(res.Body)
+	if res.StatusCode != http.StatusOK {
+		slog.Warn("http request returned non-200", "method", req.Method, "path", req.URL.Path, "status", res.Status)
+		return errors.New(res.Status)
+	}
+	if target == nil {
+		var wrapper APIResponse[json.RawMessage]
+		if err = json.NewDecoder(res.Body).Decode(&wrapper); err != nil {
+			slog.Error("http response decode failed", "method", req.Method, "path", req.URL.Path, "error", err)
+			return err
+		}
+		if !wrapper.Success {
+			slog.Warn("http api response failed", "method", req.Method, "path", req.URL.Path, "message", wrapper.Message)
+			return errors.New(wrapper.Message)
+		}
+		return nil
+	}
+	if err = json.NewDecoder(res.Body).Decode(target); err != nil {
+		slog.Error("http response decode failed", "method", req.Method, "path", req.URL.Path, "error", err)
+		return err
+	}
+	return nil
+}

openflare_relay/internal/relay/runner.go

@@ -0,0 +1,94 @@
+package relay
+
+import (
+	"context"
+	"encoding/json"
+	"log/slog"
+	"time"
+
+	"openflare-relay/internal/config"
+	"openflare-relay/internal/frps"
+	"openflare-relay/internal/heartbeat"
+	"openflare-relay/internal/httpclient"
+	"openflare-relay/internal/state"
+	"openflare-relay/internal/wsclient"
+	"openflare/service"
+)
+
+type Runner struct {
+	Config           *config.Config
+	StateStore       *state.Store
+	HeartbeatService *heartbeat.Service
+	FrpsManager      *frps.Manager
+	WebSocketService *wsclient.Client
+	HttpClient       *httpclient.Client
+}
+
+func (r *Runner) Run(ctx context.Context) error {
+	// Start heartbeat loop in background
+	go r.HeartbeatService.Run(ctx)
+
+	// WebSocket reconnection loop
+	for {
+		select {
+		case <-ctx.Done():
+			r.FrpsManager.Stop()
+			return ctx.Err()
+		default:
+		}
+
+		conn, err := r.WebSocketService.Connect(ctx)
+		if err != nil {
+			slog.Error("relay ws connect failed, will retry", "error", err)
+			r.sleepContext(ctx, 5*time.Second)
+			continue
+		}
+
+		r.handleConnection(ctx, conn)
+		_ = conn.Close()
+		slog.Info("relay ws connection closed, reconnecting...")
+		r.sleepContext(ctx, 2*time.Second)
+	}
+}
+
+func (r *Runner) handleConnection(ctx context.Context, conn *wsclient.Connection) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+
+		msg, err := conn.Receive()
+		if err != nil {
+			slog.Error("relay ws receive failed", "error", err)
+			return
+		}
+
+		switch msg.Type {
+		case "ping":
+			_ = conn.SendPong()
+		case "relay_config":
+			payloadBytes, ok := msg.Payload.(json.RawMessage)
+			if !ok {
+				slog.Error("invalid relay_config payload type")
+				continue
+			}
+			var cfg service.RelayConfig
+			if err := json.Unmarshal(payloadBytes, &cfg); err != nil {
+				slog.Error("failed to unmarshal relay_config", "error", err)
+				continue
+			}
+			r.FrpsManager.UpdateConfig(&cfg)
+		default:
+			slog.Debug("ignored unknown ws message type", "type", msg.Type)
+		}
+	}
+}
+
+func (r *Runner) sleepContext(ctx context.Context, d time.Duration) {
+	select {
+	case <-ctx.Done():
+	case <-time.After(d):
+	}
+}

openflare_relay/internal/state/store.go

@@ -0,0 +1,55 @@
+package state
+
+import (
+	"encoding/json"
+	"log/slog"
+	"os"
+	"sync"
+)
+
+type Store struct {
+	path string
+	mu   sync.RWMutex
+}
+
+type State struct {
+	LastAuthToken string `json:"last_auth_token"`
+}
+
+func NewStore(path string) *Store {
+	return &Store{
+		path: path,
+	}
+}
+
+func (s *Store) Load() (*State, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	data, err := os.ReadFile(s.path)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return &State{}, nil
+		}
+		return nil, err
+	}
+
+	var state State
+	if err := json.Unmarshal(data, &state); err != nil {
+		return &State{}, nil // Return empty state on corrupted file
+	}
+	return &state, nil
+}
+
+func (s *Store) Save(state *State) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	data, err := json.MarshalIndent(state, "", "  ")
+	if err != nil {
+		return err
+	}
+
+	slog.Debug("saving relay state")
+	return os.WriteFile(s.path, data, 0644)
+}

openflare_relay/internal/wsclient/client.go

@@ -0,0 +1,143 @@
+package wsclient
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"log/slog"
+	"net"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+
+	"golang.org/x/net/websocket"
+	"openflare/service"
+)
+
+type Client struct {
+	baseURL string
+	token   string
+	timeout time.Duration
+}
+
+type Connection struct {
+	conn        *websocket.Conn
+	url         string
+	readTimeout time.Duration
+}
+
+func New(baseURL string, token string, timeout time.Duration) *Client {
+	return &Client{
+		baseURL: strings.TrimRight(baseURL, "/"),
+		token:   strings.TrimSpace(token),
+		timeout: timeout,
+	}
+}
+
+func (c *Client) SetToken(token string) {
+	c.token = strings.TrimSpace(token)
+	slog.Debug("relay ws client token updated")
+}
+
+func (c *Client) Connect(ctx context.Context) (*Connection, error) {
+	wsURL, err := buildWebsocketURL(c.baseURL)
+	if err != nil {
+		return nil, err
+	}
+	if strings.TrimSpace(c.token) == "" {
+		return nil, errors.New("relay ws token is empty")
+	}
+	origin := strings.TrimSpace(c.baseURL)
+	if origin == "" {
+		origin = "http://localhost"
+	}
+	config, err := websocket.NewConfig(wsURL, origin)
+	if err != nil {
+		return nil, err
+	}
+	config.Header = http.Header{}
+	config.Header.Set("X-Agent-Token", c.token)
+	if c.timeout > 0 {
+		config.Dialer = &net.Dialer{Timeout: c.timeout}
+	}
+	slog.Debug("relay ws dialing server", "url", wsURL)
+	conn, err := config.DialContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+	slog.Debug("relay ws dial succeeded", "url", wsURL)
+	return &Connection{conn: conn, url: wsURL, readTimeout: websocketReadTimeout(c.timeout)}, nil
+}
+
+func buildWebsocketURL(baseURL string) (string, error) {
+	parsed, err := url.Parse(strings.TrimRight(baseURL, "/"))
+	if err != nil {
+		return "", err
+	}
+	switch parsed.Scheme {
+	case "http":
+		parsed.Scheme = "ws"
+	case "https":
+		parsed.Scheme = "wss"
+	case "ws", "wss":
+	default:
+		return "", errors.New("server_url scheme must be http, https, ws, or wss")
+	}
+	parsed.Path = strings.TrimRight(parsed.Path, "/") + "/api/relay/ws"
+	parsed.RawQuery = ""
+	parsed.Fragment = ""
+	return parsed.String(), nil
+}
+
+func (conn *Connection) SendPong() error {
+	if conn == nil || conn.conn == nil {
+		return errors.New("relay ws connection is nil")
+	}
+	slog.Debug("relay ws sending pong")
+	return websocket.JSON.Send(conn.conn, service.WSMessage{
+		Type: "pong",
+	})
+}
+
+func (conn *Connection) Receive() (service.WSMessage, error) {
+	var message service.WSMessage
+	if conn == nil || conn.conn == nil {
+		return message, errors.New("relay ws connection is nil")
+	}
+	if conn.readTimeout > 0 {
+		_ = conn.conn.SetReadDeadline(time.Now().Add(conn.readTimeout))
+	}
+	// Use custom json unmarshaling to handle any type
+	var raw struct {
+		Type    string          `json:"type"`
+		Payload json.RawMessage `json:"payload,omitempty"`
+	}
+	err := websocket.JSON.Receive(conn.conn, &raw)
+	if err != nil {
+		var netErr net.Error
+		if errors.As(err, &netErr) && netErr.Timeout() {
+			slog.Debug("relay ws receive timeout waiting for server message", "timeout", conn.readTimeout)
+		}
+		return message, err
+	}
+	message.Type = raw.Type
+	message.Payload = raw.Payload
+	slog.Debug("relay ws received message", "type", message.Type)
+	return message, nil
+}
+
+func websocketReadTimeout(requestTimeout time.Duration) time.Duration {
+	timeout := requestTimeout * 6
+	if timeout < 75*time.Second {
+		return 75 * time.Second
+	}
+	return timeout
+}
+
+func (conn *Connection) Close() error {
+	if conn == nil || conn.conn == nil {
+		return nil
+	}
+	return conn.conn.Close()
+}