creator-hub/cmd/docker-gateway/main.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/gofiber/fiber/v3"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	browserImage       = "git.ipao.vip/rogee/fingerprint-chromium@sha256:b9f23b8e3ac640174db0dfa49e9095fe7eb06f5db55a4e7550d979b35ff3a1b7"
	browserUser        = "1000:1000"
	browserEntrypoint  = "/usr/local/bin/docker-entrypoint.sh"
	managedLabel       = "io.creatorhub.managed"
	idLabel            = "io.creatorhub.runtime-id"
	networkRoleLabel   = "io.creatorhub.network-role"
	browserNetworkRole = "browser"
	controlNetworkName = "creatorhub_control"
	namePrefix         = "creatorhub-browser-"
)

var runtimeIDPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,31}$`)
var networkNamePattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$`)
var errInvalidRuntimeID = errors.New("invalid runtime id")
var errUnmanagedContainer = errors.New("refusing to operate on a container not owned by CreatorHub")

type browserConfig struct {
	image   string
	network string
}

type serviceConfig struct {
	listenAddr string
	dockerSock string
	network    string
	logLevel   logrus.Level
}

type dockerClient struct {
	baseURL string
	client  *http.Client
}

type gateway struct {
	docker dockerClient
	config browserConfig
}

type createRequest struct {
	Name string `json:"name"`
	Seed int64  `json:"seed"`
}

type browser struct {
	ID       string `json:"id"`
	Name     string `json:"name"`
	State    string `json:"state"`
	Status   string `json:"status"`
	Profile  string `json:"profile"`
	Endpoint string `json:"endpoint"`
}

func main() {
	logrus.SetFormatter(&logrus.JSONFormatter{})
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	if err := newCommand().ExecuteContext(ctx); err != nil {
		logrus.WithField("service", "docker-gateway").WithError(err).Error("service stopped")
		os.Exit(1)
	}
}

func newCommand() *cobra.Command {
	command := &cobra.Command{
		Use:           "docker-gateway",
		Short:         "Run the restricted CreatorHub Docker gateway",
		Args:          cobra.NoArgs,
		SilenceErrors: true,
		SilenceUsage:  true,
		RunE: func(command *cobra.Command, _ []string) error {
			cfg, err := loadConfig()
			if err != nil {
				return err
			}
			logrus.SetLevel(cfg.logLevel)
			return run(command, cfg)
		},
	}
	return command
}

func loadConfig() (serviceConfig, error) {
	v := viper.New()
	v.SetDefault("listen_addr", ":8081")
	v.SetDefault("docker_socket", "/var/run/docker.sock")
	v.SetDefault("browser_network", "creatorhub_browser")
	v.SetDefault("log_level", "info")
	_ = v.BindEnv("listen_addr", "LISTEN_ADDR")
	_ = v.BindEnv("docker_socket", "DOCKER_SOCKET")
	_ = v.BindEnv("browser_network", "BROWSER_NETWORK")
	_ = v.BindEnv("log_level", "LOG_LEVEL")

	level, err := logrus.ParseLevel(v.GetString("log_level"))
	if err != nil {
		return serviceConfig{}, errors.New("LOG_LEVEL must be panic, fatal, error, warn, info, debug, or trace")
	}
	cfg := serviceConfig{
		listenAddr: strings.TrimSpace(v.GetString("listen_addr")),
		dockerSock: strings.TrimSpace(v.GetString("docker_socket")),
		network:    strings.TrimSpace(v.GetString("browser_network")),
		logLevel:   level,
	}
	if cfg.listenAddr == "" {
		return serviceConfig{}, errors.New("LISTEN_ADDR must not be empty")
	}
	if err := validateListenAddr(cfg.listenAddr); err != nil {
		return serviceConfig{}, err
	}
	if cfg.dockerSock == "" {
		return serviceConfig{}, errors.New("DOCKER_SOCKET must not be empty")
	}
	if err := validateBrowserNetwork(cfg.network); err != nil {
		return serviceConfig{}, err
	}
	return cfg, nil
}

func validateListenAddr(addr string) error {
	_, port, err := net.SplitHostPort(addr)
	if err != nil {
		return errors.New("LISTEN_ADDR must be a host:port address")
	}
	number, err := strconv.Atoi(port)
	if err != nil || number < 1 || number > 65535 {
		return errors.New("LISTEN_ADDR port must be 1..65535")
	}
	return nil
}

