fix(controller): refresh keep-alive on every authenticated request

controller/controller.go

@@ -112,6 +112,12 @@ func (c *Controller) NewRequest() {
 	c.lastRequest = time.Now()
 }
 
+func (c *Controller) LastRequest() time.Time {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.lastRequest
+}
+
 // sameBackendConfig reports whether two Backend descriptors describe the same
 // running configuration. Users are intentionally excluded — they are managed
 // through their own sync RPCs, so a Start that only differs in the user list

controller/rest/middleware.go

@@ -61,15 +61,15 @@ func LogRequest(next http.Handler) http.Handler {
 	})
 }
 
-func (s *Service) trackSuccessfulRequest(next http.Handler) http.Handler {
+// trackRequest refreshes keep-alive for every authenticated request when it
+// arrives and again when it ends, whatever the status: a request that fails,
+// times out on the panel side or is still running still proves the panel is
+// alive. It must run after validateApiKey.
+func (s *Service) trackRequest(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ww := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
-
-		next.ServeHTTP(ww, r)
+		s.NewRequest()
+		defer s.NewRequest()
 
-		// Only track successful requests (status codes 200-299)
-		if status := ww.Status(); status >= 200 && status < 300 {
-			s.NewRequest()
-		}
+		next.ServeHTTP(w, r)
 	})
 }

controller/rest/middleware_test.go

@@ -0,0 +1,118 @@
+package rest
+
+import (
+	"errors"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+
+	"github.com/pasarguard/node/config"
+)
+
+func newKeepAliveTestService(t *testing.T) (*Service, uuid.UUID) {
+	t.Helper()
+	key := uuid.New()
+	s := New(config.NewTestConfig(t.TempDir(), key))
+	t.Cleanup(s.Disconnect)
+	return s, key
+}
+
+// Any request the panel makes proves the panel is alive, even one that fails
+// or is still running, so keep-alive counts it when it arrives.
+func TestKeepAliveRefreshedWhenAuthenticatedRequestArrives(t *testing.T) {
+	s, key := newKeepAliveTestService(t)
+
+	// A Start whose body is still uploading.
+	body, upload := io.Pipe()
+	defer upload.Close()
+	req := httptest.NewRequest(http.MethodPost, "/start", body)
+	req.Header.Set("x-api-key", key.String())
+	rec := httptest.NewRecorder()
+	done := make(chan struct{})
+	go func() {
+		s.Router.ServeHTTP(rec, req)
+		close(done)
+	}()
+
+	deadline := time.Now().Add(2 * time.Second)
+	for s.LastRequest().IsZero() && time.Now().Before(deadline) {
+		time.Sleep(5 * time.Millisecond)
+	}
+	refreshed := !s.LastRequest().IsZero()
+	upload.CloseWithError(errors.New("panel went away"))
+	<-done
+
+	if !refreshed {
+		t.Fatal("keep-alive was not refreshed while the request was in flight")
+	}
+	// The upload ended with an error ("panel went away"), so the handler
+	// cannot read a body and answers 400.
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+}
+
+func TestKeepAliveRefreshedForFailedRequests(t *testing.T) {
+	for _, path := range []string{"/stats/backend", "/no/such/route"} {
+		s, key := newKeepAliveTestService(t)
+
+		req := httptest.NewRequest(http.MethodGet, path, nil)
+		req.Header.Set("x-api-key", key.String())
+		rec := httptest.NewRecorder()
+		s.Router.ServeHTTP(rec, req)
+
+		if rec.Code < 400 {
+			t.Fatalf("%s: status = %d, want an error", path, rec.Code)
+		}
+		if s.LastRequest().IsZero() {
+			t.Fatalf("%s: keep-alive was not refreshed for an authenticated request that failed (%d)", path, rec.Code)
+		}
+	}
+}
+
+func TestKeepAliveNotRefreshedWithoutValidKey(t *testing.T) {
+	s, _ := newKeepAliveTestService(t)
+
+	for name, key := range map[string]string{
+		"no key":     "",
+		"not a uuid": "not-a-uuid",
+		"wrong key":  uuid.NewString(),
+	} {
+		req := httptest.NewRequest(http.MethodGet, "/info", nil)
+		if key != "" {
+			req.Header.Set("x-api-key", key)
+		}
+		rec := httptest.NewRecorder()
+		s.Router.ServeHTTP(rec, req)
+
+		if rec.Code < 400 {
+			t.Fatalf("%s: status = %d, want an auth error", name, rec.Code)
+		}
+		if !s.LastRequest().IsZero() {
+			t.Fatalf("%s: an unauthenticated request refreshed keep-alive", name)
+		}
+	}
+}
+
+func TestKeepAliveRefreshedWhenLongRequestEnds(t *testing.T) {
+	s, _ := newKeepAliveTestService(t)
+
+	var atArrival time.Time
+	handler := s.trackRequest(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		atArrival = s.LastRequest()
+		time.Sleep(10 * time.Millisecond)
+		http.Error(w, "client went away", http.StatusInternalServerError)
+	}))
+	handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/logs", nil))
+
+	if atArrival.IsZero() {
+		t.Fatal("keep-alive was not refreshed when the request arrived")
+	}
+	if !s.LastRequest().After(atArrival) {
+		t.Fatal("keep-alive was not refreshed when the request ended")
+	}
+}

