diff --git a/server/recovery.go b/server/recovery.go new file mode 100644 index 00000000..24e8a7df --- /dev/null +++ b/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) + } +} diff --git a/server/recovery_test.go b/server/recovery_test.go new file mode 100644 index 00000000..b5c5c1e0 --- /dev/null +++ b/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) + } +} diff --git a/server/server.go b/server/server.go index cbe75314..a23a37ed 100644 --- a/server/server.go +++ b/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), diff --git a/services/exec/server/exec.go b/services/exec/server/exec.go index 0ed9c400..d1a42253 100644 --- a/services/exec/server/exec.go +++ b/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 diff --git a/services/mpa/server/server.go b/services/mpa/server/server.go index 93b03719..418c045f 100644 --- a/services/mpa/server/server.go +++ b/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{ diff --git a/services/mpa/server/server_test.go b/services/mpa/server/server_test.go index 82e05fc8..8e175388 100644 --- a/services/mpa/server/server_test.go +++ b/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() diff --git a/services/sysinfo/server/sysinfo.go b/services/sysinfo/server/sysinfo.go index 0f4cc182..0e15314d 100644 --- a/services/sysinfo/server/sysinfo.go +++ b/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) diff --git a/services/sysinfo/server/sysinfo_linux_test.go b/services/sysinfo/server/sysinfo_linux_test.go index 5ca8a31f..882f575f 100644 --- a/services/sysinfo/server/sysinfo_linux_test.go +++ b/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) { diff --git a/services/util/util.go b/services/util/util.go index 1e1676f2..9ddfaa83 100644 --- a/services/util/util.go +++ b/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. diff --git a/services/util/util_test.go b/services/util/util_test.go index 13dd8ea6..d08541b0 100644 --- a/services/util/util_test.go +++ b/services/util/util_test.go @@ -19,11 +19,14 @@ package util import ( "bytes" "context" + "errors" "fmt" "os" + "os/exec" "path" "reflect" "testing" + "time" "github.com/Snowflake-Labs/sansshell/testing/testutil" "golang.org/x/sys/unix" @@ -188,6 +191,51 @@ func TestRunCommand(t *testing.T) { } } +func TestRunCommandProcessGroup(t *testing.T) { + for _, tc := range []struct { + name string + script string + }{ + { + name: "parent still running", + script: "(sleep 2; touch %s) & wait", + }, + { + name: "parent already exited", + script: "(sleep 2; touch %s) &", + }, + } { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + marker := path.Join(t.TempDir(), "marker") + ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + + start := time.Now() + _, err := RunCommand(ctx, "/bin/sh", []string{"-c", fmt.Sprintf(tc.script, marker)}, ProcessGroup()) + testutil.FatalOnErr("RunCommand", err, t) + if elapsed := time.Since(start); elapsed > 1500*time.Millisecond { + t.Fatalf("RunCommand returned after %v, want it to return soon after cancellation", elapsed) + } + + time.Sleep(3 * time.Second) + if _, err := os.Stat(marker); err == nil { + t.Fatal("background process kept running after cancellation") + } + }) + } +} + +func TestKillProcessGroupOnCancelAfterExit(t *testing.T) { + cmd := exec.CommandContext(context.Background(), "/bin/sh", "-c", "exit 0") + KillProcessGroupOnCancel(cmd) + testutil.FatalOnErr("Run", cmd.Run(), t) + if err := cmd.Cancel(); !errors.Is(err, os.ErrProcessDone) { + t.Fatalf("Cancel after exit = %v, want %v", err, os.ErrProcessDone) + } +} + func TestTrimString(t *testing.T) { b := &bytes.Buffer{} for i := 0; i < 2*MaxBuf; i++ {