Add DeregisterExtension method (#92)

When the extension manager shuts down, osquery states it should
deregisters. When the extension is started again, it can avoid errors
when osquery starts the new process before the register watcher has had
time to clean up the old one.

Author: a-cognet

client.go

@@ -17,6 +17,7 @@ type ExtensionManager interface {
 	Call(registry, item string, req osquery.ExtensionPluginRequest) (*osquery.ExtensionResponse, error)
 	Extensions() (osquery.InternalExtensionList, error)
 	RegisterExtension(info *osquery.InternalExtensionInfo, registry osquery.ExtensionRegistry) (*osquery.ExtensionStatus, error)
+	DeregisterExtension(uuid osquery.ExtensionRouteUUID) (*osquery.ExtensionStatus, error)
 	Options() (osquery.InternalOptionList, error)
 	Query(sql string) (*osquery.ExtensionResponse, error)
 	GetQueryColumns(sql string) (*osquery.ExtensionResponse, error)
@@ -73,6 +74,11 @@ func (c *ExtensionManagerClient) RegisterExtension(info *osquery.InternalExtensi
 	return c.Client.RegisterExtension(context.Background(), info, registry)
 }
 
+// DeregisterExtension de-registers the extension plugins with the osquery process.
+func (c *ExtensionManagerClient) DeregisterExtension(uuid osquery.ExtensionRouteUUID) (*osquery.ExtensionStatus, error) {
+	return c.Client.DeregisterExtension(context.Background(), uuid)
+}
+
 // Options requests the list of bootstrap or configuration options.
 func (c *ExtensionManagerClient) Options() (osquery.InternalOptionList, error) {
 	return c.Client.Options(context.Background())

mock_manager.go

@@ -16,6 +16,8 @@ type ExtensionsFunc func() (osquery.InternalExtensionList, error)
 
 type RegisterExtensionFunc func(info *osquery.InternalExtensionInfo, registry osquery.ExtensionRegistry) (*osquery.ExtensionStatus, error)
 
+type DeregisterExtensionFunc func(uuid osquery.ExtensionRouteUUID) (*osquery.ExtensionStatus, error)
+
 type OptionsFunc func() (osquery.InternalOptionList, error)
 
 type QueryFunc func(sql string) (*osquery.ExtensionResponse, error)
@@ -38,6 +40,9 @@ type MockExtensionManager struct {
 	RegisterExtensionFunc        RegisterExtensionFunc
 	RegisterExtensionFuncInvoked bool
 
+	DeRegisterExtensionFunc        DeregisterExtensionFunc
+	DeRegisterExtensionFuncInvoked bool
+
 	OptionsFunc        OptionsFunc
 	OptionsFuncInvoked bool
 
@@ -73,6 +78,11 @@ func (m *MockExtensionManager) RegisterExtension(info *osquery.InternalExtension
 	return m.RegisterExtensionFunc(info, registry)
 }
 
+func (m *MockExtensionManager) DeregisterExtension(uuid osquery.ExtensionRouteUUID) (*osquery.ExtensionStatus, error) {
+	m.DeRegisterExtensionFuncInvoked = true
+	return m.DeRegisterExtensionFunc(uuid)
+}
+
 func (m *MockExtensionManager) Options() (osquery.InternalOptionList, error) {
 	m.OptionsFuncInvoked = true
 	return m.OptionsFunc()

server.go

@@ -49,6 +49,7 @@ type ExtensionManagerServer struct {
 	timeout      time.Duration
 	pingInterval time.Duration // How often to ping osquery server
 	mutex        sync.Mutex
+	uuid         osquery.ExtensionRouteUUID
 	started      bool // Used to ensure tests wait until the server is actually started
 }
 
@@ -153,14 +154,20 @@ func (s *ExtensionManagerServer) Start() error {
 		if stat.Code != 0 {
 			return errors.Errorf("status %d registering extension: %s", stat.Code, stat.Message)
 		}
+		s.uuid = stat.UUID
 
 		listenPath := fmt.Sprintf("%s.%d", s.sockPath, stat.UUID)
 
 		processor := osquery.NewExtensionProcessor(s)
 
 		s.transport, err = transport.OpenServer(listenPath, s.timeout)
 		if err != nil {
-			return errors.Wrapf(err, "opening server socket (%s)", listenPath)
+			openError := errors.Wrapf(err, "opening server socket (%s)", listenPath)
+			_, err = s.serverClient.DeregisterExtension(stat.UUID)
+			if err != nil {
+				return errors.Wrapf(err, "deregistering extension - follows %s", openError.Error())
+			}
+			return openError
 		}
 
 		s.server = thrift.NewTSimpleServer2(processor, s.transport)
@@ -242,10 +249,16 @@ func (s *ExtensionManagerServer) Call(ctx context.Context, registry string, item
 	return &response, nil
 }
 
-// Shutdown stops the server and closes the listening socket.
-func (s *ExtensionManagerServer) Shutdown(ctx context.Context) error {
+// Shutdown deregisters the extension, stops the server and closes all sockets.
+func (s *ExtensionManagerServer) Shutdown(ctx context.Context) (err error) {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
+	stat, err := s.serverClient.DeregisterExtension(s.uuid)
+	err = errors.Wrap(err, "deregistering extension")
+	if err == nil && stat.Code != 0 {
+		err = errors.Errorf("status %d deregistering extension: %s", stat.Code, stat.Message)
+	}
+	s.serverClient.Close()
 	if s.server != nil {
 		server := s.server
 		s.server = nil
@@ -258,7 +271,7 @@ func (s *ExtensionManagerServer) Shutdown(ctx context.Context) error {
 		}()
 	}
 
-	return nil
+	return
 }
 
 // Useful for testing

server_test.go

@@ -36,6 +36,10 @@ func TestNoDeadlockOnError(t *testing.T) {
 		PingFunc: func() (*osquery.ExtensionStatus, error) {
 			return &osquery.ExtensionStatus{}, nil
 		},
+		DeRegisterExtensionFunc: func(uuid osquery.ExtensionRouteUUID) (*osquery.ExtensionStatus, error) {
+			return &osquery.ExtensionStatus{}, nil
+		},
+		CloseFunc: func() {},
 	}
 	server := &ExtensionManagerServer{
 		serverClient: mock,
@@ -70,6 +74,10 @@ func TestShutdownWhenPingFails(t *testing.T) {
 			// As if the socket was closed
 			return nil, syscall.EPIPE
 		},
+		DeRegisterExtensionFunc: func(uuid osquery.ExtensionRouteUUID) (*osquery.ExtensionStatus, error) {
+			return &osquery.ExtensionStatus{}, nil
+		},
+		CloseFunc: func() {},
 	}
 	server := &ExtensionManagerServer{
 		serverClient: mock,
@@ -79,6 +87,8 @@ func TestShutdownWhenPingFails(t *testing.T) {
 	err := server.Run()
 	assert.Error(t, err)
 	assert.Contains(t, err.Error(), "broken pipe")
+	assert.True(t, mock.DeRegisterExtensionFuncInvoked)
+	assert.True(t, mock.CloseFuncInvoked)
 }
 
 // How many parallel tests to run (because sync issues do not occur on every
@@ -104,6 +114,10 @@ func testShutdownDeadlock(t *testing.T) {
 		RegisterExtensionFunc: func(info *osquery.InternalExtensionInfo, registry osquery.ExtensionRegistry) (*osquery.ExtensionStatus, error) {
 			return &osquery.ExtensionStatus{Code: 0, UUID: retUUID}, nil
 		},
+		DeRegisterExtensionFunc: func(uuid osquery.ExtensionRouteUUID) (*osquery.ExtensionStatus, error) {
+			return &osquery.ExtensionStatus{}, nil
+		},
+		CloseFunc: func() {},
 	}
 	server := ExtensionManagerServer{serverClient: mock, sockPath: tempPath.Name()}
 
@@ -172,6 +186,10 @@ func TestShutdownBasic(t *testing.T) {
 		RegisterExtensionFunc: func(info *osquery.InternalExtensionInfo, registry osquery.ExtensionRegistry) (*osquery.ExtensionStatus, error) {
 			return &osquery.ExtensionStatus{Code: 0, UUID: retUUID}, nil
 		},
+		DeRegisterExtensionFunc: func(uuid osquery.ExtensionRouteUUID) (*osquery.ExtensionStatus, error) {
+			return &osquery.ExtensionStatus{}, nil
+		},
+		CloseFunc: func() {},
 	}
 	server := ExtensionManagerServer{serverClient: mock, sockPath: tempPath.Name()}