pkg/peertls: extension handling refactor (#2831)

pkg/peertls/extensions/extensions.go

@@ -20,8 +20,9 @@ const (
 )
 
 var (
-	// AllHandlers holds all registered extension handlers
-	AllHandlers HandlerFactories
+	// DefaultHandlers is a slice of handlers that we use by default.
+	//   - IDVersionHandler
+	DefaultHandlers HandlerFactories
 
 	// CAWhitelistSignedLeafHandler verifies that the leaf cert of the remote peer's
 	// identity was signed by one of the CA certs in the whitelist.
@@ -101,13 +102,6 @@ type HandlerFunc func(pkix.Extension, [][]*x509.Certificate) error
 // underlying extension id value.
 type HandlerFuncMap map[*ExtensionID]HandlerFunc
 
-func init() {
-	// NB: register all handlers defined in this file.
-	AllHandlers.Register(
-		CAWhitelistSignedLeafHandler,
-	)
-}
-
 // NewHandlerFactory builds a `HandlerFactory` pointer from an `ExtensionID` and a `HandlerFactoryFunc`.
 func NewHandlerFactory(id *ExtensionID, handlerFactory HandlerFactoryFunc) *HandlerFactory {
 	return &HandlerFactory{

pkg/peertls/extensions/extensions_test.go

@@ -63,13 +63,6 @@ func TestHandlers_Register(t *testing.T) {
 		err = handlers[i].NewHandlerFunc(opts[i])(exts[i], chains[i])
 		assert.Errorf(t, err, strconv.Itoa(i))
 	}
-
-	{ // test `extensions.AllHandlers`
-		for _, handler := range extensions.AllHandlers {
-			assert.NotNil(t, handler.ID())
-			assert.NotNil(t, handler.NewHandlerFunc(nil))
-		}
-	}
 }
 
 func TestHandlers_WithOptions(t *testing.T) {
@@ -122,16 +115,4 @@ func TestHandlers_WithOptions(t *testing.T) {
 		assert.True(t, ok)
 		assert.NotNil(t, handleFunc)
 	}
-
-	{ // test `extensions.AllHandlers`
-		handlerFuncMap := extensions.AllHandlers.WithOptions(&extensions.Options{})
-		for _, handler := range extensions.AllHandlers {
-			id := handler.ID()
-			require.NotNil(t, id)
-
-			handleFunc, ok := handlerFuncMap[id]
-			assert.True(t, ok)
-			assert.NotNil(t, handleFunc)
-		}
-	}
 }

pkg/peertls/extensions/revocations.go

@@ -56,14 +56,6 @@ type RevocationDB interface {
 	Close() error
 }
 
-func init() {
-	// NB: register all handlers defined in this file.
-	AllHandlers.Register(
-		RevocationCheckHandler,
-		RevocationUpdateHandler,
-	)
-}
-
 // NewRevocationExt generates a revocation extension for a certificate.
 func NewRevocationExt(key crypto.PrivateKey, revokedCert *x509.Certificate) (pkix.Extension, error) {
 	nowUnix := time.Now().Unix()

pkg/peertls/extensions/revocations_test.go

@@ -17,6 +17,7 @@ import (
 	"storj.io/storj/pkg/identity"
 	"storj.io/storj/pkg/peertls"
 	"storj.io/storj/pkg/peertls/extensions"
+	"storj.io/storj/pkg/peertls/tlsopts"
 	"storj.io/storj/pkg/storj"
 	"storj.io/storj/storage"
 )
@@ -153,3 +154,15 @@ func TestRevocationUpdateHandler(t *testing.T) {
 		}
 	})
 }
+
+func TestWithOptions_NilRevocationDB(t *testing.T) {
+	_, chain, err := testpeertls.NewCertChain(2, storj.LatestIDVersion().Number)
+	require.NoError(t, err)
+
+	opts := &extensions.Options{RevDB: nil}
+	handlerFuncMap := extensions.DefaultHandlers.WithOptions(opts)
+
+	extMap := tlsopts.NewExtensionsMap(chain[peertls.LeafIndex])
+	err = extMap.HandleExtensions(handlerFuncMap, identity.ToChains(chain))
+	require.NoError(t, err)
+}

pkg/peertls/tlsopts/options.go

@@ -10,7 +10,7 @@ import (
 	"io/ioutil"
 
 	"github.com/zeebo/errs"
-	monkit "gopkg.in/spacemonkeygo/monkit.v2"
+	"gopkg.in/spacemonkeygo/monkit.v2"
 
 	"storj.io/storj/pkg/identity"
 	"storj.io/storj/pkg/peertls"
@@ -101,7 +101,21 @@ func (opts *Options) configure() (err error) {
 		opts.VerificationFuncs.ClientAdd(peertls.VerifyCAWhitelist(opts.PeerCAWhitelist))
 	}
 
-	opts.handleExtensions(extensions.AllHandlers)
+	handlers := make(extensions.HandlerFactories, len(extensions.DefaultHandlers))
+	copy(handlers, extensions.DefaultHandlers)
+
+	if opts.Config.Extensions.Revocation {
+		handlers.Register(
+			extensions.RevocationCheckHandler,
+			extensions.RevocationUpdateHandler,
+		)
+	}
+
+	if opts.Config.Extensions.WhitelistSignedLeaf {
+		handlers.Register(extensions.CAWhitelistSignedLeafHandler)
+	}
+
+	opts.handleExtensions(handlers)
 
 	opts.Cert, err = peertls.TLSCert(opts.Ident.RawChain(), opts.Ident.Leaf, opts.Ident.Key)
 	return err

pkg/peertls/tlsopts/tls_test.go

@@ -112,7 +112,7 @@ func TestExtensionMap_HandleExtensions(t *testing.T) {
 			}
 
 			{
-				handlerFuncMap := extensions.AllHandlers.WithOptions(opts)
+				handlerFuncMap := extensions.DefaultHandlers.WithOptions(opts)
 				for _, testcase := range testcases {
 					t.Log(testcase.name)
 					extensionsMap := tlsopts.NewExtensionsMap(testcase.chain...)

pkg/storj/identity_version.go

@@ -48,7 +48,7 @@ type IDVersion struct {
 }
 
 func init() {
-	extensions.AllHandlers.Register(IDVersionHandler)
+	extensions.DefaultHandlers.Register(IDVersionHandler)
 }
 
 // GetIDVersion looks up the given version number in the map of registered

pkg/storj/identity_version_test.go

@@ -99,7 +99,7 @@ func TestIDVersionExtensionHandler_success(t *testing.T) {
 		assert.NoError(t, err)
 
 		extensionMap := tlsopts.NewExtensionsMap(testcase.chain...)
-		handlerFuncMap := extensions.AllHandlers.WithOptions(opts)
+		handlerFuncMap := extensions.DefaultHandlers.WithOptions(opts)
 
 		err = extensionMap.HandleExtensions(handlerFuncMap, identity.ToChains(testcase.chain))
 		assert.NoError(t, err)