gochat/backend/internal/middleware/upload_security.go

package middleware

import (
	"fmt"
	"net/http"
	"path/filepath"
	"strings"

	"github.com/gin-gonic/gin"
	applogger "github.com/gochat/gochat/pkg/logger"
)

// UploadSecurityConfig holds configuration for upload security middleware.
type UploadSecurityConfig struct {
	// MaxFileSize is the maximum allowed file size in bytes (default 20MB).
	MaxFileSize int64
	// AllowedMIMETypes is the whitelist of allowed MIME types.
	AllowedMIMETypes map[string]bool
	// AllowedExtensions is the whitelist of allowed file extensions.
	AllowedExtensions map[string]bool
	// MaxZipCompressionRatio is the maximum compression ratio for ZIP files
	// to prevent zip bomb attacks (default 10.0, meaning compressed size * 10 < uncompressed size).
	MaxZipCompressionRatio float64
	// MaxZipEntries is the maximum number of files inside a ZIP archive (default 1000).
	MaxZipEntries int64
}

// DefaultUploadSecurityConfig returns a sensible default configuration.
func DefaultUploadSecurityConfig() *UploadSecurityConfig {
	return &UploadSecurityConfig{
		MaxFileSize: 20 * 1024 * 1024, // 20MB
		AllowedMIMETypes: map[string]bool{
			"image/jpeg":              true,
			"image/png":               true,
			"image/gif":               true,
			"application/pdf":         true,
			"text/plain":              true,
			"audio/mpeg":              true,
			"audio/ogg":               true,
			"video/mp4":               true,
			"application/zip":         true,
			"application/x-zip-compressed": true,
		},
		AllowedExtensions: map[string]bool{
			".jpeg": true,
			".jpg":  true,
			".png":  true,
			".gif":  true,
			".pdf":  true,
			".txt":  true,
			".mp3":  true,
			".ogg":  true,
			".mp4":  true,
			".zip":  true,
		},
		MaxZipCompressionRatio: 10.0,
		MaxZipEntries:          1000,
	}
}

// sanitizeFilename removes dangerous characters from uploaded filenames.
// Removes path components, null bytes, and special characters that could
// be used for directory traversal or other attacks.
func sanitizeFilename(filename string) string {
	// Remove any directory path components
	filename = filepath.Base(filename)

	// Remove null bytes
	filename = strings.ReplaceAll(filename, "\x00", "")

	// Remove potentially dangerous characters: backslash, control chars
	sanitized := strings.Map(func(r rune) rune {
		if r < 32 && r != '\t' && r != '\n' && r != '\r' {
			return -1 // drop control characters
		}
		switch r {
		case '\\', '/', ':', '*', '?', '"', '<', '>', '|':
			return -1 // drop path/dangerous characters
		}
		return r
	}, filename)

	// Trim whitespace and dots from the edges (prevent hidden files or extension tricks)
	sanitized = strings.Trim(sanitized, " .")

	// Ensure we have a valid filename
	if sanitized == "" {
		sanitized = "upload"
	}

	return sanitized
}

// validateFileExtension checks if the file extension is in the whitelist.
func validateFileExtension(filename string, cfg *UploadSecurityConfig) bool {
	ext := strings.ToLower(filepath.Ext(filename))
	return cfg.AllowedExtensions[ext]
}

// validateMIMEType checks if the detected MIME type is in the whitelist.
func validateMIMEType(mimeType string, cfg *UploadSecurityConfig) bool {
	// Normalize: strip parameters like charset
	mimeType = strings.Split(mimeType, ";")[0]
	mimeType = strings.TrimSpace(strings.ToLower(mimeType))
	return cfg.AllowedMIMETypes[mimeType]
}

// detectMIMEType reads the first 512 bytes to detect the actual content type.
func detectMIMEType(data []byte) string {
	return http.DetectContentType(data)
}

