Luminary/internal/middleware/middleware.go

package middleware

import (
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/store/cache"
)

func IPFilter(c *cache.Store, scope model.IPRuleScope) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		rules := c.ListIPRules()
		if len(rules) == 0 {
			ctx.Next()
			return
		}
		clientIP := clientIP(ctx)
		ip := net.ParseIP(clientIP)
		if ip == nil {
			ctx.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "invalid client ip"})
			return
		}
		var allowRules, denyRules []model.IPRule
		for _, r := range rules {
			if !r.Enabled {
				continue
			}
			if r.Scope != scope && r.Scope != model.IPScopeAll {
				continue
			}
			if r.Type == model.IPRuleAllow {
				allowRules = append(allowRules, r)
			} else {
				denyRules = append(denyRules, r)
			}
		}
		for _, r := range denyRules {
			if ipMatchesRule(r.CIDR, ip) {
				ctx.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "ip denied"})
				return
			}
		}
		if len(allowRules) > 0 {
			allowed := false
			for _, r := range allowRules {
				if ipMatchesRule(r.CIDR, ip) {
					allowed = true
					break
				}
			}
			if !allowed {
				ctx.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "ip not allowed"})
				return
			}
		}
		ctx.Next()
	}
}

func cidrContains(cidr string, ip net.IP) bool {
	if !strings.Contains(cidr, "/") {
		return ip.String() == cidr
	}
	_, network, err := net.ParseCIDR(cidr)
	if err != nil {
		return false
	}
	return network.Contains(ip)
}

// ipMatchesRule checks CIDR/IP match, treating IPv4/IPv6 loopback as equivalent.
func ipMatchesRule(cidr string, ip net.IP) bool {
	if cidrContains(cidr, ip) {
		return true
	}
	if !ip.IsLoopback() {
		return false
	}
	alt := loopbackAlt(ip)
	return alt != nil && cidrContains(cidr, alt)
}

func loopbackAlt(ip net.IP) net.IP {
	if v4 := ip.To4(); v4 != nil && v4.IsLoopback() {
		return net.ParseIP("::1")
	}
	if ip.Equal(net.ParseIP("::1")) {
		return net.ParseIP("127.0.0.1")
	}
	return nil
}

func clientIP(ctx *gin.Context) string {
	return ClientIP(ctx)
}

const sessionCookie = "luminary_session"

type SessionValidator func(token string) (bool, time.Time)

func AdminAuth(validate SessionValidator) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		token := extractToken(ctx)
		if token == "" {
			ctx.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"})
			return
		}
		ok, expires := validate(token)
		if !ok || time.Now().After(expires) {
			ctx.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "session expired"})
			return
		}
		ctx.Set("session_token", token)
		ctx.Next()
	}
}

func extractToken(ctx *gin.Context) string {
	if auth := ctx.GetHeader("Authorization"); strings.HasPrefix(auth, "Bearer ") {
		return strings.TrimPrefix(auth, "Bearer ")
	}
	if cookie, err := ctx.Cookie(sessionCookie); err == nil {
		return cookie
	}
	return ""
}

func SessionCookieName() string { return sessionCookie }

func IngressAuth(gw IngressKeyResolver) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		key := extractIngressKey(ctx)
		if key == "" {
			ctx.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "missing api key"})
			return
		}
		ingress, err := gw.ResolveIngressKey(key)
		if err != nil {
			ctx.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid api key"})
			return
		}
		ctx.Set("ingress_key", ingress)
		ctx.Next()
	}
}

func extractIngressKey(ctx *gin.Context) string {
	if auth := ctx.GetHeader("Authorization"); strings.HasPrefix(auth, "Bearer ") {
		return strings.TrimPrefix(auth, "Bearer ")
	}
	return ctx.GetHeader("x-api-key")
}

type IngressKeyResolver interface {
	ResolveIngressKey(key string) (*model.IngressKey, error)
}

// RequireIngressProtocol restricts routes to ingress keys of the given protocol family.
func RequireIngressProtocol(want model.IngressKeyProtocol) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		v, ok := ctx.Get("ingress_key")
		if !ok {
			ctx.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "missing api key"})
			return
		}
		k, ok := v.(*model.IngressKey)
		if !ok {
			ctx.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid api key"})
			return
		}
		if k.EffectiveProtocol() != want {
			abortWrongProtocol(ctx, want, protocolMismatchMessage(k.EffectiveProtocol(), want))
			return
		}
		ctx.Next()
	}
}

func protocolMismatchMessage(got, want model.IngressKeyProtocol) string {
	if want == model.IngressProtocolImage {
		return "this api key is for text (OpenAI) API only; use an image (Replicate) ingress key"
	}
	return "this api key is for image (Replicate) API only; use a text (OpenAI) ingress key"
}

func abortWrongProtocol(ctx *gin.Context, want model.IngressKeyProtocol, msg string) {
	if want == model.IngressProtocolImage {
		ctx.AbortWithStatusJSON(http.StatusForbidden, gin.H{"detail": msg})
		return
	}
	ctx.AbortWithStatusJSON(http.StatusForbidden, gin.H{
		"error": gin.H{"message": msg, "type": "forbidden"},
	})
}

func Logger() gin.HandlerFunc {
	return gin.LoggerWithFormatter(func(param gin.LogFormatterParams) string {
		return param.TimeStamp.Format(time.RFC3339) + " " + param.Method + " " + param.Path + " " + param.ClientIP + "\n"
	})
}