Merge pull request #123 from gliderlabs/optimize-add-host-key

Update AddHostKey to avoid always appending

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,19 +324,42 @@ 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.
+
+	// This check is based on the AddHostKey method from the x/crypto/ssh
+	// library. This allows us to only keep one active key for each type on a
+	// server at once. So, if you're dynamically updating keys at runtime, this
+	// list will not keep growing.
+	for i, k := range srv.HostSigners {
+		if k.PublicKey().Type() == key.PublicKey().Type() {
+			srv.HostSigners[i] = key
+			return
+		}
+	}
+
 	srv.HostSigners = append(srv.HostSigners, key)
 }
 
 // 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()
 }
 
@@ -362,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{})
 	}
@@ -382,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{})
 	}

server_test.go

@@ -8,6 +8,26 @@ import (
 	"time"
 )
 
+func TestAddHostKey(t *testing.T) {
+	s := Server{}
+	signer, err := generateSigner()
+	if err != nil {
+		t.Fatal(err)
+	}
+	s.AddHostKey(signer)
+	if len(s.HostSigners) != 1 {
+		t.Fatal("Key was not properly added")
+	}
+	signer, err = generateSigner()
+	if err != nil {
+		t.Fatal(err)
+	}
+	s.AddHostKey(signer)
+	if len(s.HostSigners) != 1 {
+		t.Fatal("Key was not properly replaced")
+	}
+}
+
 func TestServerShutdown(t *testing.T) {
 	l := newLocalListener()
 	testBytes := []byte("Hello world\n")

session_test.go

@@ -289,20 +289,40 @@ func TestPtyResize(t *testing.T) {
 func TestSignals(t *testing.T) {
 	t.Parallel()
 
+	// errChan lets us get errors back from the session
+	errChan := make(chan error, 5)
+
+	// doneChan lets us specify that we should exit.
+	doneChan := make(chan interface{})
+
 	session, _, cleanup := newTestSession(t, &Server{
 		Handler: func(s Session) {
-			signals := make(chan Signal)
+			// We need to use a buffered channel here, otherwise it's possible for the
+			// second call to Signal to get discarded.
+			signals := make(chan Signal, 2)
 			s.Signals(signals)
-			if sig := <-signals; sig != SIGINT {
-				t.Fatalf("expected signal %v but got %v", SIGINT, sig)
+
+			select {
+			case sig := <-signals:
+				if sig != SIGINT {
+					errChan <- fmt.Errorf("expected signal %v but got %v", SIGINT, sig)
+					return
+				}
+			case <-doneChan:
+				errChan <- fmt.Errorf("Unexpected done")
+				return
 			}
-			exiter := make(chan bool)
-			go func() {
-				if sig := <-signals; sig == SIGKILL {
-					close(exiter)
+
+			select {
+			case sig := <-signals:
+				if sig != SIGKILL {
+					errChan <- fmt.Errorf("expected signal %v but got %v", SIGKILL, sig)
+					return
 				}
-			}()
-			<-exiter
+			case <-doneChan:
+				errChan <- fmt.Errorf("Unexpected done")
+				return
+			}
 		},
 	}, nil)
 	defer cleanup()
@@ -312,7 +332,13 @@ func TestSignals(t *testing.T) {
 		session.Signal(gossh.SIGKILL)
 	}()
 
-	err := session.Run("")
+	go func() {
+		errChan <- session.Run("")
+	}()
+
+	err := <-errChan
+	close(doneChan)
+
 	if err != nil {
 		t.Fatalf("expected nil but got %v", err)
 	}