shark expose: put a local port on the internet through Hammerhead
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
VeillaxandClaude Sonnet 5.5 committed
commit
61143d107c
15 files changed
+1538
No files matched your search
@@ -0,0 +1,146 @@
|
||||
// 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
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package expose
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.veillax.com/Veillax/shark-expose/internal/api"
|
||||
)
|
||||
|
||||
// fakeHammerhead implements just enough of the expose API to watch what the client does.
|
||||
type fakeHammerhead struct {
|
||||
mu sync.Mutex
|
||||
calls []string
|
||||
renewErr int // status to answer renewals with; 0 = ok
|
||||
gone atomic.Bool
|
||||
srv *httptest.Server
|
||||
created api.Request
|
||||
}
|
||||
|
||||
func newFake(t *testing.T) *fakeHammerhead {
|
||||
f := &fakeHammerhead{}
|
||||
mux := http.NewServeMux()
|
||||
reply := func(w http.ResponseWriter, code int, v any) {
|
||||
w.WriteHeader(code)
|
||||
json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
exp := api.Exposure{ID: "abc", Proto: "tcp", Host: "edge.test", Port: 20001, SiteToken: "tok", Live: true}
|
||||
mux.HandleFunc("POST /api/orgs/-/expose", func(w http.ResponseWriter, r *http.Request) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.calls = append(f.calls, "create")
|
||||
if r.Header.Get("Authorization") != "Bearer hhk_x_y" {
|
||||
reply(w, 401, map[string]string{"error": "Sign in to continue."})
|
||||
return
|
||||
}
|
||||
json.NewDecoder(r.Body).Decode(&f.created)
|
||||
reply(w, 200, exp)
|
||||
})
|
||||
mux.HandleFunc("GET /api/orgs/-/expose/abc", func(w http.ResponseWriter, r *http.Request) { reply(w, 200, exp) })
|
||||
mux.HandleFunc("POST /api/orgs/-/expose/abc/renew", func(w http.ResponseWriter, r *http.Request) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.calls = append(f.calls, "renew")
|
||||
if f.renewErr != 0 {
|
||||
reply(w, f.renewErr, map[string]string{"error": "no such exposure"})
|
||||
return
|
||||
}
|
||||
reply(w, 200, exp)
|
||||
})
|
||||
mux.HandleFunc("DELETE /api/orgs/-/expose/abc", func(w http.ResponseWriter, r *http.Request) {
|
||||
f.mu.Lock()
|
||||
f.calls = append(f.calls, "delete")
|
||||
f.mu.Unlock()
|
||||
f.gone.Store(true)
|
||||
reply(w, 200, map[string]bool{"ok": true})
|
||||
})
|
||||
// the tunnel agent long-polls this; there is nothing for it to do
|
||||
mux.HandleFunc("/site/", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(204) })
|
||||
f.srv = httptest.NewServer(mux)
|
||||
t.Cleanup(f.srv.Close)
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *fakeHammerhead) count(call string) int {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
n := 0
|
||||
for _, c := range f.calls {
|
||||
if c == call {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func TestRunRenewsAndDeletesOnStop(t *testing.T) {
|
||||
f := newFake(t)
|
||||
c := api.New(f.srv.URL, "hhk_x_y")
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
ready := make(chan api.Exposure, 1)
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- Run(ctx, c, Options{Request: api.Request{Proto: "tcp", TargetPort: 3000}, TTL: 15 * time.Second}, func(e api.Exposure) { ready <- e })
|
||||
}()
|
||||
e := <-ready
|
||||
if e.Address() != "tcp://edge.test:20001" {
|
||||
t.Errorf("address %q", e.Address())
|
||||
}
|
||||
if f.created.TTLSeconds != 15 || f.created.TargetPort != 3000 {
|
||||
t.Errorf("request %+v", f.created)
|
||||
}
|
||||
cancel()
|
||||
if err := <-done; err != nil {
|
||||
t.Fatalf("clean stop returned %v", err)
|
||||
}
|
||||
if !f.gone.Load() {
|
||||
t.Fatal("exposure was not deleted on stop")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunStopsWhenHammerheadForgetsTheExposure(t *testing.T) {
|
||||
f := newFake(t)
|
||||
f.renewErr = 404
|
||||
c := api.New(f.srv.URL, "hhk_x_y")
|
||||
// renewals are never sooner than 5s; a short TTL is enough to reach one
|
||||
err := Run(context.Background(), c, Options{Request: api.Request{Proto: "tcp", TargetPort: 1}, TTL: 3 * time.Second}, nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected an error once the exposure expired")
|
||||
}
|
||||
if !f.gone.Load() {
|
||||
t.Error("cleanup should still ask Hammerhead to delete")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunRejectedToken(t *testing.T) {
|
||||
f := newFake(t)
|
||||
err := Run(context.Background(), api.New(f.srv.URL, "wrong"), Options{Request: api.Request{Proto: "tcp", TargetPort: 1}}, nil)
|
||||
if !api.IsStatus(err, 401) {
|
||||
t.Fatalf("got %v", err)
|
||||
}
|
||||
if f.count("delete") != 0 {
|
||||
t.Error("deleted something that was never created")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunLoopSurvivesFlakyRenewals(t *testing.T) {
|
||||
f := newFake(t)
|
||||
f.renewErr = 500
|
||||
c := api.New(f.srv.URL, "hhk_x_y")
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 7*time.Second)
|
||||
defer cancel()
|
||||
err := keepAlive(ctx, c, "abc", Options{TTL: time.Minute})
|
||||
if err != nil {
|
||||
t.Fatalf("a failing renewal inside the TTL should not end the run: %v", err)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user