Merge pull request #6918 from skorhone/fix/hijacked_connection_handling

Fix: Correct connection counters for hijacked connections

5 files changed, 148 insertions, 92 deletions

diff --git a/pkg/api/handlers/compat/containers_attach.go b/pkg/api/handlers/compat/containers_attach.go
index 325f96b40..71586fca4 100644
--- a/pkg/api/handlers/compat/containers_attach.go
+++ b/pkg/api/handlers/compat/containers_attach.go
@@ -1,23 +1,22 @@
 package compat
 
 import (
+	"bufio"
 	"fmt"
+	"io"
+	"net"
 	"net/http"
+	"strings"
 
 	"github.com/containers/libpod/v2/libpod"
 	"github.com/containers/libpod/v2/libpod/define"
 	"github.com/containers/libpod/v2/pkg/api/handlers/utils"
+	"github.com/containers/libpod/v2/pkg/api/server/idletracker"
 	"github.com/gorilla/schema"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 )
 
-// AttachHeader is the literal header sent for upgraded/hijacked connections for
-// attach, sourced from Docker at:
-// https://raw.githubusercontent.com/moby/moby/b95fad8e51bd064be4f4e58a996924f343846c85/api/server/router/container/container_routes.go
-// Using literally to ensure compatibility with existing clients.
-const AttachHeader = "HTTP/1.1 101 UPGRADED\r\nContent-Type: application/vnd.docker.raw-stream\r\nConnection: Upgrade\r\nUpgrade: tcp\r\n\r\n"
-
 func AttachContainer(w http.ResponseWriter, r *http.Request) {
 	runtime := r.Context().Value("runtime").(*libpod.Runtime)
 	decoder := r.Context().Value("decoder").(*schema.Decoder)
@@ -98,21 +97,11 @@ func AttachContainer(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Hijack the connection
-	hijacker, ok := w.(http.Hijacker)
-	if !ok {
-		utils.InternalServerError(w, errors.Errorf("unable to hijack connection"))
-		return
-	}
-
-	connection, buffer, err := hijacker.Hijack()
+	connection, buffer, err := AttachConnection(w, r)
 	if err != nil {
-		utils.InternalServerError(w, errors.Wrapf(err, "error hijacking connection"))
+		utils.InternalServerError(w, err)
 		return
 	}
-
-	fmt.Fprintf(connection, AttachHeader)
-
 	logrus.Debugf("Hijack for attach of container %s successful", ctr.ID())
 
 	// Perform HTTP attach.
@@ -126,3 +115,57 @@ func AttachContainer(w http.ResponseWriter, r *http.Request) {
 
 	logrus.Debugf("Attach for container %s completed successfully", ctr.ID())
 }
+
+type HijackedConnection struct {
+	net.Conn                             // Connection
+	idleTracker *idletracker.IdleTracker // Connection tracker
+}
+
+func (c HijackedConnection) Close() error {
+	logrus.Debugf("Hijacked connection closed")
+
+	c.idleTracker.TrackHijackedClosed()
+	return c.Conn.Close()
+}
+
+func AttachConnection(w http.ResponseWriter, r *http.Request) (net.Conn, *bufio.ReadWriter, error) {
+	idleTracker := r.Context().Value("idletracker").(*idletracker.IdleTracker)
+
+	// Hijack the connection
+	hijacker, ok := w.(http.Hijacker)
+	if !ok {
+		return nil, nil, errors.Errorf("unable to hijack connection")
+	}
+
+	connection, buffer, err := hijacker.Hijack()
+	if err != nil {
+		return nil, nil, errors.Wrapf(err, "error hijacking connection")
+	}
+	trackedConnection := HijackedConnection{
+		Conn:        connection,
+		idleTracker: idleTracker,
+	}
+
+	WriteAttachHeaders(r, trackedConnection)
+
+	return trackedConnection, buffer, nil
+}
+
+func WriteAttachHeaders(r *http.Request, connection io.Writer) {
+	// AttachHeader is the literal header sent for upgraded/hijacked connections for
+	// attach, sourced from Docker at:
+	// https://raw.githubusercontent.com/moby/moby/b95fad8e51bd064be4f4e58a996924f343846c85/api/server/router/container/container_routes.go
+	// Using literally to ensure compatibility with existing clients.
+	c := r.Header.Get("Connection")
+	proto := r.Header.Get("Upgrade")
+	if len(proto) == 0 || !strings.EqualFold(c, "Upgrade") {
+		// OK - can't upgrade if not requested or protocol is not specified
+		fmt.Fprintf(connection,
+			"HTTP/1.1 200 OK\r\nContent-Type: application/vnd.docker.raw-stream\r\n\r\n")
+	} else {
+		// Upraded
+		fmt.Fprintf(connection,
+			"HTTP/1.1 101 UPGRADED\r\nContent-Type: application/vnd.docker.raw-stream\r\nConnection: Upgrade\r\nUpgrade: %s\r\n\r\n",
+			proto)
+	}
+}
diff --git a/pkg/api/handlers/compat/exec.go b/pkg/api/handlers/compat/exec.go
index aee4196dd..a3b8cb573 100644
--- a/pkg/api/handlers/compat/exec.go
+++ b/pkg/api/handlers/compat/exec.go
@@ -173,21 +173,11 @@ func ExecStartHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Hijack the connection
-	hijacker, ok := w.(http.Hijacker)
-	if !ok {
-		utils.InternalServerError(w, errors.Errorf("unable to hijack connection"))
-		return
-	}
-
-	connection, buffer, err := hijacker.Hijack()
+	connection, buffer, err := AttachConnection(w, r)
 	if err != nil {
-		utils.InternalServerError(w, errors.Wrapf(err, "error hijacking connection"))
+		utils.InternalServerError(w, err)
 		return
 	}
-
-	fmt.Fprintf(connection, AttachHeader)
-
 	logrus.Debugf("Hijack for attach of container %s exec session %s successful", sessionCtr.ID(), sessionID)
 
 	if err := sessionCtr.ExecHTTPStartAndAttach(sessionID, connection, buffer, nil, nil, nil); err != nil {
diff --git a/pkg/api/server/handler_api.go b/pkg/api/server/handler_api.go
index b0fd932ba..53fe8952b 100644
--- a/pkg/api/server/handler_api.go
+++ b/pkg/api/server/handler_api.go
@@ -37,6 +37,7 @@ func (s *APIServer) APIHandler(h http.HandlerFunc) http.HandlerFunc {
 			c := context.WithValue(r.Context(), "decoder", s.Decoder) //nolint
 			c = context.WithValue(c, "runtime", s.Runtime)            //nolint
 			c = context.WithValue(c, "shutdownFunc", s.Shutdown)      //nolint
+			c = context.WithValue(c, "idletracker", s.idleTracker)    //nolint
 			r = r.WithContext(c)
 
 			h(w, r)
diff --git a/pkg/api/server/idletracker/idletracker.go b/pkg/api/server/idletracker/idletracker.go
new file mode 100644
index 000000000..1ee905a99
--- /dev/null
+++ b/pkg/api/server/idletracker/idletracker.go
@@ -0,0 +1,74 @@
+package idletracker
+
+import (
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/sirupsen/logrus"
+)
+
+type IdleTracker struct {
+	http     map[net.Conn]struct{}
+	hijacked int
+	total    int
+	mux      sync.Mutex
+	timer    *time.Timer
+	Duration time.Duration
+}
+
+func NewIdleTracker(idle time.Duration) *IdleTracker {
+	return &IdleTracker{
+		http:     make(map[net.Conn]struct{}),
+		Duration: idle,
+		timer:    time.NewTimer(idle),
+	}
+}
+
+func (t *IdleTracker) ConnState(conn net.Conn, state http.ConnState) {
+	t.mux.Lock()
+	defer t.mux.Unlock()
+
+	oldActive := t.ActiveConnections()
+	logrus.Debugf("IdleTracker %p:%v %d/%d connection(s)", conn, state, oldActive, t.TotalConnections())
+	switch state {
+	case http.StateNew, http.StateActive:
+		t.http[conn] = struct{}{}
+		// stop the timer if we transitioned from idle
+		if oldActive == 0 {
+			t.timer.Stop()
+		}
+		t.total++
+	case http.StateHijacked:
+		// hijacked connections are handled elsewhere
+		delete(t.http, conn)
+		t.hijacked++
+	case http.StateIdle, http.StateClosed:
+		delete(t.http, conn)
+		// Restart the timer if we've become idle
+		if oldActive > 0 && len(t.http) == 0 {
+			t.timer.Stop()
+			t.timer.Reset(t.Duration)
+		}
+	}
+}
+
+func (t *IdleTracker) TrackHijackedClosed() {
+	t.mux.Lock()
+	defer t.mux.Unlock()
+
+	t.hijacked--
+}
+
+func (t *IdleTracker) ActiveConnections() int {
+	return len(t.http) + t.hijacked
+}
+
+func (t *IdleTracker) TotalConnections() int {
+	return t.total
+}
+
+func (t *IdleTracker) Done() <-chan time.Time {
+	return t.timer.C
+}
diff --git a/pkg/api/server/server.go b/pkg/api/server/server.go
index 8af6d3186..1c6007745 100644
--- a/pkg/api/server/server.go
+++ b/pkg/api/server/server.go
@@ -10,12 +10,12 @@ import (
 	"runtime"
 	goRuntime "runtime"
 	"strings"
-	"sync"
 	"syscall"
 	"time"
 
 	"github.com/containers/libpod/v2/libpod"
 	"github.com/containers/libpod/v2/pkg/api/handlers"
+	"github.com/containers/libpod/v2/pkg/api/server/idletracker"
 	"github.com/coreos/go-systemd/v22/activation"
 	"github.com/gorilla/mux"
 	"github.com/gorilla/schema"
@@ -24,14 +24,14 @@ import (
 )
 
 type APIServer struct {
-	http.Server                     // The  HTTP work happens here
-	*schema.Decoder                 // Decoder for Query parameters to structs
-	context.Context                 // Context to carry objects to handlers
-	*libpod.Runtime                 // Where the real work happens
-	net.Listener                    // mux for routing HTTP API calls to libpod routines
-	context.CancelFunc              // Stop APIServer
-	idleTracker        *IdleTracker // Track connections to support idle shutdown
-	pprof              *http.Server // Sidecar http server for providing performance data
+	http.Server                                 // The  HTTP work happens here
+	*schema.Decoder                             // Decoder for Query parameters to structs
+	context.Context                             // Context to carry objects to handlers
+	*libpod.Runtime                             // Where the real work happens
+	net.Listener                                // mux for routing HTTP API calls to libpod routines
+	context.CancelFunc                          // Stop APIServer
+	idleTracker        *idletracker.IdleTracker // Track connections to support idle shutdown
+	pprof              *http.Server             // Sidecar http server for providing performance data
 }
 
 // Number of seconds to wait for next request, if exceeded shutdown server
@@ -68,7 +68,7 @@ func newServer(runtime *libpod.Runtime, duration time.Duration, listener *net.Li
 	}
 
 	router := mux.NewRouter().UseEncodedPath()
-	idle := NewIdleTracker(duration)
+	idle := idletracker.NewIdleTracker(duration)
 
 	server := APIServer{
 		Server: http.Server{
@@ -231,55 +231,3 @@ func (s *APIServer) Shutdown() error {
 func (s *APIServer) Close() error {
 	return s.Server.Close()
 }
-
-type IdleTracker struct {
-	active   map[net.Conn]struct{}
-	total    int
-	mux      sync.Mutex
-	timer    *time.Timer
-	Duration time.Duration
-}
-
-func NewIdleTracker(idle time.Duration) *IdleTracker {
-	return &IdleTracker{
-		active:   make(map[net.Conn]struct{}),
-		Duration: idle,
-		timer:    time.NewTimer(idle),
-	}
-}
-
-func (t *IdleTracker) ConnState(conn net.Conn, state http.ConnState) {
-	t.mux.Lock()
-	defer t.mux.Unlock()
-
-	oldActive := len(t.active)
-	logrus.Debugf("IdleTracker %p:%v %d/%d connection(s)", conn, state, t.ActiveConnections(), t.TotalConnections())
-	switch state {
-	case http.StateNew, http.StateActive, http.StateHijacked:
-		t.active[conn] = struct{}{}
-		// stop the timer if we transitioned from idle
-		if oldActive == 0 {
-			t.timer.Stop()
-		}
-		t.total++
-	case http.StateIdle, http.StateClosed:
-		delete(t.active, conn)
-		// Restart the timer if we've become idle
-		if oldActive > 0 && len(t.active) == 0 {
-			t.timer.Stop()
-			t.timer.Reset(t.Duration)
-		}
-	}
-}
-
-func (t *IdleTracker) ActiveConnections() int {
-	return len(t.active)
-}
-
-func (t *IdleTracker) TotalConnections() int {
-	return t.total
-}
-
-func (t *IdleTracker) Done() <-chan time.Time {
-	return t.timer.C
-}