Rework the AutoRegister API and add comprehensive integration tests!

I worked on this in pieces over several weekends. First I got carried
away rewriting things and then later I added the tests.

For AutoRegister changes, I didn't like passing in WaitGroup as a
parameter in the API and I felt like it added too much boilerplate to
call the function. I also didn't like breaking standard go conventions
and returning the error in the struct. Later on after I wrote the tests
I also realised the "return only the first registration" API is nice to
use without the default client as well so I made it a client-specific
version.

In order to make testing work easily I exposed the LeaseTTL as a global
variable. I'm now thinking I want to change it to a flag so each
instance of embedportal.Run can have a different TTL. I already had too
many changes going on and I had to stop and commit.

Tools changes: With the AutoRegister changes I ended up making
tools.AutoTLSConf, but I'm not entirely happy with the API yet. To
tools.HTTPServer I added the Listener option mainly so that I could use
a file descriptor listener in the tests. This also allows binding to a
specific IP. Kinda makes the port argument useless if you use it but
it's nice to have the extra option and get the nice shutdown code.

The host, assimilate, and spawn changes are for the most part to use the
new API so this commit compiles. Also in spawn I ended up adding a
context and WaitGroup to the dashboard. I think I need to rework this
some too.

When I wrote the tests I also found and fixed two bugs:

 1. In gate.ResolveFlags() the conditions weren't quite right for people
    setting the pointers directly in code
 2. There was a subtle threading issue in clientLeasor that I discovered
    when I made the tests parallel. Because I was trying to cache the
    heap pointers myself setting the nextLeasor variable wasn't a thread
    safe operation and it was possible to have race issues. The proper
    fix is to use sync.Pool, it was a bit complicated to get exactly
    right.

Last and most importantly, I did a lot of refinement on the tests before
committing.

It runs one portal server for all of the parallel tests and uses open
ports automatically selected by the OS and even passes those open
sockets to portal via FDs so that there's no thread race issue when we
close it and portal reopens. Lots of interesting careful use of
t.Cleanup.

I also wrote some cool test helpers for subtests using reflection. So
you don't have to manually name the subtests and update the list of
t.Run calls when you add new test functions. It's almost like it's built
in syntax for go test. Also the FreePort function is kind of reusable,
unfortunately I had to make it ugly to support windows.

At first I wrote it to get the portal token by writing the state file to
the test temp dir and then repeatedly reading the state file and parsing
the proto until the token appeared. I ended up changing it over to
reading the logs like portal does because I didn't really like polling
the filesystem and using it for communication like that. Also I decided
I want to be able to test the way spawn actually works and I will write
a portal restart test which proves the state file works without actually
parsing it. Unfortunately it ended up longer code. I think I need to
make extracting the token from portal better but I'm not sure how yet.

As for the actual tests: I was careful to cover all the branches of
AutoRegister although it's not perfect yet. I left lots of TODO
comments. I'm really happy with the cleanup sequences and for example
the extra check that unregister happened at the end of the HTTP test. I
really like that I was able to stick to only the public interface and
make such a nice setup.

Author: fsmv

.github/workflows/tests.yml

@@ -18,12 +18,15 @@ jobs:
     - uses: actions/checkout@v4
 
     - name: Set up Go
-      uses: actions/setup-go@v5
+      uses: actions/setup-go@v5.5.0
       with:
         go-version-file: go.mod
 
     - name: Build daemon
-      run: go build -race -v ./...
+      shell: bash
+      run: |
+        GOTOOLCHAIN=local go version
+        GOTOOLCHAIN=local go build -race -v ./...
 
     - name: Run daemon unit tests
       run: go test -race -v ./...
@@ -64,7 +67,7 @@ jobs:
         path: 'pebble'
 
     - name: Set up Go
-      uses: actions/setup-go@v5
+      uses: actions/setup-go@v5.5.0
       with:
         go-version-file: pebble/go.mod
         cache-dependency-path: |
@@ -137,7 +140,7 @@ jobs:
         path: 'pebble'
 
     - name: Set up Go
-      uses: actions/setup-go@v5
+      uses: actions/setup-go@v5.5.0
       with:
         go-version-file: pebble/go.mod
         cache-dependency-path: |

assimilate/embedassimilate/assimilate.go

