zoraxy-headers/server.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"sort"
	"strings"
)

// ZoraxyClient wraps the Zoraxy management API with a scoped API key.
type ZoraxyClient struct {
	BaseURL string // e.g. http://127.0.0.1:8000
	APIKey  string
	HTTP    *http.Client
}

func NewZoraxyClient(baseURL, apiKey string) *ZoraxyClient {
	return &ZoraxyClient{BaseURL: baseURL, APIKey: apiKey, HTTP: &http.Client{}}
}

func (c *ZoraxyClient) do(method, path string, form url.Values) ([]byte, error) {
	var body io.Reader
	if form != nil {
		body = strings.NewReader(form.Encode())
	}
	// Zoraxy plugin API endpoints are mounted under /plugin/.
	if !strings.HasPrefix(path, "/plugin/") {
		path = "/plugin" + path
	}
	req, err := http.NewRequest(method, c.BaseURL+path, body)
	if err != nil {
		return nil, err
	}
	if c.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.APIKey)
	}
	if form != nil {
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	}
	resp, err := c.HTTP.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("zoraxy API %s %s: HTTP %d — %s", method, path, resp.StatusCode, string(data))
	}
	return data, nil
}

// listProxyRules returns all host-type proxy rules.
func (c *ZoraxyClient) listProxyRules() ([]ProxyRule, error) {
	data, err := c.do("GET", "/api/proxy/list?type=host", nil)
	if err != nil {
		return nil, err
	}
	var rules []ProxyRule
	if err := json.Unmarshal(data, &rules); err != nil {
		return nil, fmt.Errorf("parse proxy list: %w", err)
	}
	return rules, nil
}

// listHeaders returns the user-defined headers for one domain.
func (c *ZoraxyClient) listHeaders(domain string) ([]HeaderRule, error) {
	q := url.Values{}
	q.Set("type", "host")
	q.Set("domain", domain)
	data, err := c.do("GET", "/api/proxy/header/list?"+q.Encode(), nil)
	if err != nil {
		return nil, err
	}
	var headers []HeaderRule
	if err := json.Unmarshal(data, &headers); err != nil {
		return nil, fmt.Errorf("parse header list for %s: %w", domain, err)
	}
	return headers, nil
}

// addHeader adds (or overwrites) a header on a domain.
func (c *ZoraxyClient) addHeader(domain string, direction int, name, value string) error {
	form := url.Values{}
	form.Set("type", "add")
	form.Set("domain", domain)
	form.Set("direction", fmt.Sprintf("%d", direction))
	form.Set("name", name)
	form.Set("value", value)
	_, err := c.do("POST", "/api/proxy/header/add", form)
	return err
}

// removeHeader removes a header by name from a domain.
func (c *ZoraxyClient) removeHeader(domain, name string) error {
	form := url.Values{}
	form.Set("domain", domain)
	form.Set("name", name)
	_, err := c.do("POST", "/api/proxy/header/remove", form)
	return err
}

// buildTagGroups groups domains by tag and identifies the source (a domain that
// already has user-defined headers) for each tag.
func (c *ZoraxyClient) buildTagGroups() ([]TagGroup, error) {
	rules, err := c.listProxyRules()
	if err != nil {
		return nil, err
	}

	// Collect domains per tag.
	tagDomains := map[string][]string{}
	for _, r := range rules {
		for _, tag := range r.Tags {
			tag = strings.TrimSpace(tag)
			if tag == "" {
				continue
			}
			tagDomains[tag] = append(tagDomains[tag], r.RootOrMatchingDomain)
		}
	}

	groups := make([]TagGroup, 0, len(tagDomains))
	for tag, domains := range tagDomains {
		sort.Strings(domains)
		group := TagGroup{Tag: tag, Domains: domains, Source: "", HasSource: false}

		// Identify source: the first domain that already has headers.
		for _, d := range domains {
			headers, err := c.listHeaders(d)
			if err != nil {
				log.Printf("WARNING: list headers for %s: %v", d, err)
				continue
			}
			if len(headers) > 0 {
				group.Source = d
				group.HasSource = true
				break
			}
		}
		groups = append(groups, group)
	}

	sort.Slice(groups, func(i, j int) bool { return groups[i].Tag < groups[j].Tag })
	return groups, nil
}

// applyHeaders copies the source domain's headers to all target domains.
func (c *ZoraxyClient) applyHeaders(tag, source string, targets []string) []ActionResult {
	results := []ActionResult{}

	// Read the source headers once.
	srcHeaders, err := c.listHeaders(source)
	if err != nil {
		results = append(results, ActionResult{Domain: source, Success: false, Error: err.Error()})
		return results
	}
	if len(srcHeaders) == 0 {
		results = append(results, ActionResult{Domain: source, Success: false, Error: "source has no headers to copy"})
		return results
	}

	for _, t := range targets {
		if t == source {
			continue
		}
		var firstErr error
		for _, h := range srcHeaders {
			if h.IsRemove {
				if err := c.removeHeader(t, h.Key); err != nil {
					firstErr = err
				}
			} else {
				if err := c.addHeader(t, h.Direction, h.Key, h.Value); err != nil {
					firstErr = err
				}
			}
		}
		if firstErr != nil {
			results = append(results, ActionResult{Domain: t, Success: false, Error: firstErr.Error()})
		} else {
			results = append(results, ActionResult{Domain: t, Success: true})
		}
	}
	return results
}

// removeHeaders removes the source's header keys from all target domains.
func (c *ZoraxyClient) removeHeaders(source string, targets []string) []ActionResult {
	results := []ActionResult{}

	srcHeaders, err := c.listHeaders(source)
	if err != nil {
		results = append(results, ActionResult{Domain: source, Success: false, Error: err.Error()})
		return results
	}

	for _, t := range targets {
		if t == source {
			continue
		}
		var firstErr error
		for _, h := range srcHeaders {
			if err := c.removeHeader(t, h.Key); err != nil {
				firstErr = err
			}
		}
		if firstErr != nil {
			results = append(results, ActionResult{Domain: t, Success: false, Error: firstErr.Error()})
		} else {
			results = append(results, ActionResult{Domain: t, Success: true})
		}
	}
	return results
}

// ---------------------------------------------------------------------------
// HTTP handlers
// ---------------------------------------------------------------------------

type Handler struct {
	client *ZoraxyClient
}

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

func (h *Handler) handleGroups(w http.ResponseWriter, r *http.Request) {
	groups, err := h.client.buildTagGroups()
	if err != nil {
		log.Printf("ERROR build tag groups: %v", err)
		writeJSON(w, http.StatusInternalServerError, ErrorResponse{Error: "failed to load tags"})
		return
	}
	writeJSON(w, http.StatusOK, TagGroupsResponse{Groups: groups})
}

func (h *Handler) handleApply(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "invalid form"})
		return
	}
	tag := r.FormValue("tag")
	source := r.FormValue("source")
	targets := r.Form["targets"]

	if tag == "" || source == "" {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "tag and source required"})
		return
	}
	if len(targets) == 0 {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "no target domains"})
		return
	}

	results := h.client.applyHeaders(tag, source, targets)
	writeJSON(w, http.StatusOK, ApplyResponse{Results: results})
}

func (h *Handler) handleRemove(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "invalid form"})
		return
	}
	source := r.FormValue("source")
	targets := r.Form["targets"]

	if source == "" {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "source required"})
		return
	}
	if len(targets) == 0 {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "no target domains"})
		return
	}

	results := h.client.removeHeaders(source, targets)
	writeJSON(w, http.StatusOK, ApplyResponse{Results: results})
}