Properly close connections in the proxy.

README.md

@@ -65,6 +65,32 @@ the project root directory.
 Building SansShell requires a recent version of Go (check the go.mod file for
 the current version).
 
+## Build and run
+
+You need to populate ~/.sansshell with certificates before running.
+
+```
+$ cp -r auth/mtls/testdata ~/.sanshell
+```
+
+Or copy the test certificates from auth/mtls/testdata to ~/.sanshell
+
+Then you can build and run the server, in separate terminal windows:
+```
+$ go run ./cmd/sansshell-server
+$ go run ./cmd/sanssh --targets=localhost file read /etc/hosts
+```
+
+You can also run the proxy to try the full flow:
+
+```
+$ go run ./cmd/sansshell-server
+$ go run ./cmd/proxy-server
+$ go run ./cmd/sanssh --proxy=localhost:50043 --targets=localhost:50042 file read /etc/hosts
+```
+
+Minimal debugging UIs are available at http://localhost:50044 for the server and http://localhost:50046 for the proxy by default.
+
 ## Environment setup : protoc
 
 When making any change to the protocol buffers, you'll also need the protocol
@@ -100,8 +126,9 @@ do this for you, as well as re-generating the service proto files.
 $ go generate tools.go
 ```
 
-## Build and run
-You only need to do these steps once to configure example mTLS certs:
+## Creating your own certificates
+
+As an alternative to copying auth/mtls/testdata, you can create your onwn example mTLS certs:
 ```
 $ go install github.com/meterup/generate-cert@latest
 $ mkdir -m 0700 certs
@@ -111,14 +138,6 @@ $ cd ../
 $ ln -s $(pwd)/certs ~/.sansshell
 ```
 
-Or copy the test certificates from auth/mtls/testdata to ~/.sanshell
-
-Then you can build and run the server, in separate terminal windows:
-```
-$ cd cmd/sansshell-server && go build && ./sansshell-server
-$ cd cmd/sanssh && go build && ./sanssh --targets=localhost file read /etc/hosts
-```
-
 ## Debugging
 Reflection is included in the RPC servers (proxy and sansshell-server)
 allowing for the use of [grpc_cli](https://github.com/grpc/grpc/blob/master/doc/command_line_tool.md).

proxy/server/server.go

@@ -45,7 +45,13 @@ var (
 // connections (such as client credentials, deadlines, etc) which
 // the proxy can use without needing to understand them.
 type TargetDialer interface {
-	DialContext(ctx context.Context, target string, dialOpts ...grpc.DialOption) (grpc.ClientConnInterface, error)
+	DialContext(ctx context.Context, target string, dialOpts ...grpc.DialOption) (ClientConnCloser, error)
+}
+
+// ClientConnCloser is a closeable grpc.ClientConnInterface
+type ClientConnCloser interface {
+	grpc.ClientConnInterface
+	Close() error
 }
 
 // an optionsDialer implements TargetDialer using native grpc.Dial
@@ -54,7 +60,7 @@ type optionsDialer struct {
 }
 
 // See TargetDialer.DialContext
-func (o *optionsDialer) DialContext(ctx context.Context, target string, dialOpts ...grpc.DialOption) (grpc.ClientConnInterface, error) {
+func (o *optionsDialer) DialContext(ctx context.Context, target string, dialOpts ...grpc.DialOption) (ClientConnCloser, error) {
 	opts := o.opts
 	opts = append(opts, dialOpts...)
 	return grpc.DialContext(ctx, target, opts...)

proxy/server/target.go

@@ -65,7 +65,7 @@ type TargetStream struct {
 	serviceMethod *ServiceMethod
 
 	// The underlying grpc.ClientConnInterface to the target server
-	grpcConn grpc.ClientConnInterface
+	grpcConn ClientConnCloser
 
 	// The underlying grpc.ClientStream to the target server.
 	grpcStream grpc.ClientStream
@@ -198,13 +198,14 @@ func (s *TargetStream) Run(nonce uint32, replyChan chan *pb.ProxyReply) {
 		}
 		var err error
 		defer cancel()
-		s.grpcConn, err = s.dialer.DialContext(dialCtx, s.target, opts...)
+		grpcConn, err := s.dialer.DialContext(dialCtx, s.target, opts...)
 		if err != nil {
 			// We cannot create a new stream to the target. So we need to cancel this stream.
 			s.logger.Info("unable to create stream", "status", err)
 			s.cancelFunc()
 			return err
 		}
+		s.grpcConn = grpcConn
 		grpcStream, err := s.grpcConn.NewStream(s.ctx, s.serviceMethod.StreamDesc(), s.serviceMethod.FullName())
 		if err != nil {
 			// We cannot create a new stream to the target. So we need to cancel this stream.
@@ -323,6 +324,14 @@ func (s *TargetStream) Run(nonce uint32, replyChan chan *pb.ProxyReply) {
 	// a server-close call
 	err := group.Wait()
 
+	// Once all calls are complete, we need to close our network connection
+	// to the server.
+	if s.grpcConn != nil {
+		if closeErr := s.grpcConn.Close(); err == nil && closeErr != nil {
+			err = closeErr
+		}
+	}
+
 	// The error status may by set/overidden if CloseWith was used to
 	// terminate the stream.
 	select {

proxy/server/target_test.go

@@ -35,7 +35,7 @@ import (
 // A TargetDialer than returns an error for all Dials
 type dialErrTargetDialer codes.Code
 
-func (e dialErrTargetDialer) DialContext(ctx context.Context, target string, dialOpts ...grpc.DialOption) (grpc.ClientConnInterface, error) {
+func (e dialErrTargetDialer) DialContext(ctx context.Context, target string, dialOpts ...grpc.DialOption) (ClientConnCloser, error) {
 	return nil, status.Error(codes.Code(e), "")
 }
 
@@ -150,10 +150,14 @@ func (b blockingClientConn) NewStream(ctx context.Context, desc *grpc.StreamDesc
 	return nil, ctx.Err()
 }
 
+func (b blockingClientConn) Close() error {
+	return nil
+}
+
 // a context dialer that returns blockingClientConn
 type blockingClientDialer struct{}
 
-func (b blockingClientDialer) DialContext(ctx context.Context, target string, dialOpts ...grpc.DialOption) (grpc.ClientConnInterface, error) {
+func (b blockingClientDialer) DialContext(ctx context.Context, target string, dialOpts ...grpc.DialOption) (ClientConnCloser, error) {
 	return blockingClientConn{}, nil
 }