Tutorial: Build a File Upload Service

package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"sync"
	"syscall"
	"time"
)

// ===========================================
// 1. TYPES & CONFIG
// ===========================================
const (
	ChunkSize    = 5 * 1024 * 1024             // 5MB
	MaxFileSize  = 5 * 1024 * 1024 * 1024      // 5GB
	SigningSecret = "super-secret-key"
)

var allowedTypes = map[string]bool{
	"image/jpeg": true, "image/png": true, "image/gif": true,
	"video/mp4": true, "video/webm": true, "application/pdf": true,
	"text/plain": true, "application/zip": true,
}

type UploadSession struct {
	ID             string         `json:"id"`
	FileName       string         `json:"fileName"`
	FileSize       int64          `json:"fileSize"`
	ContentType    string         `json:"contentType"`
	TotalChunks    int            `json:"totalChunks"`
	UploadedChunks map[int]bool   `json:"-"`
	Uploaded       []int          `json:"uploadedChunks"`
	Checksum       string         `json:"checksum"`
	Status         string         `json:"status"`
	CreatedAt      time.Time      `json:"createdAt"`
	OwnerID        string         `json:"ownerId"`
}

type FileRecord struct {
	ID          string    `json:"id"`
	UploadID    string    `json:"uploadId"`
	FileName    string    `json:"fileName"`
	FileSize    int64     `json:"fileSize"`
	ContentType string    `json:"contentType"`
	Checksum    string    `json:"checksum"`
	Status      string    `json:"status"`
	ScanResult  *string   `json:"scanResult"`
	CreatedAt   time.Time `json:"createdAt"`
}

// ===========================================
// 2. UPLOAD MANAGER
// ===========================================
type UploadManager struct {
	mu       sync.Mutex
	sessions map[string]*UploadSession
	chunks   map[string]map[int][]byte
	files    map[string]*FileRecord
	counter  int64
}

func NewUploadManager() *UploadManager {
	return &UploadManager{
		sessions: make(map[string]*UploadSession),
		chunks:   make(map[string]map[int][]byte),
		files:    make(map[string]*FileRecord),
	}
}

func (um *UploadManager) Initiate(fileName string, fileSize int64, contentType, checksum, ownerID string) (*UploadSession, error) {
	if !allowedTypes[contentType] {
		return nil, fmt.Errorf("file type %s not allowed", contentType)
	}
	if fileSize > MaxFileSize {
		return nil, fmt.Errorf("file size exceeds maximum")
	}

	um.mu.Lock()
	defer um.mu.Unlock()
	um.counter++
	totalChunks := int(math.Ceil(float64(fileSize) / float64(ChunkSize)))

	session := &UploadSession{
		ID: fmt.Sprintf("upload_%d", um.counter),
		FileName: fileName, FileSize: fileSize, ContentType: contentType,
		TotalChunks: totalChunks, UploadedChunks: make(map[int]bool),
		Checksum: checksum, Status: "uploading",
		CreatedAt: time.Now().UTC(), OwnerID: ownerID,
	}
	um.sessions[session.ID] = session
	um.chunks[session.ID] = make(map[int][]byte)
	return session, nil
}

func (um *UploadManager) UploadChunk(uploadID string, chunkNum int, data []byte, chunkChecksum string) error {
	um.mu.Lock()
	defer um.mu.Unlock()

	session := um.sessions[uploadID]
	if session == nil { return fmt.Errorf("upload not found") }
	if session.Status != "uploading" { return fmt.Errorf("upload not in uploading state") }
	if chunkNum < 0 || chunkNum >= session.TotalChunks { return fmt.Errorf("invalid chunk number") }

	hash := sha256.Sum256(data)
	actual := hex.EncodeToString(hash[:])
	if actual != chunkChecksum { return fmt.Errorf("chunk checksum mismatch") }

	um.chunks[uploadID][chunkNum] = data
	session.UploadedChunks[chunkNum] = true
	return nil
}

