cloak-docker/main.go

package main

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"sync"
	"syscall"
	"time"
)

type Config struct {
	Servers []json.RawMessage `json:"servers"`
	Clients []ClientEntry     `json:"clients"`
}

type ClientEntry struct {
	Listen  ListenConfig    `json:"listen"`
	Connect ConnectConfig   `json:"connect"`
	Config  json.RawMessage `json:"config"`
}

type ListenConfig struct {
	IP    string `json:"ip"`
	Port  int    `json:"port"`
	Proto string `json:"proto"`
}

type ConnectConfig struct {
	IP   string `json:"ip"`
	Port int    `json:"port"`
}

type managedProcess struct {
	name    string
	cmdLine []string
	cmd     *exec.Cmd
	mu      sync.Mutex
}

func newGUID() string {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		log.Fatalf("failed to generate GUID: %v", err)
	}
	return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x",
		b[0:4], b[4:6], b[6:8], b[8:10], b[10:16])
}

func writeTempFile(tmpDir string, data []byte) (string, error) {
	name := newGUID() + ".json"
	path := filepath.Join(tmpDir, name)
	if err := os.WriteFile(path, data, 0644); err != nil {
		return "", fmt.Errorf("write temp file: %w", err)
	}
	return path, nil
}

func (mp *managedProcess) start(ctx context.Context) {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	cmd := exec.CommandContext(ctx, mp.cmdLine[0], mp.cmdLine[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Start(); err != nil {
		log.Printf("failed to start %s: %v", mp.name, err)
		return
	}

	mp.cmd = cmd
	log.Printf("started %s (pid %d)", mp.name, cmd.Process.Pid)
}

func (mp *managedProcess) stop() {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	if mp.cmd != nil && mp.cmd.Process != nil {
		log.Printf("stopping %s (pid %d)", mp.name, mp.cmd.Process.Pid)
		mp.cmd.Process.Signal(syscall.SIGTERM)
	}
}

func (mp *managedProcess) wait() error {
	mp.mu.Lock()
	cmd := mp.cmd
	mp.mu.Unlock()

	if cmd == nil {
		return nil
	}

	return cmd.Wait()
}

func runLauncher(ctx context.Context, configPath string) error {
	data, err := os.ReadFile(configPath)
	if err != nil {
		return fmt.Errorf("read config: %w", err)
	}

	var cfg Config
	if err := json.Unmarshal(data, &cfg); err != nil {
		return fmt.Errorf("parse config: %w", err)
	}

	tmpDir := filepath.Join(".", "tmp")
	if err := os.MkdirAll(tmpDir, 0755); err != nil {
		return fmt.Errorf("create tmp dir: %w", err)
	}

	var processes []*managedProcess

	for i, raw := range cfg.Servers {
		path, err := writeTempFile(tmpDir, []byte(raw))
		if err != nil {
			return fmt.Errorf("server[%d]: %w", i, err)
		}
		log.Printf("server[%d] config -> %s", i, path)

		mp := &managedProcess{
			name:    fmt.Sprintf("server[%d]", i),
			cmdLine: []string{"/usr/bin/ck-server", "-c", path},
		}
		processes = append(processes, mp)
	}

	for i, ce := range cfg.Clients {
		path, err := writeTempFile(tmpDir, []byte(ce.Config))
		if err != nil {
			return fmt.Errorf("client[%d]: %w", i, err)
		}
		log.Printf("client[%d] config -> %s", i, path)

		args := []string{
			"-s", ce.Connect.IP,
			"-p", fmt.Sprintf("%d", ce.Connect.Port),
			"-l", fmt.Sprintf("%d", ce.Listen.Port),
			"-i", ce.Listen.IP,
		}
		if ce.Listen.Proto == "udp" {
			args = append(args, "-u")
		}
		args = append(args, "-c", path)

		mp := &managedProcess{
			name:    fmt.Sprintf("client[%d]", i),
			cmdLine: append([]string{"/usr/bin/ck-client"}, args...),
		}
		processes = append(processes, mp)
	}

	if len(processes) == 0 {
		log.Println("no servers or clients to launch")
		return nil
	}

	for _, mp := range processes {
		mp.start(ctx)
	}

	var wg sync.WaitGroup
	for _, mp := range processes {
		wg.Add(1)
		go func(p *managedProcess) {
			defer wg.Done()
			for {
				err := p.wait()
				if err != nil && ctx.Err() != nil {
					return
				}

				select {
				case <-ctx.Done():
					return
				default:
				}

				if err != nil {
					log.Printf("%s exited with error: %v, restarting in 15s", p.name, err)
				} else {
					log.Printf("%s exited, restarting in 15s", p.name)
				}

				select {
				case <-time.After(15 * time.Second):
				case <-ctx.Done():
					return
				}

				select {
				case <-ctx.Done():
					return
				default:
					p.start(ctx)
				}
			}
		}(mp)
	}

	<-ctx.Done()
	log.Println("shutting down...")

	var stopWg sync.WaitGroup
	for _, mp := range processes {
		stopWg.Add(1)
		go func(p *managedProcess) {
			defer stopWg.Done()
			p.stop()
		}(mp)
	}
	stopWg.Wait()

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		log.Println("all processes stopped")
	case <-time.After(10 * time.Second):
		log.Println("force killing remaining processes")
		for _, mp := range processes {
			mp.mu.Lock()
			if mp.cmd != nil && mp.cmd.Process != nil {
				mp.cmd.Process.Kill()
			}
			mp.mu.Unlock()
		}
	}

	return nil
}

func main() {
	if len(os.Args) < 2 {
		log.Fatalf("usage: %s <config.json>", os.Args[0])
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigCh
		cancel()
	}()

	if err := runLauncher(ctx, os.Args[1]); err != nil {
		log.Fatalf("error: %v", err)
	}
}