IdempotencyIdempotency

@Component
@RequiredArgsConstructor
public class IdempotencyFilter extends OncePerRequestFilter {
    private final JdbcTemplate jdbc;

    @Override
    protected boolean shouldNotFilter(HttpServletRequest req) {
        return !"POST".equals(req.getMethod()) || req.getHeader("Idempotency-Key") == null;
    }

    @Override
    protected void doFilterInternal(HttpServletRequest req, HttpServletResponse res,
                                    FilterChain chain) throws ServletException, IOException {
        String key = req.getHeader("Idempotency-Key");
        String user = req.getUserPrincipal().getName();
        CachedBodyRequest body = new CachedBodyRequest(req);   // body read once, re-readable
        String hash = DigestUtils.sha256Hex(body.bytes());

        int claimed = jdbc.update("""1
            INSERT INTO idempotency_keys(user_id, key, request_hash, status, locked_at, expires_at)
            VALUES (?, ?, ?, 'IN_PROGRESS', now(), now() + interval '24 hours')
            ON CONFLICT (user_id, key) DO NOTHING""", user, key, hash);2

        if (claimed == 0) {
            StoredKey k = jdbc.queryForObject("""
                SELECT request_hash, status, response_code, response_body
                FROM idempotency_keys WHERE user_id = ? AND key = ?""", StoredKey.MAPPER, user, key);
            if (!k.requestHash().equals(hash)) { res.sendError(422, "Key reused with another body"); return; }3
            if (k.inProgress()) { res.setHeader("Retry-After", "1"); res.sendError(409); return; }
            res.setStatus(k.responseCode());
            res.setHeader("Idempotent-Replayed", "true");
            res.getOutputStream().write(k.responseBody());
            return;
        }

        ContentCachingResponseWrapper out = new ContentCachingResponseWrapper(res);
        try {
            chain.doFilter(body, out);
        } catch (ServletException | IOException | RuntimeException e) {
            release(user, key);
            throw e;
        }
        if (out.getStatus() >= 500) release(user, key);4
        else jdbc.update("""
            UPDATE idempotency_keys SET status = 'COMPLETED', response_code = ?, response_body = ?
            WHERE user_id = ? AND key = ?""", out.getStatus(), out.getContentAsByteArray(), user, key);
        out.copyBodyToResponse();
    }

    private void release(String user, String key) {
        jdbc.update("DELETE FROM idempotency_keys WHERE user_id = ? AND key = ?", user, key);
    }
}

type recorder struct {1
	http.ResponseWriter
	status int
	buf    bytes.Buffer
}

func (r *recorder) WriteHeader(c int)           { r.status = c; r.ResponseWriter.WriteHeader(c) }
func (r *recorder) Write(b []byte) (int, error) { r.buf.Write(b); return r.ResponseWriter.Write(b) }

func Idempotency(db *sql.DB, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		key := r.Header.Get("Idempotency-Key")
		if r.Method != http.MethodPost || key == "" {
			next.ServeHTTP(w, r)
			return
		}
		ctx, user := r.Context(), auth.UserID(r.Context())
		body, err := io.ReadAll(io.LimitReader(r.Body, 1<<20))2
		if err != nil {
			http.Error(w, "bad body", http.StatusBadRequest)
			return
		}
		r.Body = io.NopCloser(bytes.NewReader(body))
		sum := sha256.Sum256(body)
		hash := hex.EncodeToString(sum[:])

		res, err := db.ExecContext(ctx, `INSERT INTO idempotency_keys3
			(user_id, key, request_hash, status, locked_at, expires_at)
			VALUES ($1, $2, $3, 'IN_PROGRESS', now(), now() + interval '24 hours')
			ON CONFLICT (user_id, key) DO NOTHING`, user, key, hash)
		if err != nil {
			http.Error(w, "store unavailable", http.StatusServiceUnavailable)
			return
		}
		if n, _ := res.RowsAffected(); n == 0 {
			replay(ctx, db, w, user, key, hash) // 422, 409 or the stored response
			return
		}

		rec := &recorder{ResponseWriter: w, status: http.StatusOK}
		next.ServeHTTP(rec, r)4
		if rec.status >= 500 {
			db.ExecContext(ctx, `DELETE FROM idempotency_keys WHERE user_id = $1 AND key = $2`, user, key)
			return
		}
		db.ExecContext(ctx, `UPDATE idempotency_keys
			SET status = 'COMPLETED', response_code = $1, response_body = $2
			WHERE user_id = $3 AND key = $4`, rec.status, rec.buf.Bytes(), user, key)
	})
}