func (um *UploadManager) GetMissing(uploadID string) ([]int, error) {
	um.mu.Lock()
	defer um.mu.Unlock()
	session := um.sessions[uploadID]
	if session == nil { return nil, fmt.Errorf("upload not found") }
	var missing []int
	for i := 0; i < session.TotalChunks; i++ {
		if !session.UploadedChunks[i] { missing = append(missing, i) }
	}
	return missing, nil
}

func (um *UploadManager) Complete(uploadID string) (*FileRecord, error) {
	um.mu.Lock()
	session := um.sessions[uploadID]
	if session == nil { um.mu.Unlock(); return nil, fmt.Errorf("upload not found") }

	for i := 0; i < session.TotalChunks; i++ {
		if !session.UploadedChunks[i] {
			um.mu.Unlock()
			return nil, fmt.Errorf("missing chunk %d", i)
		}
	}

	session.Status = "assembling"
	chunkData := um.chunks[uploadID]

	h := sha256.New()
	var totalSize int64
	for i := 0; i < session.TotalChunks; i++ {
		h.Write(chunkData[i])
		totalSize += int64(len(chunkData[i]))
	}
	actual := hex.EncodeToString(h.Sum(nil))
	if actual != session.Checksum {
		session.Status = "failed"
		um.mu.Unlock()
		return nil, fmt.Errorf("file checksum mismatch")
	}

	um.counter++
	file := &FileRecord{
		ID: fmt.Sprintf("file_%d", um.counter), UploadID: uploadID,
		FileName: session.FileName, FileSize: totalSize,
		ContentType: session.ContentType, Checksum: actual,
		Status: "processing", CreatedAt: time.Now().UTC(),
	}
	um.files[file.ID] = file
	session.Status = "processing"
	delete(um.chunks, uploadID)
	um.mu.Unlock()

	// Async post-processing
	go func() {
		log.Printf("[SCAN] Scanning %s...", file.FileName)
		time.Sleep(2 * time.Second)
		um.mu.Lock()
		clean := "clean"
		file.ScanResult = &clean
		file.Status = "ready"
		um.mu.Unlock()
		log.Printf("[SCAN] %s is clean — file ready", file.FileName)
	}()

	return file, nil
}

func (um *UploadManager) GetFile(fileID string) *FileRecord {
	um.mu.Lock()
	defer um.mu.Unlock()
	return um.files[fileID]
}

func (um *UploadManager) GetSession(uploadID string) *UploadSession {
	um.mu.Lock()
	defer um.mu.Unlock()
	return um.sessions[uploadID]
}

func (um *UploadManager) GenerateSignedURL(fileID string, expiresIn int64) (string, error) {
	um.mu.Lock()
	file := um.files[fileID]
	um.mu.Unlock()
	if file == nil { return "", fmt.Errorf("file not found") }
	if file.Status != "ready" { return "", fmt.Errorf("file not ready") }

	expires := time.Now().Unix() + expiresIn
	payload := fmt.Sprintf("%s:%d", fileID, expires)
	mac := hmac.New(sha256.New, []byte(SigningSecret))
	mac.Write([]byte(payload))
	sig := hex.EncodeToString(mac.Sum(nil))

	return fmt.Sprintf("/files/%s?expires=%d&sig=%s", fileID, expires, sig), nil
}

// ===========================================
// 3. HTTP SERVER
// ===========================================
func writeJSON(w http.ResponseWriter, status int, data interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(data)
}

