johnandersen777.bsky.social/socialweb-computer-kcp

Switch repository

Loading…

package iroh

import (
	"bufio"
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/johnandersen777/socialweb-computer-kcp/internal/clusterlink"
	"github.com/johnandersen777/socialweb-computer-kcp/internal/contract/component"
	"github.com/johnandersen777/socialweb-computer-kcp/internal/contract/ingress"
	"github.com/johnandersen777/socialweb-computer-kcp/internal/irohdial"
	"github.com/johnandersen777/socialweb-computer-kcp/internal/kit/childproc"
	"github.com/johnandersen777/socialweb-computer-kcp/internal/kit/clock"
	"github.com/johnandersen777/socialweb-computer-kcp/internal/kit/loopback"
)

const (
	URLScheme = "iroh"

	DefaultListenHost = "127.0.0.1"

	TicketTimeout = 20 * time.Second

	MinRestartBackoff = time.Second

	MaxRestartBackoff = 30 * time.Second
)

type Config struct {
	DumbpipePath string

	RelayURL string

	ListenHost string

	StateDir string

	Log *slog.Logger
}

var (
	_ ingress.Provider = (*Provider)(nil)

	_ ingress.Registrar = (*Provider)(nil)

	_ ingress.Registration = (*registration)(nil)
)

func (c Config) Check() error {
	path := strings.TrimSpace(c.DumbpipePath)
	switch {
	case path == "":
		return errors.New("ingressproxy iroh: spec.iroh.dumbpipePath is required")
	case !filepath.IsAbs(path):
		return fmt.Errorf("ingressproxy iroh: dumbpipe %q is not an absolute path", path)
	}
	st, err := os.Stat(path)
	if err != nil {
		return fmt.Errorf("ingressproxy iroh: dumbpipe: %w", err)
	}
	if st.IsDir() || st.Mode().Perm()&0o111 == 0 {
		return fmt.Errorf("ingressproxy iroh: dumbpipe %s is not an executable file", path)
	}
	if relay := strings.TrimSpace(c.RelayURL); relay != "" {
		parsed, err := url.Parse(relay)
		if err != nil || parsed.Scheme == "" || parsed.Host == "" {
			return fmt.Errorf("ingressproxy iroh: relayURL %q is not an absolute URL", relay)
		}
	}
	host := strings.TrimSpace(c.ListenHost)
	if host != "" && strings.ContainsAny(host, ":/[]") {
		return fmt.Errorf("ingressproxy iroh: listenHost %q is a host without a port", host)
	}
	return nil
}

func (c Config) listenHost() string {
	host := strings.TrimSpace(c.ListenHost)
	if host == "" {
		return DefaultListenHost
	}
	return host
}

type Provider struct {
	cfg Config

	mu sync.Mutex

	log *slog.Logger

	ctx context.Context

	cancel context.CancelFunc

	stopCh chan struct{}

	stopOnce sync.Once

	ready bool

	readiness component.ReadySignal

	registrations map[*registration]struct{}

	wg sync.WaitGroup
}

func New(cfg Config) *Provider {
	return &Provider{cfg: cfg}
}

func (p *Provider) Start(ctx context.Context) error {
	if err := p.cfg.Check(); err != nil {
		return err
	}
	log := p.cfg.Log
	if log == nil {
		log = slog.New(slog.DiscardHandler)
	}
	runCtx, cancel := context.WithCancel(ctx)
	p.mu.Lock()
	if p.ctx != nil {
		p.mu.Unlock()
		cancel()
		return errors.New("ingressproxy iroh: the provider already started")
	}
	p.log = log
	p.ctx = runCtx
	p.cancel = cancel
	p.stopCh = make(chan struct{})
	p.stopOnce = sync.Once{}
	p.ready = true
	p.registrations = map[*registration]struct{}{}
	stopCh := p.stopCh
	p.mu.Unlock()

	log.Info("ingressproxy_iroh_started", "provider", "iroh", "dumbpipe", p.cfg.DumbpipePath,
		"listenHost", p.cfg.listenHost(), "relay", p.cfg.RelayURL)
	p.readiness.Signal()

	select {
	case <-ctx.Done():
	case <-stopCh:
	}
	cancel()
	p.mu.Lock()
	p.ready = false
	p.ctx = nil
	p.cancel = nil
	p.mu.Unlock()
	p.closeAll()
	p.wg.Wait()
	return nil
}

func (p *Provider) Ready() component.Readiness {
	return &p.readiness
}

func (p *Provider) Stop(ctx context.Context) error {
	p.mu.Lock()
	stopCh := p.stopCh
	p.mu.Unlock()
	if stopCh == nil {
		return nil
	}
	p.stopOnce.Do(func() { close(stopCh) })
	p.closeAll()
	return nil
}

func (p *Provider) Endpoints() []component.Endpoint {
	return nil
}

