EvoBGP/internal/agentserver/firewall.go

package agentserver

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"evobgp/internal/authkey"
	"evobgp/internal/firewall"
)

type firewallReplicatePayload struct {
	TenantID            string                `json:"tenant_id"`
	RevisionID          string                `json:"revision_id"`
	GeneratedAt         string                `json:"generated_at"`
	UpdatedAt           string                `json:"updated_at"`
	ClientsByHash       map[string]clientMeta `json:"clients_by_hash"`
	Rules               []firewallRuleJSON    `json:"rules"`
	PrefixesByCommunity map[string][]string   `json:"prefixes_by_community"`
}

type clientMeta struct {
	ClientID string `json:"client_id"`
	Name     string `json:"name"`
}

type firewallRuleJSON struct {
	ClientID    *string `json:"client_id"`
	Priority    int     `json:"priority"`
	Action      string  `json:"action"`
	CommunityID *string `json:"community_id"`
}

type firewallAgent struct {
	mu        sync.RWMutex
	stateFile string
	cpURL     string
	enabled   bool
	state     firewallReplicatePayload
}

func newFirewallAgent(cfg Config) *firewallAgent {
	enabled := strings.TrimSpace(os.Getenv("EVOBGP_FIREWALL_FAILOVER_ENABLED")) == "1"
	stateFile := envOr("EVOBGP_FIREWALL_STATE_FILE", "/var/lib/evobgp-agent/firewall-state.json")
	fa := &firewallAgent{
		stateFile: stateFile,
		cpURL:     strings.TrimSpace(cfg.ControlPlaneURL),
		enabled:   enabled,
	}
	if enabled {
		_ = fa.load()
	}
	return fa
}

func (s *Server) handleFirewallReplicate(w http.ResponseWriter, r *http.Request) {
	if !s.authorize(r) {
		writeProblem(w, http.StatusUnauthorized, "missing or invalid Authorization")
		return
	}
	var raw struct {
		TenantID    string `json:"tenant_id"`
		RevisionID  string `json:"revision_id"`
		GeneratedAt string `json:"generated_at"`
		Clients     []struct {
			TokenHashHex string `json:"token_hash_hex"`
			ClientID     string `json:"client_id"`
			Name         string `json:"name"`
		} `json:"clients"`
		Rules               []firewallRuleJSON  `json:"rules"`
		PrefixesByCommunity map[string][]string `json:"prefixes_by_community"`
	}
	if err := json.NewDecoder(r.Body).Decode(&raw); err != nil {
		writeProblem(w, http.StatusBadRequest, "invalid JSON body")
		return
	}
	byHash := make(map[string]clientMeta, len(raw.Clients))
	for _, c := range raw.Clients {
		byHash[strings.ToLower(c.TokenHashHex)] = clientMeta{ClientID: c.ClientID, Name: c.Name}
	}
	payload := firewallReplicatePayload{
		TenantID:            raw.TenantID,
		RevisionID:          raw.RevisionID,
		GeneratedAt:         raw.GeneratedAt,
		UpdatedAt:           time.Now().UTC().Format(time.RFC3339),
		ClientsByHash:       byHash,
		Rules:               raw.Rules,
		PrefixesByCommunity: raw.PrefixesByCommunity,
	}
	s.firewall.mu.Lock()
	s.firewall.state = payload
	s.firewall.mu.Unlock()
	if err := s.firewall.save(); err != nil {
		log.Printf("agentserver: firewall state save: %v", err)
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"ok":            true,
		"clients_count": len(byHash),
		"rules_count":   len(raw.Rules),
		"prefix_count":  len(raw.PrefixesByCommunity),
	})
}

