Skip to content

Commit

Permalink
Merge pull request from GHSA-c8xw-vjgf-94hr
Browse files Browse the repository at this point in the history
Signed-off-by: pashakostohrys <pavel@codefresh.io>
  • Loading branch information
pasha-codefresh committed Aug 21, 2023
1 parent f978e04 commit e047efa
Show file tree
Hide file tree
Showing 5 changed files with 132 additions and 15 deletions.
12 changes: 10 additions & 2 deletions server/application/terminal.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"net/http"
"time"

util_session "github.com/argoproj/argo-cd/v2/util/session"
"github.com/argoproj/gitops-engine/pkg/utils/kube"
log "github.com/sirupsen/logrus"
v1 "k8s.io/api/core/v1"
Expand Down Expand Up @@ -37,11 +38,12 @@ type terminalHandler struct {
allowedShells []string
namespace string
enabledNamespaces []string
sessionManager util_session.SessionManager
}

// NewHandler returns a new terminal handler.
func NewHandler(appLister applisters.ApplicationLister, namespace string, enabledNamespaces []string, db db.ArgoDB, enf *rbac.Enforcer, cache *servercache.Cache,
appResourceTree AppResourceTreeFn, allowedShells []string) *terminalHandler {
appResourceTree AppResourceTreeFn, allowedShells []string, sessionManager util_session.SessionManager) *terminalHandler {
return &terminalHandler{
appLister: appLister,
db: db,
Expand All @@ -51,6 +53,7 @@ func NewHandler(appLister applisters.ApplicationLister, namespace string, enable
allowedShells: allowedShells,
namespace: namespace,
enabledNamespaces: enabledNamespaces,
sessionManager: sessionManager,
}
}

Expand Down Expand Up @@ -222,7 +225,7 @@ func (s *terminalHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {

fieldLog.Info("terminal session starting")

session, err := newTerminalSession(w, r, nil)
session, err := newTerminalSession(w, r, nil, s.sessionManager)
if err != nil {
http.Error(w, "Failed to start terminal session", http.StatusBadRequest)
return
Expand Down Expand Up @@ -282,6 +285,11 @@ type TerminalMessage struct {
Cols uint16 `json:"cols"`
}

// TerminalCommand is the struct for websocket commands,For example you need ask client to reconnect
type TerminalCommand struct {
Code int
}

// startProcess executes specified commands in the container and connects it up with the ptyHandler (a session)
func startProcess(k8sClient kubernetes.Interface, cfg *rest.Config, namespace, podName, containerName string, cmd []string, ptyHandler PtyHandler) error {
req := k8sClient.CoreV1().RESTClient().Post().
Expand Down
77 changes: 66 additions & 11 deletions server/application/websocket.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@ package application
import (
"encoding/json"
"fmt"
"github.com/argoproj/argo-cd/v2/common"
httputil "github.com/argoproj/argo-cd/v2/util/http"
util_session "github.com/argoproj/argo-cd/v2/util/session"
"net/http"
"sync"
"time"
Expand All @@ -12,6 +15,11 @@ import (
"k8s.io/client-go/tools/remotecommand"
)

const (
ReconnectCode = 1
ReconnectMessage = "\nReconnect because the token was refreshed...\n"
)

var upgrader = func() websocket.Upgrader {
upgrader := websocket.Upgrader{}
upgrader.HandshakeTimeout = time.Second * 2
Expand All @@ -23,25 +31,40 @@ var upgrader = func() websocket.Upgrader {

// terminalSession implements PtyHandler
type terminalSession struct {
wsConn *websocket.Conn
sizeChan chan remotecommand.TerminalSize
doneChan chan struct{}
tty bool
readLock sync.Mutex
writeLock sync.Mutex
wsConn *websocket.Conn
sizeChan chan remotecommand.TerminalSize
doneChan chan struct{}
tty bool
readLock sync.Mutex
writeLock sync.Mutex
sessionManager util_session.SessionManager
token *string
}

// getToken get auth token from web socket request
func getToken(r *http.Request) (string, error) {
cookies := r.Cookies()
return httputil.JoinCookies(common.AuthCookieName, cookies)
}

// newTerminalSession create terminalSession
func newTerminalSession(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*terminalSession, error) {
func newTerminalSession(w http.ResponseWriter, r *http.Request, responseHeader http.Header, sessionManager util_session.SessionManager) (*terminalSession, error) {
token, err := getToken(r)
if err != nil {
return nil, err
}

conn, err := upgrader.Upgrade(w, r, responseHeader)
if err != nil {
return nil, err
}
session := &terminalSession{
wsConn: conn,
tty: true,
sizeChan: make(chan remotecommand.TerminalSize),
doneChan: make(chan struct{}),
wsConn: conn,
tty: true,
sizeChan: make(chan remotecommand.TerminalSize),
doneChan: make(chan struct{}),
sessionManager: sessionManager,
token: &token,
}
return session, nil
}
Expand Down Expand Up @@ -78,8 +101,40 @@ func (t *terminalSession) Next() *remotecommand.TerminalSize {
}
}

// reconnect send reconnect code to client and ask them init new ws session
func (t *terminalSession) reconnect() (int, error) {
reconnectCommand, _ := json.Marshal(TerminalCommand{
Code: ReconnectCode,
})
reconnectMessage, _ := json.Marshal(TerminalMessage{
Operation: "stdout",
Data: ReconnectMessage,
})
t.writeLock.Lock()
err := t.wsConn.WriteMessage(websocket.TextMessage, reconnectMessage)
if err != nil {
log.Errorf("write message err: %v", err)
return 0, err
}
err = t.wsConn.WriteMessage(websocket.TextMessage, reconnectCommand)
if err != nil {
log.Errorf("write message err: %v", err)
return 0, err
}
t.writeLock.Unlock()
return 0, nil
}

// Read called in a loop from remotecommand as long as the process is running
func (t *terminalSession) Read(p []byte) (int, error) {
// check if token still valid
_, newToken, err := t.sessionManager.VerifyToken(*t.token)
// err in case if token is revoked, newToken in case if refresh happened
if err != nil || newToken != "" {
// need to send reconnect code in case if token was refreshed
return t.reconnect()
}

t.readLock.Lock()
_, message, err := t.wsConn.ReadMessage()
t.readLock.Unlock()
Expand Down
46 changes: 46 additions & 0 deletions server/application/websocket_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
package application

import (
"encoding/json"
"github.com/gorilla/websocket"
"github.com/stretchr/testify/assert"
"net/http"
"net/http/httptest"
"strings"
"testing"
)

func reconnect(w http.ResponseWriter, r *http.Request) {
var upgrader = websocket.Upgrader{}
c, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}

ts := terminalSession{wsConn: c}
_, _ = ts.reconnect()
}

func TestReconnect(t *testing.T) {

s := httptest.NewServer(http.HandlerFunc(reconnect))
defer s.Close()

u := "ws" + strings.TrimPrefix(s.URL, "http")

// Connect to the server
ws, _, err := websocket.DefaultDialer.Dial(u, nil)
assert.NoError(t, err)

defer ws.Close()

_, p, _ := ws.ReadMessage()

var message TerminalMessage

err = json.Unmarshal(p, &message)

assert.NoError(t, err)
assert.Equal(t, message.Data, ReconnectMessage)

}
3 changes: 2 additions & 1 deletion server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -975,7 +975,7 @@ func (a *ArgoCDServer) newHTTPServer(ctx context.Context, port int, grpcWebHandl
}
mux.Handle("/api/", handler)

terminal := application.NewHandler(a.appLister, a.Namespace, a.ApplicationNamespaces, a.db, a.enf, a.Cache, appResourceTreeFn, a.settings.ExecShells).
terminal := application.NewHandler(a.appLister, a.Namespace, a.ApplicationNamespaces, a.db, a.enf, a.Cache, appResourceTreeFn, a.settings.ExecShells, *a.sessionMgr).
WithFeatureFlagMiddleware(a.settingsMgr.GetSettings)
th := util_session.WithAuthMiddleware(a.DisableAuth, a.sessionMgr, terminal)
mux.Handle("/terminal", th)
Expand All @@ -988,6 +988,7 @@ func (a *ArgoCDServer) newHTTPServer(ctx context.Context, port int, grpcWebHandl
// will be added in mux.
registerExtensions(mux, a)
}

mustRegisterGWHandler(versionpkg.RegisterVersionServiceHandler, ctx, gwmux, conn)
mustRegisterGWHandler(clusterpkg.RegisterClusterServiceHandler, ctx, gwmux, conn)
mustRegisterGWHandler(applicationpkg.RegisterApplicationServiceHandler, ctx, gwmux, conn)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,14 @@ export const PodTerminalViewer: React.FC<PodTerminalViewerProps> = ({

const onConnectionMessage = (e: MessageEvent) => {
const msg = JSON.parse(e.data);
connSubject.next(msg);
if (!msg?.Code) {
connSubject.next(msg);
} else {
// Do reconnect due to refresh token event
onConnectionClose();
setupConnection()
}

};

const onConnectionOpen = () => {
Expand Down

0 comments on commit e047efa

Please sign in to comment.