shark-expose/internal/expose/expose_test.go

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)
	}
}