`scmr`: `change`: Handle timeouts properly

1 files changed, 29 insertions, 14 deletions

diff --git a/pkg/goexec/scmr/change.go b/pkg/goexec/scmr/change.go
index 3dd506e..bf2fc4d 100644
--- a/pkg/goexec/scmr/change.go
+++ b/pkg/goexec/scmr/change.go
@@ -2,7 +2,6 @@ package scmrexec
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"github.com/FalconOpsLLC/goexec/pkg/goexec"
 	"github.com/oiweiwei/go-msrpc/msrpc/scmr/svcctl/v2"
@@ -24,6 +23,7 @@ type ScmrChange struct {
 	IO goexec.ExecutionIO
 
 	NoStart     bool
+	NoRevert    bool
 	ServiceName string
 }
 
@@ -110,6 +110,8 @@ func (m *ScmrChange) Execute(ctx context.Context, in *goexec.ExecutionIO) (err e
 		//Dependencies:     []byte(svc.originalConfig.Dependencies), // TODO
 	}
 
+	bpn := svc.originalConfig.BinaryPathName
+
 	_, err = m.ctl.ChangeServiceConfigW(ctx, req)
 
 	if err != nil {
@@ -117,26 +119,39 @@ func (m *ScmrChange) Execute(ctx context.Context, in *goexec.ExecutionIO) (err e
 		return fmt.Errorf("change service config request: %w", err)
 	}
 
-	m.AddCleaners(func(ctxInner context.Context) error {
-		req.BinaryPathName = svc.originalConfig.BinaryPathName
-
-		if ctxInner.Err() != nil && errors.Is(ctxInner.Err(), context.DeadlineExceeded) {
-			ctxInner = context.WithoutCancel(context.Background())
-		}
-		_, err := m.ctl.ChangeServiceConfigW(ctxInner, req)
+	if !m.NoStart {
+		stopResponse, err = m.ctl.ControlService(ctx, &svcctl.ControlServiceRequest{
+			Service: svc.handle,
+			Control: 0,
+		})
 
+		err = m.startService(ctx, svc)
 		if err != nil {
-			return fmt.Errorf("restore service config: %w", err)
+			log.Error().Err(err).Msg("Failed to start service")
 		}
-		return nil
-	})
+	}
 
-	if !m.NoStart {
+	if !m.NoRevert {
+		if svc.handle == nil {
+
+			if err = m.Reconnect(ctx); err != nil {
+				return err
+			}
+			svc, err = m.openService(ctx, svc.name)
+
+			if err != nil {
+				log.Error().Err(err).Msg("Failed to reopen service handle")
+				return fmt.Errorf("reopen service: %w", err)
+			}
+		}
+		req.BinaryPathName = bpn
+		req.Service = svc.handle
+		_, err := m.ctl.ChangeServiceConfigW(ctx, req)
 
-		err = m.startService(ctx, svc)
 		if err != nil {
-			log.Error().Err(err).Msg("Failed to start service")
+			return fmt.Errorf("restore service config: %w", err)
 		}
+		return nil
 	}
 
 	return