feat: implement ssh forwarding

internal/docker/docker.go

@@ -0,0 +1,21 @@
+package docker
+
+import (
+	"context"
+	"github.com/docker/docker/api/types"
+	"github.com/docker/go-connections/nat"
+	"strconv"
+)
+
+// ContainerSSHHostPort returns the port on the host that is exposing the internal ssh port of the given container info
+func ContainerSSHHostPort(ctx context.Context, container types.ContainerJSON) int {
+	ports := container.NetworkSettings.Ports[nat.Port("22/tcp")]
+	if len(ports) == 0 {
+		return -1
+	}
+	port, err := strconv.Atoi(ports[0].HostPort)
+	if err != nil {
+		return -1
+	}
+	return port
+}

internal/service/config.go

@@ -9,6 +9,7 @@ import (
 type Config struct {
 	DatabasePath          string `json:"databasePath"`
 	TemplateDirectoryPath string `json:"templateDirectoryPath"`
+	HostKeyDirectoryPath  string `json:"hostKeyDirectoryPath"`
 	HostName              string `json:"hostName"`
 }
 
@@ -29,5 +30,10 @@ func ReadConfigFrom(reader io.Reader) (Config, error) {
 		return Config{}, err
 	}
 
+	config.HostKeyDirectoryPath, err = filepath.Abs(config.HostKeyDirectoryPath)
+	if err != nil {
+		return Config{}, err
+	}
+
 	return config, nil
 }

internal/service/service.go

@@ -12,6 +12,7 @@ import (
 	"github.com/uptrace/bun/extra/bundebug"
 	_ "modernc.org/sqlite"
 	"net/http"
+	"tesseract/internal/sshproxy"
 )
 
 const (
@@ -19,7 +20,7 @@ const (
 	keyDockerClient = "dockerClient"
 	keyDB           = "db"
 	keyConfig       = "config"
-	keyMelody       = "melody"
+	keySSHProxy     = "sshProxy"
 )
 
 type Services struct {
@@ -27,6 +28,7 @@ type Services struct {
 	DockerClient *client.Client
 	Database     *bun.DB
 	Config       Config
+	SSHProxy     *sshproxy.SSHProxy
 	Melody       *melody.Melody
 }
 
@@ -42,6 +44,10 @@ func Database(c echo.Context) *bun.DB {
 	return c.Get(keyDB).(*bun.DB)
 }
 
+func SSHProxy(c echo.Context) *sshproxy.SSHProxy {
+	return c.Get(keySSHProxy).(*sshproxy.SSHProxy)
+}
+
 func Initialize(config Config) (Services, error) {
 	hc := &http.Client{}
 
@@ -57,12 +63,15 @@ func Initialize(config Config) (Services, error) {
 	bundb := bun.NewDB(db, sqlitedialect.New())
 	bundb.AddQueryHook(bundebug.NewQueryHook(bundebug.WithVerbose(true)))
 
+	sshProxy := sshproxy.New()
+
 	return Services{
 		HTTPClient:   hc,
 		DockerClient: docker,
 		Database:     bundb,
 		Config:       config,
 		Melody:       melody.New(),
+		SSHProxy:     sshProxy,
 	}, nil
 }
 
@@ -73,6 +82,7 @@ func (s Services) Middleware() echo.MiddlewareFunc {
 			c.Set(keyDockerClient, s.DockerClient)
 			c.Set(keyDB, s.Database)
 			c.Set(keyConfig, s.Config)
+			c.Set(keySSHProxy, s.SSHProxy)
 			return next(c)
 		}
 	}

internal/sshproxy/proxy_connection.go

@@ -0,0 +1,77 @@
+package sshproxy
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"sync"
+)
+
+type proxyConnection struct {
+	internalPort int
+	externalPort int
+	listener     net.Listener
+}
+
+func newProxyConnection(toPort int) (*proxyConnection, error) {
+	l, err := net.Listen("tcp", ":0")
+	if err != nil {
+		return nil, err
+	}
+
+	externalPort := l.Addr().(*net.TCPAddr).Port
+
+	return &proxyConnection{
+		internalPort: toPort,
+		externalPort: externalPort,
+		listener:     l,
+	}, nil
+}
+
+func (c *proxyConnection) start() {
+	for {
+		conn, err := c.listener.Accept()
+		if err != nil {
+			fmt.Printf("error accepting connection at %v: %v\n", c.listener.Addr(), err)
+		}
+		go c.forwardConnectionToSSH(conn)
+	}
+}
+
+func (c *proxyConnection) forwardConnectionToSSH(conn net.Conn) {
+	containerConn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", c.internalPort))
+	if err != nil {
+		fmt.Printf("error connecting to container ssh at port %d\n", c.internalPort)
+		return
+	}
+	defer containerConn.Close()
+
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		defer conn.Close()
+		for {
+			_, err := io.Copy(conn, containerConn)
+			if err != nil {
+				fmt.Println("read remote conn err", err)
+				break
+			}
+		}
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		defer conn.Close()
+		for {
+			_, err := io.Copy(containerConn, conn)
+			if err != nil {
+				fmt.Println("write remote conn err", err)
+				break
+			}
+		}
+	}()
+}