func (p *Provider) Register(ctx context.Context, spec ingress.Spec, target ingress.Target) (ingress.Registration, error) {
	p.mu.Lock()
	ready := p.ready
	root := p.ctx
	log := p.log
	p.mu.Unlock()
	if !ready || root == nil {
		return nil, errors.New("ingressproxy iroh: the provider is not started")
	}
	if target == nil {
		return nil, errors.New("ingressproxy iroh: the handler registry is not configured")
	}
	bridge, err := loopback.Start(root, loopback.Config{
		Host:     p.cfg.listenHost(),
		Protocol: spec.Protocol,
	}, target)
	if err != nil {
		return nil, err
	}
	secret, err := p.secret(spec.Workspace, spec.Exposure)
	if err != nil {
		_ = bridge.Close()
		return nil, err
	}
	runCtx, cancel := context.WithCancel(root)
	reg := &registration{
		provider: p,
		spec:     spec,
		bridge:   bridge,
		secret:   secret,
		log:      log.With("workspace", spec.Workspace, "exposure", spec.Exposure),
		cancel:   cancel,
		done:     make(chan struct{}),
	}
	child, err := reg.startChild(runCtx)
	if err != nil {
		cancel()
		_ = bridge.Close()
		return nil, err
	}
	reg.mu.Lock()
	reg.cmd = child.cmd
	reg.url = URLScheme + "://" + child.ticket
	reg.mu.Unlock()
	p.mu.Lock()
	running := p.ctx == root
	if running {
		if p.registrations == nil {
			p.registrations = map[*registration]struct{}{}
		}
		p.registrations[reg] = struct{}{}
		p.wg.Go(func() { reg.supervise(runCtx, child) })
	}
	p.mu.Unlock()
	if !running {
		cancel()
		child.stop()
		_ = bridge.Close()
		return nil, errors.New("ingressproxy iroh: the provider stopped")
	}
	return reg, nil
}

func (p *Provider) closeAll() {
	p.mu.Lock()
	registrations := p.registrations
	p.registrations = nil
	p.mu.Unlock()
	for registration := range registrations {
		_ = registration.Close()
	}
}

func (p *Provider) forget(registration *registration) {
	p.mu.Lock()
	if p.registrations != nil {
		delete(p.registrations, registration)
	}
	p.mu.Unlock()
}

func (p *Provider) secret(workspace, exposure string) (string, error) {
	dir := filepath.Join(p.cfg.StateDir, "iroh")
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return "", fmt.Errorf("ingressproxy iroh: the secret dir: %w", err)
	}
	sum := sha256.Sum256([]byte(workspace + "/" + exposure))
	path := filepath.Join(dir, hex.EncodeToString(sum[:8])+".secret")
	data, err := os.ReadFile(path)
	switch {
	case err == nil:
		secret := strings.TrimSpace(string(data))
		if !validSecret(secret) {
			return "", fmt.Errorf("ingressproxy iroh: %s holds no 32-byte hex iroh secret", path)
		}
		return secret, nil
	case !errors.Is(err, os.ErrNotExist):
		return "", fmt.Errorf("ingressproxy iroh: reading %s: %w", path, err)
	}
	seed := make([]byte, 32)
	if _, err := rand.Read(seed); err != nil {
		return "", fmt.Errorf("ingressproxy iroh: the secret: %w", err)
	}
	secret := hex.EncodeToString(seed)
	if err := os.WriteFile(path, []byte(secret), 0o600); err != nil {
		return "", fmt.Errorf("ingressproxy iroh: writing %s: %w", path, err)
	}
	return secret, nil
}

func validSecret(secret string) bool {
	if len(secret) != 64 {
		return false
	}
	_, err := hex.DecodeString(secret)
	return err == nil
}

type registration struct {
	provider *Provider

	spec ingress.Spec

	bridge *loopback.Bridge

	secret string

	log *slog.Logger

	cancel context.CancelFunc

	done chan struct{}

	closeOnce sync.Once

	mu sync.Mutex

	cmd *exec.Cmd

	url string
}

func (r *registration) URL() string {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.url
}

func (r *registration) Provider() string {
	return "iroh (" + URLScheme + "://<ticket>)"
}

func (r *registration) Close() error {
	r.closeOnce.Do(func() {
		r.cancel()
		r.mu.Lock()
		cmd := r.cmd
		r.mu.Unlock()
		if cmd != nil && cmd.Process != nil {
			_ = cmd.Process.Kill()
		}
		_ = r.bridge.Close()
		<-r.done
		r.provider.forget(r)
	})
	return nil
}

type child struct {
	cmd *exec.Cmd

	ticket string

	tail *childproc.Tail

	started time.Time

	output *sync.WaitGroup
}

func (c *child) stop() {
	childproc.Kill(c.cmd)
	c.output.Wait()
}

