short-link/backend/main.go

package main

import (
	"context"
	"flag"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

func main() {
	port := flag.String("port", "8819", "listen port")
	dbPath := flag.String("db", "shortlink.db", "sqlite database path")
	baseURL := flag.String("base-url", "", "base url for short links (e.g. https://xieyaxin.top:8899)")
	flag.Parse()

	store, err := NewStore(*dbPath)
	if err != nil {
		log.Fatalf("failed to open database: %v", err)
	}
	defer store.Close()

	adminToken := strings.TrimSpace(os.Getenv("ADMIN_TOKEN"))
	h := NewHandler(store, *baseURL, adminToken)

	mux := http.NewServeMux()
	mux.HandleFunc("POST /api/shorten", h.HandleCreate)
	mux.HandleFunc("GET /{code}", h.HandleRedirect)

	// Admin routes
	mux.HandleFunc("GET /api/admin/links", h.requireAdmin(h.HandleListLinks))
	mux.HandleFunc("PUT /api/admin/links/{code}", h.requireAdmin(h.HandleUpdateLink))
	mux.HandleFunc("DELETE /api/admin/links/{code}", h.requireAdmin(h.HandleDeleteLink))
	mux.HandleFunc("GET /api/admin/export", h.requireAdmin(h.HandleExport))

	srv := &http.Server{
		Addr:         ":" + *port,
		Handler:      mux,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 10 * time.Second,
		IdleTimeout:  60 * time.Second,
	}

	// Graceful shutdown — wait for SIGINT/SIGTERM, then drain connections
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		log.Printf("listening on :%s", *port)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("server error: %v", err)
		}
	}()

	<-quit
	log.Println("shutting down...")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := srv.Shutdown(ctx); err != nil {
		log.Fatalf("forced shutdown: %v", err)
	}

	log.Println("bye")
}