OpenGFW/cmd/web.go

package cmd

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sync"

	"github.com/apernet/OpenGFW/analyzer"
	"github.com/apernet/OpenGFW/engine"
	"github.com/apernet/OpenGFW/modifier"
	"github.com/apernet/OpenGFW/ruleset"
	"github.com/apernet/OpenGFW/web"

	"go.uber.org/zap"
	"gopkg.in/yaml.v3"
)

// hub is the (optional) statistics collector feeding the web UI. It is nil
// when the web UI is disabled.
var hub *web.Hub

// ruleManager owns the rule file and knows how to hot reload the engine.
// It is shared between the SIGHUP handler and the web UI.
type ruleManager struct {
	path      string
	analyzers []analyzer.Analyzer
	modifiers []modifier.Modifier
	rsConfig  *ruleset.BuiltinConfig

	mu     sync.Mutex
	engine engine.Engine
}

var _ web.RuleManager = (*ruleManager)(nil)

func newRuleManager(path string, ans []analyzer.Analyzer, mods []modifier.Modifier, rsConfig *ruleset.BuiltinConfig) *ruleManager {
	return &ruleManager{path: path, analyzers: ans, modifiers: mods, rsConfig: rsConfig}
}

func (m *ruleManager) SetEngine(en engine.Engine) {
	m.mu.Lock()
	m.engine = en
	m.mu.Unlock()
}

func (m *ruleManager) Path() string { return m.path }

// Compile parses and compiles a rule file content, without applying it.
func (m *ruleManager) Compile(raw string) (ruleset.Ruleset, []ruleset.ExprRule, error) {
	rawRs, err := ruleset.ExprRulesFromYAMLBytes([]byte(raw))
	if err != nil {
		return nil, nil, fmt.Errorf("failed to parse rules: %w", err)
	}
	rs, err := ruleset.CompileExprRules(rawRs, m.analyzers, m.modifiers, m.rsConfig)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to compile rules: %w", err)
	}
	return rs, rawRs, nil
}

// Reload re-reads the rule file from disk and applies it to the engine.
func (m *ruleManager) Reload() error {
	bs, err := os.ReadFile(m.path)
	if err != nil {
		return err
	}
	rs, _, err := m.Compile(string(bs))
	if err != nil {
		return err
	}
	return m.update(rs)
}

func (m *ruleManager) update(rs ruleset.Ruleset) error {
	m.mu.Lock()
	en := m.engine
	m.mu.Unlock()
	if en == nil {
		return errors.New("engine is not running")
	}
	return en.UpdateRuleset(rs)
}

func (m *ruleManager) Load() (string, []web.Rule, error) {
	bs, err := os.ReadFile(m.path)
	if err != nil {
		return "", nil, err
	}
	rules, err := ruleset.ExprRulesFromYAMLBytes(bs)
	if err != nil {
		// The file is still shown as-is so that the user can fix it in the editor.
		return string(bs), nil, nil
	}
	return string(bs), toWebRules(rules), nil
}

func (m *ruleManager) Validate(raw string) ([]web.Rule, error) {
	_, rules, err := m.Compile(raw)
	if err != nil {
		return nil, err
	}
	return toWebRules(rules), nil
}

func (m *ruleManager) Marshal(rules []web.Rule) (string, error) {
	out := make([]web.Rule, 0, len(rules))
	for _, r := range rules {
		if r.Modifier != nil && r.Modifier.Name == "" {
			r.Modifier = nil // Leftover from switching a rule away from `modify`
		}
		out = append(out, r)
	}
	bs, err := yaml.Marshal(out)
	if err != nil {
		return "", err
	}
	return string(bs), nil
}

// Apply compiles the given rules, persists them to the rule file and hot
// reloads the engine. The file is only written once the rules compile.
func (m *ruleManager) Apply(raw string) ([]web.Rule, error) {
	rs, rules, err := m.Compile(raw)
	if err != nil {
		return nil, err
	}
	if err := m.writeFile(raw); err != nil {
		return nil, fmt.Errorf("failed to write rule file: %w", err)
	}
	if err := m.update(rs); err != nil {
		return nil, fmt.Errorf("failed to update ruleset: %w", err)
	}
	return toWebRules(rules), nil
}