@@ -70,7 +70,7 @@ func Run(ctx context.Context, flags *flag.FlagSet, args []string) {
 		idx := i
 		wg.Add(1)
 		go func() {
-			err := client.AutoRegister(ctx, registration, nil)
+			err := client.AutoRegisterChan(ctx, registration, nil)
 			wg.Done()
 			if err != nil && !errors.Is(err, context.Cause(ctx)) {
 				log.Printf("Error for registration #%v: %v", idx, err)

go.mod

@@ -1,12 +1,14 @@
 module ask.systems/daemon
 
-go 1.20
+go 1.21.13
+
+toolchain go1.24.0
 
 require (
 	golang.org/x/crypto v0.33.0
 	golang.org/x/sys v0.30.0
 	google.golang.org/grpc v1.64.1
-	google.golang.org/protobuf v1.34.2
+	google.golang.org/protobuf v1.36.5
 )
 
 require (

go.sum

@@ -1,4 +1,5 @@
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 golang.org/x/crypto v0.33.0 h1:IOBPskki6Lysi0lo9qQvbxiQ+FvsCC/YWOecCHAixus=
 golang.org/x/crypto v0.33.0/go.mod h1:bVdXmD7IV/4GdElGPozy6U7lWdRXA4qyRVGJV57uQ5M=
 golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8=
@@ -11,5 +12,5 @@ google.golang.org/genproto/googleapis/rpc v0.0.0-20240318140521-94a12d6c2237 h1:
 google.golang.org/genproto/googleapis/rpc v0.0.0-20240318140521-94a12d6c2237/go.mod h1:WtryC6hu0hhx87FDGxWCDptyssuo68sk10vYjF+T9fY=
 google.golang.org/grpc v1.64.1 h1:LKtvyfbX3UGVPFcGqJ9ItpVWW6oN/2XqTxfAnwRRXiA=
 google.golang.org/grpc v1.64.1/go.mod h1:hiQF4LFZelK2WKaP6W0L92zGHtiQdZxk8CrSdvyjeP0=
-google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
-google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
+google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM=
+google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=

host/embedhost/host.go

@@ -13,7 +13,6 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
-	"sync"
 
 	_ "ask.systems/daemon/portal/flags"
 	_ "ask.systems/daemon/tools/flags"
@@ -47,24 +46,6 @@ func Run(ctx context.Context, flags *flag.FlagSet, args []string) {
 		"instead of hosting a server.")
 	flags.Parse(args[1:])
 
-	// Setup graceful stopping
-	ctx, cancel := context.WithCancel(context.Background())
-	ctx = tools.ContextWithQuitSignals(ctx)
-	wg := &sync.WaitGroup{}
-	defer func() {
-		cancel()
-		wg.Wait()
-		log.Print("Goodbye.")
-	}()
-
-	reg, err := gate.AutoRegister(ctx, &gate.RegisterRequest{
-		Pattern: *urlPath,
-	}, wg)
-	if err != nil {
-		log.Print("Fatal error registering with portal:", err)
-		return
-	}
-
 	// Extract the path part of the pattern and the prefix to remove
 	_, servePath := gate.ParsePattern(*urlPath)
 	prefix := servePath
@@ -110,7 +91,23 @@ func Run(ctx context.Context, flags *flag.FlagSet, args []string) {
 		log.Print("WARNING: Failed to open and stat web_root directory, we probably can't serve anything. Error: ", err)
 	}
 
+	// Setup graceful stopping
+	ctx = tools.ContextWithQuitSignals(context.Background())
+
+	// Register the reverse proxy pattern with portal.
+	// Only blocks until the initial registration is done, then keeps renewing.
+	reg, waitForUnregister, err := gate.AutoRegister(ctx, &gate.RegisterRequest{
+		Pattern: *urlPath,
+	})
+	if err != nil {
+		log.Print("Fatal error registering with portal:", err)
+		return
+	}
+	// Start serving files (blocks until graceful stop is done)
 	tools.HTTPServer(ctx.Done(), reg.Lease.Port, reg.TLSConfig, nil)
+	// Wait for the AutoRegister background goroutine to gracefully stop
+	<-waitForUnregister
+	log.Print("Goodbye.")
 }
 
 // The -hash_pasword utility

portal/embedportal/http_proxyserv.go

@@ -50,7 +50,6 @@ type httpProxy struct {
 type forwarder struct {
 	Handler   http.Handler
 	Lease     *gate.Lease
-	AdminOnly bool
 	AllowHTTP bool
 }
 

portal/embedportal/port_leasor.go

@@ -14,7 +14,7 @@ import (
 )
 
 // How often to look through the Leases and unregister those past TTL
-const (
+var (
 	ttlCheckFreq = leaseTTL / 100
 )
 
@@ -27,8 +27,8 @@ var (
 type onCancelFunc func(*gate.Lease)
 
 type clientLeasor struct {
-	leasors    *sync.Map
-	nextLeasor *portLeasor
+	leasors    *sync.Map  // Key: string Value: *portLeasor
+	leasorPool *sync.Pool // Contains *portLeasor
 
 	startPort     uint16
 	endPort       uint16
@@ -43,8 +43,7 @@ func makeClientLeasor(startPort, endPort uint16, reservedPorts map[uint16]bool,
 	if endPort < startPort {
 		startPort, endPort = endPort, startPort
 	}
-	return &clientLeasor{
-		nextLeasor:  &portLeasor{},
+	c := &clientLeasor{
 		leasors:     &sync.Map{},
 		onCancelMut: &sync.Mutex{},
 
@@ -53,6 +52,27 @@ func makeClientLeasor(startPort, endPort uint16, reservedPorts map[uint16]bool,
 		reservedPorts: reservedPorts,
 		quit:          quit,
 	}
+	c.leasorPool = &sync.Pool{
+		New: func() any {
+			return &portLeasor{
+				mut:       &sync.Mutex{},
+				startPort: startPort,
+				endPort:   endPort,
+				// TODO: clean up this weird system with copyOnCancel and and c.OnCancel
+				// calling OnCancel on all the portLeasors. I think the onCancel list
+				// should not be copied to each portLeasor and instead they should
+				// report cancels back to the clientLeasor which should keep the only
+				// list.
+				//
+				// I think leasorPool.New is garunteed to only be called when we call
+				// Get and that is the only reason this system is safe
+				onCancel:    c.copyOnCancel(),
+				leases:      make(map[uint32][]*gate.Lease),
+				unusedPorts: makeUnusedPorts(startPort, endPort, reservedPorts),
+			}
+		},
+	}
+	return c
 }
 
 func leaseString(lease *gate.Lease) string {
@@ -62,22 +82,32 @@ func leaseString(lease *gate.Lease) string {
 }
 
 func (c *clientLeasor) PortLeasorForClient(clientAddr string) *portLeasor {
-	leasor, loaded := c.leasors.LoadOrStore(clientAddr, c.nextLeasor)
+	// We use a sync.Pool because we have to call LoadOrStore on the leasors map,
+	// we need to have a pointer to space on the heap every time we lookup a
+	// client so that we can start a new leasor if one doesn't exist.
+	//
+	// The pool lets us save and re-use the heap space here when we find a client
+	// we already have a leasor for.
+	nextLeasor := c.leasorPool.Get()
+	// This thread exclusively owns nextLeasor right now becasue we called Get and
+	// have not stored it in the map yet.
+	nextLeasor.(*portLeasor).clientAddr = clientAddr
+	leasor, loaded := c.leasors.LoadOrStore(clientAddr, nextLeasor)
 	if !loaded {
-		// Because we have to call LoadOrStore on the leasors map, we need to have a
-		// pointer to space on the heap every time we lookup a client so that we can
-		// start a new leasor if one doesn't exist.
-		//
-		// So, save and re-use the heap space here when we find a client we already
-		// have a leasor for.
-		c.nextLeasor.Start(clientAddr, c.startPort, c.endPort,
-			makeUnusedPorts(c.startPort, c.endPort, c.reservedPorts),
-			c.quit, c.copyOnCancel())
-		c.nextLeasor = &portLeasor{}
+		// This same pointer might have been returned from a concurrent call before
+		// we start this goroutine, but this will still only happen once and it's
+		// okay if we start using leasor before this.
+		go leasor.(*portLeasor).monitorTTLs(c.quit)
+	} else {
+		c.leasorPool.Put(nextLeasor)
 	}
 	return leasor.(*portLeasor)
 }
 
+func (c *clientLeasor) makePortLeasor() any {
+	return &portLeasor{}
+}
+
 func (c *clientLeasor) copyOnCancel() []onCancelFunc {
 	c.onCancelMut.Lock()
 	defer c.onCancelMut.Unlock()
@@ -90,9 +120,9 @@ func (c *clientLeasor) OnCancel(cancelFunc func(*gate.Lease)) {
 	c.onCancelMut.Lock()
 	defer c.onCancelMut.Unlock()
 	c.onCancel = append(c.onCancel, cancelFunc)
-	// Since we still have the mutex lock, this covers all existing leasors. New
-	// leasors can't be created until the mutex is released and they'll get the
-	// new list.
+	// Since we still have the mutex lock and copyOnCancel also uses it, this
+	// covers all existing leasors. New leasors can't be created until the mutex
+	// is released and they'll get the new list.
 	c.leasors.Range(func(key, value interface{}) bool {
 		l := value.(*portLeasor)
 		l.OnCancel(cancelFunc)
@@ -117,19 +147,6 @@ type portLeasor struct {
 	clientAddr  string
 }
 
-func (l *portLeasor) Start(clientAddr string, startPort, endPort uint16, unusedPorts []uint16, quit chan struct{}, onCancel []onCancelFunc) {
-	*l = portLeasor{
-		mut:         &sync.Mutex{},
-		clientAddr:  clientAddr,
-		startPort:   startPort,
-		endPort:     endPort,
-		onCancel:    onCancel,
-		leases:      make(map[uint32][]*gate.Lease),
-		unusedPorts: unusedPorts,
-	}
-	go l.monitorTTLs(quit)
-}
-
 func (l *portLeasor) OnCancel(cancelFunc func(*gate.Lease)) {
 	l.mut.Lock()
 	defer l.mut.Unlock()

portal/embedportal/portal.go

@@ -27,15 +27,17 @@ import (
 //go:generate protoc -I ../ ../embedportal/storage.proto --go_out ../ --go_opt=paths=source_relative
 //go:generate protoc -I ../ ../gate/service.proto --go_out ../ --go-grpc_out ../ --go_opt=paths=source_relative --go-grpc_opt=paths=source_relative
 
-const (
-	leaseTTL         = 24 * time.Hour
+var (
+	LeaseTTL         = 24 * time.Hour
+	leaseTTL         = LeaseTTL // TODO
 	ttlRandomStagger = 0.05
 )
 
 var kACMEAddress string
 
 // TODO: actually use ctx
 func Run(ctx context.Context, flags *flag.FlagSet, args []string) {
+	leaseTTL = LeaseTTL // TODO
 	flags.Usage = func() {
 		fmt.Fprintf(flags.Output(), ""+
 			"Usage: %s [flags]\n"+
@@ -63,7 +65,7 @@ func Run(ctx context.Context, flags *flag.FlagSet, args []string) {
 			flags.Name(), flags.Name())
 		flags.PrintDefaults()
 	}
-	rpcPort := flags.Uint("rpc_port", 2048, ""+
+	rpcPort := flags.Int("rpc_port", 2048, ""+
 		"The port to bind for the portal RPC server that clients use to register\n"+
 		"with. You shouldn't need to change this unless there's a conflict or you\n"+
 		"run multiple instances of portal.")
@@ -198,7 +200,7 @@ func Run(ctx context.Context, flags *flag.FlagSet, args []string) {
 	// Starts serving the rpc server port.
 	// First loads the registrations from the state into the two proxy servers.
 	_, err = startRPCServer(leasor,
-		tcpProxy, httpProxy, uint16(*rpcPort),
+		tcpProxy, httpProxy, *rpcPort,
 		rootCert, state, quit)
 	if err != nil {
 		log.Fatal("Failed to start RPC server:", err)

portal/embedportal/rpcserv.go

@@ -8,7 +8,6 @@ import (
 	"fmt"
 	"log"
 	"net"
-	"strconv"
 	"strings"
 	"time"
 
@@ -220,7 +219,7 @@ func (s *rpcServ) Renew(ctx context.Context, lease *gate.Lease) (*gate.Lease, er
 // StartNew creates a new RPCServ and starts it
 func startRPCServer(clientLeasor *clientLeasor,
 	tcpProxy *tcpProxy, httpProxy *httpProxy,
-	port uint16, rootCert *tls.Config,
+	port int, rootCert *tls.Config,
 	state *stateManager, quit chan struct{}) (*rpcServ, error) {
 
 	s := &rpcServ{
@@ -259,7 +258,7 @@ func startRPCServer(clientLeasor *clientLeasor,
 		}),
 	)
 	gate.RegisterPortalServer(server, s)
-	l, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
+	l, err := listenerFromPortOrFD(port)
 	if err != nil {
 		return nil, fmt.Errorf("Failed to start listener: %v", err)
 	}

portal/embedportal/tls_cert.go

@@ -194,7 +194,7 @@ func (t *tlsRefresher) refreshCertFile(idx int, cert, key *os.File) (*tls.Certif
 		newCertFile, errCert := os.Open(cert.Name())
 		newKeyFile, errKey := os.Open(key.Name())
 		if errCert != nil || errKey != nil {
-			err := fmt.Errorf("Failed to reopen tls cert (%w) or key (%w) files", errCert, errKey)
+			err := fmt.Errorf("Failed to reopen TLS cert (%w) or key (%w) files", errCert, errKey)
 			log.Print(err.Error())
 			newCertFile.Close()
 			newKeyFile.Close()