func run(command *cobra.Command, cfg serviceConfig) error {
	transport := &http.Transport{
		DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, "unix", cfg.dockerSock)
		},
	}
	defer transport.CloseIdleConnections()
	docker := dockerClient{
		baseURL: "http://docker/v1.43",
		client:  &http.Client{Transport: transport, Timeout: 30 * time.Second},
	}
	if err := docker.ensureBrowserNetwork(cfg.network); err != nil {
		return err
	}
	logrus.WithFields(logrus.Fields{
		"service":     "docker-gateway",
		"listen_addr": cfg.listenAddr,
		"network":     cfg.network,
	}).Info("service starting")
	return newGateway(docker, browserConfig{image: browserImage, network: cfg.network}).Listen(cfg.listenAddr, fiber.ListenConfig{
		GracefulContext:       command.Context(),
		DisableStartupMessage: true,
	})
}

func newGateway(client dockerClient, cfg browserConfig) *fiber.App {
	api := gateway{docker: client, config: cfg}
	app := fiber.New(fiber.Config{
		AppName:     "CreatorHub Docker gateway",
		BodyLimit:   1 << 20,
		ReadTimeout: 5 * time.Second,
		IdleTimeout: 60 * time.Second,
		ErrorHandler: func(c fiber.Ctx, err error) error {
			status := http.StatusInternalServerError
			var fiberError *fiber.Error
			if errors.As(err, &fiberError) && fiberError != nil {
				status = fiberError.Code
			}
			return writeError(c, status, err)
		},
	})
	app.Get("/healthz", func(c fiber.Ctx) error {
		c.Status(fiber.StatusNoContent)
		return nil
	})
	app.Get("/v1/browsers", api.list)
	app.Post("/v1/browsers", api.create)
	app.Post("/v1/browsers/:id/:action", api.changeState)
	app.Delete("/v1/browsers/:id", api.remove)
	return app
}

func (api gateway) list(c fiber.Ctx) error {
	filters, _ := json.Marshal(map[string][]string{"label": {managedLabel + "=true"}})
	result, err := api.docker.request(http.MethodGet, "/containers/json?all=1&filters="+url.QueryEscape(string(filters)), nil)
	if err != nil {
		return writeError(c, http.StatusBadGateway, err)
	}
	defer result.Body.Close()
	if result.StatusCode != http.StatusOK {
		return forwardDockerError(c, result)
	}

	var containers []struct {
		ID     string            `json:"Id"`
		State  string            `json:"State"`
		Status string            `json:"Status"`
		Labels map[string]string `json:"Labels"`
	}
	if err := json.NewDecoder(result.Body).Decode(&containers); err != nil {
		return writeError(c, http.StatusBadGateway, fmt.Errorf("decode Docker response: %w", err))
	}

	browsers := make([]browser, 0, len(containers))
	for _, container := range containers {
		id := container.Labels[idLabel]
		if !runtimeIDPattern.MatchString(id) {
			continue
		}
		browsers = append(browsers, browser{
			ID:       container.ID,
			Name:     id,
			State:    container.State,
			Status:   container.Status,
			Profile:  "creatorhub-profile-" + id,
			Endpoint: "http://" + namePrefix + id + ":9222",
		})
	}
	return writeJSON(c, http.StatusOK, browsers)
}

