at-container-registry/pkg/labeler/server.go

package labeler

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"atcr.io/pkg/atproto"
	indigooauth "github.com/bluesky-social/indigo/atproto/auth/oauth"
	"github.com/go-chi/chi/v5"
)

// Server is the labeler HTTP server.
type Server struct {
	config    *Config
	db        *sql.DB
	router    chi.Router
	clientApp *indigooauth.ClientApp
	auth      *Auth
}

// NewServer creates a new labeler server.
func NewServer(cfg *Config) (*Server, error) {
	db, err := OpenDB(cfg.Labeler.DBPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}

	publicURL := cfg.PublicURL()

	// Set up OAuth client for admin login
	oauthStore := indigooauth.NewMemStore()
	scopes := []string{"atproto"}

	var oauthConfig indigooauth.ClientConfig
	var redirectURI string

	u, err := url.Parse(publicURL)
	if err != nil {
		return nil, fmt.Errorf("invalid public URL: %w", err)
	}

	host := u.Hostname()
	if isLocalhost(host) {
		port := u.Port()
		if port == "" {
			port = "5002"
		}
		oauthBaseURL := "http://127.0.0.1:" + port
		redirectURI = oauthBaseURL + "/auth/oauth/callback"
		oauthConfig = indigooauth.NewLocalhostConfig(redirectURI, scopes)
	} else {
		clientID := publicURL + "/oauth-client-metadata.json"
		redirectURI = publicURL + "/auth/oauth/callback"
		oauthConfig = indigooauth.NewPublicConfig(clientID, redirectURI, scopes)
	}

	clientApp := indigooauth.NewClientApp(&oauthConfig, oauthStore)
	clientApp.Dir = atproto.GetDirectory()

	auth := NewAuth(cfg.Labeler.OwnerDID)

	s := &Server{
		config:    cfg,
		db:        db,
		clientApp: clientApp,
		auth:      auth,
	}

	s.setupRoutes()
	return s, nil
}

func (s *Server) setupRoutes() {
	r := chi.NewRouter()

	// DID document
	r.Get("/.well-known/did.json", s.handleDIDDocument)

	// OAuth client metadata
	r.Get("/oauth-client-metadata.json", s.handleClientMetadata)

	// Auth routes (public)
	r.Get("/auth/login", s.handleLogin)
	r.Get("/auth/oauth/authorize", s.handleAuthorize)
	r.Get("/auth/oauth/callback", s.handleCallback)
	r.Get("/auth/logout", s.handleLogout)

	// XRPC endpoints (public)
	r.Get("/xrpc/com.atproto.label.subscribeLabels", s.handleSubscribeLabels)
	r.Get("/xrpc/com.atproto.label.queryLabels", s.handleQueryLabels)

	// Protected routes (require owner)
	r.Group(func(r chi.Router) {
		r.Use(s.auth.RequireOwner)

		r.Get("/", s.handleDashboard)
		r.Get("/takedown", s.handleTakedownForm)
		r.Post("/takedown", s.handleTakedownSubmit)
		r.Post("/reverse", s.handleReverse)
	})

	s.router = r
}

// Serve starts the HTTP server with graceful shutdown.
func (s *Server) Serve() error {
	slog.Info("Starting labeler service",
		"addr", s.config.Labeler.Addr,
		"public_url", s.config.PublicURL(),
		"did", s.config.DID(),
		"owner", s.config.Labeler.OwnerDID,
	)

	srv := &http.Server{
		Addr:    s.config.Labeler.Addr,
		Handler: s.router,
	}

	// Graceful shutdown
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	errCh := make(chan error, 1)
	go func() {
		errCh <- srv.ListenAndServe()
	}()

	select {
	case err := <-errCh:
		if err != http.ErrServerClosed {
			return err
		}
	case <-ctx.Done():
		slog.Info("Shutting down labeler service")
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5000000000) // 5s
		defer cancel()
		if err := srv.Shutdown(shutdownCtx); err != nil {
			return fmt.Errorf("shutdown error: %w", err)
		}
	}

	s.db.Close()
	return nil
}

func isLocalhost(host string) bool {
	return host == "localhost" || host == "127.0.0.1" || strings.HasPrefix(host, "192.168.")
}