package middleware

import (
	"context"
	"net"
	"net/http"
)

type peerAddressKey struct{}

// CapturePeerAddress 在 RealIP 中间件改写 RemoteAddr 之前保存真实 TCP 对端。
//
// 不能在 LocalOnly 里直接读 r.RemoteAddr：chi 的 RealIP 会信任
// X-Forwarded-For，外部调用者可伪造该头冒充 127.0.0.1。这个中间件必须注册在
// RealIP 之前，后续安全判断只读这里保存的原始连接地址。
func CapturePeerAddress(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := context.WithValue(r.Context(), peerAddressKey{}, r.RemoteAddr)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// LocalOnly 只允许真实 TCP 对端为回环地址的请求。
// 未经过 CapturePeerAddress、地址格式错误或非 IP 地址时一律拒绝（fail closed）。
func LocalOnly(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		peer, ok := r.Context().Value(peerAddressKey{}).(string)
		if !ok || !isLoopbackAddress(peer) {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusForbidden)
			_, _ = w.Write([]byte(`{"error":"该操作只允许从 Gateway 本机执行"}`))
			return
		}
		next.ServeHTTP(w, r)
	})
}

func isLoopbackAddress(address string) bool {
	host, _, err := net.SplitHostPort(address)
	if err != nil {
		return false
	}
	ip := net.ParseIP(host)
	return ip != nil && ip.IsLoopback()
}