DockMV/internal/api/sources.go

package api

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strconv"
	"time"

	"github.com/arescom/dockmv/internal/dkr"
	"github.com/arescom/dockmv/internal/sshx"
	"github.com/arescom/dockmv/internal/store"
)

// sourceConn is a live connection to one source daemon.
//
// It is reference counted rather than closed eagerly: a migration holds its
// source for as long as it runs, so switching source in the UI half way through
// a transfer must not pull the socket out from under it. The connection is
// closed once it is both replaced and unused.
type sourceConn struct {
	src     store.Source
	docker  *dkr.Client
	ssh     *sshx.Client
	version string

	refs  int
	stale bool
}

func (c *sourceConn) close() {
	if c.docker != nil {
		_ = c.docker.Close()
	}
	if c.ssh != nil {
		_ = c.ssh.Close()
	}
}

// sourceStatus is what the UI shows about the source it is looking at.
type sourceStatus struct {
	ID            string           `json:"id"`
	Name          string           `json:"name"`
	Kind          store.SourceKind `json:"kind"`
	Endpoint      string           `json:"endpoint"`
	DockerVersion string           `json:"dockerVersion,omitempty"`
	Connected     bool             `json:"connected"`
	Error         string           `json:"error,omitempty"`
}

func (c *sourceConn) status() sourceStatus {
	return sourceStatus{
		ID: c.src.ID, Name: c.src.Name, Kind: c.src.Kind,
		Endpoint: c.docker.Endpoint, DockerVersion: c.version, Connected: true,
	}
}

// source returns the current source, connecting to it on first use. The
// returned release function must be called when the caller is done with it —
// for a job, when the job finishes.
func (s *Server) source(ctx context.Context) (*sourceConn, func(), error) {
	s.srcMu.Lock()
	if cur := s.cur; cur != nil {
		cur.refs++
		s.srcMu.Unlock()
		return cur, func() { s.releaseSource(cur) }, nil
	}
	id := s.sources.Selected()
	s.srcMu.Unlock()

	conn, err := s.dialSource(ctx, id)
	if err != nil {
		return nil, nil, err
	}

	s.srcMu.Lock()
	defer s.srcMu.Unlock()
	// Another request may have connected while this one was dialling; one
	// connection is enough, so the loser is dropped.
	if s.cur != nil {
		conn.close()
		conn = s.cur
	} else {
		s.cur = conn
	}
	conn.refs++
	return conn, func() { s.releaseSource(conn) }, nil
}

func (s *Server) releaseSource(c *sourceConn) {
	s.srcMu.Lock()
	defer s.srcMu.Unlock()
	c.refs--
	if c.refs <= 0 && c.stale {
		c.close()
	}
}

// selectSource connects to a source and, once that worked, makes it the current
// one and records the choice for the next run.
func (s *Server) selectSource(ctx context.Context, id string) (sourceStatus, error) {
	conn, err := s.dialSource(ctx, id)
	if err != nil {
		return sourceStatus{}, err
	}
	st := conn.status()

	s.srcMu.Lock()
	old := s.cur
	s.cur = conn
	if old != nil {
		old.stale = true
		if old.refs <= 0 {
			old.close()
		}
	}
	s.srcMu.Unlock()

	if err := s.sources.Select(conn.src.ID); err != nil {
		return st, fmt.Errorf("remember the selected source: %w", err)
	}
	s.log.Info("source selected", "id", conn.src.ID, "endpoint", conn.docker.Endpoint)
	return st, nil
}

// invalidateSource drops the cached connection when the source behind it has
// been edited or removed. An empty id invalidates whatever is current.
func (s *Server) invalidateSource(id string) {
	s.srcMu.Lock()
	defer s.srcMu.Unlock()
	if s.cur == nil {
		return
	}
	if id != "" && s.cur.src.ID != id {
		return
	}
	s.cur.stale = true
	if s.cur.refs <= 0 {
		s.cur.close()
	}
	s.cur = nil
}

// dialSource opens a connection to one source and verifies the daemon answers.
func (s *Server) dialSource(ctx context.Context, id string) (*sourceConn, error) {
	src, err := s.sources.Get(id)
	if err != nil {
		return nil, err
	}

	conn := &sourceConn{src: src}
	switch src.Kind {
	case store.SourceLocal, store.SourceDocker:
		c, err := dkr.New(src.DockerHost)
		if err != nil {
			return nil, err
		}
		conn.docker = c
	case store.SourceSSH:
		if src.SSH == nil {
			return nil, errors.New("source has no ssh configuration")
		}
		client, err := sshx.Dial(ctx, *src.SSH, s.hosts)
		if err != nil {
			return nil, err
		}
		rd := sshx.NewRemoteDocker(client)
		// The CLI is checked first: a missing binary or a user outside the
		// docker group is a readable error here, and an unexplained broken
		// socket if it is left to the tunnel.
		if _, err := rd.ProbeCLI(ctx); err != nil {
			client.Close()
			return nil, err
		}
		c, err := dkr.NewTunnel(sourceEndpoint(src), func(ctx context.Context, _, _ string) (net.Conn, error) {
			return rd.DialAPI(ctx)
		})
		if err != nil {
			client.Close()
			return nil, err
		}
		conn.ssh, conn.docker = client, c
	default:
		return nil, fmt.Errorf("unknown source kind %q", src.Kind)
	}

	pingCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	version, err := conn.docker.Ping(pingCtx)
	if err != nil {
		conn.close()
		return nil, fmt.Errorf("connect to docker at %s: %w", conn.docker.Endpoint, err)
	}
	conn.version = version
	return conn, nil
}

