shark-expose/internal/expose/expose.go

// Package expose runs one exposure: ask Hammerhead for it, bring up the tunnel, keep it alive,
// and take everything down again when told to stop.
package expose

import (
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strconv"
	"time"

	"git.veillax.com/Veillax/bamboo"
	"git.veillax.com/Veillax/shark-expose/internal/api"
)

type Options struct {
	Request api.Request
	TTL     time.Duration
	// LiveWait is how long to wait for the edge to start listening before giving up on a "live" confirmation.
	LiveWait time.Duration
	Out      io.Writer // where status lines go; nil = discard
	Verbose  bool      // let the tunnel agent log
}

func (o *Options) say(format string, a ...any) {
	if o.Out != nil {
		fmt.Fprintf(o.Out, format+"\n", a...)
	}
}

// Run blocks until ctx is cancelled or the exposure is lost, and always removes what it created.
// onReady is called once with the exposure after the edge is listening (or LiveWait has passed).
func Run(ctx context.Context, c *api.Client, o Options, onReady func(api.Exposure)) error {
	if o.TTL <= 0 {
		o.TTL = 5 * time.Minute
	}
	if o.LiveWait <= 0 {
		o.LiveWait = 20 * time.Second
	}
	o.Request.TTLSeconds = int(o.TTL.Seconds())

	e, err := c.Create(ctx, o.Request)
	if err != nil {
		return err
	}
	// Whatever happens next, the exposure must not outlive us. The caller's ctx is probably
	// cancelled by then, so cleanup gets its own.
	defer func() {
		cctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
		defer cancel()
		if err := c.Delete(cctx, e.ID); err != nil {
			o.say("could not remove the exposure (it will expire by itself): %v", err)
		}
	}()

	dir, err := os.MkdirTemp("", "shark-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(dir)
	if !o.Verbose {
		log.SetOutput(io.Discard)
	}
	agent, err := bamboo.New(bamboo.Config{Hammerhead: c.Base, Token: e.SiteToken, DataDir: dir})
	if err != nil {
		return err
	}
	actx, stopAgent := context.WithCancel(ctx)
	agentDone := make(chan struct{})
	go func() { defer close(agentDone); agent.Run(actx) }()
	defer func() {
		stopAgent()
		select {
		case <-agentDone:
		case <-time.After(5 * time.Second):
		}
	}()

	// Hammerhead tells us when the edge has the port open.
	deadline := time.Now().Add(o.LiveWait)
	for !e.Live && time.Now().Before(deadline) {
		select {
		case <-ctx.Done():
			return nil
		case <-time.After(300 * time.Millisecond):
		}
		if got, err := c.Get(ctx, e.ID); err == nil {
			e.Live, e.Error, e.ExpiresAt = got.Live, got.Error, got.ExpiresAt
		}
	}
	if e.Error != "" {
		return fmt.Errorf("the edge could not open the port: %s", e.Error)
	}
	if onReady != nil {
		onReady(e)
	}
	return keepAlive(ctx, c, e.ID, o)
}

// keepAlive renews the exposure until ctx ends. It gives up once Hammerhead has forgotten the
// exposure, or when renewals have been failing for longer than the TTL (it has expired by then).
func keepAlive(ctx context.Context, c *api.Client, id string, o Options) error {
	every := o.TTL / 3
	if every < 5*time.Second {
		every = 5 * time.Second
	}
	if every > time.Minute {
		every = time.Minute
	}
	t := time.NewTicker(every)
	defer t.Stop()
	lastOK := time.Now()
	for {
		select {
		case <-ctx.Done():
			return nil
		case <-t.C:
		}
		_, err := c.Renew(ctx, id, o.TTL)
		switch {
		case err == nil:
			lastOK = time.Now()
		case ctx.Err() != nil:
			return nil
		case api.IsStatus(err, 404):
			return fmt.Errorf("Hammerhead no longer has this exposure (it expired or was removed)")
		case time.Since(lastOK) > o.TTL:
			return fmt.Errorf("lost contact with Hammerhead: %v", err)
		default:
			o.say("could not renew, will retry: %v", err)
		}
	}
}

// Listening reports whether something accepts connections at host:port.
func Listening(host string, port int) bool {
	c, err := net.DialTimeout("tcp", net.JoinHostPort(host, strconv.Itoa(port)), time.Second)
	if err != nil {
		return false
	}
	c.Close()
	return true
}