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

package main

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

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

const (
	browserUser        = "1000:1000"
	browserEntrypoint  = "/usr/local/bin/docker-entrypoint.sh"
	managedLabel       = "io.creatorhub.managed"
	idLabel            = "io.creatorhub.runtime-id"
	nameLabel          = "io.creatorhub.display-name"
	networkRoleLabel   = "io.creatorhub.network-role"
	browserNetworkRole = "browser"
	controlNetworkName = "creatorhub_control"
	namePrefix         = "creatorhub-browser-"
	pullTimeout        = 10 * time.Minute
)

var (
	runtimeIDPattern  = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,31}$`)
	networkNamePattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$`)
	imageRefPattern   = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._:/@-]{0,300}$`)
	volumePattern     = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$`)
)

var (
	errInvalidRuntimeID     = errors.New("invalid runtime id")
	errUnmanagedContainer   = errors.New("refusing to operate on a container not owned by CreatorHub")
	errUnauthorized         = errors.New("gateway token rejected")
)

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

type dockerClient struct {
	baseURL string
	client  *http.Client // 常规 Docker 调用
	slow    *http.Client // 镜像拉取等长操作，不设整体超时
}

type gateway struct {
	docker  dockerClient
	network string
	token   string
}

// createRequest 全量字段由平台下发；网关不做业务决策，只做输入合法性校验。
type createRequest struct {
	Alias  string   `json:"alias"`
	Name   string   `json:"name"`
	Image  string   `json:"image"`
	Cmd    []string `json:"cmd"`
	Volume string   `json:"volume"`
}

type browser struct {
	ID       string `json:"id"`
	Alias    string `json:"alias"`
	Name     string `json:"name"`
	State    string `json:"state"`
	Status   string `json:"status"`
	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("gateway_token", "GATEWAY_TOKEN")
	_ = 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")),
		token:      strings.TrimSpace(v.GetString("gateway_token")),
		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 len(cfg.token) < 16 {
		return serviceConfig{}, errors.New("GATEWAY_TOKEN must be at least 16 characters")
	}
	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},
		slow:    &http.Client{Transport: transport},
	}
	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, cfg.network, cfg.token).Listen(cfg.listenAddr, fiber.ListenConfig{
		GracefulContext:       command.Context(),
		DisableStartupMessage: true,
	})
}

func newGateway(client dockerClient, network, token string) *fiber.App {
	api := gateway{docker: client, network: network, token: token}
	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.Use("/v1", api.authorize)
	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) authorize(c fiber.Ctx) error {
	expected := "Bearer " + api.token
	if subtle.ConstantTimeCompare([]byte(c.Get(fiber.HeaderAuthorization)), []byte(expected)) != 1 {
		return writeError(c, http.StatusUnauthorized, errUnauthorized)
	}
	return c.Next()
}

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 {
		alias := container.Labels[idLabel]
		if !runtimeIDPattern.MatchString(alias) {
			continue
		}
		name := container.Labels[nameLabel]
		if name == "" {
			name = alias
		}
		browsers = append(browsers, browser{
			ID:       container.ID,
			Alias:    alias,
			Name:     name,
			State:    container.State,
			Status:   container.Status,
			Endpoint: "http://" + namePrefix + alias + ":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 alias, name, image, cmd and volume"))
	}
	if err := validateCreate(input); err != nil {
		return writeError(c, http.StatusBadRequest, err)
	}
	if err := api.docker.pullIfMissing(c.Context(), input.Image); err != nil {
		return writeError(c, http.StatusBadGateway, err)
	}

	pidsLimit := int64(512)
	payload := map[string]any{
		"Image":      input.Image,
		"User":       browserUser,
		"Entrypoint": []string{browserEntrypoint},
		"Cmd":        input.Cmd,
		"Env":        []string{"REMOTE_DEBUGGING_PORT=9222"},
		"Labels": map[string]string{
			managedLabel: "true",
			idLabel:      input.Alias,
			nameLabel:    input.Name,
		},
		"ExposedPorts": map[string]any{"9222/tcp": map[string]any{}},
		"HostConfig": map[string]any{
			"NetworkMode":    api.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": input.Volume,
				"Target": "/data",
			}},
		},
	}
	result, err := api.docker.request(http.MethodPost, "/containers/create?name="+url.QueryEscape(namePrefix+input.Alias), 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.Alias)+"?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, "alias": input.Alias})
}

