[ws-proxy] close connection if workspace or port stopped share (#20248)

* remove useless InfoProvider

* close connection if visibility change

* nit: fix typo

* fn name typo fix

---------

Co-authored-by: Filip Troníček <[email protected]>

Author: iQQBot

components/ws-proxy/cmd/run.go

@@ -78,15 +78,13 @@ var runCmd = &cobra.Command{
 			log.WithError(err).Fatal(err, "unable to start manager")
 		}
 
-		var infoprov proxy.CompositeInfoProvider
-		crdInfoProv, err := proxy.NewCRDWorkspaceInfoProvider(mgr.GetClient(), mgr.GetScheme())
+		infoprov, err := proxy.NewCRDWorkspaceInfoProvider(mgr.GetClient(), mgr.GetScheme())
 		if err != nil {
 			log.WithError(err).Fatal("cannot create CRD-based info provider")
 		}
-		if err = crdInfoProv.SetupWithManager(mgr); err != nil {
+		if err = infoprov.SetupWithManager(mgr); err != nil {
 			log.WithError(err).Fatal(err, "unable to create CRD-based info provider", "controller", "Workspace")
 		}
-		infoprov = append(infoprov, crdInfoProv)
 
 		if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
 			log.WithError(err).Fatal("unable to set up health check")

components/ws-proxy/pkg/common/infoprovider.go

@@ -5,6 +5,7 @@
 package common
 
 import (
+	"context"
 	"time"
 
 	"github.com/gitpod-io/gitpod/ws-manager/api"
@@ -43,6 +44,9 @@ type WorkspaceCoords struct {
 type WorkspaceInfoProvider interface {
 	// WorkspaceInfo returns the workspace information of a workspace using it's workspace ID
 	WorkspaceInfo(workspaceID string) *WorkspaceInfo
+
+	AcquireContext(ctx context.Context, workspaceID, port string) (context.Context, string, error)
+	ReleaseContext(id string)
 }
 
 // WorkspaceInfo is all the infos ws-proxy needs to know about a workspace.

components/ws-proxy/pkg/proxy/auth.go

@@ -5,6 +5,7 @@
 package proxy
 
 import (
+	"errors"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -17,6 +18,12 @@ import (
 	"github.com/gitpod-io/gitpod/ws-proxy/pkg/common"
 )
 
+var (
+	ErrTokenNotFound = fmt.Errorf("no owner cookie present")
+	ErrTokenMismatch = fmt.Errorf("owner token mismatch")
+	ErrTokenDecode   = fmt.Errorf("cannot decode owner token")
+)
+
 // WorkspaceAuthHandler rejects requests which are not authenticated or authorized to access a workspace.
 func WorkspaceAuthHandler(domain string, info common.WorkspaceInfoProvider) mux.MiddlewareFunc {
 	return func(h http.Handler) http.Handler {
@@ -47,68 +54,76 @@ func WorkspaceAuthHandler(domain string, info common.WorkspaceInfoProvider) mux.
 				return
 			}
 
+			isPublic := false
 			if ws.Auth != nil && ws.Auth.Admission == api.AdmissionLevel_ADMIT_EVERYONE {
-				// workspace is free for all - no tokens or cookies matter
-				h.ServeHTTP(resp, req)
-
-				return
-			}
-
-			if port != "" {
-				// this is a workspace port request and ports can be public or private.
-				// For public ports no tokens or cookies matter, private ports are subject
-				// to the same access policies as the workspace itself is.
-				var isPublic bool
-
+				isPublic = true
+			} else if port != "" {
 				prt, err := strconv.ParseUint(port, 10, 16)
 				if err != nil {
 					log.WithField("port", port).WithError(err).Error("cannot convert port to int")
 				} else {
 					for _, p := range ws.Ports {
 						if p.Port == uint32(prt) {
 							isPublic = p.Visibility == api.PortVisibility_PORT_VISIBILITY_PUBLIC
-
 							break
 						}
 					}
 				}
-
-				if isPublic {
-					// workspace port is free for all - no tokens or cookies matter
-					h.ServeHTTP(resp, req)
-
-					return
-				}
-
-				// port seems to be private - subject it to the same access policy as the workspace itself
 			}
 
-			tkn := req.Header.Get("x-gitpod-owner-token")
-			if tkn == "" {
-				cn := fmt.Sprintf("%s%s_owner_", cookiePrefix, ws.InstanceID)
-				c, err := req.Cookie(cn)
+			authenticate := func() (bool, error) {
+				tkn := req.Header.Get("x-gitpod-owner-token")
+				if tkn == "" {
+					cn := fmt.Sprintf("%s%s_owner_", cookiePrefix, ws.InstanceID)
+					c, err := req.Cookie(cn)
+					if err != nil {
+						return false, ErrTokenNotFound
+					}
+					tkn = c.Value
+				}
+				tkn, err := url.QueryUnescape(tkn)
 				if err != nil {
-					log.WithField("cookieName", cn).Debug("no owner cookie present")
-					resp.WriteHeader(http.StatusUnauthorized)
-
-					return
+					return false, ErrTokenDecode
 				}
 
-				tkn = c.Value
+				if tkn != ws.Auth.OwnerToken {
+					return false, ErrTokenMismatch
+				}
+				return true, nil
 			}
-			tkn, err := url.QueryUnescape(tkn)
-			if err != nil {
-				log.WithError(err).Warn("cannot decode owner token")
-				resp.WriteHeader(http.StatusBadRequest)
 
+			authenticated, err := authenticate()
+			if !authenticated && !isPublic {
+				if err != nil {
+					if errors.Is(err, ErrTokenNotFound) {
+						resp.WriteHeader(http.StatusUnauthorized)
+						return
+					}
+					if errors.Is(err, ErrTokenMismatch) {
+						log.Warn("owner token mismatch")
+						resp.WriteHeader(http.StatusForbidden)
+						return
+					}
+					if errors.Is(err, ErrTokenDecode) {
+						log.Warn("cannot decode owner token")
+						resp.WriteHeader(http.StatusBadRequest)
+						return
+					}
+				}
+				log.WithError(err).Error("cannot authenticate")
+				resp.WriteHeader(http.StatusInternalServerError)
 				return
 			}
 
-			if tkn != ws.Auth.OwnerToken {
-				log.Warn("owner token mismatch")
-				resp.WriteHeader(http.StatusForbidden)
-
-				return
+			if !authenticated && isPublic {
+				ctx, id, err := info.AcquireContext(req.Context(), wsID, port)
+				if err != nil {
+					log.WithError(err).Error("cannot acquire context")
+					resp.WriteHeader(http.StatusInternalServerError)
+					return
+				}
+				defer info.ReleaseContext(id)
+				req = req.WithContext(ctx)
 			}
 
 			h.ServeHTTP(resp, req)

components/ws-proxy/pkg/proxy/auth_test.go

@@ -36,8 +36,8 @@ func TestWorkspaceAuthHandler(t *testing.T) {
 		testPort    = 8080
 	)
 	var (
-		ownerOnlyInfos = map[string]*common.WorkspaceInfo{
-			workspaceID: {
+		ownerOnlyInfos = []common.WorkspaceInfo{
+			{
 				WorkspaceID: workspaceID,
 				InstanceID:  instanceID,
 				Auth: &api.WorkspaceAuthentication{
@@ -47,8 +47,8 @@ func TestWorkspaceAuthHandler(t *testing.T) {
 				Ports: []*api.PortSpec{{Port: testPort, Visibility: api.PortVisibility_PORT_VISIBILITY_PRIVATE}},
 			},
 		}
-		publicPortInfos = map[string]*common.WorkspaceInfo{
-			workspaceID: {
+		publicPortInfos = []common.WorkspaceInfo{
+			{
 				WorkspaceID: workspaceID,
 				InstanceID:  instanceID,
 				Auth: &api.WorkspaceAuthentication{
@@ -58,8 +58,8 @@ func TestWorkspaceAuthHandler(t *testing.T) {
 				Ports: []*api.PortSpec{{Port: testPort, Visibility: api.PortVisibility_PORT_VISIBILITY_PUBLIC}},
 			},
 		}
-		admitEveryoneInfos = map[string]*common.WorkspaceInfo{
-			workspaceID: {
+		admitEveryoneInfos = []common.WorkspaceInfo{
+			{
 				WorkspaceID: workspaceID,
 				InstanceID:  instanceID,
 				Auth:        &api.WorkspaceAuthentication{Admission: api.AdmissionLevel_ADMIT_EVERYONE},
@@ -68,7 +68,7 @@ func TestWorkspaceAuthHandler(t *testing.T) {
 	)
 	tests := []struct {
 		Name        string
-		Infos       map[string]*common.WorkspaceInfo
+		Infos       []common.WorkspaceInfo
 		OwnerCookie string
 		WorkspaceID string
 		Port        string
@@ -225,7 +225,7 @@ func TestWorkspaceAuthHandler(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.Name, func(t *testing.T) {
 			var res testResult
-			handler := WorkspaceAuthHandler(domain, &fixedInfoProvider{Infos: test.Infos})(http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
+			handler := WorkspaceAuthHandler(domain, &fakeWsInfoProvider{infos: test.Infos})(http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
 				res.HandlerCalled = true
 				resp.WriteHeader(http.StatusOK)
 			}))

components/ws-proxy/pkg/proxy/infoprovider.go

@@ -8,10 +8,12 @@ import (
 	"context"
 	"net/url"
 	"sort"
+	"strconv"
 
 	"golang.org/x/xerrors"
 	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/util/uuid"
 	"k8s.io/client-go/tools/cache"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -20,6 +22,7 @@ import (
 
 	wsk8s "github.com/gitpod-io/gitpod/common-go/kubernetes"
 	"github.com/gitpod-io/gitpod/common-go/log"
+	"github.com/gitpod-io/gitpod/ws-manager/api"
 	wsapi "github.com/gitpod-io/gitpod/ws-manager/api"
 	workspacev1 "github.com/gitpod-io/gitpod/ws-manager/api/crd/v1"
 	"github.com/gitpod-io/gitpod/ws-proxy/pkg/common"
@@ -48,11 +51,19 @@ func getPortStr(urlStr string) string {
 	return portURL.Port()
 }
 
+type ConnectionContext struct {
+	WorkspaceID string
+	Port        string
+	UUID        string
+	CancelFunc  context.CancelCauseFunc
+}
+
 type CRDWorkspaceInfoProvider struct {
 	client.Client
 	Scheme *runtime.Scheme
 
-	store cache.ThreadSafeStore
+	store        cache.ThreadSafeStore
+	contextStore cache.ThreadSafeStore
 }
 
 // NewCRDWorkspaceInfoProvider creates a fresh WorkspaceInfoProvider.
@@ -67,12 +78,21 @@ func NewCRDWorkspaceInfoProvider(client client.Client, scheme *runtime.Scheme) (
 			return nil, xerrors.Errorf("object is not a WorkspaceInfo")
 		},
 	}
+	contextIndexers := cache.Indexers{
+		workspaceIndex: func(obj interface{}) ([]string, error) {
+			if connCtx, ok := obj.(*ConnectionContext); ok {
+				return []string{connCtx.WorkspaceID}, nil
+			}
+			return nil, xerrors.Errorf("object is not a ConnectionContext")
+		},
+	}
 
 	return &CRDWorkspaceInfoProvider{
 		Client: client,
 		Scheme: scheme,
 
-		store: cache.NewThreadSafeStore(indexers, cache.Indices{}),
+		store:        cache.NewThreadSafeStore(indexers, cache.Indices{}),
+		contextStore: cache.NewThreadSafeStore(contextIndexers, cache.Indices{}),
 	}, nil
 }
 
@@ -101,6 +121,28 @@ func (r *CRDWorkspaceInfoProvider) WorkspaceInfo(workspaceID string) *common.Wor
 	return nil
 }
 
+func (r *CRDWorkspaceInfoProvider) AcquireContext(ctx context.Context, workspaceID string, port string) (context.Context, string, error) {
+	ws := r.WorkspaceInfo(workspaceID)
+	if ws == nil {
+		return ctx, "", xerrors.Errorf("workspace %s not found", workspaceID)
+	}
+	id := string(uuid.NewUUID())
+	ctx, cancel := context.WithCancelCause(ctx)
+	connCtx := &ConnectionContext{
+		WorkspaceID: workspaceID,
+		Port:        port,
+		CancelFunc:  cancel,
+		UUID:        id,
+	}
+
+	r.contextStore.Add(id, connCtx)
+	return ctx, id, nil
+}
+
+func (r *CRDWorkspaceInfoProvider) ReleaseContext(id string) {
+	r.contextStore.Delete(id)
+}
+
 func (r *CRDWorkspaceInfoProvider) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
 	var ws workspacev1.Workspace
 	err := r.Client.Get(context.Background(), req.NamespacedName, &ws)
@@ -162,11 +204,44 @@ func (r *CRDWorkspaceInfoProvider) Reconcile(ctx context.Context, req ctrl.Reque
 	}
 
 	r.store.Update(req.Name, wsinfo)
+	r.invalidateConnectionContext(wsinfo)
 	log.WithField("workspace", req.Name).WithField("details", wsinfo).Debug("adding/updating workspace details")
 
 	return ctrl.Result{}, nil
 }
 
+func (r *CRDWorkspaceInfoProvider) invalidateConnectionContext(ws *common.WorkspaceInfo) {
+	connCtxs, err := r.contextStore.ByIndex(workspaceIndex, ws.WorkspaceID)
+	if err != nil {
+		return
+	}
+	if len(connCtxs) == 0 {
+		return
+	}
+
+	if ws.Auth != nil && ws.Auth.Admission == wsapi.AdmissionLevel_ADMIT_EVERYONE {
+		return
+	}
+	publicPorts := make(map[string]struct{})
+	for _, p := range ws.Ports {
+		if p.Visibility == api.PortVisibility_PORT_VISIBILITY_PUBLIC {
+			publicPorts[strconv.FormatUint(uint64(p.Port), 10)] = struct{}{}
+		}
+	}
+
+	for _, _connCtx := range connCtxs {
+		connCtx, ok := _connCtx.(*ConnectionContext)
+		if !ok {
+			continue
+		}
+		if _, ok := publicPorts[connCtx.Port]; ok {
+			continue
+		}
+		connCtx.CancelFunc(xerrors.Errorf("workspace %s is no longer public", ws.WorkspaceID))
+		r.contextStore.Delete(connCtx.UUID)
+	}
+}
+
 // SetupWithManager sets up the controller with the Manager.
 func (r *CRDWorkspaceInfoProvider) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).
@@ -177,28 +252,3 @@ func (r *CRDWorkspaceInfoProvider) SetupWithManager(mgr ctrl.Manager) error {
 		).
 		Complete(r)
 }
-
-// CompositeInfoProvider checks each of its info providers and returns the first info found.
-type CompositeInfoProvider []common.WorkspaceInfoProvider
-
-func (c CompositeInfoProvider) WorkspaceInfo(workspaceID string) *common.WorkspaceInfo {
-	for _, ip := range c {
-		res := ip.WorkspaceInfo(workspaceID)
-		if res != nil {
-			return res
-		}
-	}
-	return nil
-}
-
-type fixedInfoProvider struct {
-	Infos map[string]*common.WorkspaceInfo
-}
-
-// WorkspaceInfo returns the workspace information of a workspace using it's workspace ID.
-func (fp *fixedInfoProvider) WorkspaceInfo(workspaceID string) *common.WorkspaceInfo {
-	if fp.Infos == nil {
-		return nil
-	}
-	return fp.Infos[workspaceID]
-}