a301_server/pkg/middleware/idempotency.go

package middleware

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"time"

	"a301_server/pkg/database"
	"github.com/gofiber/fiber/v2"
)

const idempotencyTTL = 10 * time.Minute
const redisTimeout = 5 * time.Second

type cachedResponse struct {
	StatusCode int             `json:"s"`
	Body       json.RawMessage `json:"b"`
}

// IdempotencyRequired rejects requests without an Idempotency-Key header,
// then delegates to Idempotency for cache/replay logic.
func IdempotencyRequired(c *fiber.Ctx) error {
	if c.Get("Idempotency-Key") == "" {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{
			"error": "Idempotency-Key 헤더가 필요합니다",
		})
	}
	return Idempotency(c)
}

// Idempotency checks the Idempotency-Key header to prevent duplicate transactions.
// If the same key is seen again within the TTL, the cached response is returned.
func Idempotency(c *fiber.Ctx) error {
	key := c.Get("Idempotency-Key")
	if key == "" {
		return c.Next()
	}
	if len(key) > 256 {
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "Idempotency-Key가 너무 깁니다"})
	}

	// userID가 있으면 키에 포함하여 사용자 간 캐시 충돌 방지
	redisKey := "idempotency:"
	if uid, ok := c.Locals("userID").(uint); ok {
		redisKey += fmt.Sprintf("u%d:", uid)
	}
	redisKey += key

	ctx, cancel := context.WithTimeout(context.Background(), redisTimeout)
	defer cancel()

	// Atomically claim the key using SET NX (only succeeds if key doesn't exist)
	set, err := database.RDB.SetNX(ctx, redisKey, "processing", idempotencyTTL).Result()
	if err != nil {
		// Redis error — let the request through rather than blocking
		log.Printf("WARNING: idempotency SetNX failed (key=%s): %v", key, err)
		return c.Next()
	}

	if !set {
		// Key already exists — either processing or completed
		getCtx, getCancel := context.WithTimeout(context.Background(), redisTimeout)
		defer getCancel()

		cached, err := database.RDB.Get(getCtx, redisKey).Bytes()
		if err != nil {
			return c.Status(fiber.StatusConflict).JSON(fiber.Map{"error": "요청이 처리 중입니다"})
		}
		if string(cached) == "processing" {
			return c.Status(fiber.StatusConflict).JSON(fiber.Map{"error": "요청이 처리 중입니다"})
		}
		var cr cachedResponse
		if json.Unmarshal(cached, &cr) == nil {
			c.Set("Content-Type", "application/json")
			c.Set("X-Idempotent-Replay", "true")
			return c.Status(cr.StatusCode).Send(cr.Body)
		}
		return c.Status(fiber.StatusConflict).JSON(fiber.Map{"error": "요청이 처리 중입니다"})
	}

	// We claimed the key — process the request
	if err := c.Next(); err != nil {
		// Processing failed — remove the key so it can be retried
		delCtx, delCancel := context.WithTimeout(context.Background(), redisTimeout)
		defer delCancel()
		if delErr := database.RDB.Del(delCtx, redisKey).Err(); delErr != nil {
			log.Printf("WARNING: idempotency cache delete failed (key=%s): %v", key, delErr)
		}
		return err
	}

	// Cache successful responses (2xx), otherwise remove the key for retry
	status := c.Response().StatusCode()
	if status >= 200 && status < 300 {
		cr := cachedResponse{StatusCode: status, Body: c.Response().Body()}
		if data, err := json.Marshal(cr); err == nil {
			writeCtx, writeCancel := context.WithTimeout(context.Background(), redisTimeout)
			defer writeCancel()
			if err := database.RDB.Set(writeCtx, redisKey, data, idempotencyTTL).Err(); err != nil {
				log.Printf("WARNING: idempotency cache write failed (key=%s): %v", key, err)
			}
		}
	} else {
		// Non-success — allow retry by removing the key
		delCtx, delCancel := context.WithTimeout(context.Background(), redisTimeout)
		defer delCancel()
		if delErr := database.RDB.Del(delCtx, redisKey).Err(); delErr != nil {
			log.Printf("WARNING: idempotency cache delete failed (key=%s): %v", key, delErr)
		}
	}

	return nil
}