Fatal when calling GlobalProtocolRegister with server(s) already star…

…ted (#373)

* Added flag to ProtocolStorage and fatal when calling GlobalProtocolRegister while at least a server already started.

* Good style for comments.

* Panic instead of logging when calling 'GlobalProtocolRegister' after a server has been started.

* Format protocol_test.go.

* Deferred Unlock was causing a deadlock in GlobalProtocolRegister, should not be deferred.

local.go

@@ -176,6 +176,7 @@ func (l *LocalTest) panicClosed() {
 // CloseAll closes all the servers.
 func (l *LocalTest) CloseAll() {
 	log.Lvl3("Stopping all")
+	InformAllServersStopped()
 	// If the debug-level is 0, we copy all errors to a buffer that
 	// will be discarded at the end.
 	if log.DebugVisible() == 0 {

protocol.go

@@ -55,10 +55,16 @@ var protocols = newProtocolStorage()
 
 // protocolStorage holds all protocols either globally or per-Server.
 type protocolStorage struct {
+	// Lock used because of the 'serverStarted' flag: it can be changed from a
+	// call to 'Server.Start' and is checked when calling
+	// 'GlobalProtocolRegister'.
+	sync.Mutex
 	// Instantiators maps the name of the protocols to the `NewProtocol`-
 	// methods.
 	instantiators map[string]NewProtocol
-	sync.Mutex
+	// Flag indicating if a server has already started; here to avoid calls
+	// to 'GlobalProtocolRegister' when a server has already started.
+	serverStarted bool
 }
 
 // newProtocolStorage returns an initialized ProtocolStorage-struct.
@@ -117,9 +123,30 @@ func ProtocolNameToID(name string) ProtocolID {
 // All registered protocols will be copied to every instantiated Server. If a
 // protocol is tied to a service, use `Server.ProtocolRegisterName`
 func GlobalProtocolRegister(name string, protocol NewProtocol) (ProtocolID, error) {
+	protocols.Lock()
+	// Cannot defer the "Unlock" because "Register" is using the lock too.
+	if protocols.serverStarted {
+		protocols.Unlock()
+		panic("Cannot call 'GlobalProtocolRegister' when a server has already started.")
+	}
+	protocols.Unlock()
 	return protocols.Register(name, protocol)
 }
 
+// InformServerStarted allows to set the 'serverStarted' flag to true.
+func InformServerStarted() {
+	protocols.Lock()
+	defer protocols.Unlock()
+	protocols.serverStarted = true
+}
+
+// InformAllServersStopped allows to set the 'serverStarted' flag to false.
+func InformAllServersStopped() {
+	protocols.Lock()
+	defer protocols.Unlock()
+	protocols.serverStarted = false
+}
+
 // MessageProxy is an interface that allows one protocol to completely define its
 // wire protocol format while still using the Overlay.
 // Cothority sends different messages dynamically as slices of bytes, whereas

protocol_test.go

@@ -137,7 +137,8 @@ func TestProtocolAutomaticInstantiation(t *testing.T) {
 		return &ps, nil
 	}
 
-	GlobalProtocolRegister(simpleProto, fn)
+	_, err := GlobalProtocolRegister(simpleProto, fn)
+	require.Nil(t, err)
 	local := NewLocalTest(tSuite)
 	defer local.CloseAll()
 	h, _, tree := local.GenTree(2, true)
@@ -177,7 +178,8 @@ func TestProtocolError(t *testing.T) {
 		return &ps, nil
 	}
 
-	GlobalProtocolRegister(simpleProto, fn)
+	_, err := GlobalProtocolRegister(simpleProto, fn)
+	require.Nil(t, err)
 	local := NewLocalTest(tSuite)
 	h, _, tree := local.GenTree(2, true)
 	h1 := h[0]
@@ -228,6 +230,27 @@ func TestProtocolError(t *testing.T) {
 
 }
 
+func TestGlobalProtocolRegisterTooLate(t *testing.T) {
+	var simpleProto = "simplePE"
+	done := make(chan bool)
+	fn := func(n *TreeNodeInstance) (ProtocolInstance, error) {
+		ps := SimpleProtocol{
+			TreeNodeInstance: n,
+			Chan:             done,
+		}
+		log.ErrFatal(ps.RegisterHandler(ps.ReturnError))
+		return &ps, nil
+	}
+
+	local := NewLocalTest(tSuite)
+	defer local.CloseAll()
+	local.GenTree(2, true)
+	fnShouldPanic := func() {
+		GlobalProtocolRegister(simpleProto, fn)
+	}
+	assert.Panics(t, fnShouldPanic)
+}
+
 func TestMessageProxyFactory(t *testing.T) {
 	defer eraseAllMessageProxy()
 	RegisterMessageProxy(NewTestMessageProxyChan)
@@ -240,7 +263,8 @@ func TestMessageProxyStore(t *testing.T) {
 	defer local.CloseAll()
 
 	RegisterMessageProxy(NewTestMessageProxy)
-	GlobalProtocolRegister(testProtoIOName, newTestProtocolInstance)
+	_, err := GlobalProtocolRegister(testProtoIOName, newTestProtocolInstance)
+	require.Nil(t, err)
 	h, _, tree := local.GenTree(2, true)
 
 	go func() {
@@ -252,7 +276,7 @@ func TestMessageProxyStore(t *testing.T) {
 		require.Equal(t, "", res)
 
 	}()
-	_, err := h[0].StartProtocol(testProtoIOName, tree)
+	_, err = h[0].StartProtocol(testProtoIOName, tree)
 	require.Nil(t, err)
 
 	res := <-chanTestProtoInstance

server.go

@@ -162,6 +162,7 @@ func (c *Server) protocolInstantiate(protoID ProtocolID, tni *TreeNodeInstance)
 // Start makes the router and the websocket listen on their respective
 // ports.
 func (c *Server) Start() {
+	InformServerStarted()
 	c.started = time.Now()
 	log.Lvlf1("Starting server at %s on address %s with public key %s",
 		c.started.Format("2006-01-02 15:04:05"),