package controller
import (
"net/http"
"strconv"
"strings"
"github.com/go-chi/chi"
httputil "github.com/goodrain/rainbond/util/http"
"github.com/sirupsen/logrus"
)
type ChunkUploadController struct{}
func (c *ChunkUploadController) InitUpload(w http.ResponseWriter, r *http.Request) {
SetPackageBuildCORSHeaders(w, r)
eventID := strings.TrimSpace(chi.URLParam(r, "eventID"))
if eventID == "" {
httputil.ReturnError(r, w, 400, "eventID is required")
return
}
var req struct {
FileName string `json:"file_name" validate:"required"`
FileSize int64 `json:"file_size" validate:"required"`
FileMD5 string `json:"file_md5"`
ChunkSize int `json:"chunk_size"`
}
if !httputil.ValidatorRequestStructAndErrorResponse(r, w, &req, nil) {
return
}
manager := GetChunkUploadManager()
session, err := manager.InitUploadSession(eventID, req.FileName, req.FileSize, req.FileMD5, req.ChunkSize)
if err != nil {
logrus.Errorf("Failed to init upload session: %v", err)
httputil.ReturnError(r, w, 500, "Failed to initialize upload: "+err.Error())
return
}
uploadedChunks := manager.parseUploadedChunks(session.UploadedChunks)
response := map[string]interface{}{
"session_id": session.ID,
"chunk_size": session.ChunkSize,
"total_chunks": session.TotalChunks,
"uploaded_chunks": uploadedChunks,
"status": session.Status,
}
logrus.Infof("Initialized upload session: %s for event: %s, file: %s", session.ID, eventID, req.FileName)
httputil.ReturnSuccess(r, w, response)
}
func (c *ChunkUploadController) UploadChunk(w http.ResponseWriter, r *http.Request) {
SetPackageBuildCORSHeaders(w, r)
eventID := strings.TrimSpace(chi.URLParam(r, "eventID"))
if eventID == "" {
httputil.ReturnError(r, w, 400, "eventID is required")
return
}
sessionID := r.FormValue("session_id")
chunkIndexStr := r.FormValue("chunk_index")
if sessionID == "" || chunkIndexStr == "" {
httputil.ReturnError(r, w, 400, "session_id and chunk_index are required")
return
}
chunkIndex, err := strconv.Atoi(chunkIndexStr)
if err != nil {
httputil.ReturnError(r, w, 400, "invalid chunk_index")
return
}
reader, header, err := r.FormFile("file")
if err != nil {
logrus.Errorf("Failed to parse chunk file: %v", err)
httputil.ReturnError(r, w, 400, "Failed to parse chunk file")
return
}
defer reader.Close()
logrus.Debugf("Receiving chunk %d for session %s, size: %d", chunkIndex, sessionID, header.Size)
manager := GetChunkUploadManager()
if err := manager.SaveChunk(sessionID, chunkIndex, reader); err != nil {
logrus.Errorf("Failed to save chunk: %v", err)
httputil.ReturnError(r, w, 500, "Failed to save chunk: "+err.Error())
return
}
status, err := manager.GetUploadStatus(sessionID)
if err != nil {
logrus.Errorf("Failed to get upload status: %v", err)
httputil.ReturnError(r, w, 500, "Failed to get status")
return
}
response := map[string]interface{}{
"session_id": sessionID,
"chunk_index": chunkIndex,
"received_size": header.Size,
"uploaded_chunks": status.UploadedChunks,
"progress": status.Progress,
}
httputil.ReturnSuccess(r, w, response)
}
func (c *ChunkUploadController) CompleteUpload(w http.ResponseWriter, r *http.Request) {
SetPackageBuildCORSHeaders(w, r)
eventID := strings.TrimSpace(chi.URLParam(r, "eventID"))
if eventID == "" {
httputil.ReturnError(r, w, 400, "eventID is required")
return
}
var req struct {
SessionID string `json:"session_id" validate:"required"`
}
if !httputil.ValidatorRequestStructAndErrorResponse(r, w, &req, nil) {
return
}
manager := GetChunkUploadManager()
filePath, err := manager.CompleteUpload(req.SessionID)
if err != nil {
logrus.Errorf("Failed to complete upload: %v", err)
httputil.ReturnError(r, w, 500, "Failed to complete upload: "+err.Error())
return
}
response := map[string]interface{}{
"file_path": filePath,
"status": "completed",
}
logrus.Infof("Upload completed for session: %s, file: %s", req.SessionID, filePath)
httputil.ReturnSuccess(r, w, response)
}
func (c *ChunkUploadController) GetUploadStatus(w http.ResponseWriter, r *http.Request) {
SetPackageBuildCORSHeaders(w, r)
sessionID := strings.TrimSpace(chi.URLParam(r, "sessionID"))
if sessionID == "" {
httputil.ReturnError(r, w, 400, "sessionID is required")
return
}
manager := GetChunkUploadManager()
status, err := manager.GetUploadStatus(sessionID)
if err != nil {
logrus.Errorf("Failed to get upload status: %v", err)
httputil.ReturnError(r, w, 404, "Session not found")
return
}
httputil.ReturnSuccess(r, w, status)
}
func (c *ChunkUploadController) CancelUpload(w http.ResponseWriter, r *http.Request) {
SetPackageBuildCORSHeaders(w, r)
sessionID := strings.TrimSpace(chi.URLParam(r, "sessionID"))
if sessionID == "" {
httputil.ReturnError(r, w, 400, "sessionID is required")
return
}
manager := GetChunkUploadManager()
if err := manager.CancelUpload(sessionID); err != nil {
logrus.Errorf("Failed to cancel upload: %v", err)
httputil.ReturnError(r, w, 500, "Failed to cancel upload: "+err.Error())
return
}
logrus.Infof("Cancelled upload session: %s", sessionID)
httputil.ReturnSuccess(r, w, map[string]string{"message": "Upload cancelled"})
}
func (c *ChunkUploadController) HandleOptions(w http.ResponseWriter, r *http.Request) {
SetPackageBuildCORSHeaders(w, r)
w.WriteHeader(http.StatusOK)
}