// writeFile replaces the rule file atomically so that a crash in the middle of
// a save cannot leave a truncated ruleset behind.
func (m *ruleManager) writeFile(raw string) error {
	mode := os.FileMode(0o644)
	if fi, err := os.Stat(m.path); err == nil {
		mode = fi.Mode().Perm()
	}
	dir := filepath.Dir(m.path)
	tmp, err := os.CreateTemp(dir, ".rules-*.yaml")
	if err != nil {
		return err
	}
	tmpName := tmp.Name()
	defer os.Remove(tmpName) // No-op once the rename succeeded
	if _, err := tmp.WriteString(raw); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Sync(); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	if err := os.Chmod(tmpName, mode); err != nil {
		return err
	}
	return os.Rename(tmpName, m.path)
}

func toWebRules(rules []ruleset.ExprRule) []web.Rule {
	out := make([]web.Rule, 0, len(rules))
	for _, r := range rules {
		wr := web.Rule{Name: r.Name, Action: r.Action, Log: r.Log, Expr: r.Expr}
		if r.Modifier.Name != "" {
			wr.Modifier = &web.RuleModifier{Name: r.Modifier.Name, Args: r.Modifier.Args}
		}
		out = append(out, wr)
	}
	return out
}

// startWebServer starts the web UI. It returns nil when the UI is disabled.
func startWebServer(ctx context.Context, config *cliConfig, rm *ruleManager) error {
	secret := config.Web.Secret
	generated := false
	if secret == "" {
		secret = web.RandomSecret()
		generated = true
	}
	srv, err := web.NewServer(web.Config{
		Listen:   config.Web.Listen,
		Secret:   secret,
		CertFile: config.Web.Cert,
		KeyFile:  config.Web.Key,
		Hub:      hub,
		Rules:    rm,
		Meta:     webMeta(),
		Info:     func() web.Info { return webInfo(config, rm) },
		Logf: func(format string, args ...interface{}) {
			logger.Info(fmt.Sprintf(format, args...))
		},
	})
	if err != nil {
		return configError{Field: "web", Err: err}
	}
	scheme := "http"
	if srv.TLS() {
		scheme = "https"
	}
	fields := []zap.Field{
		zap.String("listen", srv.Addr()),
		zap.String("scheme", scheme),
	}
	if generated {
		fields = append(fields, zap.String("password", secret))
		logger.Warn("web UI password was not set, using a generated one", fields...)
	} else {
		logger.Info("web UI started", fields...)
	}
	go func() {
		if err := srv.Run(ctx); err != nil {
			logger.Error("web UI stopped", zap.Error(err))
		}
	}()
	return nil
}

func webMeta() web.MetaInfo {
	meta := web.MetaInfo{
		Actions:   []string{"allow", "block", "drop", "modify"},
		Functions: []string{"geoip", "geosite", "cidr", "lookup"},
	}
	for _, a := range analyzers {
		proto := "tcp"
		if _, ok := a.(analyzer.UDPAnalyzer); ok {
			proto = "udp"
		}
		meta.Analyzers = append(meta.Analyzers, web.AnalyzerInfo{Name: a.Name(), Proto: proto})
	}
	for _, m := range modifiers {
		meta.Modifiers = append(meta.Modifiers, m.Name())
	}
	return meta
}

func webInfo(config *cliConfig, rm *ruleManager) web.Info {
	hostname, _ := os.Hostname()
	return web.Info{
		Version:   appVersion,
		Commit:    appCommit,
		Platform:  runtime.GOOS + "/" + runtime.GOARCH,
		GoVersion: runtime.Version(),
		Hostname:  hostname,
		RuleFile:  rm.Path(),
		Config: web.ConfigDigest{
			IOQueueSize:   config.IO.QueueSize,
			IOLocal:       config.IO.Local,
			IORST:         config.IO.RST,
			Workers:       config.Workers.Count,
			WorkerQueue:   config.Workers.QueueSize,
			UDPMaxStreams: config.Workers.UDPMaxStreams,
			GeoIP:         config.Ruleset.GeoIp,
			GeoSite:       config.Ruleset.GeoSite,
		},
	}
}

// hubInfo converts engine stream info into the shape the hub understands.
func hubInfo(info ruleset.StreamInfo) web.StreamInfo {
	return web.StreamInfo{
		ID:      info.ID,
		Proto:   info.Protocol.String(),
		SrcIP:   info.SrcIP.String(),
		SrcPort: info.SrcPort,
		DstIP:   info.DstIP.String(),
		DstPort: info.DstPort,
		Props:   toWebProps(info.Props),
	}
}

func toWebProps(props analyzer.CombinedPropMap) web.Props {
	if len(props) == 0 {
		return nil
	}
	out := make(web.Props, len(props))
	for name, p := range props {
		out[name] = p
	}
	return out
}