fix(inference): forbid invalid endpoint and body content (#352)

Author: Ravenyjh

api/inference/const/const.go

@@ -18,6 +18,14 @@ var (
 		"/audio/transcriptions": {},
 	}
 
+	// FreePrefixes defines path prefixes that can be accessed without charging
+	// These are typically metadata or system endpoints that don't consume GPU resources
+	// Note: Paths here should NOT include /v1/proxy prefix (it's already stripped)
+	FreePrefixes = []string{
+		"/attestation", // TEE attestation endpoints (e.g., /attestation/report)
+		"/signature",   // TEE signature endpoints (e.g., /signature/{chatID})
+	}
+
 	// Keep this as to remove duplicate headers from incoming request
 	RequestMetaDataDuplicate = map[string]struct{}{
 		"Address":           {},

api/inference/internal/ctrl/proxy.go

@@ -3,6 +3,7 @@ package ctrl
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
 	"net/url"
@@ -24,6 +25,19 @@ func (c *Ctrl) PrepareHTTPRequest(ctx *gin.Context, targetURL string, reqBody []
 			return nil, errors.Wrap(err, "ensure stream options")
 		}
 		reqBody = modifiedBody
+
+		// Enforce configured model to prevent users from requesting more expensive models
+		// Pass user address from context for rate limiting
+		userAddr, _ := ctx.Get("userAddress")
+		userAddrStr, _ := userAddr.(string)
+		modifiedBody, err = c.EnforceConfiguredModel(reqBody, userAddrStr)
+		if err != nil {
+			// Model validation failure is a user input error (similar to invalid token, bad request body)
+			// Mark as expected error to prevent polluting error monitoring
+			ctx.Set("ignoreError", true)
+			return nil, errors.Wrap(err, "enforce configured model")
+		}
+		reqBody = modifiedBody
 	}
 
 	// For text-to-image requests, ensure wait=true query parameter is set
@@ -310,3 +324,85 @@ func (c *Ctrl) EnsureStreamOptions(body []byte) ([]byte, error) {
 
 	return modifiedBody, nil
 }
+
+// EnforceConfiguredModel ensures that requests use the configured model from the service config.
+// This prevents users from requesting more expensive models while paying for cheaper ones.
+//
+// Security rationale:
+// - Provider advertises a specific model in the service configuration
+// - Pricing is based on that specific model
+// - Allowing users to change the model could result in:
+//   1. Provider paying more to backend service than they charge users
+//   2. Users getting access to premium models at cheaper prices
+//
+// This function forcibly overwrites any "model" field in the request body with the
+// configured model from c.Service.ModelType.
+func (c *Ctrl) EnforceConfiguredModel(body []byte, userAddr string) ([]byte, error) {
+	// Return original body if empty (e.g., GET requests)
+	if len(body) == 0 {
+		return body, nil
+	}
+
+	// Return original body if no model is configured
+	if c.Service.ModelType == "" {
+		c.logger.Warnf("Model enforcement skipped: c.Service.ModelType is empty (Type=%s)", c.Service.Type)
+		return body, nil
+	}
+
+	// Debug log to verify configuration
+	c.logger.Debugf("EnforceConfiguredModel: Service.Type=%s, Service.ModelType=%s",
+		c.Service.Type, c.Service.ModelType)
+
+	var bodyMap map[string]interface{}
+
+	err := json.Unmarshal(body, &bodyMap)
+	if err != nil {
+		// Return original body for non-JSON requests
+		return body, nil
+	}
+
+	// Check if request contains a model field
+	requestModel, hasModel := bodyMap["model"]
+	if !hasModel {
+		// No model specified, add the configured model
+		c.logger.Infof("No model specified in request, adding configured model: %s", c.Service.ModelType)
+		bodyMap["model"] = c.Service.ModelType
+	} else {
+		// Model specified in request, check if it matches configured model
+		requestModelStr, ok := requestModel.(string)
+		if !ok {
+			// Invalid model type, reject request
+			return nil, errors.New(fmt.Sprintf("invalid model type in request (expected string), configured model is: %s", c.Service.ModelType))
+		}
+
+		if requestModelStr != c.Service.ModelType {
+			// Model mismatch detected - record in rate limiter and REJECT
+			c.logger.Warnf("Model mismatch detected and REJECTED: user=%s, requested=%s, configured=%s",
+				userAddr, requestModelStr, c.Service.ModelType)
+
+			// Record this attempt in rate limiter if user address is available
+			if userAddr != "" {
+				rateLimiter := GetRateLimiter()
+				shouldBlock, blockedUntil := rateLimiter.RecordModelMismatch(userAddr)
+				if shouldBlock {
+					c.logger.Warnf("User will be blocked due to excessive model mismatch: user=%s, blocked_until=%s",
+						userAddr, blockedUntil.Format("2006-01-02 15:04:05"))
+				}
+			}
+
+			return nil, errors.New(fmt.Sprintf("model not supported: requested '%s', only '%s' is available for this service",
+				requestModelStr, c.Service.ModelType))
+		}
+
+		// Model matches - log for audit
+		c.logger.Debugf("Model validation passed: requested=%s matches configured=%s", requestModelStr, c.Service.ModelType)
+	}
+
+	// Marshal back to JSON
+	modifiedBody, err := json.Marshal(bodyMap)
+	if err != nil {
+		return body, errors.Wrap(err, "failed to marshal modified JSON body after enforcing model")
+	}
+
+	return modifiedBody, nil
+}

api/inference/internal/ctrl/rate_limiter.go

@@ -0,0 +1,118 @@
+package ctrl
+
+import (
+	"sync"
+	"time"
+)
+
+// ModelMismatchLimiter tracks model mismatch attempts for users
+// Similar to common/middleware/RateLimiter but with count-based blocking instead of token bucket
+type ModelMismatchLimiter struct {
+	mu    sync.RWMutex
+	users map[string]*UserMismatchInfo
+	// Configuration
+	limit  int           // Max mismatches allowed
+	window time.Duration // Time window for counting mismatches
+	block  time.Duration // Block duration after exceeding limit
+}
+
+// UserMismatchInfo stores model mismatch information for a user
+type UserMismatchInfo struct {
+	Count        int       // Number of model mismatch attempts
+	LastAttempt  time.Time // Last time model mismatch occurred
+	BlockedUntil time.Time // Time until which user is blocked
+}
+
+var (
+	globalMismatchLimiter *ModelMismatchLimiter
+	mismatchLimiterOnce   sync.Once
+)
+
+// GetRateLimiter returns the global model mismatch limiter instance
+func GetRateLimiter() *ModelMismatchLimiter {
+	mismatchLimiterOnce.Do(func() {
+		globalMismatchLimiter = &ModelMismatchLimiter{
+			users:  make(map[string]*UserMismatchInfo),
+			limit:  5,                // Max 5 mismatches
+			window: 5 * time.Minute,  // Within 5 minutes
+			block:  1 * time.Hour,    // Block for 1 hour
+		}
+		// Start cleanup goroutine (similar to common/middleware/RateLimiter)
+		go globalMismatchLimiter.cleanup()
+	})
+	return globalMismatchLimiter
+}
+
+// RecordModelMismatch records a model mismatch attempt for a user
+// Returns true if the user should be blocked
+func (ml *ModelMismatchLimiter) RecordModelMismatch(userAddr string) (shouldBlock bool, blockedUntil time.Time) {
+	ml.mu.Lock()
+	defer ml.mu.Unlock()
+
+	now := time.Now()
+
+	// Get or create user info
+	info, exists := ml.users[userAddr]
+	if !exists {
+		info = &UserMismatchInfo{}
+		ml.users[userAddr] = info
+	}
+
+	// Check if user is already blocked
+	if now.Before(info.BlockedUntil) {
+		return true, info.BlockedUntil
+	}
+
+	// Reset count if outside the time window
+	if now.Sub(info.LastAttempt) > ml.window {
+		info.Count = 0
+	}
+
+	// Increment count
+	info.Count++
+	info.LastAttempt = now
+
+	// Check if user should be blocked
+	if info.Count >= ml.limit {
+		info.BlockedUntil = now.Add(ml.block)
+		return true, info.BlockedUntil
+	}
+
+	return false, time.Time{}
+}
+
+// IsBlocked checks if a user is currently blocked
+func (ml *ModelMismatchLimiter) IsBlocked(userAddr string) (blocked bool, blockedUntil time.Time) {
+	ml.mu.RLock()
+	defer ml.mu.RUnlock()
+
+	info, exists := ml.users[userAddr]
+	if !exists {
+		return false, time.Time{}
+	}
+
+	now := time.Now()
+	if now.Before(info.BlockedUntil) {
+		return true, info.BlockedUntil
+	}
+
+	return false, time.Time{}
+}
+
+// cleanup removes old entries periodically (similar to common/middleware/RateLimiter.cleanupVisitors)
+func (ml *ModelMismatchLimiter) cleanup() {
+	ticker := time.NewTicker(10 * time.Minute)
+	defer ticker.Stop()
+
+	for range ticker.C {
+		ml.mu.Lock()
+		now := time.Now()
+		for addr, info := range ml.users {
+			// Remove entries that are old (> 24 hours) and not blocked
+			if now.Sub(info.LastAttempt) > 24*time.Hour && now.After(info.BlockedUntil) {
+				delete(ml.users, addr)
+			}
+		}
+		ml.mu.Unlock()
+	}
+}

api/inference/internal/proxy/proxy.go

@@ -1,10 +1,12 @@
 package proxy
 
 import (
+	"fmt"
 	"io"
 	"net/http"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/gin-contrib/cors"
 	"github.com/google/uuid"
@@ -172,16 +174,42 @@ func (p *Proxy) proxyHTTPRequest(ctx *gin.Context) {
 
 	// handle endpoints not need to be charged
 	if _, ok := constant.TargetRoute[targetPath]; !ok {
+		// Check if this is a signature endpoint with special handling (targetSeparated=false)
+		// handleSignatureRoute returns true if it handled the request from broker cache
+		// returns false if it should be forwarded to backend (targetSeparated=true)
 		if p.handleSignatureRoute(ctx, targetPath) {
 			return
 		}
 
-		httpReq, err := p.ctrl.PrepareHTTPRequest(ctx, targetURL, reqBody, svcType)
-		if err != nil {
-			p.handleBrokerError(ctx, err, "prepare HTTP request")
+		// Check if this path matches any free prefixes (attestation, signature, etc.)
+		isFree := false
+		for _, prefix := range constant.FreePrefixes {
+			if strings.HasPrefix(strings.ToLower(targetPath), prefix) {
+				isFree = true
+				break
+			}
+		}
+
+		if isFree {
+			// Log free endpoint access for audit purposes
+			p.logger.Infof("Free endpoint access: path=%s, method=%s, remote=%s, user_agent=%s",
+				targetPath, ctx.Request.Method, ctx.Request.RemoteAddr, ctx.Request.UserAgent())
+
+			httpReq, err := p.ctrl.PrepareHTTPRequest(ctx, targetURL, reqBody, svcType)
+			if err != nil {
+				p.handleBrokerError(ctx, err, "prepare HTTP request")
+				return
+			}
+			p.ctrl.ProcessHTTPRequest(ctx, svcType, httpReq, model.Request{}, "0", false)
 			return
 		}
-		p.ctrl.ProcessHTTPRequest(ctx, svcType, httpReq, model.Request{}, "0", false)
+
+		// Reject all other endpoints that are not in TargetRoute or FreePrefixes
+		// This prevents unauthorized access to unknown endpoints
+		p.logger.Warnf("Blocked unsupported endpoint: path=%s, method=%s, remote=%s, user_agent=%s",
+			targetPath, ctx.Request.Method, ctx.Request.RemoteAddr, ctx.Request.UserAgent())
+		ctx.Set("ignoreError", true)
+		p.handleBrokerError(ctx, errors.New("endpoint not supported"), "unsupported endpoint")
 		return
 	}
 
@@ -193,6 +221,24 @@ func (p *Proxy) proxyHTTPRequest(ctx *gin.Context) {
 		return
 	}
 
+	// Store user address in context for rate limiting
+	ctx.Set("userAddress", userAddress)
+
+	// Check if user is rate-limited due to excessive model mismatch attempts
+	rateLimiter := ctrl.GetRateLimiter()
+	if blocked, blockedUntil := rateLimiter.IsBlocked(userAddress); blocked {
+		// User is blocked - return error immediately without processing
+		ctx.Set("ignoreError", true)
+		remainingTime := blockedUntil.Sub(time.Now())
+		p.logger.Warnf("User blocked due to excessive model mismatch attempts: user=%s, blocked_until=%s, remaining=%v",
+			userAddress, blockedUntil.Format(time.RFC3339), remainingTime)
+
+		ctx.JSON(http.StatusTooManyRequests, gin.H{
+			"error": fmt.Sprintf("Rate limit exceeded: too many invalid model requests. Please try again in %v", remainingTime.Round(time.Minute)),
+		})
+		return
+	}
+
 	// Record unique user for DAU tracking
 	monitor.RecordUniqueUser(userAddress)