diff --git a/client/x/dkg/keeper/dkg_svc.go b/client/x/dkg/keeper/dkg_svc.go index fff42351..2b323546 100644 --- a/client/x/dkg/keeper/dkg_svc.go +++ b/client/x/dkg/keeper/dkg_svc.go @@ -90,6 +90,14 @@ func releaseDKGSvc(round uint32) { // is cut short — so this value must stay above 60s to be meaningful. const dkgAsyncTimeout = 2 * time.Minute +// decryptKernelCallTimeout bounds a single kernel RPC made by the decrypt worker. +// The worker runs on a process-lifetime context.Background with no deadline of its +// own, so without a per-call timeout a wedged kernel (accepts the connection but +// never responds) blocks the worker loop forever and decryptWorkerRunning is never +// reset. 60s matches the kernel-operation budget in dkgAsyncTimeout and the +// contract-call timeout in contract_client.go. +const decryptKernelCallTimeout = 60 * time.Second + // dkgAsyncContext creates a new context for async DKG service goroutines with a timeout. // This replaces the consensus context that would otherwise be canceled after block processing. func dkgAsyncContext() (context.Context, context.CancelFunc) { @@ -100,6 +108,18 @@ var decryptWorkerRunning atomic.Bool // ResumeDKGService reloads unfinished DKG sessions and resumes their execution safely without spawning duplicate goroutines. func (k *Keeper) ResumeDKGService(ctx context.Context, dkgNetwork *types.DKGNetwork) { + // The decrypt worker serves the latest ACTIVE round, which may differ from the + // latest round once the next round has opened. StartDecryptWorker is idempotent. + if active, err := k.getLatestActiveDKGNetwork(ctx); err != nil { + log.Warn(ctx, "Failed to get latest active DKG round while resuming decrypt worker", err) + } else if active != nil { + k.StartDecryptWorker() + } else { + // Debug (not Info): this runs every block; a missing active round is the + // normal steady state and Info would be too noisy. + log.Debug(ctx, "No active DKG round; decrypt worker not started") + } + session, err := k.stateManager.GetSession(dkgNetwork.Round) if err != nil { log.Error(ctx, "Failed to get DKG session while resuming the DKG service", err) @@ -279,7 +299,10 @@ func (k *Keeper) resumeFailedSession(ctx context.Context, session *types.DKGSess // goroutine exits. The decrypt worker must run for the lifetime of the process. func (k *Keeper) StartDecryptWorker() { if !decryptWorkerRunning.CompareAndSwap(false, true) { - // already running + // Debug (not Info): this runs every block; skipping because the worker is + // already running is the normal steady state and Info would be too noisy. + log.Debug(context.Background(), "Decrypt worker already running; skipping start") + return } @@ -620,8 +643,13 @@ func (k *Keeper) computePartialDecrypt(ctx context.Context, session *types.DKGSe return } + // Bound the kernel RPC so a wedged kernel cannot block the worker loop + // indefinitely. The worker's context has no deadline of its own. + callCtx, cancel := context.WithTimeout(ctx, decryptKernelCallTimeout) + defer cancel() + kernelStart := time.Now() - resp, err := client.PartialDecryptTDH2(ctx, &types.PartialDecryptTDH2Request{ + resp, err := client.PartialDecryptTDH2(callCtx, &types.PartialDecryptTDH2Request{ CodeCommitment: session.CodeCommitment, Round: session.Round, Ciphertext: req.Ciphertext, diff --git a/client/x/dkg/keeper/dkg_svc_test.go b/client/x/dkg/keeper/dkg_svc_test.go index dcaaa189..30b816f8 100644 --- a/client/x/dkg/keeper/dkg_svc_test.go +++ b/client/x/dkg/keeper/dkg_svc_test.go @@ -1871,6 +1871,133 @@ func TestResumeDKGService_NoSession(t *testing.T) { k.ResumeDKGService(ctx, dkgNetwork) } +// TestResumeDKGService_ActiveRoundStartsWorker verifies that the decrypt worker is +// restarted whenever an active DKG round exists, even if the latest round has already +// advanced to a pre-active stage (e.g. the next round opened registration). This is +// the restart-window hole from issue piplabs/story#854. +func TestResumeDKGService_ActiveRoundStartsWorker(t *testing.T) { + if testing.Short() { + t.Skip("skipping worker test in short mode") + } + + k, _, _, ctx := setupDKGKeeperWithMocks(t) + + sm, err := NewStateManager(t.TempDir()) + require.NoError(t, err) + k.stateManager = sm + + resetDKGSvcRound() + defer resetDKGSvcRound() + // Force a fresh worker spawn so the assertion exercises the new branch, and + // leave decryptWorkerRunning true afterwards. The worker is a process-lifetime + // singleton; resetting the flag to false here would let a later mock-backed test + // spawn a real worker that calls its gomock client after that test completes. + decryptWorkerRunning.Store(false) + + // Detach the mock contract client so the process-lifetime worker goroutine + // returns early at its nil guard instead of issuing unexpected mock calls + // after the test completes. + k.contractClient = nil + + // An earlier round is still the active (completed) round. + activeNetwork := &types.DKGNetwork{ + Round: 5, + Stage: types.DKGStageActive, + } + require.NoError(t, k.setDKGNetwork(ctx, activeNetwork)) + require.NoError(t, k.setLatestActiveRound(ctx, activeNetwork)) + + // The latest round has opened registration for the NEXT round; its session is in a + // valid (non-stuck) phase, so none of the legacy branches would start the worker. + session := &types.DKGSession{ + Round: 6, + Phase: types.PhaseInitialized, + } + require.NoError(t, sm.CreateSession(ctx, session)) + + dkgNetwork := &types.DKGNetwork{ + Round: 6, + Stage: types.DKGStageRegistration, + } + + k.ResumeDKGService(ctx, dkgNetwork) + + require.True(t, decryptWorkerRunning.Load(), "worker should run whenever an active round exists") +} + +// TestResumeDKGService_NoActiveRoundNoWorker verifies that on a fresh chain with no +// active round, ResumeDKGService does not start the decrypt worker via the new branch. +func TestResumeDKGService_NoActiveRoundNoWorker(t *testing.T) { + k, _, _, ctx := setupDKGKeeperWithMocks(t) + + sm, err := NewStateManager(t.TempDir()) + require.NoError(t, err) + k.stateManager = sm + + resetDKGSvcRound() + defer resetDKGSvcRound() + decryptWorkerRunning.Store(false) + // Restore the singleton flag to true on exit so a later mock-backed test does + // not spawn a real worker that outlives it (see the note above). + defer decryptWorkerRunning.Store(true) + + // Latest round session is in a valid (non-stuck) phase and no active round exists. + session := &types.DKGSession{ + Round: 1, + Phase: types.PhaseInitialized, + } + require.NoError(t, sm.CreateSession(ctx, session)) + + dkgNetwork := &types.DKGNetwork{ + Round: 1, + Stage: types.DKGStageRegistration, + } + + k.ResumeDKGService(ctx, dkgNetwork) + + require.False(t, decryptWorkerRunning.Load(), "worker must not start when no active round exists") +} + +// TestResumeDKGService_CompletedActiveStartsWorker verifies the pre-existing behavior: +// when the latest round is itself active with a completed session, the worker is started. +func TestResumeDKGService_CompletedActiveStartsWorker(t *testing.T) { + if testing.Short() { + t.Skip("skipping worker test in short mode") + } + + k, _, _, ctx := setupDKGKeeperWithMocks(t) + + sm, err := NewStateManager(t.TempDir()) + require.NoError(t, err) + k.stateManager = sm + + resetDKGSvcRound() + defer resetDKGSvcRound() + // Force a fresh spawn for the assertion; leave the flag true on exit (see the + // note in TestResumeDKGService_ActiveRoundStartsWorker). + decryptWorkerRunning.Store(false) + + // Detach the mock contract client so the process-lifetime worker goroutine + // returns early at its nil guard instead of issuing unexpected mock calls + // after the test completes. + k.contractClient = nil + + session := &types.DKGSession{ + Round: 7, + Phase: types.PhaseCompleted, + } + require.NoError(t, sm.CreateSession(ctx, session)) + + dkgNetwork := &types.DKGNetwork{ + Round: 7, + Stage: types.DKGStageActive, + } + + k.ResumeDKGService(ctx, dkgNetwork) + + require.True(t, decryptWorkerRunning.Load(), "worker should run when latest active round session is completed") +} + // --- resumeFailedSession --- func TestResumeFailedSession_DealingStage(t *testing.T) { @@ -2089,3 +2216,102 @@ func TestResumeFailedSession_ActiveStage(t *testing.T) { require.Equal(t, types.PhaseCompleted, got.Phase, "active stage: handleDKGComplete should advance phase to PhaseCompleted") } + +// --- decrypt worker kernel-call deadline (issue piplabs/story#854 family) --- + +// TestComputePartialDecrypt_DerivesCallDeadline verifies that the kernel RPC is +// invoked with a bounded context even when the worker passes context.Background +// (which has no deadline). This proves decryptKernelCallTimeout is wired so a +// wedged kernel cannot block the worker loop indefinitely. +func TestComputePartialDecrypt_DerivesCallDeadline(t *testing.T) { + t.Parallel() + + ctx := context.Background() + ctrl := gomock.NewController(t) + + cc := []byte("cc-deadline") + mockKernel := dkgtestutil.NewMockKernelServiceClient(ctrl) + + router := NewKernelRouter(nil, nil) + router.RegisterClient(cc, mockKernel) + + k := &Keeper{kernelRouter: router} + + session := &types.DKGSession{ + Round: 1, Index: 1, GlobalPubKey: []byte("pub"), CodeCommitment: cc, + } + req := types.DecryptRequest{Ciphertext: []byte("ct"), Label: make([]byte, 32)} + + var ( + gotDeadline bool + remaining time.Duration + ) + mockKernel.EXPECT().PartialDecryptTDH2(gomock.Any(), gomock.Any()). + DoAndReturn(func(callCtx context.Context, _ *types.PartialDecryptTDH2Request, _ ...interface{}) (*types.PartialDecryptTDH2Response, error) { + deadline, ok := callCtx.Deadline() + gotDeadline = ok + if ok { + remaining = time.Until(deadline) + } + return &types.PartialDecryptTDH2Response{ + EncryptedPartialDecryption: []byte("p"), EphemeralPubKey: []byte("e"), + PubShare: []byte("s"), Signature: []byte("sig"), + }, nil + }).Times(1) + + // Parent context has NO deadline; the deadline must come from the fix. + result := k.computePartialDecrypt(ctx, session, req) + require.NoError(t, result.err) + require.True(t, gotDeadline, "kernel call must receive a bounded context even from a deadline-less parent") + require.Greater(t, remaining, time.Duration(0), "derived deadline must be in the future") + require.LessOrEqual(t, remaining, decryptKernelCallTimeout, "derived deadline must not exceed the configured timeout") +} + +// TestComputePartialDecrypt_WedgedKernelReturns simulates a kernel that accepts the +// call but never responds on its own. The worker call must still return (via context +// timeout) rather than block forever, so the worker loop can complete and +// decryptWorkerRunning can be reset. Uses a short-deadline parent context to keep the +// test fast; the production timeout is decryptKernelCallTimeout. +func TestComputePartialDecrypt_WedgedKernelReturns(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + + cc := []byte("cc-wedged") + mockKernel := dkgtestutil.NewMockKernelServiceClient(ctrl) + + router := NewKernelRouter(nil, nil) + router.RegisterClient(cc, mockKernel) + + k := &Keeper{kernelRouter: router} + + session := &types.DKGSession{ + Round: 1, Index: 1, GlobalPubKey: []byte("pub"), CodeCommitment: cc, + } + req := types.DecryptRequest{Ciphertext: []byte("ct"), Label: make([]byte, 32)} + + // Kernel never returns on its own: it blocks until the call context is done, + // then surfaces the context error (mirrors a real gRPC deadline cancellation). + mockKernel.EXPECT().PartialDecryptTDH2(gomock.Any(), gomock.Any()). + DoAndReturn(func(callCtx context.Context, _ *types.PartialDecryptTDH2Request, _ ...interface{}) (*types.PartialDecryptTDH2Response, error) { + <-callCtx.Done() + return nil, callCtx.Err() + }).Times(1) + + // Short-deadline parent stands in for decryptKernelCallTimeout to keep the test fast. + parentCtx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + done := make(chan decryptComputeResult, 1) + go func() { + done <- k.computePartialDecrypt(parentCtx, session, req) + }() + + select { + case result := <-done: + require.Error(t, result.err, "wedged kernel call must return an error, not a response") + require.Contains(t, result.err.Error(), "generating partial decrypt failed") + case <-time.After(5 * time.Second): + t.Fatal("computePartialDecrypt blocked on a wedged kernel; worker would never reset decryptWorkerRunning") + } +} diff --git a/client/x/dkg/keeper/kernel_client.go b/client/x/dkg/keeper/kernel_client.go index b54ae6b1..c9f56efb 100644 --- a/client/x/dkg/keeper/kernel_client.go +++ b/client/x/dkg/keeper/kernel_client.go @@ -6,6 +6,7 @@ import ( "io" "os" "strings" + "time" "github.com/piplabs/story/client/x/dkg/types" "github.com/piplabs/story/lib/errors" @@ -13,8 +14,19 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/keepalive" ) +// kernelKeepalive detects a wedged kernel connection (accepts the socket but stops +// responding) at the transport layer so a dead connection surfaces as a call error +// instead of hanging a caller forever. PermitWithoutStream keeps pinging even when +// no RPC is in flight, since the decrypt worker is idle between ticks. +var kernelKeepalive = keepalive.ClientParameters{ + Time: 30 * time.Second, + Timeout: 10 * time.Second, + PermitWithoutStream: true, +} + // TLSConfig holds the TLS certificate paths for gRPC client connections. // When CAFile is set, TLS is enabled with server certificate verification. // When CertFile and KeyFile are also set, mutual TLS (mTLS) is used. @@ -56,7 +68,10 @@ func CreateKernelClient(endpoint string, tlsCfg *TLSConfig) (types.KernelService target = "passthrough:///" + target } - conn, err := grpc.NewClient(target, grpc.WithTransportCredentials(creds)) + conn, err := grpc.NewClient(target, + grpc.WithTransportCredentials(creds), + grpc.WithKeepaliveParams(kernelKeepalive), + ) if err != nil { return nil, nil, errors.Wrap(err, "failed to connect to story-kernel client") }