main.go

// Copyright 2025 Edgeless Systems GmbH
// SPDX-License-Identifier: BUSL-1.1

package main

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/containerd/ttrpc"
	"github.com/edgelesssys/contrast/imagestore/internal/service"
	"github.com/edgelesssys/contrast/internal/katacomponents"
	"github.com/spf13/cobra"
	"golang.org/x/sync/errgroup"
)

var version = "0.0.0-dev"

func main() {
	if err := newRootCmd().Execute(); err != nil {
		os.Exit(1)
	}
}

func newRootCmd() *cobra.Command {
	return &cobra.Command{
		Use:          "imagestore",
		Short:        "securely mount a block device",
		Version:      version,
		SilenceUsage: true,
		RunE:         run,
	}
}

func run(cmd *cobra.Command, _ []string) error {
	log := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
		Level: slog.LevelDebug,
	}))
	ctx, cancel := signal.NotifyContext(cmd.Context(), syscall.SIGTERM, syscall.SIGINT)
	defer cancel()

	fmt.Fprintf(os.Stderr, "Contrast imagestore %s\n", version)
	fmt.Fprintln(os.Stderr, "Report issues at https://github.com/edgelesssys/contrast/issues")

	if err := os.MkdirAll(filepath.Dir(katacomponents.SecuremountSocket), os.ModePerm); err != nil {
		return fmt.Errorf("creating directory for socket: %w", err)
	}
	if err := os.Remove(katacomponents.SecuremountSocket); err != nil && !errors.Is(err, os.ErrNotExist) {
		return fmt.Errorf("removing existing socket: %w", err)
	}

	l, err := (&net.ListenConfig{}).Listen(ctx, "unix", katacomponents.SecuremountSocket)
	if err != nil {
		return fmt.Errorf("listening on socket: %w", err)
	}
	defer l.Close()
	defer os.RemoveAll(katacomponents.SecuremountSocket)

	s, err := ttrpc.NewServer()
	if err != nil {
		return fmt.Errorf("creating ttRPC server: %w", err)
	}
	defer s.Close()

	katacomponents.RegisterSecureMountServiceService(s, &service.SecureImageStoreService{Logger: log})
	eg, ctxEg := errgroup.WithContext(ctx)

	eg.Go(func() error {
		log.Info("Started imagestore", "socket", katacomponents.SecuremountSocket)
		log.Info("Waiting for imagestore request...")
		if err := s.Serve(ctxEg, l); err != nil {
			return fmt.Errorf("starting the ttRPC server: %w", err)
		}
		return nil
	})

	eg.Go(func() error {
		<-ctxEg.Done()
		if ctx.Err() != nil {
			log.Info("Received signal, shutting down.")
		} else {
			log.Info("Unexpected shutdown", "err", ctxEg.Err())
		}
		ctxCleanup, cancel := context.WithTimeout(context.Background(), 20*time.Second)
		defer cancel()

		if err := s.Shutdown(ctxCleanup); err != nil { //nolint:contextcheck
			return fmt.Errorf("shutting down the ttRPC server: %w", err)
		}
		return nil
	})

	return eg.Wait()
}