From 114c5f8ac3fbb9f30f3eec5b40a01703c3fbca38 Mon Sep 17 00:00:00 2001
From: Sanskar Jaiswal <jaiswalsanskar078@gmail.com>
Date: Fri, 6 May 2022 15:07:20 +0530
Subject: [PATCH] modify ssh server to have artifical latency

Signed-off-by: Sanskar Jaiswal <jaiswalsanskar078@gmail.com>
---
 ssh.go      |  11 +++---
 ssh_test.go | 100 ++++++++++++++++++++++++++++++++++++++++------------
 2 files changed, 84 insertions(+), 27 deletions(-)

diff --git a/ssh.go b/ssh.go
index 6d765f6..93bbd16 100644
--- a/ssh.go
+++ b/ssh.go
@@ -46,7 +46,9 @@ type SSH struct {
 	DisableConnReuse bool
 	// DisableSimultaneousConns, if true will disable simultaneous conns from the same host.
 	DisableSimultaneousConns bool
-	PublicKeyLookupFunc      func(string) (*PublicKey, error)
+	// Latency, if set will introduce some synthetic latency to the ssh server operations.
+	Latency             *time.Duration
+	PublicKeyLookupFunc func(string) (*PublicKey, error)
 }
 
 func NewSSH(config Config) *SSH {
@@ -121,7 +123,6 @@ func (s *SSH) handleConnection(keyID string, chans <-chan ssh.NewChannel, sConn
 					host, _ := getHost(sConn.RemoteAddr().String())
 					mux.Lock()
 					defer mux.Unlock()
-					log.Println("disable simultaneous conns")
 					for i, connHost := range connHosts {
 						if host == connHost {
 							connHosts[i] = connHosts[len(connHosts)-1]
@@ -213,6 +214,10 @@ func (s *SSH) handleConnection(keyID string, chans <-chan ssh.NewChannel, sConn
 					io.Copy(ch, stdout)
 					io.Copy(ch.Stderr(), stderr)
 
+					if s.Latency != nil {
+						time.Sleep(*s.Latency)
+					}
+
 					if err = cmd.Wait(); err != nil {
 						log.Printf("ssh: command failed: %v", err)
 						return
@@ -395,8 +400,6 @@ func (s *SSH) Serve() error {
 			}
 			if !matched {
 				connHosts = append(connHosts, host)
-			} else {
-				continue
 			}
 		}
 
diff --git a/ssh_test.go b/ssh_test.go
index 97da587..74950ae 100644
--- a/ssh_test.go
+++ b/ssh_test.go
@@ -39,28 +39,13 @@ func TestListenAndServe(t *testing.T) {
 		{
 			name: "default ssh server",
 			serverFunc: func(repo, keyDir string) *SSH {
-				server := NewSSH(Config{
-					Dir:    filepath.Dir(repo),
-					KeyDir: keyDir,
-				})
-
-				server.PublicKeyLookupFunc = func(s string) (*PublicKey, error) {
-					return &PublicKey{Id: "12345"}, nil
-				}
-				return server
+				return setupSSHServer(repo, keyDir)
 			},
 		},
 		{
 			name: "ssh server times out",
 			serverFunc: func(repo, keyDir string) *SSH {
-				server := NewSSH(Config{
-					Dir:    filepath.Dir(repo),
-					KeyDir: keyDir,
-				})
-
-				server.PublicKeyLookupFunc = func(s string) (*PublicKey, error) {
-					return &PublicKey{Id: "12345"}, nil
-				}
+				server := setupSSHServer(repo, keyDir)
 				timeout := time.Nanosecond * 1
 				server.Timeout = &timeout
 				return server
@@ -103,18 +88,15 @@ func TestListenAndServe(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			cmd := exec.Command("git", "clone", "ssh://git@localhost:2222/"+filepath.Base(repo))
-			cmd.Dir = cloned
-			cmd.Env = []string{"GIT_SSH_COMMAND=ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no"}
-
+			cmd := getCloneCommand(filepath.Base(repo), cloned)
 			e := new(strings.Builder)
 			cmd.Stderr = e
 			err = cmd.Start()
 			if err != nil {
-				panic(err)
+				t.Fatalf("faile to start git clone command: %s", e.String())
 			}
-			err = cmd.Wait()
 
+			err = cmd.Wait()
 			g.Expect(err != nil).To(Equal(tt.err))
 			_, err = os.Stat(filepath.Join(cloned, filepath.Base(repo)))
 			if !tt.err {
@@ -125,6 +107,78 @@ func TestListenAndServe(t *testing.T) {
 
 }
 
+func TestSshServerLatency(t *testing.T) {
+	repo, err := createRepo()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(repo)
+	keyDir, err := os.MkdirTemp("", "key-dir")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(keyDir)
+
+	server := setupSSHServer(repo, keyDir)
+	latency := time.Second * 5
+	server.Latency = &latency
+
+	defer server.Stop()
+
+	go func() {
+		server.ListenAndServe(":2222")
+	}()
+
+	cloned, err := os.MkdirTemp("", "cloned")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(cloned)
+
+	if err = retry(10, time.Second*1, func() error {
+		_, err := net.Dial("tcp", "localhost:2222")
+		return err
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	cmd := getCloneCommand(filepath.Base(repo), cloned)
+	e := new(strings.Builder)
+	cmd.Stderr = e
+
+	err = cmd.Start()
+	if err != nil {
+		t.Fatalf("faile to start git clone command: %s", e.String())
+	}
+
+	g := NewWithT(t)
+	start := time.Now()
+	err = cmd.Wait()
+	g.Expect(err).ToNot(HaveOccurred())
+	g.Expect(time.Now()).ShouldNot(BeTemporally("~", start, *server.Latency))
+	_, err = os.Stat(filepath.Join(cloned, filepath.Base(repo)))
+	g.Expect(err).ToNot(HaveOccurred())
+}
+
+func getCloneCommand(repoName, cmdDir string) *exec.Cmd {
+	cmd := exec.Command("git", "clone", "ssh://git@localhost:2222/"+repoName)
+	cmd.Dir = cmdDir
+	cmd.Env = []string{"GIT_SSH_COMMAND=ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no"}
+	return cmd
+}
+
+func setupSSHServer(repo, keyDir string) *SSH {
+	server := NewSSH(Config{
+		Dir:    filepath.Dir(repo),
+		KeyDir: keyDir,
+	})
+
+	server.PublicKeyLookupFunc = func(s string) (*PublicKey, error) {
+		return &PublicKey{Id: "12345"}, nil
+	}
+	return server
+}
+
 func createRepo() (string, error) {
 	repo, err := os.MkdirTemp("", "ssh-test")
 	if err != nil {