Wormhole/internal/proxy/headers.go

package proxy

import (
	"encoding/json"
	"net"
	"net/http"
	"strings"

	"github.com/wormhole/wormhole/internal/auth"
	"github.com/wormhole/wormhole/internal/store"
)

const (
	// IPForwardReplace 覆盖 X-Forwarded-For / X-Real-IP 为访客真实 IP。
	IPForwardReplace = "replace"
	// IPForwardAppend 按反向代理规范追加 X-Forwarded-For（Nginx proxy_add_x_forwarded_for）。
	IPForwardAppend = "append"
)

func normalizeIPForwardMode(mode string) string {
	switch mode {
	case IPForwardReplace, IPForwardAppend:
		return mode
	case "real": // 兼容旧值
		return IPForwardReplace
	case "proxy": // 兼容旧值
		return IPForwardAppend
	default:
		return IPForwardReplace
	}
}

func parseHeaderMap(raw string) map[string]string {
	raw = strings.TrimSpace(raw)
	if raw == "" || raw == "{}" {
		return nil
	}
	var m map[string]string
	if err := json.Unmarshal([]byte(raw), &m); err != nil {
		return nil
	}
	return m
}

func mergeHeaderMaps(base, override map[string]string) map[string]string {
	out := make(map[string]string)
	for k, v := range base {
		out[k] = v
	}
	for k, v := range override {
		out[k] = v
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

func effectiveIPMode(resourceMode, clientMode string) string {
	if resourceMode != "" {
		return normalizeIPForwardMode(resourceMode)
	}
	if clientMode != "" {
		return normalizeIPForwardMode(clientMode)
	}
	return IPForwardReplace
}

func resolveProxyConfig(st *store.Store, clientID uint, resourceMode, resourceHeaders string) (mode string, headers map[string]string) {
	client, _ := st.GetClientByID(clientID)
	clientMode := ""
	clientHeaders := ""
	if client != nil {
		clientMode = client.IPForwardMode
		clientHeaders = client.CustomHeaders
	}
	mode = effectiveIPMode(resourceMode, clientMode)
	headers = mergeHeaderMaps(parseHeaderMap(clientHeaders), parseHeaderMap(resourceHeaders))
	return mode, headers
}

func hostOnly(addr string) string {
	if addr == "" {
		return ""
	}
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return strings.Trim(addr, "[]")
	}
	return host
}

func directClientIP(r *http.Request) string {
	return hostOnly(r.RemoteAddr)
}

func forwardIPForHTTP(r *http.Request, mode string) string {
	mode = normalizeIPForwardMode(mode)
	if mode == IPForwardAppend {
		return directClientIP(r)
	}
	return auth.ExtractRemoteIP(r)
}

func forwardIPForTCP(clientAddr string, _ string) string {
	return hostOnly(clientAddr)
}

func isForbiddenHeader(name string) bool {
	switch strings.ToLower(name) {
	case "connection", "transfer-encoding", "keep-alive", "proxy-connection",
		"te", "trailer", "upgrade", "content-length", "host":
		return true
	default:
		return false
	}
}

func applyOutboundHeaders(req *http.Request, r *http.Request, mode string, custom map[string]string) {
	mode = normalizeIPForwardMode(mode)
	req.Header.Del("X-Forwarded-For")
	req.Header.Del("X-Real-IP")
	req.Header.Del("Forwarded")

	switch mode {
	case IPForwardAppend:
		clientIP := directClientIP(r)
		if clientIP == "" {
			break
		}
		existing := strings.TrimSpace(r.Header.Get("X-Forwarded-For"))
		xff := clientIP
		if existing != "" {
			xff = existing + ", " + clientIP
		}
		req.Header.Set("X-Forwarded-For", xff)
		req.Header.Set("X-Real-IP", clientIP)
	default:
		visitorIP := auth.ExtractRemoteIP(r)
		if visitorIP != "" {
			req.Header.Set("X-Forwarded-For", visitorIP)
			req.Header.Set("X-Real-IP", visitorIP)
		}
	}

	for k, v := range custom {
		if k == "" || isForbiddenHeader(k) {
			continue
		}
		req.Header.Set(k, v)
	}
}