Merge pull request #225 from aojea/router_reenroll

Router reenroll

Author: aojea

internal/node/enroll.go

@@ -150,6 +150,10 @@ func (n *SamNode) Enroll(ctx context.Context, hubURL string, jwt string) error {
 	n.keysMu.Unlock()
 
 	// Connect and Auth to hub after enrollment to join the mesh
+	if len(enrollResp.HubAddresses) == 0 {
+		return fmt.Errorf("failed to connect and authenticate after HTTP enrollment: control plane returned no hub addresses")
+	}
+
 	var lastAuthErr error
 	var authed bool
 	for _, addrStr := range enrollResp.HubAddresses {
@@ -329,6 +333,10 @@ func (n *SamNode) EnrollBootstrap(ctx context.Context, hubURL string, bootstrapT
 	n.keysMu.Unlock()
 
 	// Connect and Auth to hub after enrollment to join the mesh
+	if len(enrollResp.HubAddresses) == 0 {
+		return fmt.Errorf("failed to connect and authenticate after bootstrap enrollment: control plane returned no hub addresses")
+	}
+
 	var lastAuthErr error
 	var authed bool
 	for _, addrStr := range enrollResp.HubAddresses {

internal/node/node.go

@@ -638,13 +638,13 @@ func (n *SamNode) startConnectionMonitor(ctx context.Context, bootstrapDuration,
 }
 
 func (n *SamNode) RegisterStaticServices(ctx context.Context, services []api.ServiceConfig) error {
-	// Wait for DHT to be ready (size > 0)
+	// Wait for node to be connected to a hub or DHT to be ready
 	// This avoids failure if we try to register immediately after enrollment
-	// before the DHT has discovered peers.
+	// before the connection is established.
 	ticker := time.NewTicker(100 * time.Millisecond)
 	defer ticker.Stop()
 
-	timeout := time.After(10 * time.Second) // 10 seconds timeout for DHT readiness
+	timeout := time.After(10 * time.Second)
 
 dhtLoop:
 	for {
@@ -654,7 +654,7 @@ dhtLoop:
 		case <-timeout:
 			return fmt.Errorf("timeout waiting for DHT to be ready before registering static services")
 		case <-ticker.C:
-			if n.DHT.RoutingTable().Size() > 0 {
+			if n.IsConnected() || (n.DHT != nil && n.DHT.RoutingTable().Size() > 0) {
 				break dhtLoop
 			}
 		}

internal/router/router.go

@@ -94,6 +94,7 @@ type Router struct {
 	biscuitExpiration time.Time
 	trustedPublicKeys []ed25519.PublicKey
 	keysMu            sync.RWMutex
+	enrollMu          sync.Mutex
 	privKey           crypto.PrivKey
 
 	// Control contexts
@@ -320,9 +321,9 @@ func (r *Router) enroll(peerID peer.ID) error {
 		return err
 	}
 
+	r.keysMu.Lock()
 	r.biscuitToken = enrollResp.BiscuitToken
 	r.biscuitExpiration = time.Unix(enrollResp.Expiration, 0)
-	r.keysMu.Lock()
 	r.trustedPublicKeys = []ed25519.PublicKey{enrollResp.HubPublicKey}
 	r.keysMu.Unlock()
 
@@ -442,9 +443,9 @@ func (r *Router) enrollBootstrap(peerID peer.ID) error {
 		return fmt.Errorf("enrollment not approved (status: %v)", enrollResp.Status)
 	}
 
+	r.keysMu.Lock()
 	r.biscuitToken = enrollResp.BiscuitToken
 	r.biscuitExpiration = time.Unix(enrollResp.Expiration, 0)
-	r.keysMu.Lock()
 	r.trustedPublicKeys = []ed25519.PublicKey{enrollResp.HubPublicKey}
 	r.keysMu.Unlock()
 
@@ -456,6 +457,22 @@ func (r *Router) enrollBootstrap(peerID peer.ID) error {
 	return nil
 }
 
+func (r *Router) reEnroll() error {
+	r.enrollMu.Lock()
+	defer r.enrollMu.Unlock()
+
+	if r.Host == nil {
+		return fmt.Errorf("cannot re-enroll: router host is not initialized")
+	}
+
+	if r.config.BootstrapToken != "" {
+		return r.enrollBootstrap(r.Host.ID())
+	} else if r.config.OIDCToken != "" {
+		return r.enroll(r.Host.ID())
+	}
+	return fmt.Errorf("no enrollment token available for re-enrollment")
+}
+
 func (r *Router) syncKeys() error {
 	client := &http.Client{Timeout: 10 * time.Second}
 	resp, err := client.Get(r.config.ControlPlaneURL + "/keys")
@@ -534,7 +551,11 @@ func (r *Router) runLeaseRenewalLoop() {
 }
 
 func (r *Router) renewLease() {
-	if len(r.biscuitToken) == 0 {
+	r.keysMu.RLock()
+	biscuit := r.biscuitToken
+	r.keysMu.RUnlock()
+
+	if len(biscuit) == 0 {
 		logger.Warn("Cannot renew lease: router is not enrolled (no biscuit)")
 		return
 	}
@@ -564,7 +585,7 @@ func (r *Router) renewLease() {
 	req := &api.RouterLeaseRequest{
 		PeerId:         r.Host.ID().String(),
 		Addresses:      addrs,
-		Biscuit:        r.biscuitToken,
+		Biscuit:        biscuit,
 		ConnectedPeers: connectedPeers,
 		DhtSize:        dhtSize,
 	}
@@ -578,6 +599,18 @@ func (r *Router) renewLease() {
 	}
 	defer func() { _ = resp.Body.Close() }()
 
+	if resp.StatusCode == http.StatusUnauthorized {
+		body, _ := io.ReadAll(resp.Body)
+		logger.Warnf("Control plane lease renewal rejected (401 Unauthorized: %s), attempting re-enrollment...", string(body))
+		if err := r.reEnroll(); err != nil {
+			logger.Errorf("Re-enrollment failed after 401 Unauthorized lease renewal: %v", err)
+			return
+		}
+		logger.Info("Successfully re-enrolled after 401 Unauthorized lease renewal, retrying lease renewal...")
+		r.renewLease()
+		return
+	}
+
 	if resp.StatusCode != http.StatusOK {
 		body, _ := io.ReadAll(resp.Body)
 		logger.Errorf("Control plane lease renewal rejected, status %s: %s", resp.Status, string(body))
@@ -718,10 +751,14 @@ func (r *Router) HandleAuthHandshake(s network.Stream) {
 	logger.Infof("[AuthN] Successfully authenticated peer %s", remotePeer)
 
 	// Send mutual response (our biscuit)
+	r.keysMu.RLock()
+	ourBiscuit := r.biscuitToken
+	r.keysMu.RUnlock()
+
 	writer := msgio.NewVarintWriter(s)
 	resp := &api.AuthResponse{
 		Success: true,
-		Biscuit: r.biscuitToken,
+		Biscuit: ourBiscuit,
 	}
 	respBytes, _ := proto.Marshal(resp)
 	if err := writer.WriteMsg(respBytes); err != nil {
@@ -859,10 +896,11 @@ func isLoopbackOrLinkLocal(addr multiaddr.Multiaddr) bool {
 
 // RefreshEnrollment trades the expiring biscuit token for a new one using a cryptographic challenge.
 func (r *Router) RefreshEnrollment(ctx context.Context) error {
-	r.keysMu.Lock()
-	defer r.keysMu.Unlock()
+	r.keysMu.RLock()
+	currentBiscuit := r.biscuitToken
+	r.keysMu.RUnlock()
 
-	if len(r.biscuitToken) == 0 {
+	if len(currentBiscuit) == 0 {
 		return fmt.Errorf("router not enrolled (no biscuit)")
 	}
 
@@ -890,8 +928,7 @@ func (r *Router) RefreshEnrollment(ctx context.Context) error {
 		return fmt.Errorf("failed to create http request: %w", err)
 	}
 	httpReq.Header.Set("Content-Type", "application/x-protobuf")
-	// Set current biscuit in authorization header
-	b64Biscuit := base64.StdEncoding.EncodeToString(r.biscuitToken)
+	b64Biscuit := base64.StdEncoding.EncodeToString(currentBiscuit)
 	httpReq.Header.Set("Authorization", "Bearer "+b64Biscuit)
 
 	client := &http.Client{Timeout: 10 * time.Second}
@@ -901,6 +938,12 @@ func (r *Router) RefreshEnrollment(ctx context.Context) error {
 	}
 	defer func() { _ = resp.Body.Close() }()
 
+	if resp.StatusCode == http.StatusUnauthorized {
+		body, _ := io.ReadAll(resp.Body)
+		logger.Warnf("Proactive biscuit refresh rejected (401 Unauthorized: %s), attempting re-enrollment...", string(body))
+		return r.reEnroll()
+	}
+
 	if resp.StatusCode == http.StatusForbidden {
 		logger.Errorf("Refresh rejected: Router is banned (403 Forbidden). Hard-killing router.")
 		if r.Host != nil {
@@ -928,9 +971,12 @@ func (r *Router) RefreshEnrollment(ctx context.Context) error {
 		return fmt.Errorf("refresh error: %s", refreshResp.ErrorMessage)
 	}
 
-	// Update local biscuit token and expiration
+	// Update local biscuit token and expiration under lock
+	r.keysMu.Lock()
 	r.biscuitToken = refreshResp.BiscuitToken
 	r.biscuitExpiration = time.Unix(refreshResp.ExpiresAt, 0)
+	r.keysMu.Unlock()
+
 	logger.Infof("Router biscuit token refreshed successfully.")
 	return nil
 }

internal/router/router_test.go

@@ -457,3 +457,130 @@ func TestRouterProactiveTokenRefresh(t *testing.T) {
 		t.Errorf("expected refreshed expiration %v to be after initial expiration %v", r.biscuitExpiration, initialExpiration)
 	}
 }
+
+func TestRouterLeaseRenewalReEnrollOn401(t *testing.T) {
+	issuer, mintToken := startCustomMockOIDC(t)
+	cp, cpStore, cpURL := setupControlPlane(t, issuer)
+	defer func() {
+		_ = cp.Close()
+		_ = cpStore.Close()
+	}()
+
+	tempDir := t.TempDir()
+	routerKeyPath := filepath.Join(tempDir, "router.key")
+
+	routerJWT := mintToken(map[string]interface{}{
+		"sub":    "router-reenroll-test",
+		"groups": []string{"routers"},
+	})
+
+	rOpts := Options{
+		ControlPlaneURL:    cpURL,
+		ListenAddrs:        []string{"/ip4/127.0.0.1/tcp/0"},
+		KeysSyncInterval:   20 * time.Second,
+		LeaseRenewInterval: 20 * time.Second,
+		OIDCToken:          routerJWT,
+		KeysDBPath:         routerKeyPath,
+		AllowLoopback:      true,
+		BiscuitTimeout:     1 * time.Second,
+	}
+
+	r, err := NewRouter(context.Background(), rOpts)
+	if err != nil {
+		t.Fatalf("failed to create router: %v", err)
+	}
+
+	if err := r.Start(); err != nil {
+		t.Fatalf("failed to start router: %v", err)
+	}
+	defer func() { _ = r.Close() }()
+
+	// Corrupt router's biscuit token to simulate key rotation/expiration
+	r.keysMu.Lock()
+	r.biscuitToken = []byte("invalid-corrupted-biscuit")
+	r.keysMu.Unlock()
+
+	// Call renewLease which should fail with 401, trigger reEnroll, and succeed
+	r.renewLease()
+
+	// Verify that biscuitToken was replaced with a valid non-corrupted token
+	r.keysMu.RLock()
+	currentToken := r.biscuitToken
+	r.keysMu.RUnlock()
+
+	if bytes.Equal(currentToken, []byte("invalid-corrupted-biscuit")) {
+		t.Error("expected biscuitToken to be updated after 401 re-enrollment, but it remained corrupted")
+	}
+}
+
+func TestRouterReEnrollUninitializedHostGuard(t *testing.T) {
+	r := &Router{}
+	err := r.reEnroll()
+	if err == nil {
+		t.Fatal("expected reEnroll to return error when r.Host is nil, got nil")
+	}
+	expectedMsg := "cannot re-enroll: router host is not initialized"
+	if err.Error() != expectedMsg {
+		t.Fatalf("expected error %q, got %q", expectedMsg, err.Error())
+	}
+}
+
+func TestRouterProactiveRefreshReEnrollOn401(t *testing.T) {
+	issuer, mintToken := startCustomMockOIDC(t)
+	cp, cpStore, cpURL := setupControlPlane(t, issuer)
+	defer func() {
+		_ = cp.Close()
+		_ = cpStore.Close()
+	}()
+
+	tempDir := t.TempDir()
+	routerKeyPath := filepath.Join(tempDir, "router.key")
+
+	routerJWT := mintToken(map[string]interface{}{
+		"sub":    "router-refresh-401-test",
+		"groups": []string{"routers"},
+	})
+
+	rOpts := Options{
+		ControlPlaneURL:    cpURL,
+		ListenAddrs:        []string{"/ip4/127.0.0.1/tcp/0"},
+		KeysSyncInterval:   20 * time.Second,
+		LeaseRenewInterval: 20 * time.Second,
+		OIDCToken:          routerJWT,
+		KeysDBPath:         routerKeyPath,
+		AllowLoopback:      true,
+		BiscuitTimeout:     1 * time.Second,
+	}
+
+	r, err := NewRouter(context.Background(), rOpts)
+	if err != nil {
+		t.Fatalf("failed to create router: %v", err)
+	}
+
+	if err := r.Start(); err != nil {
+		t.Fatalf("failed to start router: %v", err)
+	}
+	defer func() { _ = r.Close() }()
+
+	// Corrupt router's biscuit token to simulate key rotation
+	r.keysMu.Lock()
+	r.biscuitToken = []byte("stale-biscuit-signed-by-purged-key")
+	r.keysMu.Unlock()
+
+	// Trigger proactive refresh which should fail with 401, trigger reEnroll, and succeed
+	err = r.RefreshEnrollment(context.Background())
+	if err != nil {
+		t.Fatalf("expected RefreshEnrollment to recover via reEnroll, but got error: %v", err)
+	}
+
+	// Verify that biscuitToken was replaced with a valid fresh token
+	r.keysMu.RLock()
+	currentToken := r.biscuitToken
+	r.keysMu.RUnlock()
+
+	if bytes.Equal(currentToken, []byte("stale-biscuit-signed-by-purged-key")) {
+		t.Error("expected biscuitToken to be updated after 401 refresh fallback, but it remained stale")
+	}
+}
+
+