func (api gateway) create(c fiber.Ctx) error {
	var input createRequest
	decoder := json.NewDecoder(bytes.NewReader(c.Body()))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&input); err != nil {
		return writeError(c, http.StatusBadRequest, errors.New("body must contain only name and numeric seed"))
	}
	if !runtimeIDPattern.MatchString(input.Name) || input.Seed < 1 || input.Seed > 2147483647 {
		return writeError(c, http.StatusBadRequest, errors.New("name must match [a-z0-9][a-z0-9-]{0,31} and seed must be 1..2147483647"))
	}

	pidsLimit := int64(512)
	payload := map[string]any{
		"Image":      api.config.image,
		"User":       browserUser,
		"Entrypoint": []string{browserEntrypoint},
		"Cmd":        []string{"--fingerprint=" + strconv.FormatInt(input.Seed, 10), "about:blank"},
		"Env":        []string{"REMOTE_DEBUGGING_PORT=9222"},
		"Labels": map[string]string{
			managedLabel: "true",
			idLabel:      input.Name,
		},
		"ExposedPorts": map[string]any{"9222/tcp": map[string]any{}},
		"HostConfig": map[string]any{
			"NetworkMode":    api.config.network,
			"ReadonlyRootfs": true,
			"CapDrop":        []string{"ALL"},
			"SecurityOpt":    []string{"no-new-privileges"},
			"PidsLimit":      &pidsLimit,
			"Memory":         int64(1 << 30),
			"NanoCpus":       int64(2_000_000_000),
			"Tmpfs": map[string]string{
				"/tmp":           "rw,nosuid,nodev,noexec,mode=1777,size=256m",
				"/tmp/.X11-unix": "rw,nosuid,nodev,noexec,mode=1777,size=1m",
				"/dev/shm":       "rw,nosuid,nodev,noexec,size=256m",
				"/home/ubuntu":   "rw,nosuid,nodev,noexec,uid=1000,gid=1000,mode=700,size=64m",
			},
			"Mounts": []map[string]any{{
				"Type":   "volume",
				"Source": "creatorhub-profile-" + input.Name,
				"Target": "/data",
			}},
		},
	}
	result, err := api.docker.request(http.MethodPost, "/containers/create?name="+url.QueryEscape(namePrefix+input.Name), payload)
	if err != nil {
		return writeError(c, http.StatusBadGateway, err)
	}
	defer result.Body.Close()
	if result.StatusCode != http.StatusCreated {
		return forwardDockerError(c, result)
	}
	var created struct {
		ID string `json:"Id"`
	}
	if err := json.NewDecoder(result.Body).Decode(&created); err != nil || created.ID == "" {
		cleanupErr := api.docker.expect(http.MethodDelete, "/containers/"+url.PathEscape(namePrefix+input.Name)+"?force=1&v=0", nil, http.StatusNoContent, http.StatusNotFound)
		if cleanupErr != nil {
			return writeError(c, http.StatusBadGateway, fmt.Errorf("Docker returned an invalid container id and cleanup failed: %w", cleanupErr))
		}
		return writeError(c, http.StatusBadGateway, errors.New("Docker returned an invalid container id; container was removed while preserving its Profile volume"))
	}
	if err := api.docker.expect(http.MethodPost, "/containers/"+url.PathEscape(created.ID)+"/start", nil, http.StatusNoContent, http.StatusNotModified); err != nil {
		cleanupErr := api.docker.expect(http.MethodDelete, "/containers/"+url.PathEscape(created.ID)+"?force=1&v=0", nil, http.StatusNoContent)
		if cleanupErr != nil {
			return writeError(c, http.StatusBadGateway, fmt.Errorf("container did not start and cleanup failed: %w; cleanup: %v", err, cleanupErr))
		}
		return writeError(c, http.StatusBadGateway, fmt.Errorf("container did not start and was removed while preserving its Profile volume: %w", err))
	}
	return writeJSON(c, http.StatusCreated, map[string]string{"id": created.ID, "name": input.Name})
}

func (api gateway) changeState(c fiber.Ctx) error {
	id := c.Params("id")
	action := c.Params("action")
	if err := api.requireManaged(id); err != nil {
		return writeError(c, statusFor(err), err)
	}
	var path string
	switch action {
	case "start":
		path = "/containers/" + url.PathEscape(namePrefix+id) + "/start"
	case "stop":
		path = "/containers/" + url.PathEscape(namePrefix+id) + "/stop?t=10"
	default:
		return writeError(c, http.StatusNotFound, errors.New("unknown action"))
	}
	if err := api.docker.expect(http.MethodPost, path, nil, http.StatusNoContent, http.StatusNotModified); err != nil {
		return writeError(c, http.StatusBadGateway, err)
	}
	c.Status(http.StatusNoContent)
	return nil
}

func (api gateway) remove(c fiber.Ctx) error {
	id := c.Params("id")
	if err := api.requireManaged(id); err != nil {
		return writeError(c, statusFor(err), err)
	}
	path := "/containers/" + url.PathEscape(namePrefix+id) + "?force=1&v=0"
	if err := api.docker.expect(http.MethodDelete, path, nil, http.StatusNoContent); err != nil {
		return writeError(c, http.StatusBadGateway, err)
	}
	c.Status(http.StatusNoContent)
	return nil
}