func (fa *firewallAgent) handleBlocklist(w http.ResponseWriter, r *http.Request) {
	token := bearerToken(r)
	if token == "" {
		writeProblem(w, http.StatusUnauthorized, "missing bearer token")
		return
	}
	hash := hex.EncodeToString(authkey.HashToken(token))
	fa.mu.RLock()
	st := fa.state
	meta, ok := st.ClientsByHash[hash]
	fa.mu.RUnlock()
	if !ok {
		writeProblem(w, http.StatusUnauthorized, "unknown firewall client")
		return
	}
	rules := make([]firewall.Rule, 0, len(st.Rules))
	for _, r := range st.Rules {
		rules = append(rules, firewall.Rule{
			ClientID:    r.ClientID,
			Priority:    r.Priority,
			Action:      r.Action,
			CommunityID: r.CommunityID,
		})
	}
	prefixes := firewall.Evaluate(meta.ClientID, rules, st.PrefixesByCommunity)
	updatedAt, _ := time.Parse(time.RFC3339, st.UpdatedAt)
	age := int(time.Since(updatedAt).Seconds())
	if updatedAt.IsZero() {
		age = 0
	}
	w.Header().Set("X-EvoBGP-Source", "speaker")
	w.Header().Set("X-EvoBGP-Revision-ID", st.RevisionID)
	w.Header().Set("X-EvoBGP-Generated-At", st.GeneratedAt)
	w.Header().Set("X-EvoBGP-Rules-Version", firewall.RulesVersionHash(rules))
	if age > 3600 {
		w.Header().Set("X-EvoBGP-Stale", "true")
	}
	if age > 0 {
		w.Header().Set("Age", strconvItoa(age))
	}
	resp := map[string]any{
		"client_id":             meta.ClientID,
		"revision_id":           st.RevisionID,
		"generated_at":          st.GeneratedAt,
		"source":                "speaker",
		"rules_applied":         len(rules),
		"communities_evaluated": len(st.PrefixesByCommunity),
		"prefixes":              prefixes,
		"total":                 len(prefixes),
		"hash":                  prefixHash(prefixes),
	}
	if strings.Contains(r.Header.Get("Accept"), "text/plain") {
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		for _, p := range prefixes {
			_, _ = w.Write([]byte(p + "\n"))
		}
		return
	}
	writeJSON(w, http.StatusOK, resp)
}

func (fa *firewallAgent) handleApplyReportForward(w http.ResponseWriter, r *http.Request) {
	body, _ := io.ReadAll(r.Body)
	forwarded := false
	if fa.cpURL != "" {
		req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, strings.TrimSuffix(fa.cpURL, "/")+"/v1/firewall/apply-report", strings.NewReader(string(body)))
		if err == nil {
			req.Header.Set("Authorization", r.Header.Get("Authorization"))
			req.Header.Set("Content-Type", "application/json")
			resp, err := http.DefaultClient.Do(req)
			if err == nil {
				forwarded = resp.StatusCode >= 200 && resp.StatusCode < 300
				_ = resp.Body.Close()
			}
		}
	}
	writeJSON(w, http.StatusOK, map[string]any{"ok": true, "forwarded": forwarded})
}

func (fa *firewallAgent) load() error {
	b, err := os.ReadFile(fa.stateFile)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	var st firewallReplicatePayload
	if err := json.Unmarshal(b, &st); err != nil {
		return err
	}
	fa.mu.Lock()
	fa.state = st
	fa.mu.Unlock()
	return nil
}

func (fa *firewallAgent) save() error {
	fa.mu.RLock()
	b, err := json.MarshalIndent(fa.state, "", "  ")
	fa.mu.RUnlock()
	if err != nil {
		return err
	}
	dir := filepath.Dir(fa.stateFile)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	tmp := fa.stateFile + ".tmp"
	if err := os.WriteFile(tmp, b, 0o600); err != nil {
		return err
	}
	return os.Rename(tmp, fa.stateFile)
}

func bearerToken(r *http.Request) string {
	h := r.Header.Get("Authorization")
	const p = "Bearer "
	if !strings.HasPrefix(h, p) {
		return ""
	}
	return strings.TrimSpace(h[len(p):])
}

func prefixHash(prefixes []string) string {
	cp := append([]string(nil), prefixes...)
	sort.Strings(cp)
	sum := sha256.Sum256([]byte(strings.Join(cp, "\n")))
	return "sha256:" + hex.EncodeToString(sum[:])
}

func strconvItoa(n int) string {
	if n == 0 {
		return "0"
	}
	var b [20]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	return string(b[i:])
}