func validateCreate(input createRequest) error {
	if !runtimeIDPattern.MatchString(input.Alias) {
		return errors.New("alias must match [a-z0-9][a-z0-9-]{0,31}")
	}
	if input.Name == "" || utf8.RuneCountInString(input.Name) > 64 || hasControlRunes(input.Name) {
		return errors.New("name must be 1..64 visible characters")
	}
	if !imageRefPattern.MatchString(input.Image) {
		return errors.New("image must be a valid image reference")
	}
	if !volumePattern.MatchString(input.Volume) {
		return errors.New("volume must be a valid volume name")
	}
	if len(input.Cmd) == 0 || len(input.Cmd) > 64 {
		return errors.New("cmd must contain 1..64 arguments")
	}
	total := 0
	for _, arg := range input.Cmd {
		if arg == "" || hasControlRunes(arg) {
			return errors.New("cmd arguments must be non-empty visible strings")
		}
		total += len(arg)
	}
	if total > 4096 {
		return errors.New("cmd arguments exceed 4096 characters")
	}
	return nil
}

func hasControlRunes(value string) bool {
	for _, r := range value {
		if r < 0x20 || r == 0x7f {
			return true
		}
	}
	return false
}

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
}

// pullIfMissing 在镜像不在本地时从远端仓库拉取；镜像缺失属于可恢复错误，调用方可直接重试。
func (docker dockerClient) pullIfMissing(ctx context.Context, ref string) error {
	inspect, err := docker.request(http.MethodGet, "/images/"+url.PathEscape(ref)+"/json", nil)
	if err != nil {
		return fmt.Errorf("inspect image %s: %w", ref, err)
	}
	_, _ = io.Copy(io.Discard, inspect.Body)
	_ = inspect.Body.Close()
	switch inspect.StatusCode {
	case http.StatusOK:
		return nil
	case http.StatusNotFound:
		// 本地无此镜像，继续拉取
	default:
		return fmt.Errorf("inspect image %s returned %s", ref, inspect.Status)
	}

	pullCtx, cancel := context.WithTimeout(ctx, pullTimeout)
	defer cancel()
	query := url.Values{"fromImage": {ref}}
	if !strings.Contains(ref, "@") {
		if repository, tag := splitImageRef(ref); tag != "" {
			query = url.Values{"fromImage": {repository}, "tag": {tag}}
		}
	}
	request, err := http.NewRequestWithContext(pullCtx, http.MethodPost, docker.baseURL+"/images/create?"+query.Encode(), nil)
	if err != nil {
		return fmt.Errorf("build image pull request: %w", err)
	}
	response, err := docker.slow.Do(request)
	if err != nil {
		return fmt.Errorf("pull image %s: %w", ref, err)
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		message, _ := io.ReadAll(io.LimitReader(response.Body, 4096))
		return fmt.Errorf("pull image %s returned %s: %s", ref, response.Status, strings.TrimSpace(string(message)))
	}
	_, _ = io.Copy(io.Discard, response.Body)
	return nil
}

func splitImageRef(ref string) (repository, tag string) {
	if at := strings.Index(ref, "@"); at >= 0 {
		return ref[:at], ref[at+1:]
	}
	if colon := strings.LastIndex(ref, ":"); colon > strings.LastIndex(ref, "/") {
		return ref[:colon], ref[colon+1:]
	}
	return ref, ""
}

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)
}