Merge branch '18633-install-sudo-passwordless-in-shell'

[arvados.git] / sdk / go / arvados / container_gateway.go

diff --git a/sdk/go/arvados/container_gateway.go b/sdk/go/arvados/container_gateway.go
index 07f8c0793caacf0ae728e798a4e195deba1e59a7..00c98d572ea010e2db43ca1c62bfa0eb341b1e7d 100644 (file)
--- a/sdk/go/arvados/container_gateway.go
+++ b/sdk/go/arvados/container_gateway.go
@@ -8,8 +8,10 @@ import (
        "context"
        "io"
        "net/http"
+       "sync"
 
        "git.arvados.org/arvados.git/sdk/go/ctxlog"
+       "github.com/sirupsen/logrus"
 )
 
 func (sshconn ContainerSSHConnection) ServeHTTP(w http.ResponseWriter, req *http.Request) {
@@ -28,26 +30,45 @@ func (sshconn ContainerSSHConnection) ServeHTTP(w http.ResponseWriter, req *http
        }
        defer conn.Close()
 
+       var bytesIn, bytesOut int64
+       var wg sync.WaitGroup
        ctx, cancel := context.WithCancel(context.Background())
+       wg.Add(1)
        go func() {
+               defer wg.Done()
                defer cancel()
-               _, err := io.CopyN(conn, sshconn.Bufrw, int64(sshconn.Bufrw.Reader.Buffered()))
+               n, err := io.CopyN(conn, sshconn.Bufrw, int64(sshconn.Bufrw.Reader.Buffered()))
+               bytesOut += n
                if err == nil {
-                       _, err = io.Copy(conn, sshconn.Conn)
+                       n, err = io.Copy(conn, sshconn.Conn)
+                       bytesOut += n
                }
                if err != nil {
                        ctxlog.FromContext(req.Context()).WithError(err).Error("error copying downstream")
                }
        }()
+       wg.Add(1)
        go func() {
+               defer wg.Done()
                defer cancel()
-               _, err := io.CopyN(sshconn.Conn, bufrw, int64(bufrw.Reader.Buffered()))
+               n, err := io.CopyN(sshconn.Conn, bufrw, int64(bufrw.Reader.Buffered()))
+               bytesIn += n
                if err == nil {
-                       _, err = io.Copy(sshconn.Conn, conn)
+                       n, err = io.Copy(sshconn.Conn, conn)
+                       bytesIn += n
                }
                if err != nil {
                        ctxlog.FromContext(req.Context()).WithError(err).Error("error copying upstream")
                }
        }()
        <-ctx.Done()
+       if sshconn.Logger != nil {
+               go func() {
+                       wg.Wait()
+                       sshconn.Logger.WithFields(logrus.Fields{
+                               "bytesIn":  bytesIn,
+                               "bytesOut": bytesOut,
+                       }).Info("closed connection")
+               }()
+       }
 }