torrent-gateway/internal/blossom/server.go

package blossom

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"git.sovbit.dev/enki/torrentGateway/internal/config"
	"git.sovbit.dev/enki/torrentGateway/internal/proxy"
	"git.sovbit.dev/enki/torrentGateway/internal/storage"
	"golang.org/x/time/rate"
)

// Server implements a Blossom server
type Server struct {
	storage     *storage.Backend
	config      *config.BlossomServerConfig
	rateLimiter *rate.Limiter
	mux         *http.ServeMux
	smartProxy  *proxy.SmartProxy
	fullConfig  *config.Config
}

// BlobUploadResponse represents the response for blob uploads
type BlobUploadResponse struct {
	Hash      string    `json:"hash"`
	Size      int64     `json:"size"`
	Type      string    `json:"type"`
	Timestamp time.Time `json:"timestamp"`
	Message   string    `json:"message,omitempty"`
}

// ErrorResponse represents an error response
type ErrorResponse struct {
	Error   string `json:"error"`
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// NewServer creates a new Blossom server
func NewServer(storage *storage.Backend, config *config.BlossomServerConfig, fullConfig *config.Config) *Server {
	// Create rate limiter
	limiter := rate.NewLimiter(
		rate.Limit(config.RateLimit.RequestsPerMinute)/60, // requests per second
		config.RateLimit.BurstSize,
	)

	var smartProxy *proxy.SmartProxy
	if fullConfig.Proxy.Enabled {
		smartProxy = proxy.NewSmartProxy(storage, fullConfig)
	}

	server := &Server{
		storage:     storage,
		config:      config,
		rateLimiter: limiter,
		mux:         http.NewServeMux(),
		smartProxy:  smartProxy,
		fullConfig:  fullConfig,
	}

	server.setupRoutes()
	return server
}

// setupRoutes configures the HTTP routes
func (s *Server) setupRoutes() {
	// Blob download endpoint: GET /{hash}
	s.mux.HandleFunc("/", s.handleBlobRequest)
	
	// Upload endpoint: PUT /upload
	s.mux.HandleFunc("/upload", s.handleUpload)
	
	// Server info endpoint
	s.mux.HandleFunc("/info", s.handleInfo)
	
	// Health check
	s.mux.HandleFunc("/health", s.handleHealth)
}

// ServeHTTP implements http.Handler
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Apply rate limiting
	if !s.rateLimiter.Allow() {
		s.writeError(w, http.StatusTooManyRequests, "rate limit exceeded")
		return
	}

	// Add CORS headers for web clients
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "GET, PUT, POST, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")

	if r.Method == http.MethodOptions {
		w.WriteHeader(http.StatusOK)
		return
	}

	s.mux.ServeHTTP(w, r)
}

// handleBlobRequest handles GET requests for blobs
func (s *Server) handleBlobRequest(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		if r.URL.Path == "/" {
			s.handleRoot(w, r)
			return
		}
		s.writeError(w, http.StatusMethodNotAllowed, "method not allowed")
		return
	}

	// Extract hash from path
	path := strings.TrimPrefix(r.URL.Path, "/")
	if path == "" {
		s.handleRoot(w, r)
		return
	}

	// Validate hash format (should be 64 character hex)
	if len(path) != 64 || !isValidHash(path) {
		s.writeError(w, http.StatusBadRequest, "invalid hash format")
		return
	}

	// Get blob from storage
	reader, info, err := s.storage.GetBlobData(path)
	if err != nil {
		log.Printf("Error retrieving blob %s: %v", path, err)
		s.writeError(w, http.StatusInternalServerError, "internal server error")
		return
	}

	if reader == nil {
		// Try smart proxy if enabled and configured
		if s.smartProxy != nil && s.fullConfig.Proxy.Enabled {
			log.Printf("Blob %s not found in storage, trying smart proxy for chunked file", path)
			if err := s.smartProxy.ServeBlob(w, path); err != nil {
				log.Printf("Smart proxy failed for hash %s: %v", path, err)
				s.writeError(w, http.StatusNotFound, "blob not found")
				return
			}
			log.Printf("Successfully served chunked file via smart proxy: %s", path)
			return
		}
		s.writeError(w, http.StatusNotFound, "blob not found")
		return
	}
	defer reader.Close()

	// Set appropriate headers
	if info.MimeType != "" {
		w.Header().Set("Content-Type", info.MimeType)
	} else {
		w.Header().Set("Content-Type", "application/octet-stream")
	}
	
	w.Header().Set("Content-Length", strconv.FormatInt(info.Size, 10))
	w.Header().Set("Cache-Control", "public, max-age=31536000") // Cache for 1 year
	w.Header().Set("ETag", fmt.Sprintf(`"%s"`, path))

	// Check for conditional requests
	if match := r.Header.Get("If-None-Match"); match != "" {
		if strings.Contains(match, path) {
			w.WriteHeader(http.StatusNotModified)
			return
		}
	}

	// Stream the blob
	if _, err := io.Copy(w, reader); err != nil {
		log.Printf("Error streaming blob %s: %v", path, err)
		return
	}
}