internal/sshproxy/sshproxy.go

@@ -0,0 +1,37 @@
+package sshproxy
+
+type SSHProxy struct {
+	// internalPorts maps internal docker ssh ports to the corresponding external ssh ports
+	// that users use to ssh into workspaces
+	internalPorts map[int]int
+
+	connections map[int]*proxyConnection
+}
+
+func New() *SSHProxy {
+	return &SSHProxy{
+		internalPorts: map[int]int{},
+		connections:   map[int]*proxyConnection{},
+	}
+}
+
+func (p *SSHProxy) NewProxyEntryTo(toPort int) error {
+	c, err := newProxyConnection(toPort)
+	if err != nil {
+		return err
+	}
+
+	go c.start()
+
+	p.connections[toPort] = c
+	p.internalPorts[toPort] = c.externalPort
+
+	return nil
+}
+
+func (p *SSHProxy) FindExternalPort(internalPort int) int {
+	if port, ok := p.internalPorts[internalPort]; ok {
+		return port
+	}
+	return -1
+}

internal/workspace/docker.go

@@ -0,0 +1,55 @@
+package workspace
+
+import (
+	"context"
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/client"
+	"io"
+)
+
+type spawnedShell struct {
+	io.Reader
+	io.Writer
+
+	execID string
+}
+
+func stopContainer(ctx context.Context, docker *client.Client, containerID string) error {
+	return docker.ContainerStop(ctx, containerID, container.StopOptions{})
+}
+
+func startContainer(ctx context.Context, docker *client.Client, containerID string) error {
+	return docker.ContainerStart(ctx, containerID, container.StartOptions{})
+}
+
+func deleteContainer(ctx context.Context, docker *client.Client, containerID string) error {
+	return docker.ContainerRemove(ctx, containerID, container.RemoveOptions{
+		RemoveVolumes: true,
+	})
+}
+
+func inspectContainer(ctx context.Context, docker *client.Client, containerID string) (types.ContainerJSON, error) {
+	return docker.ContainerInspect(ctx, containerID)
+}
+
+func spawnNewShell(ctx context.Context, docker *client.Client, containerID string) (*spawnedShell, error) {
+	res, err := docker.ContainerExecCreate(ctx, containerID, container.ExecOptions{
+		Tty:    true,
+		Detach: true,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	attached, err := docker.ContainerExecAttach(ctx, res.ID, container.ExecAttachOptions{})
+	if err != nil {
+		return nil, err
+	}
+
+	return &spawnedShell{
+		Reader: attached.Reader,
+		Writer: attached.Conn,
+		execID: res.ID,
+	}, nil
+}

internal/workspace/http_handlers.go

@@ -5,9 +5,13 @@ import (
 	"encoding/json"
 	"errors"
 	"github.com/docker/docker/api/types/container"
+	"github.com/docker/go-connections/nat"
 	"github.com/google/uuid"
 	"github.com/labstack/echo/v4"
 	"net/http"
+	"strconv"
+	"sync"
+	docker2 "tesseract/internal/docker"
 	"tesseract/internal/service"
 	"tesseract/internal/template"
 	"time"
@@ -17,8 +21,13 @@ type createWorkspaceRequestBody struct {
 	ImageID string `json:"imageId"`
 }
 
+type updateWorkspaceRequestBody struct {
+	Status string `json:"status"`
+}
+
 func fetchAllWorkspaces(c echo.Context) error {
 	db := service.Database(c)
+	ctx := c.Request().Context()
 
 	var workspaces []workspace
 	err := db.NewSelect().Model(&workspaces).Scan(c.Request().Context())
@@ -33,10 +42,56 @@ func fetchAllWorkspaces(c echo.Context) error {
 		return c.JSON(http.StatusOK, make([]workspace, 0))
 	}
 
+	docker := service.DockerClient(c)
+	sshProxy := service.SSHProxy(c)
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var errs []error
+	for i, w := range workspaces {
+		wg.Add(1)
+		i, w := i, w
+		go func() {
+			defer wg.Done()
+
+			inspect, err := docker.ContainerInspect(ctx, w.ContainerID)
+			if err != nil {
+				mu.Lock()
+				errs = append(errs, err)
+				mu.Unlock()
+			} else {
+				switch inspect.State.Status {
+				case "running":
+					workspaces[i].Status = statusRunning
+				case "exited":
+					workspaces[i].Status = statusStopped
+				case "paused":
+					workspaces[i].Status = statusPaused
+				case "restarting":
+					workspaces[i].Status = statusRestarting
+				default:
+					workspaces[i].Status = statusUnknown
+				}
+
+				if internalPort := docker2.ContainerSSHHostPort(ctx, inspect); internalPort > 0 {
+					if port := sshProxy.FindExternalPort(internalPort); port > 0 {
+						workspaces[i].SSHPort = port
+					}
+				}
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	if err = errors.Join(errs...); err != nil {
+		return err
+	}
+
 	return c.JSON(http.StatusOK, workspaces)
 }
 
-func createWorkspace(c echo.Context) error {
+func updateOrCreateWorkspace(c echo.Context) error {
 	workspaceName := c.Param("workspaceName")
 	if workspaceName == "" {
 		return echo.NewHTTPError(http.StatusNotFound)
@@ -46,15 +101,55 @@ func createWorkspace(c echo.Context) error {
 		return echo.NewHTTPError(http.StatusNotFound)
 	}
 
-	body := createWorkspaceRequestBody{}
-	err := json.NewDecoder(c.Request().Body).Decode(&body)
+	db := service.Database(c)
+	ctx := c.Request().Context()
+
+	var w workspace
+	err := db.NewSelect().Model(&w).
+		Where("name = ?", workspaceName).
+		Scan(ctx)
 	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return createWorkspace(c, workspaceName)
+		}
+		return err
+	}
+
+	var body updateWorkspaceRequestBody
+	if err = json.NewDecoder(c.Request().Body).Decode(&body); err != nil {
+		return err
+	}
+
+	docker := service.DockerClient(c)
+
+	switch status(body.Status) {
+	case statusStopped:
+		if err = stopContainer(ctx, docker, workspaceName); err != nil {
+			return err
+		}
+		w.Status = statusStopped
+		break
+	case statusRunning:
+		if err = startContainer(ctx, docker, workspaceName); err != nil {
+			return err
+		}
+		w.Status = statusRunning
+		break
+	}
+
+	return c.JSON(http.StatusOK, w)
+}
+
+func createWorkspace(c echo.Context, workspaceName string) error {
+	var body createWorkspaceRequestBody
+	if err := json.NewDecoder(c.Request().Body).Decode(&body); err != nil {
 		return echo.NewHTTPError(http.StatusBadRequest)
 	}
 
 	db := service.Database(c)
+	ctx := c.Request().Context()
 
-	tx, err := db.BeginTx(c.Request().Context(), nil)
+	tx, err := db.BeginTx(ctx, nil)
 	if err != nil {
 		return err
 	}
@@ -72,18 +167,53 @@ func createWorkspace(c echo.Context) error {
 
 	docker := service.DockerClient(c)
 
-	res, err := docker.ContainerCreate(c.Request().Context(), &container.Config{
+	containerSSHPort := nat.Port("22/tcp")
+	containerConfig := &container.Config{
+		Tty:   true,
 		Image: img.ImageID,
-	}, nil, nil, nil, workspaceName)
+		ExposedPorts: nat.PortSet{
+			containerSSHPort: {},
+		},
+	}
+
+	hostConfig := &container.HostConfig{
+		PortBindings: nat.PortMap{
+			containerSSHPort: {
+				{"127.0.0.1", ""},
+			},
+		},
+	}
+
+	res, err := docker.ContainerCreate(ctx, containerConfig, hostConfig, nil, nil, workspaceName)
 	if err != nil {
 		return err
 	}
 
-	err = docker.ContainerStart(c.Request().Context(), res.ID, container.StartOptions{})
+	err = docker.ContainerStart(ctx, res.ID, container.StartOptions{})
 	if err != nil {
 		return err
 	}
 
+	inspect, err := docker.ContainerInspect(ctx, res.ID)
+	if err != nil {
+		return err
+	}
+
+	ports := inspect.NetworkSettings.Ports[containerSSHPort]
+	if len(ports) == 0 {
+		return errors.New("failed to bind ssh port for container")
+	}
+
+	sshProxy := service.SSHProxy(c)
+	hostPort, err := strconv.Atoi(ports[0].HostPort)
+	if err != nil {
+		return err
+	}
+
+	if err = sshProxy.NewProxyEntryTo(hostPort); err != nil {
+		return err
+	}
+
 	id, err := uuid.NewV7()
 	if err != nil {
 		return err
@@ -95,6 +225,8 @@ func createWorkspace(c echo.Context) error {
 		ContainerID: res.ID,
 		ImageTag:    img.ImageTag,
 		CreatedAt:   time.Now().Format(time.RFC3339),
+		SSHPort:     hostPort,
+		Status:      statusRunning,
 	}
 	_, err = tx.NewInsert().Model(&w).Exec(c.Request().Context())
 	if err != nil {
@@ -108,3 +240,75 @@ func createWorkspace(c echo.Context) error {
 
 	return c.JSON(http.StatusOK, w)
 }
+
+func deleteWorkspace(c echo.Context) error {
+	workspaceName := c.Param("workspaceName")
+	if workspaceName == "" {
+		return echo.NewHTTPError(http.StatusNotFound)
+	}
+
+	db := service.Database(c)
+	docker := service.DockerClient(c)
+	ctx := c.Request().Context()
+
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		return echo.NewHTTPError(http.StatusNotFound)
+	}
+
+	var w workspace
+	if err = tx.NewSelect().Model(&w).Scan(ctx); err != nil {
+		_ = tx.Rollback()
+		return echo.NewHTTPError(http.StatusNotFound)
+	}
+
+	inspect, err := inspectContainer(ctx, docker, w.ContainerID)
+	if err != nil {
+		_ = tx.Rollback()
+		return err
+	}
+
+	if inspect.State.Running {
+		if err = stopContainer(ctx, docker, w.ContainerID); err != nil {
+			_ = tx.Rollback()
+			return err
+		}
+	}
+
+	if err = deleteContainer(ctx, docker, w.ContainerID); err != nil {
+		return err
+	}
+
+	res, err := tx.NewDelete().
+		Table("workspaces").
+		Where("name = ?", workspaceName).
+		Exec(ctx)
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return echo.NewHTTPError(http.StatusNotFound)
+		}
+		return err
+	}
+
+	count, err := res.RowsAffected()
+	if err != nil {
+		_ = tx.Rollback()
+		return err
+	}
+
+	if count == 0 {
+		_ = tx.Rollback()
+		return echo.NewHTTPError(http.StatusNotFound)
+	}
+	if count != 1 {
+		_ = tx.Rollback()
+		return echo.NewHTTPError(http.StatusNotFound)
+	}
+
+	if err = tx.Commit(); err != nil {
+		_ = tx.Rollback()
+		return err
+	}
+
+	return c.NoContent(http.StatusOK)
+}

internal/workspace/routes.go

@@ -6,5 +6,6 @@ import (
 
 func DefineRoutes(g *echo.Group) {
 	g.GET("/workspaces", fetchAllWorkspaces)
-	g.POST("/workspaces/:workspaceName", createWorkspace)
+	g.POST("/workspaces/:workspaceName", updateOrCreateWorkspace)
+	g.DELETE("/workspaces/:workspaceName", deleteWorkspace)
 }

internal/workspace/workspace.go

@@ -1,9 +1,17 @@
 package workspace
 
 import (
+	"context"
+	"database/sql"
+	"errors"
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/client"
 	"github.com/google/uuid"
 	"github.com/uptrace/bun"
 	"regexp"
+	"sync"
+	"tesseract/internal/docker"
+	"tesseract/internal/service"
 )
 
 type workspace struct {
@@ -19,6 +27,100 @@ type workspace struct {
 	ImageTag string `json:"imageTag"`
 
 	CreatedAt string `json:"createdAt"`
+
+	SSHPort int `bun:"-" json:"sshPort,omitempty"`
+
+	Status status `bun:"-" json:"status"`
 }
 
+// status represents the status of a workspace.
+type status string
+
+const (
+	statusRunning    status = "running"
+	statusStopped    status = "stopped"
+	statusPaused     status = "paused"
+	statusRestarting status = "restarting"
+	statusUnknown    status = "unknown"
+)
+
 var workspaceNameRegex = regexp.MustCompile("^[\\w-]+$")
+
+func SyncAll(ctx context.Context, services service.Services) error {
+	tx, err := services.Database.BeginTx(ctx, nil)
+	if err != nil {
+		return err
+	}
+
+	var workspaces []workspace
+	if err = tx.NewSelect().Model(&workspaces).
+		Column("id", "container_id").
+		Scan(ctx); err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return nil
+		}
+	}
+	if len(workspaces) == 0 {
+		return nil
+	}
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var errs []error
+
+	var deletedWorkspaces []workspace
+
+	for _, w := range workspaces {
+		w := w
+		wg.Add(1)
+		go func() {
+			var err error
+			defer wg.Done()
+			defer func() {
+				mu.Lock()
+				errs = append(errs, err)
+				mu.Unlock()
+			}()
+
+			if err = services.DockerClient.ContainerStart(ctx, w.ContainerID, container.StartOptions{}); err != nil {
+				if client.IsErrNotFound(err) {
+					err = nil
+					mu.Lock()
+					deletedWorkspaces = append(deletedWorkspaces, w)
+					mu.Unlock()
+				}
+				return
+			}
+
+			inspect, err := services.DockerClient.ContainerInspect(ctx, w.ContainerID)
+			if err != nil {
+				return
+			}
+
+			internalPort := docker.ContainerSSHHostPort(ctx, inspect)
+			if internalPort <= 0 {
+				return
+			}
+
+			err = services.SSHProxy.NewProxyEntryTo(internalPort)
+		}()
+	}
+
+	if err = errors.Join(errs...); err != nil {
+		_ = tx.Rollback()
+		return err
+	}
+
+	_, err = tx.NewDelete().Model(&deletedWorkspaces).WherePK().Exec(ctx)
+	if err != nil {
+		_ = tx.Rollback()
+		return err
+	}
+
+	if err = tx.Commit(); err != nil {
+		_ = tx.Rollback()
+		return err
+	}
+
+	return nil
+}