func main() {
	mgr := NewUploadManager()
	chunkPattern := regexp.MustCompile(`^/api/uploads/([^/]+)/chunks/(\d+)$`)
	completePattern := regexp.MustCompile(`^/api/uploads/([^/]+)/complete$`)
	statusPattern := regexp.MustCompile(`^/api/uploads/([^/]+)/status$`)
	filePattern := regexp.MustCompile(`^/api/files/([^/]+)$`)
	downloadPattern := regexp.MustCompile(`^/api/files/([^/]+)/download$`)

	mux := http.NewServeMux()

	mux.HandleFunc("/api/uploads/initiate", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost { writeJSON(w, 405, map[string]string{"error": "Method not allowed"}); return }
		var body struct {
			FileName    string `json:"fileName"`
			FileSize    int64  `json:"fileSize"`
			ContentType string `json:"contentType"`
			Checksum    string `json:"checksum"`
			OwnerID     string `json:"ownerId"`
		}
		json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body)
		if body.FileName == "" || body.FileSize == 0 || body.ContentType == "" || body.Checksum == "" {
			writeJSON(w, 400, map[string]string{"error": "fileName, fileSize, contentType, checksum required"}); return
		}
		s, err := mgr.Initiate(body.FileName, body.FileSize, body.ContentType, body.Checksum, body.OwnerID)
		if err != nil { writeJSON(w, 400, map[string]string{"error": err.Error()}); return }
		writeJSON(w, 201, map[string]interface{}{"uploadId": s.ID, "totalChunks": s.TotalChunks, "chunkSize": ChunkSize})
	})

	mux.HandleFunc("/api/uploads/", func(w http.ResponseWriter, r *http.Request) {
		if m := chunkPattern.FindStringSubmatch(r.URL.Path); m != nil && r.Method == http.MethodPut {
			chunkNum, _ := strconv.Atoi(m[2])
			data, _ := io.ReadAll(http.MaxBytesReader(w, r.Body, ChunkSize+1024))
			checksum := r.Header.Get("X-Chunk-Checksum")
			if checksum == "" { writeJSON(w, 400, map[string]string{"error": "X-Chunk-Checksum required"}); return }
			if err := mgr.UploadChunk(m[1], chunkNum, data, checksum); err != nil {
				writeJSON(w, 400, map[string]string{"error": err.Error()}); return
			}
			writeJSON(w, 200, map[string]interface{}{"status": "chunk_received", "chunkNum": chunkNum}); return
		}
		if m := completePattern.FindStringSubmatch(r.URL.Path); m != nil && r.Method == http.MethodPost {
			file, err := mgr.Complete(m[1])
			if err != nil { writeJSON(w, 400, map[string]string{"error": err.Error()}); return }
			writeJSON(w, 200, file); return
		}
		if m := statusPattern.FindStringSubmatch(r.URL.Path); m != nil && r.Method == http.MethodGet {
			s := mgr.GetSession(m[1])
			if s == nil { writeJSON(w, 404, map[string]string{"error": "Upload not found"}); return }
			missing, _ := mgr.GetMissing(m[1])
			writeJSON(w, 200, map[string]interface{}{"uploadId": s.ID, "status": s.Status, "missingChunks": missing, "totalChunks": s.TotalChunks}); return
		}
		writeJSON(w, 404, map[string]string{"error": "Not found"})
	})

	mux.HandleFunc("/api/files/", func(w http.ResponseWriter, r *http.Request) {
		if m := downloadPattern.FindStringSubmatch(r.URL.Path); m != nil {
			url, err := mgr.GenerateSignedURL(m[1], 3600)
			if err != nil { writeJSON(w, 400, map[string]string{"error": err.Error()}); return }
			writeJSON(w, 200, map[string]string{"downloadUrl": url}); return
		}
		if m := filePattern.FindStringSubmatch(r.URL.Path); m != nil {
			f := mgr.GetFile(m[1])
			if f == nil { writeJSON(w, 404, map[string]string{"error": "File not found"}); return }
			writeJSON(w, 200, f); return
		}
		writeJSON(w, 404, map[string]string{"error": "Not found"})
	})

	mux.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
		writeJSON(w, 200, map[string]string{"status": "ok"})
	})

	port := os.Getenv("PORT"); if port == "" { port = "3000" }
	srv := &http.Server{Addr: ":" + port, Handler: mux, ReadTimeout: 30 * time.Second, WriteTimeout: 30 * time.Second}

	go func() {
		log.Printf("File Upload Service on http://localhost:%s", port)
		if err := srv.ListenAndServe(); err != http.ErrServerClosed { log.Fatal(err) }
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	srv.Close()
}