cheburnet/cmd/cmd.go

package cmd

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"sync"
	"syscall"

	"github.com/nil2x/cheburnet/internal/api"
	"github.com/nil2x/cheburnet/internal/config"
	"github.com/nil2x/cheburnet/internal/handler"
	"github.com/nil2x/cheburnet/internal/imap"
	"github.com/nil2x/cheburnet/internal/ok"
	"github.com/nil2x/cheburnet/internal/session"
	"github.com/nil2x/cheburnet/internal/socks"
	"github.com/nil2x/cheburnet/internal/transform"
	"github.com/nil2x/cheburnet/internal/yadisk"
)

// Run starts the program, waits its completion and exits with appropriate code.
func Run() {
	ctx, cancel := context.WithCancel(context.Background())

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)

	go func() {
		<-sigCh
		cancel()
	}()

	errs := make(chan error, 100)
	code := 0

	go func() {
		for err := range errs {
			fmt.Fprintln(os.Stderr, err)
			code = 1
			cancel()
		}
	}()

	if err := run(ctx, errs); err != nil {
		fmt.Fprintln(os.Stderr, err)
		code = 1
	}

	close(errs)

	if runtime.GOOS == "windows" {
		fmt.Fprintln(os.Stdout, "\nPress Enter to exit...")
		fmt.Scanln()
	}

	os.Exit(code)
}

func run(ctx context.Context, errs chan<- error) error {
	flags := config.ParseFlags()

	if flags.PrintVersion {
		fmt.Fprintln(os.Stdout, config.BuildInfo())

		return nil
	}

	if flags.GenerateSecret {
		secret, err := transform.GenerateSecret()

		if err != nil {
			return fmt.Errorf("generate secret: %v", err)
		}

		fmt.Fprintln(os.Stdout, secret)

		return nil
	}

	cfg, err := config.Parse(flags.ConfigPath)

	if err != nil {
		return fmt.Errorf("parse config: %v", err)
	}

	if err := transform.Init(&cfg); err != nil {
		return fmt.Errorf("init transform: %v", err)
	}

	if err := config.Validate(cfg); err != nil {
		return fmt.Errorf("validate config: %v", err)
	}

	if err := config.SetupLog(cfg.Log); err != nil {
		return fmt.Errorf("setup log: %v", err)
	}

	if err := config.SetupDNS(cfg.DNS); err != nil {
		return fmt.Errorf("setup dns: %v", err)
	}

	if err := transform.ValidateQR(cfg.QR); err != nil {
		return fmt.Errorf("validate qr: %v", err)
	}

	if err := imap.Init(cfg.IMAP); err != nil {
		return fmt.Errorf("init imap: %v", err)
	}

	if err := yadisk.Init(cfg.API, cfg.YaDisk); err != nil {
		return fmt.Errorf("init yadisk: %v", err)
	}

	if err := ok.Init(cfg.API, cfg.OK); err != nil {
		return fmt.Errorf("init ok: %v", err)
	}

	vkClient := api.NewVKClient(cfg.API)
	storageClient := api.NewStorageClient()

	if !cfg.API.SkipValidation {
		for _, club := range cfg.Clubs {
			if err := api.ValidateClub(vkClient, club); err != nil {
				return fmt.Errorf("validate club: %v: %v", club.Name, err)
			}

			if err := api.ValidateLongPoll(vkClient, club); err != nil {
				return fmt.Errorf("validate long poll: %v: %v", club.Name, err)
			}
		}

		for _, user := range cfg.Users {
			if err := api.ValidateUser(vkClient, user); err != nil {
				return fmt.Errorf("validate user: %v: %v", user.Name, err)
			}
		}

		for _, client := range imap.GetClients() {
			if err := imap.Validate(client); err != nil {
				return fmt.Errorf("validate imap: %v: %v", client.Name, err)
			}
		}

		for _, client := range yadisk.GetClients() {
			if err := yadisk.Validate(client); err != nil {
				return fmt.Errorf("validate yadisk: %v: %v", client.Name, err)
			}
		}

		for _, client := range ok.GetClients() {
			if err := ok.Validate(client); err != nil {
				return fmt.Errorf("validate ok: %v: %v", client.Name, err)
			}
		}
	}

	if err := session.Init(cfg, vkClient, storageClient); err != nil {
		return fmt.Errorf("init session: %v", err)
	}

	var wg sync.WaitGroup

	if cfg.Socks.Port > 0 {
		wg.Add(1)
		go func() {
			defer wg.Done()

			if err := socks.Listen(ctx, cfg, vkClient, storageClient); err != nil {
				errs <- fmt.Errorf("listen socks: %v", err)
			}
		}()
	}

	for _, club := range cfg.Clubs {
		wg.Add(1)
		go func(club config.Club) {
			defer wg.Done()

			if err := handler.ListenLongPoll(ctx, cfg, vkClient, storageClient, club); err != nil {
				errs <- fmt.Errorf("listen long poll: %v", err)
			}
		}(club)

		wg.Add(1)
		go func(club config.Club) {
			defer wg.Done()

			if err := handler.ListenStorage(ctx, cfg, vkClient, storageClient, club); err != nil {
				errs <- fmt.Errorf("listen storage: %v", err)
			}
		}(club)
	}

	for _, client := range imap.GetClients() {
		wg.Add(1)
		go func(client *imap.Client) {
			defer wg.Done()

			if err := handler.ListenIMAP(ctx, cfg, vkClient, storageClient, client); err != nil {
				errs <- fmt.Errorf("listen imap: %v: %v", client.Name, err)
			}
		}(client)
	}

	for _, client := range yadisk.GetClients() {
		wg.Add(1)
		go func(client *yadisk.Client) {
			defer wg.Done()

			if err := handler.ListenYaDisk(ctx, cfg, vkClient, storageClient, client); err != nil {
				errs <- fmt.Errorf("listen yadisk: %v: %v", client.Name, err)
			}
		}(client)
	}

	for _, client := range ok.GetClients() {
		wg.Add(1)
		go func(client *ok.Client) {
			defer wg.Done()

			if err := handler.ListenOkStorage(ctx, cfg, vkClient, storageClient, client); err != nil {
				errs <- fmt.Errorf("listen ok storage: %v: %v", client.Name, err)
			}
		}(client)
	}

	wg.Add(1)
	go func() {
		defer wg.Done()

		if err := handler.Clear(ctx); err != nil {
			errs <- fmt.Errorf("clear handler: %v", err)
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()

		if err := session.Clear(ctx); err != nil {
			errs <- fmt.Errorf("clear session: %v", err)
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()

		if err := imap.Clear(ctx); err != nil {
			errs <- fmt.Errorf("clear imap: %v", err)
		}
	}()

	wg.Wait()

	if err := imap.Close(); err != nil {
		errs <- fmt.Errorf("close imap: %v", err)
	}

	return nil
}