fix: server robustness for dmesg grep, handler panics, exec cancellation and MPA expiry

server/recovery.go

@@ -0,0 +1,49 @@
+/* Copyright (c) 2026 Snowflake Inc. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the
+   "License"); you may not use this file except in compliance
+   with the License.  You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing,
+   software distributed under the License is distributed on an
+   "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+   KIND, either express or implied.  See the License for the
+   specific language governing permissions and limitations
+   under the License.
+*/
+
+package server
+
+import (
+	"context"
+	"fmt"
+	"runtime/debug"
+
+	"github.com/go-logr/logr"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func recoverPanic(logger logr.Logger, method string, err *error) {
+	if r := recover(); r != nil {
+		logger.Error(fmt.Errorf("%v", r), "panic in handler", "method", method, "stack", string(debug.Stack()))
+		*err = status.Errorf(codes.Internal, "internal error handling %s", method)
+	}
+}
+
+func unaryRecoveryInterceptor(logger logr.Logger) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+		defer recoverPanic(logger, info.FullMethod, &err)
+		return handler(ctx, req)
+	}
+}
+
+func streamRecoveryInterceptor(logger logr.Logger) grpc.StreamServerInterceptor {
+	return func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) (err error) {
+		defer recoverPanic(logger, info.FullMethod, &err)
+		return handler(srv, ss)
+	}
+}

server/recovery_test.go