// sourceEndpoint describes where a source lives, before and after it is
// connected to.
func sourceEndpoint(src store.Source) string {
	if src.Kind == store.SourceSSH && src.SSH != nil {
		port := src.SSH.Port
		if port == 0 {
			port = 22
		}
		return fmt.Sprintf("ssh://%s@%s", src.SSH.User, net.JoinHostPort(src.SSH.Host, strconv.Itoa(port)))
	}
	return src.DockerHost
}

// handleListSources lists the sources without connecting to any of them; the
// state of the current one comes from /api/health.
func (s *Server) handleListSources(w http.ResponseWriter, r *http.Request) {
	s.srcMu.Lock()
	var current *sourceStatus
	if s.cur != nil {
		st := s.cur.status()
		current = &st
	}
	s.srcMu.Unlock()

	writeJSON(w, http.StatusOK, map[string]any{
		"sources":  s.sources.List(),
		"selected": s.sources.Selected(),
		"current":  current,
	})
}

func (s *Server) handleSaveSource(w http.ResponseWriter, r *http.Request) {
	var src store.Source
	if err := decode(r, &src); err != nil {
		writeError(w, http.StatusBadRequest, "%v", err)
		return
	}
	saved, err := s.sources.Save(src)
	if err != nil {
		writeError(w, http.StatusBadRequest, "%v", err)
		return
	}
	// Editing the source in use means the live connection describes the old
	// settings; drop it so the next call reconnects.
	s.invalidateSource(saved.ID)
	writeJSON(w, http.StatusOK, saved)
}

func (s *Server) handleDeleteSource(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	if err := s.sources.Delete(id); err != nil {
		if errors.Is(err, store.ErrNotFound) {
			writeError(w, http.StatusNotFound, "%v", err)
			return
		}
		writeError(w, http.StatusBadRequest, "%v", err)
		return
	}
	s.invalidateSource(id)
	w.WriteHeader(http.StatusNoContent)
}

// handleSelectSource switches the source the whole UI works against.
func (s *Server) handleSelectSource(w http.ResponseWriter, r *http.Request) {
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()

	st, err := s.selectSource(ctx, r.PathValue("id"))
	if err != nil {
		s.writeDialError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, st)
}

// handleSourceProbe reads the SSH host key of a source host, so its fingerprint
// can be approved the same way a target's is.
func (s *Server) handleSourceProbe(w http.ResponseWriter, r *http.Request) {
	cfg, err := s.sourceSSH(r.PathValue("id"))
	if err != nil {
		writeError(w, http.StatusNotFound, "%v", err)
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
	defer cancel()

	info, err := sshx.Probe(ctx, cfg, s.hosts)
	if err != nil {
		writeError(w, http.StatusBadGateway, "%v", err)
		return
	}
	writeJSON(w, http.StatusOK, info)
}

func (s *Server) handleSourceTrust(w http.ResponseWriter, r *http.Request) {
	cfg, err := s.sourceSSH(r.PathValue("id"))
	if err != nil {
		writeError(w, http.StatusNotFound, "%v", err)
		return
	}
	var body struct {
		Fingerprint string `json:"fingerprint"`
	}
	if err := decode(r, &body); err != nil {
		writeError(w, http.StatusBadRequest, "%v", err)
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
	defer cancel()

	if err := sshx.TrustFromProbe(ctx, cfg, s.hosts, body.Fingerprint); err != nil {
		writeError(w, http.StatusBadRequest, "%v", err)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"trusted": true})
}

// sourceSSH returns the SSH configuration of a source, refusing the ones that
// are not reached over SSH.
func (s *Server) sourceSSH(id string) (sshx.Config, error) {
	src, err := s.sources.Get(id)
	if err != nil {
		return sshx.Config{}, err
	}
	if src.Kind != store.SourceSSH || src.SSH == nil {
		return sshx.Config{}, fmt.Errorf("source %s is not reached over ssh", src.Name)
	}
	return *src.SSH, nil
}