func (api gateway) requireManaged(id string) error {
	if !runtimeIDPattern.MatchString(id) {
		return errInvalidRuntimeID
	}
	result, err := api.docker.request(http.MethodGet, "/containers/"+url.PathEscape(namePrefix+id)+"/json", nil)
	if err != nil {
		return err
	}
	defer result.Body.Close()
	if result.StatusCode == http.StatusNotFound {
		return os.ErrNotExist
	}
	if result.StatusCode != http.StatusOK {
		return fmt.Errorf("Docker inspect returned %s", result.Status)
	}
	var inspected struct {
		Config struct {
			Labels map[string]string `json:"Labels"`
		} `json:"Config"`
	}
	if err := json.NewDecoder(result.Body).Decode(&inspected); err != nil {
		return fmt.Errorf("decode Docker inspect: %w", err)
	}
	if inspected.Config.Labels[managedLabel] != "true" || inspected.Config.Labels[idLabel] != id {
		return errUnmanagedContainer
	}
	return nil
}

func (docker dockerClient) request(method, path string, payload any) (*http.Response, error) {
	var body io.Reader
	if payload != nil {
		encoded, err := json.Marshal(payload)
		if err != nil {
			return nil, err
		}
		body = bytes.NewReader(encoded)
	}
	request, err := http.NewRequest(method, docker.baseURL+path, body)
	if err != nil {
		return nil, err
	}
	if payload != nil {
		request.Header.Set("Content-Type", "application/json")
	}
	return docker.client.Do(request)
}

func (docker dockerClient) expect(method, path string, payload any, allowed ...int) error {
	response, err := docker.request(method, path, payload)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	for _, status := range allowed {
		if response.StatusCode == status {
			return nil
		}
	}
	message, _ := io.ReadAll(io.LimitReader(response.Body, 4096))
	return fmt.Errorf("Docker returned %s: %s", response.Status, strings.TrimSpace(string(message)))
}

func (docker dockerClient) ensureBrowserNetwork(name string) error {
	if err := validateBrowserNetwork(name); err != nil {
		return err
	}
	response, err := docker.request(http.MethodGet, "/networks/"+url.PathEscape(name), nil)
	if err != nil {
		return fmt.Errorf("inspect browser network: %w", err)
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusOK {
		var network struct {
			Name       string            `json:"Name"`
			Driver     string            `json:"Driver"`
			Internal   bool              `json:"Internal"`
			Attachable bool              `json:"Attachable"`
			Ingress    bool              `json:"Ingress"`
			Labels     map[string]string `json:"Labels"`
		}
		if err := json.NewDecoder(response.Body).Decode(&network); err != nil {
			return fmt.Errorf("decode browser network: %w", err)
		}
		if network.Name != name || network.Driver != "bridge" || network.Internal || network.Attachable || network.Ingress ||
			network.Labels[managedLabel] != "true" || network.Labels[networkRoleLabel] != browserNetworkRole {
			return fmt.Errorf("browser network %q is not a CreatorHub bridge", name)
		}
		return nil
	}
	if response.StatusCode != http.StatusNotFound {
		return fmt.Errorf("inspect browser network returned %s", response.Status)
	}
	return docker.expect(http.MethodPost, "/networks/create", map[string]any{
		"Name":           name,
		"CheckDuplicate": true,
		"Driver":         "bridge",
		"Labels": map[string]string{
			managedLabel:     "true",
			networkRoleLabel: browserNetworkRole,
		},
	}, http.StatusCreated)
}

func validateBrowserNetwork(name string) error {
	if !networkNamePattern.MatchString(name) {
		return errors.New("BROWSER_NETWORK is invalid")
	}
	if name == controlNetworkName {
		return errors.New("BROWSER_NETWORK must not reuse the control network")
	}
	return nil
}

func statusFor(err error) int {
	switch {
	case errors.Is(err, errInvalidRuntimeID):
		return http.StatusBadRequest
	case errors.Is(err, os.ErrNotExist):
		return http.StatusNotFound
	case errors.Is(err, errUnmanagedContainer):
		return http.StatusForbidden
	default:
		return http.StatusBadGateway
	}
}

func forwardDockerError(c fiber.Ctx, result *http.Response) error {
	message, _ := io.ReadAll(io.LimitReader(result.Body, 4096))
	status := http.StatusBadGateway
	if result.StatusCode == http.StatusConflict {
		status = http.StatusConflict
	}
	return writeError(c, status, fmt.Errorf("Docker returned %s: %s", result.Status, strings.TrimSpace(string(message))))
}

func writeError(c fiber.Ctx, status int, err error) error {
	return writeJSON(c, status, map[string]string{"error": err.Error()})
}

func writeJSON(c fiber.Ctx, status int, value any) error {
	return c.Status(status).JSON(value)
}