zoraxy-waf/middleware.go

package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net/http"
	"sync/atomic"
	"time"
)

var breaker = newCircuitBreaker()

var (
	requestsTotal    uint64
	requestsAllowed  uint64
	requestsBlocked  uint64
	requestsBypassed uint64
)

func metricsSnapshot() map[string]interface{} {
	return map[string]interface{}{
		"total":    atomic.LoadUint64(&requestsTotal),
		"allowed":  atomic.LoadUint64(&requestsAllowed),
		"blocked":  atomic.LoadUint64(&requestsBlocked),
		"bypassed": atomic.LoadUint64(&requestsBypassed),
		"circuit":  circuitStateStr(breaker.State()),
		"enabled":  getConfig().Enabled,
	}
}

func circuitStateStr(s CircuitState) string {
	switch s {
	case StateClosed:   return "closed"
	case StateOpen:     return "open"
	case StateHalfOpen: return "half-open"
	}
	return "unknown"
}

func handleInspect(w http.ResponseWriter, r *http.Request) {
	atomic.AddUint64(&requestsTotal, 1)
	cfg := getConfig()

	if !cfg.Enabled || breaker.State() == StateOpen {
		atomic.AddUint64(&requestsBypassed, 1)
		if breaker.State() == StateOpen {
			log.Printf("Firewall: breaker OPEN — bypass %s %s", r.Method, r.URL.Path)
		}
		// Skip WAF, proxy directly to return listener → backend.
		forwardToReturn(w, r)
		return
	}

	resp, err := forwardToWAF(w, r)
	if err != nil {
		atomic.AddUint64(&requestsBypassed, 1)
		breaker.RecordFailure()
		return
	}

	if resp.StatusCode == 403 {
		breaker.RecordSuccess()
		atomic.AddUint64(&requestsBlocked, 1)
		log.Printf("Firewall: BLOCKED %s %s", r.Method, r.URL.Path)
		return
	}

	breaker.RecordSuccess()
	atomic.AddUint64(&requestsAllowed, 1)
}

// forwardToWAF proxies the request to the WAF and copies the response back.
func forwardToWAF(w http.ResponseWriter, r *http.Request) (*http.Response, error) {
	cfg := getConfig()

	originalPath := r.Header.Get("X-Zoraxy-Uri")
	if originalPath == "" {
		originalPath = r.URL.Path
	}

	var bodyReader io.Reader
	var bodyLen int64
	if r.Body != nil && (r.Method == "POST" || r.Method == "PUT" || r.Method == "PATCH") {
		bodyBytes, readErr := io.ReadAll(r.Body)
		if readErr != nil {
			http.Error(w, "body read error", http.StatusInternalServerError)
			return nil, readErr
		}
		bodyReader = bytes.NewReader(bodyBytes)
		bodyLen = int64(len(bodyBytes))
	}

	proxyReq, err := http.NewRequestWithContext(r.Context(), r.Method, cfg.WAFURL+originalPath, bodyReader)
	if err != nil {
		http.Error(w, "proxy error", http.StatusInternalServerError)
		return nil, err
	}
	if bodyLen > 0 {
		proxyReq.ContentLength = bodyLen
	}
	proxyReq.Host = r.Host
	for key, vals := range r.Header {
		for _, v := range vals {
			proxyReq.Header.Add(key, v)
		}
	}

	client := &http.Client{
		Timeout: time.Duration(cfg.TimeoutMs) * time.Millisecond,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	resp, err := client.Do(proxyReq)
	if err != nil {
		http.Error(w, "WAF unreachable", http.StatusBadGateway)
		return nil, err
	}
	defer resp.Body.Close()

	for key, vals := range resp.Header {
		for _, v := range vals {
			w.Header().Add(key, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)

	return resp, nil
}

// forwardToReturn proxies directly to the return listener, skipping the WAF.
// Used when the firewall is disabled or the circuit breaker is open.
func forwardToReturn(w http.ResponseWriter, r *http.Request) {
	cfg := getConfig()
	originalPath := r.Header.Get("X-Zoraxy-Uri")
	if originalPath == "" {
		originalPath = r.URL.Path
	}
	target := fmt.Sprintf("http://127.0.0.1:%d%s", cfg.ReturnPort, originalPath)
	log.Printf("DEBUG forwardToReturn: %s %s -> %s (Host: %s)", r.Method, originalPath, target, r.Host)
	proxyReq, _ := http.NewRequestWithContext(r.Context(), r.Method, target, r.Body)
	proxyReq.Host = r.Host
	for key, vals := range r.Header {
		for _, v := range vals {
			proxyReq.Header.Add(key, v)
		}
	}
	client := &http.Client{
		Timeout: time.Duration(cfg.TimeoutMs) * time.Millisecond,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	resp, err := client.Do(proxyReq)
	if err != nil {
		http.Error(w, "return listener unreachable", http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()
	for key, vals := range resp.Header {
		for _, v := range vals {
			w.Header().Add(key, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}