func (r *registration) startChild(ctx context.Context) (*child, error) {
	cmd := exec.CommandContext(ctx, r.provider.cfg.DumbpipePath, "listen-tcp", "--host", r.bridge.Address())
	cmd.Env = childEnv(r.secret, r.provider.cfg.RelayURL)
	cmd.Stdout = io.Discard
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("ingressproxy iroh: dumbpipe stderr: %w", err)
	}
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("ingressproxy iroh: starting dumbpipe: %w", err)
	}
	tail := &childproc.Tail{}
	found := make(chan string, 1)
	output := &sync.WaitGroup{}
	output.Go(func() { pump(stderr, tail, found) })
	stop := func() {
		childproc.Kill(cmd)
		output.Wait()
	}
	timer := time.NewTimer(TicketTimeout)
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			stop()
			return nil, ctx.Err()
		case <-timer.C:
			stop()
			return nil, fmt.Errorf("ingressproxy iroh: dumbpipe printed no ticket within %s: %s", TicketTimeout, strings.TrimSpace(tail.String()))
		case ticket, ok := <-found:
			if !ok {
				stop()
				return nil, fmt.Errorf("ingressproxy iroh: dumbpipe ended before printing a ticket: %s", strings.TrimSpace(tail.String()))
			}
			if err := irohdial.CheckTicket(ticket); err != nil {
				stop()
				return nil, fmt.Errorf("ingressproxy iroh: dumbpipe printed %q: %w", ticket, err)
			}
			if _, err := clusterlink.ParseTicket(ticket); err != nil {
				stop()
				return nil, fmt.Errorf("ingressproxy iroh: dumbpipe printed %q: %w", ticket, err)
			}
			r.log.Info("ingressproxy_iroh_dumbpipe_listening", "address", r.bridge.Address(), "ticket", ticket)
			return &child{cmd: cmd, ticket: ticket, tail: tail, started: time.Now(), output: output}, nil
		}
	}
}

func (r *registration) supervise(ctx context.Context, first *child) {
	defer close(r.done)
	current := first
	backoff := MinRestartBackoff
	for {
		if current == nil {
			next, err := r.startChild(ctx)
			if err != nil {
				if ctx.Err() != nil {
					return
				}
				r.log.Warn("ingressproxy_iroh_dumbpipe_restart_failed", "err", err.Error())
				backoff = childproc.NextBackoff(backoff, MinRestartBackoff, MaxRestartBackoff)
				if !clock.Sleep(ctx, backoff) {
					return
				}
				continue
			}
			current = next
			r.mu.Lock()
			r.cmd = next.cmd
			r.url = URLScheme + "://" + next.ticket
			r.mu.Unlock()
		}
		started := current.started
		err := current.cmd.Wait()
		current.output.Wait()
		stderr := strings.TrimSpace(current.tail.String())
		r.mu.Lock()
		if r.cmd == current.cmd {
			r.cmd = nil
		}
		r.mu.Unlock()
		current = nil
		if ctx.Err() != nil {
			return
		}
		if time.Since(started) >= time.Minute {
			backoff = MinRestartBackoff
		}
		r.log.Warn("ingressproxy_iroh_dumbpipe_exited", "err", fmt.Sprint(err), "stderr", stderr)
		if !clock.Sleep(ctx, backoff) {
			return
		}
		backoff = childproc.NextBackoff(backoff, MinRestartBackoff, MaxRestartBackoff)
	}
}

func pump(reader io.Reader, tail *childproc.Tail, found chan<- string) {
	defer close(found)
	scanner := bufio.NewScanner(reader)
	scanner.Buffer(make([]byte, 4096), 64*1024)
	for scanner.Scan() {
		line := scanner.Text()
		tail.Write([]byte(line + "\n"))
		if ticket, ok := ticketIn(line); ok {
			select {
			case found <- ticket:
			default:
			}
		}
	}
}

func ticketIn(line string) (string, bool) {
	_, rest, ok := strings.Cut(line, "dumbpipe connect-tcp ")
	if !ok {
		return "", false
	}
	fields := strings.Fields(rest)
	if len(fields) == 0 {
		return "", false
	}
	ticket := fields[len(fields)-1]
	if !strings.HasPrefix(ticket, clusterlink.TicketPrefix) {
		return "", false
	}
	return ticket, true
}

func childEnv(secret, relay string) []string {
	env := make([]string, 0, len(os.Environ())+2)
	for _, entry := range os.Environ() {
		key, _, _ := strings.Cut(entry, "=")
		if key == clusterlink.EnvSecret || key == "IROH_RELAY_URL" {
			continue
		}
		env = append(env, entry)
	}
	env = append(env, clusterlink.EnvSecret+"="+secret)
	if relay = strings.TrimSpace(relay); relay != "" {
		env = append(env, "IROH_RELAY_URL="+relay)
	}
	return env
}