// handleUpload handles PUT requests for blob uploads
func (s *Server) handleUpload(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPut && r.Method != http.MethodPost {
		s.writeError(w, http.StatusMethodNotAllowed, "method not allowed")
		return
	}

	// Check content length
	contentLength := r.ContentLength
	if contentLength <= 0 {
		s.writeError(w, http.StatusBadRequest, "content-length required")
		return
	}

	// Check max blob size
	maxSize, err := parseSize(s.config.MaxBlobSize)
	if err != nil {
		log.Printf("Error parsing max blob size: %v", err)
		maxSize = 100 * 1024 * 1024 // Default to 100MB if config invalid
	}

	if contentLength > maxSize {
		s.writeError(w, http.StatusRequestEntityTooLarge, 
			fmt.Sprintf("blob too large (max %d bytes)", maxSize))
		return
	}

	// Determine content type
	contentType := r.Header.Get("Content-Type")
	if contentType == "" {
		contentType = "application/octet-stream"
	}

	// Create a limited reader to prevent DoS
	limitedReader := io.LimitReader(r.Body, maxSize+1)

	// Store the blob using unified storage
	metadata, err := s.storage.StoreBlobAsFile(limitedReader, "blob", contentType)
	if err != nil {
		log.Printf("Error storing blob: %v", err)
		s.writeError(w, http.StatusInternalServerError, "failed to store blob")
		return
	}
	hash := metadata.Hash

	// Return success response
	response := BlobUploadResponse{
		Hash:      hash,
		Size:      contentLength,
		Type:      contentType,
		Timestamp: time.Now(),
		Message:   "blob stored successfully",
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusCreated)
	json.NewEncoder(w).Encode(response)
}

// handleInfo provides server information
func (s *Server) handleInfo(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		s.writeError(w, http.StatusMethodNotAllowed, "method not allowed")
		return
	}

	maxSize, _ := parseSize(s.config.MaxBlobSize)

	info := map[string]interface{}{
		"server":         "Blossom-BitTorrent Gateway",
		"version":        "1.0.0",
		"blossom_spec":   "draft-01",
		"max_blob_size":  maxSize,
		"supported_types": []string{"*/*"},
		"features": []string{
			"upload",
			"download",
			"rate_limiting",
			"caching",
		},
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(info)
}

// handleHealth provides health check
func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		s.writeError(w, http.StatusMethodNotAllowed, "method not allowed")
		return
	}

	health := map[string]interface{}{
		"status":    "ok",
		"timestamp": time.Now(),
		"service":   "blossom-server",
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(health)
}

// handleRoot handles requests to the root path
func (s *Server) handleRoot(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		s.writeError(w, http.StatusMethodNotAllowed, "method not allowed")
		return
	}

	info := map[string]interface{}{
		"service": "Blossom Server",
		"message": "This is a Blossom blob storage server. Use GET /{hash} to retrieve blobs or PUT /upload to store new blobs.",
		"endpoints": map[string]string{
			"upload": "PUT /upload - Upload a new blob",
			"download": "GET /{hash} - Download a blob by hash",
			"info": "GET /info - Server information",
			"health": "GET /health - Health check",
		},
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(info)
}

// writeError writes a JSON error response
func (s *Server) writeError(w http.ResponseWriter, code int, message string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	
	response := ErrorResponse{
		Error:   http.StatusText(code),
		Code:    code,
		Message: message,
	}
	
	json.NewEncoder(w).Encode(response)
}

// isValidHash checks if a string is a valid SHA-256 hash
func isValidHash(hash string) bool {
	if len(hash) != 64 {
		return false
	}
	
	for _, c := range hash {
		if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F')) {
			return false
		}
	}
	
	return true
}

// parseSize parses size strings like "100MB", "2GB", etc.
func parseSize(sizeStr string) (int64, error) {
	if sizeStr == "" {
		return 100 * 1024 * 1024, nil // Default 100MB if not configured
	}
	
	var size int64
	var unit string
	n, err := fmt.Sscanf(sizeStr, "%d%s", &size, &unit)
	if err != nil || n != 2 {
		return 0, fmt.Errorf("invalid size format: %s", sizeStr)
	}

	switch strings.ToUpper(unit) {
	case "B":
		return size, nil
	case "KB", "K":
		return size * 1024, nil
	case "MB", "M":
		return size * 1024 * 1024, nil
	case "GB", "G":
		return size * 1024 * 1024 * 1024, nil
	default:
		return 0, fmt.Errorf("unknown unit: %s", unit)
	}
}

// Start starts the Blossom server
func (s *Server) Start() error {
	addr := fmt.Sprintf(":%d", s.config.Port)
	log.Printf("Starting Blossom server on port %d", s.config.Port)
	return http.ListenAndServe(addr, s)
}