From a4b2953017643b4a9a53bc5382554477069b0976 Mon Sep 17 00:00:00 2001
From: Brian Downs <brian.downs@gmail.com>
Date: Wed, 19 Aug 2020 13:30:51 -0700
Subject: [PATCH] add setup hook capabilities for rke2

Signed-off-by: Brian Downs <brian.downs@gmail.com>
---
 pkg/cli/cmds/server.go   | 2 ++
 pkg/cli/server/server.go | 2 ++
 pkg/server/server.go     | 4 ++++
 pkg/server/types.go      | 1 +
 4 files changed, 9 insertions(+)

diff --git a/pkg/cli/cmds/server.go b/pkg/cli/cmds/server.go
index fbd3c8d8b3..077f21e320 100644
--- a/pkg/cli/cmds/server.go
+++ b/pkg/cli/cmds/server.go
@@ -1,6 +1,7 @@
 package cmds
 
 import (
+	"github.com/rancher/k3s/pkg/daemons/config"
 	"github.com/rancher/k3s/pkg/version"
 	"github.com/rancher/spur/cli"
 	"github.com/rancher/spur/cli/altsrc"
@@ -54,6 +55,7 @@ type Server struct {
 	ClusterInit              bool
 	ClusterReset             bool
 	EncryptSecrets           bool
+	SetupHooks               []func(config.Control) error
 }
 
 var ServerConfig Server
diff --git a/pkg/cli/server/server.go b/pkg/cli/server/server.go
index c911b7ed6e..b9a441c309 100644
--- a/pkg/cli/server/server.go
+++ b/pkg/cli/server/server.go
@@ -193,6 +193,8 @@ func run(app *cli.Context, cfg *cmds.Server) error {
 		return errors.Wrap(err, "Invalid tls-min-version")
 	}
 
+	serverConfig.SetupHooks = append(serverConfig.SetupHooks, cfg.SetupHooks...)
+
 	// TLS config based on mozilla ssl-config generator
 	// https://ssl-config.mozilla.org/#server=golang&version=1.13.6&config=intermediate&guideline=5.4
 	// Need to disable the TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256 Cipher for TLS1.2
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 7c397e647d..f44f9e8c21 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -60,6 +60,10 @@ func StartServer(ctx context.Context, config *Config) error {
 		return errors.Wrap(err, "starting tls server")
 	}
 
+	for _, hook := range config.SetupHooks {
+		hook(config.ControlConfig)
+	}
+
 	ip := net2.ParseIP(config.ControlConfig.BindAddress)
 	if ip == nil {
 		hostIP, err := net.ChooseHostInterface()
diff --git a/pkg/server/types.go b/pkg/server/types.go
index 9bcc2f3639..204ec93053 100644
--- a/pkg/server/types.go
+++ b/pkg/server/types.go
@@ -10,4 +10,5 @@ type Config struct {
 	ControlConfig    config.Control
 	Rootless         bool
 	SupervisorPort   int
+	SetupHooks       []func(config.Control) error
 }