From bc3a1eefdd0f66f45d7fbcea470059c09527c869 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Mon, 26 Aug 2019 15:34:55 -0700
Subject: [PATCH] session: release forwarded ssh socket connection per
 connection

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 session/sshforward/copy.go                      | 7 +++++--
 session/sshforward/ssh.go                       | 2 +-
 session/sshforward/sshprovider/agentprovider.go | 2 +-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/session/sshforward/copy.go b/session/sshforward/copy.go
index c2763fa45..85366f19a 100644
--- a/session/sshforward/copy.go
+++ b/session/sshforward/copy.go
@@ -9,17 +9,17 @@ import (
 	"google.golang.org/grpc"
 )
 
-func Copy(ctx context.Context, conn io.ReadWriteCloser, stream grpc.Stream) error {
+func Copy(ctx context.Context, conn io.ReadWriteCloser, stream grpc.Stream, closeStream func() error) error {
 	g, ctx := errgroup.WithContext(ctx)
 
 	g.Go(func() (retErr error) {
 		p := &BytesMessage{}
 		for {
 			if err := stream.RecvMsg(p); err != nil {
+				conn.Close()
 				if err == io.EOF {
 					return nil
 				}
-				conn.Close()
 				return errors.WithStack(err)
 			}
 			select {
@@ -42,6 +42,9 @@ func Copy(ctx context.Context, conn io.ReadWriteCloser, stream grpc.Stream) erro
 			n, err := conn.Read(buf)
 			switch {
 			case err == io.EOF:
+				if closeStream != nil {
+					closeStream()
+				}
 				return nil
 			case err != nil:
 				return errors.WithStack(err)
diff --git a/session/sshforward/ssh.go b/session/sshforward/ssh.go
index 660e89f7f..0001f59b5 100644
--- a/session/sshforward/ssh.go
+++ b/session/sshforward/ssh.go
@@ -49,7 +49,7 @@ func (s *server) run(ctx context.Context, l net.Listener, id string) error {
 				return err
 			}
 
-			go Copy(ctx, conn, stream)
+			go Copy(ctx, conn, stream, stream.CloseSend)
 		}
 	})
 
diff --git a/session/sshforward/sshprovider/agentprovider.go b/session/sshforward/sshprovider/agentprovider.go
index 009a91b7d..7aa3e3dfc 100644
--- a/session/sshforward/sshprovider/agentprovider.go
+++ b/session/sshforward/sshprovider/agentprovider.go
@@ -114,7 +114,7 @@ func (sp *socketProvider) ForwardAgent(stream sshforward.SSH_ForwardAgentServer)
 
 	eg.Go(func() error {
 		defer s1.Close()
-		return sshforward.Copy(ctx, s2, stream)
+		return sshforward.Copy(ctx, s2, stream, nil)
 	})
 
 	return eg.Wait()