alanyeung-s3/aws/main.go

package main

import (
	"log"
	"net/http"
	"strings"

	"aws-sts-mock/internal/config"
	"aws-sts-mock/internal/handler"
	"aws-sts-mock/internal/kvdb"
	"aws-sts-mock/internal/middleware"
	"aws-sts-mock/internal/service"
	"aws-sts-mock/internal/storage"
	"aws-sts-mock/pkg/sigv4"
)

func main() {
	// Load configuration
	cfg := config.Load()

	// Initialize BoltDB
	log.Println("Initializing BoltDB...")
	db, err := kvdb.NewBoltKVDB(cfg.DBPath)
	if err != nil {
		log.Fatalf("Failed to initialize BoltDB: %v", err)
	}
	defer db.Close()

	// Initialize default users
	log.Println("Initializing default users...")
	if err := db.InitializeDefaultUsers(); err != nil {
		log.Fatalf("Failed to initialize default users: %v", err)
	}

	log.Println("Default users initialized:")
	log.Println("  - User 1: AKIAIOSFODNN7EXAMPLE (Account: 123456789012)")
	log.Println("  - User 2: AKIAJSIE27KKMHXI3BJQ (Account: 987654321098)")

	// Initialize user-aware storage
	userStorage, err := storage.NewUserAwareStorage(cfg.UploadsDir)
	if err != nil {
		log.Fatalf("Failed to initialize storage: %v", err)
	}

	// Initialize services
	s3Service := service.NewS3Service(userStorage, db)
	stsService := service.NewSTSService()

	// Initialize handlers
	s3Handler := handler.NewS3Handler(s3Service)
	stsHandler := handler.NewSTSHandler(stsService)
	healthHandler := handler.NewHealthHandler()
	publicViewHandler := handler.NewPublicViewHandler(userStorage, db)

	// Initialize middleware
	sigv4Middleware := sigv4.ValidateSigV4Middleware
	userValidationMiddleware := middleware.UserValidationMiddleware(db)

	// Setup routes
	mux := http.NewServeMux()

	// Health check endpoint - no authentication
	mux.HandleFunc("/health", healthHandler.Handle)

	// Main endpoint - handles both authenticated and public requests
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Check if Authorization header exists (indicates authenticated request)
		authHeader := r.Header.Get("Authorization")

		if authHeader != "" {
			// Authenticated request - determine if STS or S3 based on Content-Type
			// STS requests use application/x-www-form-urlencoded
			// S3 requests typically don't, or use application/xml
			contentType := r.Header.Get("Content-Type")

			// If it's a POST with form-urlencoded content type, it's likely STS
			if r.Method == "POST" && strings.Contains(contentType, "application/x-www-form-urlencoded") {
				// STS request
				sigv4Handler := sigv4Middleware(func(w http.ResponseWriter, r *http.Request) {
					userValidationMiddleware(http.HandlerFunc(stsHandler.Handle)).ServeHTTP(w, r)
				})
				sigv4Handler(w, r)
			} else {
				// S3 request
				sigv4Handler := sigv4Middleware(func(w http.ResponseWriter, r *http.Request) {
					userValidationMiddleware(http.HandlerFunc(s3Handler.Handle)).ServeHTTP(w, r)
				})
				sigv4Handler(w, r)
			}
		} else {
			// Unauthenticated request - try public viewing
			publicViewHandler.Handle(w, r)
		}
	})

	// Start server
	log.Printf("========================================")
	log.Printf("AWS STS/S3 Mock Server with BoltDB")
	log.Printf("========================================")
	log.Printf("Server running on port: %s", cfg.Port)
	log.Printf("Database: %s", cfg.DBPath)
	log.Printf("Uploads directory: %s", cfg.UploadsDir)
	log.Printf("User isolation: ENABLED")
	log.Printf("Folder structure: ./uploads/{accountID}/{bucketID}/")
	log.Printf("")
	log.Printf("Endpoints:")
	log.Printf("  - AWS S3/STS (authenticated): http://localhost:%s/", cfg.Port)
	log.Printf("  - Public View (unauthenticated): http://localhost:%s/{bucketName}/{objectKey}", cfg.Port)
	log.Printf("  - Health Check: http://localhost:%s/health", cfg.Port)
	log.Printf("========================================")

	if err := http.ListenAndServe(":"+cfg.Port, mux); err != nil {
		log.Fatal(err)
	}
}