Start cleaning up config to fix race conditions

Author: belak

server.go

@@ -58,14 +58,17 @@ type Server struct {
 	RequestHandlers map[string]RequestHandler
 
 	listenerWg sync.WaitGroup
-	mu         sync.Mutex
+	mu         sync.RWMutex
 	listeners  map[net.Listener]struct{}
 	conns      map[*gossh.ServerConn]struct{}
 	connWg     sync.WaitGroup
 	doneChan   chan struct{}
 }
 
 func (srv *Server) ensureHostSigner() error {
+	srv.mu.Lock()
+	defer srv.mu.Unlock()
+
 	if len(srv.HostSigners) == 0 {
 		signer, err := generateSigner()
 		if err != nil {
@@ -79,6 +82,7 @@ func (srv *Server) ensureHostSigner() error {
 func (srv *Server) ensureHandlers() {
 	srv.mu.Lock()
 	defer srv.mu.Unlock()
+
 	if srv.RequestHandlers == nil {
 		srv.RequestHandlers = map[string]RequestHandler{}
 		for k, v := range DefaultRequestHandlers {
@@ -94,6 +98,9 @@ func (srv *Server) ensureHandlers() {
 }
 
 func (srv *Server) config(ctx Context) *gossh.ServerConfig {
+	srv.mu.RLock()
+	defer srv.mu.RUnlock()
+
 	var config *gossh.ServerConfig
 	if srv.ServerConfigCallback == nil {
 		config = &gossh.ServerConfig{}
@@ -142,6 +149,9 @@ func (srv *Server) config(ctx Context) *gossh.ServerConfig {
 
 // Handle sets the Handler for the server.
 func (srv *Server) Handle(fn Handler) {
+	srv.mu.Lock()
+	defer srv.mu.Unlock()
+
 	srv.Handler = fn
 }
 
@@ -153,6 +163,7 @@ func (srv *Server) Handle(fn Handler) {
 func (srv *Server) Close() error {
 	srv.mu.Lock()
 	defer srv.mu.Unlock()
+
 	srv.closeDoneChanLocked()
 	err := srv.closeListenersLocked()
 	for c := range srv.conns {
@@ -313,6 +324,9 @@ func (srv *Server) ListenAndServe() error {
 // with the same algorithm, it is overwritten. Each server config must have at
 // least one host key.
 func (srv *Server) AddHostKey(key Signer) {
+	srv.mu.Lock()
+	defer srv.mu.Unlock()
+
 	// these are later added via AddHostKey on ServerConfig, which performs the
 	// check for one of every algorithm.
 
@@ -332,12 +346,20 @@ func (srv *Server) AddHostKey(key Signer) {
 
 // SetOption runs a functional option against the server.
 func (srv *Server) SetOption(option Option) error {
+	// NOTE: there is a potential race here for any option that doesn't call an
+	// internal method. We can't actually lock here because if something calls
+	// (as an example) AddHostKey, it will deadlock.
+
+	//srv.mu.Lock()
+	//defer srv.mu.Unlock()
+
 	return option(srv)
 }
 
 func (srv *Server) getDoneChan() <-chan struct{} {
 	srv.mu.Lock()
 	defer srv.mu.Unlock()
+
 	return srv.getDoneChanLocked()
 }
 
@@ -374,6 +396,7 @@ func (srv *Server) closeListenersLocked() error {
 func (srv *Server) trackListener(ln net.Listener, add bool) {
 	srv.mu.Lock()
 	defer srv.mu.Unlock()
+
 	if srv.listeners == nil {
 		srv.listeners = make(map[net.Listener]struct{})
 	}
@@ -394,6 +417,7 @@ func (srv *Server) trackListener(ln net.Listener, add bool) {
 func (srv *Server) trackConn(c *gossh.ServerConn, add bool) {
 	srv.mu.Lock()
 	defer srv.mu.Unlock()
+
 	if srv.conns == nil {
 		srv.conns = make(map[*gossh.ServerConn]struct{})
 	}