diff --git a/controller/controller.go b/controller/controller.go index 904b36af..ee8f760f 100644 --- a/controller/controller.go +++ b/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 diff --git a/controller/rest/middleware.go b/controller/rest/middleware.go index ad742245..d36eb4f4 100644 --- a/controller/rest/middleware.go +++ b/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) }) } diff --git a/controller/rest/middleware_test.go b/controller/rest/middleware_test.go new file mode 100644 index 00000000..2f33b063 --- /dev/null +++ b/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") + } +} diff --git a/controller/rest/service.go b/controller/rest/service.go index fc9aa6db..d4867922 100644 --- a/controller/rest/service.go +++ b/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) diff --git a/controller/rpc/middleware.go b/controller/rpc/middleware.go index 208870f8..a5694281 100644 --- a/controller/rpc/middleware.go +++ b/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)) } diff --git a/controller/rpc/middleware_test.go b/controller/rpc/middleware_test.go index 37d187eb..ed6ba60f 100644 --- a/controller/rpc/middleware_test.go +++ b/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") + } +}