@@ -0,0 +1,63 @@
+/* Copyright (c) 2026 Snowflake Inc. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the
+   "License"); you may not use this file except in compliance
+   with the License.  You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing,
+   software distributed under the License is distributed on an
+   "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+   KIND, either express or implied.  See the License for the
+   specific language governing permissions and limitations
+   under the License.
+*/
+
+package server
+
+import (
+	"context"
+	"testing"
+
+	"github.com/go-logr/logr"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func TestUnaryRecoveryInterceptor(t *testing.T) {
+	interceptor := unaryRecoveryInterceptor(logr.Discard())
+	info := &grpc.UnaryServerInfo{FullMethod: "/Test.Test/Unary"}
+
+	_, err := interceptor(context.Background(), nil, info, func(context.Context, any) (any, error) {
+		panic("boom")
+	})
+	if got := status.Code(err); got != codes.Internal {
+		t.Fatalf("panicking handler: got code %v, want %v (err: %v)", got, codes.Internal, err)
+	}
+
+	resp, err := interceptor(context.Background(), nil, info, func(context.Context, any) (any, error) {
+		return "ok", nil
+	})
+	if err != nil || resp != "ok" {
+		t.Fatalf("normal handler: got (%v, %v), want (ok, nil)", resp, err)
+	}
+}
+
+func TestStreamRecoveryInterceptor(t *testing.T) {
+	interceptor := streamRecoveryInterceptor(logr.Discard())
+	info := &grpc.StreamServerInfo{FullMethod: "/Test.Test/Stream"}
+
+	err := interceptor(nil, nil, info, func(any, grpc.ServerStream) error {
+		panic("boom")
+	})
+	if got := status.Code(err); got != codes.Internal {
+		t.Fatalf("panicking handler: got code %v, want %v (err: %v)", got, codes.Internal, err)
+	}
+
+	wantErr := status.Error(codes.NotFound, "missing")
+	if err := interceptor(nil, nil, info, func(any, grpc.ServerStream) error { return wantErr }); err != wantErr {
+		t.Fatalf("erroring handler: got %v, want %v", err, wantErr)
+	}
+}

server/server.go

@@ -215,14 +215,14 @@ func BuildServer(opts ...Option) (*grpc.Server, error) {
 
 	authz := rpcauth.NewRPCAuthorizer(ss.policy, ss.authzHooks...)
 
-	unary := ss.unaryInterceptors
+	unary := append([]grpc.UnaryServerInterceptor{unaryRecoveryInterceptor(ss.logger)}, ss.unaryInterceptors...)
 	unary = append(unary,
 		// Execute log interceptor after other interceptors so that metadata gets logged
 		telemetry.UnaryServerLogInterceptor(ss.logger),
 		// Execute authz after logger is setup
 		authz.Authorize,
 	)
-	streaming := ss.streamInterceptors
+	streaming := append([]grpc.StreamServerInterceptor{streamRecoveryInterceptor(ss.logger)}, ss.streamInterceptors...)
 	streaming = append(streaming,
 		// Execute log interceptor after other interceptors so that metadata gets logged
 		telemetry.StreamServerLogInterceptor(ss.logger),

services/exec/server/exec.go

@@ -50,7 +50,7 @@ type server struct{}
 func (s *server) Run(ctx context.Context, req *pb.ExecRequest) (res *pb.ExecResponse, err error) {
 	recorder := metrics.RecorderFromContextOrNoop(ctx)
 
-	var opts []util.Option
+	opts := []util.Option{util.ProcessGroup()}
 	if req.User != "" {
 		uid, gid, err := resolveUser(req.User)
 		if err != nil {
@@ -105,6 +105,7 @@ func (s *server) StreamingRun(req *pb.ExecRequest, stream pb.Exec_StreamingRunSe
 			}
 		}
 	}
+	util.KillProcessGroupOnCancel(cmd)
 	stdout, err := cmd.StdoutPipe()
 	if err != nil {
 		return err

services/mpa/server/server.go

@@ -117,11 +117,11 @@ func actionId(action *mpa.Action) (string, error) {
 }
 
 type storedAction struct {
-	action       *mpa.Action
-	lastModified time.Time
-	approvers    []*mpa.Principal
-	approved     chan struct{}
-	sizeBytes    int
+	action    *mpa.Action
+	created   time.Time
+	approvers []*mpa.Principal
+	approved  chan struct{}
+	sizeBytes int
 }
 
 // server is used to implement the gRPC server
@@ -142,16 +142,16 @@ func (s *server) deleteActionLocked(id string) {
 	}
 }
 
-// evictOldestLocked removes the least-recently-modified action, if any.
+// evictOldestLocked removes the oldest action, if any.
 // Callers must hold s.mu.
 func (s *server) evictOldestLocked() {
 	var oldestID string
 	var oldestTime time.Time
 	found := false
 	for id, act := range s.actions {
-		if !found || act.lastModified.Before(oldestTime) {
+		if !found || act.created.Before(oldestTime) {
 			oldestID = id
-			oldestTime = act.lastModified
+			oldestTime = act.created
 			found = true
 		}
 	}
@@ -179,7 +179,7 @@ func (s *server) clearOutdatedApprovals() {
 
 	staleTime := time.Now().Add(-maxMPAApprovedAge)
 	for id, act := range s.actions {
-		if act.lastModified.Before(staleTime) {
+		if act.created.Before(staleTime) {
 			s.curBytes -= act.sizeBytes
 			delete(s.actions, id)
 		}
@@ -240,10 +240,10 @@ func (s *server) Store(ctx context.Context, in *mpa.StoreRequest) (*mpa.StoreRes
 	}
 
 	act := &storedAction{
-		action:       action,
-		approved:     make(chan struct{}),
-		lastModified: time.Now(),
-		sizeBytes:    size,
+		action:    action,
+		approved:  make(chan struct{}),
+		created:   time.Now(),
+		sizeBytes: size,
 	}
 	s.actions[id] = act
 	s.curBytes += size
@@ -281,7 +281,6 @@ func (s *server) Approve(ctx context.Context, in *mpa.ApproveRequest) (*mpa.Appr
 	if act.action.User == p.ID {
 		return nil, status.Error(codes.InvalidArgument, "MPA requests cannot be approved by their requestor")
 	}
-	act.lastModified = time.Now()
 	// Only add the approver if it's new compared to existing approvals
 	if !containsPrincipal(act.approvers, p) {
 		act.approvers = append(act.approvers, &mpa.Principal{

services/mpa/server/server_test.go

@@ -22,6 +22,7 @@ import (
 	"strconv"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/Snowflake-Labs/sansshell/auth/rpcauth"
 	"github.com/Snowflake-Labs/sansshell/services/mpa"
@@ -207,6 +208,36 @@ func TestStoreBoundsTotalBytes(t *testing.T) {
 	}
 }
 
+func TestApproveDoesNotExtendExpiry(t *testing.T) {
+	s := &server{actions: make(map[string]*storedAction)}
+	ctx := context.Background()
+	rCtx := rpcauth.AddPeerToContext(ctx, &rpcauth.PeerAuthInput{
+		Principal: &rpcauth.PrincipalAuthInput{ID: "requester"},
+	})
+	stored, err := s.Store(rCtx, &mpa.StoreRequest{
+		Method:  "foobar",
+		Message: mustAny(anypb.New(&emptypb.Empty{})),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	s.actions[stored.Id].created = time.Now().Add(-maxMPAApprovedAge + time.Minute)
+
+	aCtx := rpcauth.AddPeerToContext(ctx, &rpcauth.PeerAuthInput{
+		Principal: &rpcauth.PrincipalAuthInput{ID: "approver"},
+	})
+	for i := 0; i < 2; i++ {
+		if _, err := s.Approve(aCtx, &mpa.ApproveRequest{Action: stored.Action}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	s.actions[stored.Id].created = s.actions[stored.Id].created.Add(-2 * time.Minute)
+
+	if _, err := s.Get(ctx, &mpa.GetRequest{Id: stored.Id}); status.Code(err) != codes.NotFound {
+		t.Fatalf("Get after expiry: got %v, want NotFound", err)
+	}
+}
+
 func TestWaitForApproval(t *testing.T) {
 	ctx := context.Background()
 

services/sysinfo/server/sysinfo.go

@@ -79,21 +79,29 @@ func (s *server) Dmesg(req *pb.DmesgRequest, stream pb.SysInfo_DmesgServer) erro
 		return err
 	}
 
+	var pattern *regexp.Regexp
+	if req.Grep != "" {
+		regex := req.Grep
+		if req.IgnoreCase {
+			regex = `(?i)` + regex
+		}
+		var err error
+		pattern, err = regexp.Compile(regex)
+		if err != nil {
+			recorder.CounterOrLog(ctx, sysinfoDmesgFailureCounter, 1, attribute.String("reason", "invalid_grep_pattern"))
+			return status.Errorf(codes.InvalidArgument, "invalid grep pattern %q: %v", req.Grep, err)
+		}
+	}
+
 	records, err := getKernelMessages(timeout, ctx.Done())
 	if err != nil {
 		recorder.CounterOrLog(ctx, sysinfoDmesgFailureCounter, 1, attribute.String("reason", "get kernel message error"))
 		return status.Errorf(codes.InvalidArgument, "can't get kernel message %v", err)
 	}
 	// grep the messages we want
-	if req.Grep != "" {
+	if pattern != nil {
 		var newRecords []*pb.DmsgRecord
-		regex := req.Grep
-		if req.IgnoreCase {
-			regex = `(?i)` + regex
-		}
-
 		isInvert := req.InvertMatch
-		pattern := regexp.MustCompile(regex)
 		for _, record := range records {
 			// Two cases are what we want
 			// case 1: isInvert = false and match = true (pattern matched)

services/sysinfo/server/sysinfo_linux_test.go

@@ -392,6 +392,14 @@ func TestDmesg(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "invalid grep pattern expect an error: dmesg -grep (",
+			req: &pb.DmesgRequest{
+				TailLines: -1,
+				Grep:      "(",
+			},
+			wantErr: true,
+		},
 	} {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {

services/util/util.go

@@ -20,6 +20,7 @@ package util
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -75,6 +76,7 @@ type cmdOptions struct {
 	uid          uint32
 	gid          uint32
 	extraFiles   []*os.File
+	processGroup bool
 }
 
 // Option will run the apply operation to change required checking/state
@@ -177,6 +179,34 @@ func ExtraFiles(files []*os.File) Option {
 	})
 }
 
+// ProcessGroup is an option which runs the command in its own process group so that
+// cancelling the context kills any processes it started along with it.
+func ProcessGroup() Option {
+	return optionfunc(func(o *cmdOptions) {
+		o.processGroup = true
+	})
+}
+
+// KillProcessGroupOnCancel configures cmd to start in a new process group and
+// makes context cancellation kill the whole group rather than just cmd.
+func KillProcessGroupOnCancel(cmd *exec.Cmd) {
+	if cmd.SysProcAttr == nil {
+		cmd.SysProcAttr = &syscall.SysProcAttr{}
+	}
+	cmd.SysProcAttr.Setpgid = true
+	cmd.Cancel = func() error {
+		return killProcessGroup(cmd.Process.Pid)
+	}
+}
+
+func killProcessGroup(pgid int) error {
+	err := syscall.Kill(-pgid, syscall.SIGKILL)
+	if errors.Is(err, syscall.ESRCH) {
+		return os.ErrProcessDone
+	}
+	return err
+}
+
 // DefRunBufLimit is the default limit we'll buffer for stdout/stderr from RunCommand exec'ing
 // a process.
 const DefRunBufLimit = 10 * 1024 * 1024
@@ -298,8 +328,20 @@ func RunCommand(ctx context.Context, bin string, args []string, opts ...Option)
 		cmd.SysProcAttr.Credential.Uid = options.uid
 		cmd.SysProcAttr.Credential.Gid = options.gid
 	}
+	if options.processGroup {
+		KillProcessGroupOnCancel(cmd)
+	}
 	logger.Info("executing local command", "cmd", cmd.String())
-	run.Error = cmd.Run()
+	run.Error = cmd.Start()
+	if run.Error == nil {
+		if options.processGroup {
+			// os/exec stops watching ctx once the direct child exits, but processes it
+			// left in the background can still hold stdout/stderr open.
+			stop := context.AfterFunc(ctx, func() { _ = killProcessGroup(cmd.Process.Pid) })
+			defer stop()
+		}
+		run.Error = cmd.Wait()
+	}
 	run.ExitCode = cmd.ProcessState.ExitCode()
 	// If this was an error it could be two different things. Just exiting non-zero results in an exec.ExitError
 	// and we can suppress that as exit code is enough. Otherwise we leave run.Error for callers to use.