pipely/cmd/server/main.go

package main

import (
    "context"
    "errors"
    "log"
    "net/http"
    "os"
    "os/signal"
    "syscall"
    "time"

    "github.com/gin-gonic/gin"

    "onixbyte.com/pipely/internal/artefact"
    "onixbyte.com/pipely/internal/auth"
    "onixbyte.com/pipely/internal/config"
    "onixbyte.com/pipely/internal/database"
    "onixbyte.com/pipely/internal/delta"
    "onixbyte.com/pipely/internal/deployment"
    "onixbyte.com/pipely/internal/device"
    "onixbyte.com/pipely/internal/download"
    "onixbyte.com/pipely/internal/middleware"
)

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

    // Connect to PostgreSQL (auto-creates database if needed)
    db, err := database.Connect(cfg.DatabaseURL)
    if err != nil {
        log.Fatalf("Database connection failed: %v", err)
    }

    // Auto-migrate tables
    if err := database.AutoMigrate(db); err != nil {
        log.Fatalf("Auto-migrate failed: %v", err)
    }

    // Initialise auth service and create default admin
    authSvc := auth.NewService(db, []byte(cfg.JWTSecret), cfg.JWTExpiryHrs)
    if err := authSvc.CreateInitialAdmin("admin", "admin"); err != nil {
        log.Printf("Warning: failed to create initial admin: %v", err)
    }
    log.Println("[main] Default admin user ready (admin / admin)")

    // Ensure artefact storage directory exists
    if err := os.MkdirAll(cfg.ArtefactDir, 0755); err != nil {
        log.Fatalf("Failed to create artefact directory: %v", err)
    }

    // Global rate limiter
    tokenBucket := middleware.NewTokenBucket(cfg.RateLimitRPS, cfg.RateLimitRPS*2)
    go func() {
        ticker := time.NewTicker(10 * time.Minute)
        for range ticker.C {
            tokenBucket.Cleanup(30 * time.Minute)
        }
    }()

    // Gated rollout semaphore
    downloadGate := middleware.NewGatedSemaphore(cfg.MaxConcurrent)

    // Offline device detection
    go func() {
        ticker := time.NewTicker(5 * time.Minute)
        for range ticker.C {
            offline, err := database.GetOfflineDevices(db, 10*time.Minute)
            if err != nil {
                log.Printf("[health] Offline check failed: %v", err)
                continue
            }
            ids := make([]string, len(offline))
            for i, d := range offline {
                ids[i] = d.DeviceID
            }
            if len(ids) > 0 {
                if err := database.MarkDevicesOffline(db, ids); err != nil {
                    log.Printf("[health] Mark offline failed: %v", err)
                } else {
                    log.Printf("[health] Marked %d devices as offline", len(ids))
                }
            }
        }
    }()

    // Initialise handlers
    authHandler := auth.NewHandler(authSvc)
    deviceHandler := device.NewHandler(db)
    artefactHandler := artefact.NewHandler(db, cfg.ArtefactDir)
    deploymentHandler := deployment.NewHandler(db)
    deltaHandler := delta.NewHandler(cfg.ArtefactDir)
    downloadHandler := download.NewHandler(db, cfg.ArtefactDir, downloadGate, cfg.BandwidthBPS)

    // Setup Gin
    gin.SetMode(gin.ReleaseMode)
    r := gin.New()
    r.Use(gin.Recovery())

    // CORS — withCredentials requires explicit origin, not wildcard
    r.Use(func(c *gin.Context) {
        origin := c.GetHeader("Origin")
        if origin != "" {
            c.Header("Access-Control-Allow-Origin", origin)
            c.Header("Access-Control-Allow-Credentials", "true")
        } else {
            c.Header("Access-Control-Allow-Origin", "*")
        }
        c.Header("Vary", "Origin")
        c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
        c.Header("Access-Control-Allow-Headers", "Origin, Content-Type, Accept, Authorization, Range")
        if c.Request.Method == "OPTIONS" {
            c.AbortWithStatus(http.StatusNoContent)
            return
        }
        c.Next()
    })

    // -- Public routes (daemon-facing) --
    api := r.Group("/api")
    {
        api.POST("/devices/register", deviceHandler.Register)
        api.POST("/devices/heartbeat", deviceHandler.Heartbeat)
        api.GET("/devices/check-update", deviceHandler.CheckUpdate)
    }

    // -- Public auth routes --
    api.POST("/auth/login", authHandler.Login)

    // -- Authenticated admin routes --
    admin := api.Group("")
    admin.Use(middleware.AuthMiddleware(authSvc))
    {
        api.POST("/auth/logout", authHandler.Logout)

        admin.GET("/devices", deviceHandler.List)
        admin.GET("/devices/:deviceId", deviceHandler.Get)

        admin.POST("/artefacts", artefactHandler.Upload)
        admin.GET("/artefacts", artefactHandler.List)
        admin.GET("/artefacts/:id", artefactHandler.Get)

        admin.POST("/delta/generate", deltaHandler.GenerateDelta)

        admin.POST("/deployments", deploymentHandler.Create)
        admin.GET("/deployments", deploymentHandler.List)
        admin.GET("/deployments/:id", deploymentHandler.Get)
        admin.POST("/deployments/:id/activate", deploymentHandler.Activate)
        admin.POST("/deployments/:id/pause", deploymentHandler.Pause)
        admin.GET("/deployments/:id/devices", deploymentHandler.ListDevices)
        admin.POST("/deployments/:id/rollback", deploymentHandler.Rollback)

        admin.GET("/admin/users", authHandler.ListUsers)
        admin.POST("/admin/users", authHandler.CreateUser)
        admin.DELETE("/admin/users/:id", authHandler.DeleteUser)
        admin.PUT("/admin/users/:id/password", authHandler.ChangePassword)
    }

    // -- Download endpoint (rate-limited + gated) --
    packages := api.Group("/packages")
    packages.Use(middleware.RateLimitMiddleware(tokenBucket))
    packages.Use(middleware.GatedMiddleware(downloadGate))
    {
        packages.GET("/download/:id", downloadHandler.HandleDownload)
    }

    // -- Device status update --
    api.POST("/deployments/:id/status", deploymentHandler.UpdateStatus)

    // -- Static admin panel (SPA fallback) --
    // Enable with SERVE_SPA=true (default). Set to false when using Caddy/Nginx as reverse proxy.
    if cfg.ServeSPA {
        r.Static("/assets", "./web/dist/assets")
        r.GET("/favicon.ico", func(c *gin.Context) { c.File("./web/dist/favicon.ico") })
        r.NoRoute(func(c *gin.Context) {
            c.File("./web/dist/index.html")
        })
    }

    // Graceful shutdown
    srv := &http.Server{
        Addr:    ":" + cfg.ServerPort,
        Handler: r,
    }

    go func() {
        log.Printf("[main] Pipely listening on port :%s", cfg.ServerPort)
        log.Printf("[main] Max concurrent downloads: %d, bandwidth per stream: %d B/s", cfg.MaxConcurrent, cfg.BandwidthBPS)
        if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
            log.Fatalf("Server error: %v", err)
        }
    }()

    quit := make(chan os.Signal, 1)
    signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
    <-quit
    log.Println("[main] Shutting down...")

    ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
    defer cancel()
    if err := srv.Shutdown(ctx); err != nil {
        log.Fatalf("Forced shutdown: %v", err)
    }
    log.Println("[main] Server stopped")
}