sdk/go/arvados/container_gateway.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package arvados
   6
   7 import (
   8         "context"
   9         "io"
  10         "net/http"
  11
  12         "git.arvados.org/arvados.git/sdk/go/ctxlog"
  13 )
  14
  15 func (sshconn ContainerSSHConnection) ServeHTTP(w http.ResponseWriter, req *http.Request) {
  16         hj, ok := w.(http.Hijacker)
  17         if !ok {
  18                 http.Error(w, "ResponseWriter does not support connection upgrade", http.StatusInternalServerError)
  19                 return
  20         }
  21         w.Header().Set("Connection", "upgrade")
  22         w.Header().Set("Upgrade", "ssh")
  23         w.WriteHeader(http.StatusSwitchingProtocols)
  24         conn, bufrw, err := hj.Hijack()
  25         if err != nil {
  26                 ctxlog.FromContext(req.Context()).WithError(err).Error("error hijacking ResponseWriter")
  27                 return
  28         }
  29         defer conn.Close()
  30
  31         ctx, cancel := context.WithCancel(context.Background())
  32         go func() {
  33                 defer cancel()
  34                 _, err := io.CopyN(conn, sshconn.Bufrw, int64(sshconn.Bufrw.Reader.Buffered()))
  35                 if err == nil {
  36                         _, err = io.Copy(conn, sshconn.Conn)
  37                 }
  38                 if err != nil {
  39                         ctxlog.FromContext(req.Context()).WithError(err).Error("error copying downstream")
  40                 }
  41         }()
  42         go func() {
  43                 defer cancel()
  44                 _, err := io.CopyN(sshconn.Conn, bufrw, int64(bufrw.Reader.Buffered()))
  45                 if err == nil {
  46                         _, err = io.Copy(sshconn.Conn, conn)
  47                 }
  48                 if err != nil {
  49                         ctxlog.FromContext(req.Context()).WithError(err).Error("error copying upstream")
  50                 }
  51         }()
  52         <-ctx.Done()
  53 }