// UploadSecurityMiddleware validates uploaded files against security rules:
//   - Enforces file size limits (default 20MB)
//   - Validates file extensions against a whitelist
//   - Sanitizes filenames to prevent directory traversal
//   - Verifies actual MIME type matches declared type
//   - Inspects ZIP archives for zip bomb and zip fragment attacks
func UploadSecurityMiddleware(cfg *UploadSecurityConfig) gin.HandlerFunc {
	if cfg == nil {
		cfg = DefaultUploadSecurityConfig()
	}

	return func(c *gin.Context) {
		form, err := c.MultipartForm()
		if err != nil {
			// No multipart form — skip this middleware (not an upload request)
			c.Next()
			return
		}

		files := form.File["file"]
		if len(files) == 0 {
			// Check other common field names
			for key, fileHeaders := range form.File {
				if key == "attachment" || key == "upload" || key == "avatar" || key == "media" {
					files = fileHeaders
					break
				}
			}
			if len(files) == 0 {
				c.Next()
				return
			}
		}

		for _, fileHeader := range files {
			// 1. File size validation
			if fileHeader.Size > cfg.MaxFileSize {
				applogger.L().Errorf("Upload rejected: file %s exceeds size limit (%d > %d bytes)",
					fileHeader.Filename, fileHeader.Size, cfg.MaxFileSize)
				c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, gin.H{
					"error": fmt.Sprintf("file size exceeds maximum allowed size of %d bytes", cfg.MaxFileSize),
				})
				return
			}

			// 2. Filename sanitization
			sanitized := sanitizeFilename(fileHeader.Filename)
			if sanitized != fileHeader.Filename {
				applogger.L().Errorf("Upload rejected: filename contains dangerous characters: original=%s sanitized=%s",
					fileHeader.Filename, sanitized)
				c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
					"error": "filename contains invalid or dangerous characters",
				})
				return
			}

			// 3. File extension whitelist
			if !validateFileExtension(fileHeader.Filename, cfg) {
				applogger.L().Errorf("Upload rejected: file extension not allowed: %s", fileHeader.Filename)
				c.AbortWithStatusJSON(http.StatusUnsupportedMediaType, gin.H{
					"error": fmt.Sprintf("file type not allowed: %s", filepath.Ext(fileHeader.Filename)),
				})
				return
			}

			// 4. Open file to validate MIME type and inspect ZIP archives
			file, err := fileHeader.Open()
			if err != nil {
				applogger.L().Errorf("Upload rejected: cannot open file %s: %v", fileHeader.Filename, err)
				c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
					"error": "cannot read uploaded file",
				})
				return
			}

			// Read first 512 bytes for MIME detection
			buf := make([]byte, 512)
			n, err := file.Read(buf)
			if err != nil && n == 0 {
				file.Close()
				applogger.L().Errorf("Upload rejected: cannot read file content for %s: %v", fileHeader.Filename, err)
				c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
					"error": "cannot read file content",
				})
				return
			}

			// 5. MIME type validation — verify actual content matches whitelist
			detectedMIME := detectMIMEType(buf[:n])
			if !validateMIMEType(detectedMIME, cfg) {
				file.Close()
				applogger.L().Errorf("Upload rejected: MIME type not allowed: detected=%s file=%s",
					detectedMIME, fileHeader.Filename)
				c.AbortWithStatusJSON(http.StatusUnsupportedMediaType, gin.H{
					"error": fmt.Sprintf("file content type not allowed: %s", detectedMIME),
				})
				return
			}

			// 6. ZIP bomb / zip fragment protection
			ext := strings.ToLower(filepath.Ext(fileHeader.Filename))
			if ext == ".zip" {
				if err := inspectZipArchive(file, fileHeader.Size, cfg); err != nil {
					file.Close()
					applogger.L().Errorf("Upload rejected: ZIP inspection failed for %s: %v", fileHeader.Filename, err)
					c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
						"error": err.Error(),
					})
					return
				}
			}

			file.Close()
		}

		// Store sanitized filenames in context for downstream handlers
		c.Set("upload_verified", true)
		c.Next()
	}
}