hooks: Allow local control of OCI stages via extensionStages

This allows callers to avoid delegating to OCI runtimes for cases
where they feel that the runtime hook handling is unreliable [1].

[1]: https://github.com/projectatomic/libpod/issues/730#issuecomment-392959938

Signed-off-by: W. Trevor King <wking@tremily.us>

Closes: #855
Approved by: rhatdan

2 files changed, 37 insertions, 15 deletions

diff --git a/pkg/hooks/hooks.go b/pkg/hooks/hooks.go
index 66e97d6bd..b78ede38b 100644
--- a/pkg/hooks/hooks.go
+++ b/pkg/hooks/hooks.go
@@ -45,6 +45,11 @@ type namedHooks []*namedHook
 // increasing preference (hook configurations in later directories
 // override configurations with the same filename from earlier
 // directories).
+//
+// extensionStages allows callers to add additional stages beyond
+// those specified in the OCI Runtime Specification and to control
+// OCI-defined stages instead of delagating to the OCI runtime.  See
+// Hooks() for more information.
 func New(ctx context.Context, directories []string, extensionStages []string, lang language.Tag) (manager *Manager, err error) {
 	manager = &Manager{
 		hooks:           map[string]*current.Hook{},
@@ -82,13 +87,24 @@ func (m *Manager) namedHooks() (hooks []*namedHook) {
 }
 
 // Hooks injects OCI runtime hooks for a given container configuration.
+//
+// If the extensionStages slice was set when initializing the Manager,
+// matching hooks requesting those stages will be returned in the
+// extensionStages map.  This takes precedence over their inclusion in
+// the OCI configuration.  For example:
+//
+//   manager, err := New(ctx, []string{DefaultDir}, []string{"poststop"}, lang)
+//   extensionStages, err := manager.Hooks(config, annotations, hasBindMounts)
+//
+// will have any matching post-stop hooks in extensionStages and will
+// not insert them into config.Hooks.Poststop.
 func (m *Manager) Hooks(config *rspec.Spec, annotations map[string]string, hasBindMounts bool) (extensionStages map[string][]rspec.Hook, err error) {
 	hooks := m.namedHooks()
 	collator := collate.New(m.language, collate.IgnoreCase, collate.IgnoreWidth)
 	collator.Sort(namedHooks(hooks))
-	validStages := map[string]bool{} // beyond the OCI stages
+	localStages := map[string]bool{} // stages destined for extensionStages
 	for _, stage := range m.extensionStages {
-		validStages[stage] = true
+		localStages[stage] = true
 	}
 	for _, namedHook := range hooks {
 		match, err := namedHook.hook.When.Match(config, annotations, hasBindMounts)
@@ -100,21 +116,22 @@ func (m *Manager) Hooks(config *rspec.Spec, annotations map[string]string, hasBi
 				config.Hooks = &rspec.Hooks{}
 			}
 			for _, stage := range namedHook.hook.Stages {
-				switch stage {
-				case "prestart":
-					config.Hooks.Prestart = append(config.Hooks.Prestart, namedHook.hook.Hook)
-				case "poststart":
-					config.Hooks.Poststart = append(config.Hooks.Poststart, namedHook.hook.Hook)
-				case "poststop":
-					config.Hooks.Poststop = append(config.Hooks.Poststop, namedHook.hook.Hook)
-				default:
-					if !validStages[stage] {
-						return extensionStages, fmt.Errorf("hook %q: unknown stage %q", namedHook.name, stage)
-					}
+				if _, ok := localStages[stage]; ok {
 					if extensionStages == nil {
 						extensionStages = map[string][]rspec.Hook{}
 					}
 					extensionStages[stage] = append(extensionStages[stage], namedHook.hook.Hook)
+				} else {
+					switch stage {
+					case "prestart":
+						config.Hooks.Prestart = append(config.Hooks.Prestart, namedHook.hook.Hook)
+					case "poststart":
+						config.Hooks.Poststart = append(config.Hooks.Poststart, namedHook.hook.Hook)
+					case "poststop":
+						config.Hooks.Poststop = append(config.Hooks.Poststop, namedHook.hook.Hook)
+					default:
+						return extensionStages, fmt.Errorf("hook %q: unknown stage %q", namedHook.name, stage)
+					}
 				}
 			}
 		}
diff --git a/pkg/hooks/hooks_test.go b/pkg/hooks/hooks_test.go
index 1b0703556..8bd89f56f 100644
--- a/pkg/hooks/hooks_test.go
+++ b/pkg/hooks/hooks_test.go
@@ -190,10 +190,10 @@ func TestExtensionStage(t *testing.T) {
 				When: current.When{
 					Always: &always,
 				},
-				Stages: []string{"prestart", "a", "b"},
+				Stages: []string{"prestart", "poststop", "a", "b"},
 			},
 		},
-		extensionStages: []string{"a", "b", "c"},
+		extensionStages: []string{"poststop", "a", "b", "c"},
 	}
 
 	config := &rspec.Spec{}
@@ -211,6 +211,11 @@ func TestExtensionStage(t *testing.T) {
 	}, config.Hooks)
 
 	assert.Equal(t, map[string][]rspec.Hook{
+		"poststop": {
+			{
+				Path: "/a/b/c",
+			},
+		},
 		"a": {
 			{
 				Path: "/a/b/c",