zoraxy-waf/main.go

package main

import (
	"embed"
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	"zoraxy-firewall/zoraxy_plugin"
)

//go:embed web/*
var webFS embed.FS

func main() {
	loadConfig()

	spec := &zoraxy_plugin.IntroSpect{
		ID:            "zoraxy-firewall",
		Name:          "Firewall",
		Author:        "Zoraxy Community",
		AuthorContact: "",
		Description:   "Inline WAF proxy with health checks and fail-open bypass.",
		URL:           "",
		Type:          zoraxy_plugin.PluginType_Router,
		VersionMajor:  1,
		VersionMinor:  0,
		VersionPatch:  0,
		UIPath:        "/ui",

		StaticCapturePaths:   []zoraxy_plugin.StaticCaptureRule{{CapturePath: "/"}},
		StaticCaptureIngress: "/inspect",
	}

	config, err := zoraxy_plugin.ServeAndRecvSpec(spec)
	if err != nil {
		log.Fatalf("failed to receive config: %v", err)
	}

	// Start background health checker.
	startHealthCheck()

	// Return port listener for WAF async callbacks.
	cfg := getConfig()
	if cfg.ReturnPort > 0 {
		go func() {
			rmux := http.NewServeMux()
			rmux.HandleFunc("/verdict", handleReturnVerdict)
			addr := fmt.Sprintf(":%d", cfg.ReturnPort)
			log.Printf("Firewall: return listener on %s", addr)
			if err := http.ListenAndServe(addr, rmux); err != nil {
				log.Printf("Firewall: return listener error: %v", err)
			}
		}()
	}

	mux := http.NewServeMux()

	// Static capture handler — must match sub-paths (Zoraxy sends /inspect/some/page).
	mux.Handle("/inspect/", http.HandlerFunc(handleInspect))
	mux.HandleFunc("/inspect", handleInspect)

	// UI + config API.
	uiRouter := zoraxy_plugin.NewPluginEmbedUIRouter(
		spec.ID, &webFS, "web", spec.UIPath,
	)
	uiRouter.HandleFunc("/api/config", handleGetConfig, mux)
	uiRouter.HandleFunc("/api/config/save", handleSaveConfig, mux)
	uiRouter.HandleFunc("/api/stats", handleGetStats, mux)

	uiRouter.RegisterTerminateHandler(func() {
		log.Println("firewall shutting down")
	}, mux)
	uiRouter.AttachHandlerToMux(mux)

	addr := fmt.Sprintf("127.0.0.1:%d", config.Port)
	log.Printf("firewall v%d.%d.%d listening on %s",
		spec.VersionMajor, spec.VersionMinor, spec.VersionPatch, addr)

	if err := http.ListenAndServe(addr, mux); err != nil {
		log.Fatalf("server error: %v", err)
	}
}

// --- Config API ---

func handleGetConfig(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, getConfig())
}

func handleSaveConfig(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid form"})
		return
	}

	cfgMu.Lock()
	cfg.Enabled = r.FormValue("enabled") == "true"
	if u := r.FormValue("waf_url"); u != "" {
		cfg.WAFURL = u
	}
	if v := r.FormValue("return_port"); v != "" {
		fmt.Sscanf(v, "%d", &cfg.ReturnPort)
	}
	if v := r.FormValue("health_interval"); v != "" {
		fmt.Sscanf(v, "%d", &cfg.HealthInterval)
	}
	if v := r.FormValue("timeout_ms"); v != "" {
		fmt.Sscanf(v, "%d", &cfg.TimeoutMs)
	}
	cfgMu.Unlock()

	if err := saveConfig(); err != nil {
		log.Printf("ERROR saving config: %v", err)
		writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "save failed"})
		return
	}
	writeJSON(w, http.StatusOK, map[string]string{"status": "saved"})
}

func handleGetStats(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, metricsSnapshot())
}

func handleReturnVerdict(w http.ResponseWriter, r *http.Request) {
	var v struct {
		Verdict string `json:"verdict"`
		Reason  string `json:"reason"`
	}
	if err := json.NewDecoder(r.Body).Decode(&v); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "bad payload"})
		return
	}
	log.Printf("Firewall: async verdict %s (reason=%s)", v.Verdict, v.Reason)
	writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}

// --- Helpers ---

func writeJSON(w http.ResponseWriter, status int, v interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}