Commilitia-Drop/internal/httpapi/relay.go

package httpapi

import (
	"errors"
	"io"
	"log/slog"
	"net/http"
	"strconv"

	"github.com/go-chi/chi/v5"

	"commilitia.net/cdrop/internal/jwtauth"
	"commilitia.net/cdrop/internal/relay"
)

// Per-chunk upload limit. Senders split files into chunks of at most this size
// (plan §M5: "sender 多次 chunked POST"). Whole-chunk read into memory before
// commit lets us cleanly reject oversized chunks without partial state.
const maxRelayChunkSize = relay.DefaultChunkLimit

func (s *Server) handleRelayChunk(w http.ResponseWriter, r *http.Request) {
	claims, _ := jwtauth.ClaimsFromContext(r.Context())
	device, _ := jwtauth.DeviceNameFromContext(r.Context())
	sessionID := chi.URLParam(r, "id")
	if sessionID == "" {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "missing session id"})
		return
	}

	// Read body fully so an oversized chunk fails *before* we touch the ring.
	body, err := io.ReadAll(io.LimitReader(r.Body, maxRelayChunkSize+1))
	if err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()})
		return
	}
	if len(body) > maxRelayChunkSize {
		writeJSON(w, http.StatusRequestEntityTooLarge,
			map[string]string{"error": "chunk exceeds 1 MiB"})
		return
	}

	sess, err := s.relay.Acquire(sessionID, claims.UserID, device)
	if err != nil {
		mapRelayError(w, err)
		return
	}

	// Resume support (plan §M10): caller sends the offset of THIS chunk's
	// first byte. We CAS against current bytesWritten so a retried chunk
	// (network blip, sender app reload) returns 409 and the X-Bytes-Received
	// header tells the client where to pick up.
	if hdr := r.Header.Get("X-Resume-Offset"); hdr != "" {
		want, perr := strconv.ParseInt(hdr, 10, 64)
		if perr != nil {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid X-Resume-Offset"})
			return
		}
		got := sess.BytesWritten()
		if want != got {
			w.Header().Set("X-Bytes-Received", strconv.FormatInt(got, 10))
			writeJSON(w, http.StatusConflict, map[string]any{
				"error":    "offset mismatch",
				"expected": got,
				"received": want,
			})
			return
		}
	}

	n, err := sess.Write(r.Context(), body)
	if err != nil {
		slog.Warn("relay chunk write failed",
			"err", err, "session", sessionID, "user", claims.UserID, "n", n)
		writeJSON(w, http.StatusGone, map[string]string{
			"error":          err.Error(),
			"bytes_received": strconv.FormatInt(sess.BytesWritten(), 10),
		})
		return
	}
	if r.Header.Get("X-End") == "true" {
		sess.CloseWriter()
	}
	w.Header().Set("X-Bytes-Received", strconv.FormatInt(sess.BytesWritten(), 10))
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) handleRelayStream(w http.ResponseWriter, r *http.Request) {
	claims, _ := jwtauth.ClaimsFromContext(r.Context())
	device, _ := jwtauth.DeviceNameFromContext(r.Context())
	sessionID := chi.URLParam(r, "id")
	if sessionID == "" {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "missing session id"})
		return
	}

	sess, err := s.relay.Acquire(sessionID, claims.UserID, device)
	if err != nil {
		mapRelayError(w, err)
		return
	}

	// The relay stream is strictly single-consumer: a second concurrent reader
	// would drain bytes out of the ring and silently corrupt the receiver's copy
	// (G2). Refuse it rather than join. ReleaseReader on return lets a legitimate
	// reconnect (network blip, page reload) take over once the first has exited.
	if !sess.AcquireReader() {
		writeJSON(w, http.StatusConflict, map[string]string{"error": "stream already active"})
		return
	}
	defer sess.ReleaseReader()

	flusher, _ := w.(http.Flusher)
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("X-Accel-Buffering", "no")
	w.WriteHeader(http.StatusOK)
	if flusher != nil {
		flusher.Flush()
	}

	buf := make([]byte, 32*1024)
	for {
		n, err := sess.Read(r.Context(), buf)
		if n > 0 {
			if _, werr := w.Write(buf[:n]); werr != nil {
				return
			}
			if flusher != nil {
				flusher.Flush()
			}
		}
		if errors.Is(err, io.EOF) {
			return
		}
		if err != nil {
			return
		}
	}
}

func mapRelayError(w http.ResponseWriter, err error) {
	switch {
	case errors.Is(err, relay.ErrTooManySessions):
		writeJSON(w, http.StatusServiceUnavailable,
			map[string]string{"error": "too many active relay sessions"})
	case errors.Is(err, relay.ErrForbidden):
		writeJSON(w, http.StatusForbidden, map[string]string{"error": "forbidden"})
	case errors.Is(err, relay.ErrNotRegistered):
		// The transfer isn't in relay mode (never fell back, or the session was
		// reaped after going idle). The client should re-issue /fallback.
		writeJSON(w, http.StatusConflict, map[string]string{"error": "relay session not active"})
	default:
		writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
	}
}