controller/rest/service.go

@@ -28,7 +28,7 @@ func (s *Service) setRouter() {
 	// Api Handlers
 	router.Use(LogRequest)
 	router.Use(s.validateApiKey)
-	router.Use(s.trackSuccessfulRequest)
+	router.Use(s.trackRequest)
 	router.Use(middleware.Recoverer)
 
 	router.Post("/start", s.Start)

controller/rpc/middleware.go

@@ -75,6 +75,38 @@ func validateApiKeyStreamMiddleware(s *Service) grpc.StreamServerInterceptor {
 	}
 }
 
+// trackRequestMiddleware refreshes keep-alive for every authenticated call when
+// it arrives and again when it ends, whatever the result: a call that fails,
+// times out on the panel side or is still running still proves the panel is
+// alive. It must run after validateApiKey.
+func trackRequestMiddleware(s *Service) grpc.UnaryServerInterceptor {
+	return func(
+		ctx context.Context,
+		req any,
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (any, error) {
+		s.NewRequest()
+		defer s.NewRequest()
+
+		return handler(ctx, req)
+	}
+}
+
+func trackRequestStreamMiddleware(s *Service) grpc.StreamServerInterceptor {
+	return func(
+		srv any,
+		ss grpc.ServerStream,
+		info *grpc.StreamServerInfo,
+		handler grpc.StreamHandler,
+	) error {
+		s.NewRequest()
+		defer s.NewRequest()
+
+		return handler(srv, ss)
+	}
+}
+
 func checkBackendStatus(s *Service) error {
 	back := s.Backend()
 	if back == nil {
@@ -147,11 +179,6 @@ func LoggingInterceptor(s *Service) grpc.UnaryServerInterceptor {
 		// Log the request
 		logRequest(ctx, info.FullMethod, err)
 
-		// Track successful requests
-		if err == nil {
-			s.NewRequest()
-		}
-
 		return resp, err
 	}
 }
@@ -175,11 +202,6 @@ func LoggingStreamInterceptor(s *Service) grpc.StreamServerInterceptor {
 		// Log the request
 		logRequest(ss.Context(), info.FullMethod, err)
 
-		// Track successful requests
-		if err == nil {
-			s.NewRequest()
-		}
-
 		return err
 	}
 }
@@ -216,6 +238,8 @@ func ConditionalMiddleware(s *Service) grpc.UnaryServerInterceptor {
 
 		interceptors = append(interceptors, validateApiKeyMiddleware(s))
 
+		interceptors = append(interceptors, trackRequestMiddleware(s))
+
 		if backendMethods[info.FullMethod] {
 			interceptors = append(interceptors, CheckBackendMiddleware(s))
 		}
@@ -238,6 +262,8 @@ func ConditionalStreamMiddleware(s *Service) grpc.StreamServerInterceptor {
 
 		interceptors = append(interceptors, validateApiKeyStreamMiddleware(s))
 
+		interceptors = append(interceptors, trackRequestStreamMiddleware(s))
+
 		if backendMethods[info.FullMethod] {
 			interceptors = append(interceptors, CheckBackendStreamMiddleware(s))
 		}

controller/rpc/middleware_test.go

@@ -1,6 +1,19 @@
 package rpc
 
-import "testing"
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+
+	"github.com/pasarguard/node/common"
+	"github.com/pasarguard/node/config"
+)
 
 func TestRoutingMethodsAreBackendGated(t *testing.T) {
 	methods := []string{
@@ -17,3 +30,110 @@ func TestRoutingMethodsAreBackendGated(t *testing.T) {
 		}
 	}
 }
+
+func newKeepAliveTestService(t *testing.T) (*Service, uuid.UUID) {
+	t.Helper()
+	key := uuid.New()
+	s := New(config.NewTestConfig(t.TempDir(), key))
+	t.Cleanup(s.Disconnect)
+	return s, key
+}
+
+func withAPIKey(key string) context.Context {
+	return metadata.NewIncomingContext(context.Background(), metadata.Pairs("x-api-key", key))
+}
+
+type fakeServerStream struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+func (f *fakeServerStream) Context() context.Context { return f.ctx }
+
+// Any call the panel makes proves the panel is alive, even one that fails or
+// is still running, so keep-alive counts it when it arrives.
+func TestKeepAliveRefreshedWhenAuthenticatedCallArrives(t *testing.T) {
+	s, key := newKeepAliveTestService(t)
+
+	info := &grpc.UnaryServerInfo{FullMethod: "/service.NodeService/GetBaseInfo"}
+	var atArrival time.Time
+	_, err := ConditionalMiddleware(s)(withAPIKey(key.String()), &common.Empty{}, info, func(context.Context, any) (any, error) {
+		atArrival = s.LastRequest()
+		return nil, status.Error(codes.NotFound, "user not found.")
+	})
+	if status.Code(err) != codes.NotFound {
+		t.Fatalf("err = %v, want NotFound", err)
+	}
+	if atArrival.IsZero() {
+		t.Fatal("keep-alive was not refreshed when the call arrived")
+	}
+}
+
+func TestKeepAliveRefreshedWhenBackendGateRejects(t *testing.T) {
+	s, key := newKeepAliveTestService(t)
+
+	info := &grpc.UnaryServerInfo{FullMethod: "/service.NodeService/GetBackendStats"}
+	_, err := ConditionalMiddleware(s)(withAPIKey(key.String()), &common.Empty{}, info, func(context.Context, any) (any, error) {
+		t.Error("handler ran without a backend")
+		return nil, nil
+	})
+	if status.Code(err) != codes.Unavailable {
+		t.Fatalf("err = %v, want Unavailable", err)
+	}
+	if s.LastRequest().IsZero() {
+		t.Fatal("keep-alive was not refreshed for an authenticated call the backend gate rejected")
+	}
+}
+
+func TestKeepAliveNotRefreshedWithoutValidKey(t *testing.T) {
+	s, _ := newKeepAliveTestService(t)
+
+	cases := map[string]context.Context{
+		"no metadata":   context.Background(),
+		"no key header": metadata.NewIncomingContext(context.Background(), metadata.Pairs("other", "x")),
+		"not a uuid":    withAPIKey("not-a-uuid"),
+		"wrong key":     withAPIKey(uuid.NewString()),
+	}
+	info := &grpc.UnaryServerInfo{FullMethod: "/service.NodeService/GetBaseInfo"}
+	streamInfo := &grpc.StreamServerInfo{FullMethod: "/service.NodeService/GetLogs", IsServerStream: true}
+	for name, ctx := range cases {
+		_, err := ConditionalMiddleware(s)(ctx, &common.Empty{}, info, func(context.Context, any) (any, error) {
+			return &common.Empty{}, nil
+		})
+		if err == nil {
+			t.Fatalf("%s: unary call passed auth", name)
+		}
+		err = ConditionalStreamMiddleware(s)(nil, &fakeServerStream{ctx: ctx}, streamInfo, func(any, grpc.ServerStream) error {
+			return nil
+		})
+		if err == nil {
+			t.Fatalf("%s: stream passed auth", name)
+		}
+		if !s.LastRequest().IsZero() {
+			t.Fatalf("%s: an unauthenticated call refreshed keep-alive", name)
+		}
+	}
+}
+
+func TestKeepAliveRefreshedWhenStreamOpensAndEnds(t *testing.T) {
+	s, key := newKeepAliveTestService(t)
+
+	// An ungated method, so the handler runs without a backend.
+	info := &grpc.StreamServerInfo{FullMethod: "/service.NodeService/Ungated", IsServerStream: true}
+	var atArrival time.Time
+	err := ConditionalStreamMiddleware(s)(nil, &fakeServerStream{ctx: withAPIKey(key.String())}, info, func(any, grpc.ServerStream) error {
+		atArrival = s.LastRequest()
+		time.Sleep(10 * time.Millisecond)
+		// The panel went away mid-stream.
+		return status.Error(codes.Canceled, "context canceled")
+	})
+	if status.Code(err) != codes.Canceled {
+		t.Fatalf("err = %v, want Canceled", err)
+	}
+	if atArrival.IsZero() {
+		t.Fatal("keep-alive was not refreshed when the stream opened")
+	}
+	if !s.LastRequest().After(atArrival) {
+		t.Fatal("keep-alive was not refreshed when the stream ended")
+	}
+}