diff --git a/internal/daemon/worker/countingconn.go b/internal/daemon/worker/countingconn.go index a8c7b9bd0..d4941798b 100644 --- a/internal/daemon/worker/countingconn.go +++ b/internal/daemon/worker/countingconn.go @@ -5,6 +5,10 @@ import ( "sync" ) +// countingConn is a `net.Conn` implementation that records the bytes that go +// across Read() and Write(). All other `net.Conn` function calls are a +// pass-through to the underlying `net.Conn`, meaning it's also safe to call +// those functions directly on the underlying object, if you have access to it. type countingConn struct { net.Conn @@ -29,7 +33,8 @@ func (c *countingConn) BytesWritten() uint64 { return c.bytesWritten } -// Read wraps the embedded conn's Read and counts the number of bytes read. +// Read wraps the embedded conn's Read() and counts the number of bytes read +// (the number of bytes the client sent to us). func (c *countingConn) Read(in []byte) (int, error) { n, err := c.Conn.Read(in) c.mu.Lock() @@ -38,7 +43,8 @@ func (c *countingConn) Read(in []byte) (int, error) { return n, err } -// Write wraps the embedded conn's Write and counts the number of bytes read. +// Write wraps the embedded conn's Write() and counts the number of bytes +// written (the number of bytes we sent to the client). func (c *countingConn) Write(in []byte) (int, error) { n, err := c.Conn.Write(in) c.mu.Lock() diff --git a/internal/daemon/worker/handler.go b/internal/daemon/worker/handler.go index 7614a447d..646ff5da4 100644 --- a/internal/daemon/worker/handler.go +++ b/internal/daemon/worker/handler.go @@ -219,9 +219,28 @@ func (w *Worker) handleProxy(listenerCfg *listenerutil.ListenerConfig, sessionMa } event.WriteSysEvent(ctx, op, "connection successfully authorized", "session_id", sessionId, "connection_id", ci.Id) + // Wrapping the client websocket with a `net.Conn` implementation that + // records the bytes that go across Read() and Write(). cc := &countingConn{Conn: websocket.NetConn(connCtx, conn, websocket.MessageBinary)} + err = sess.ApplyConnectionCounterCallbacks(ci.Id, cc.BytesRead, cc.BytesWritten) + if err != nil { + event.WriteError(ctx, op, err, event.WithInfoMsg("unable to set counter callbacks for session connection")) + err = conn.Close(websocket.StatusInternalError, "unable to set counter callbacks for session connection") + if err != nil { + event.WriteError(ctx, op, err, event.WithInfoMsg("error closing client connection")) + } + return + } + defer func() { - if sessionManager.RequestCloseConnections(ctx, map[string]string{ci.Id: sess.GetId()}) { + ccd := map[string]*session.ConnectionCloseData{ + ci.Id: { + SessionId: sess.GetId(), + BytesUp: cc.BytesRead(), + BytesDown: cc.BytesWritten(), + }, + } + if sessionManager.RequestCloseConnections(ctx, ccd) { event.WriteSysEvent(ctx, op, "connection closed", "session_id", sessionId, "connection_id", ci.Id) } }() diff --git a/internal/daemon/worker/session/manager.go b/internal/daemon/worker/session/manager.go index 1322b23ca..ceb379ec2 100644 --- a/internal/daemon/worker/session/manager.go +++ b/internal/daemon/worker/session/manager.go @@ -35,15 +35,15 @@ type Manager interface { // no error is returned. DeleteLocalSession([]string) - // RequestCloseConnections sends connection close requests to the controller, - // and sets close times within the worker. It should be called during the worker - // status loop and on connection exit on the proxy. + // RequestCloseConnections sends connection close requests to the + // controller, and sets close times within the worker. It should be called + // during the worker status loop and on connection exit on the proxy. // // The boolean indicates whether the function was successful, e.g. had any // errors. Individual events will be sent for the errors if there are any. // - // closeInfo is a map of connection ids mapped to their individual session ids. - RequestCloseConnections(context.Context, map[string]string) bool + // closeInfo is a map of connection ids mapped to connection metadata. + RequestCloseConnections(context.Context, map[string]*ConnectionCloseData) bool } type manager struct { @@ -123,7 +123,7 @@ func (m *manager) DeleteLocalSession(sessIds []string) { } } -func (m *manager) RequestCloseConnections(ctx context.Context, closeInfo map[string]string) bool { +func (m *manager) RequestCloseConnections(ctx context.Context, closeInfo map[string]*ConnectionCloseData) bool { return closeConnections(ctx, m.controllerSessionConn, m, closeInfo) } diff --git a/internal/daemon/worker/session/manager_test.go b/internal/daemon/worker/session/manager_test.go index 24ad73fbc..2f846bcca 100644 --- a/internal/daemon/worker/session/manager_test.go +++ b/internal/daemon/worker/session/manager_test.go @@ -75,7 +75,7 @@ func TestManager_RequestCloseConnections(t *testing.T) { manager, err := NewManager(mockSessionClient) require.NoError(t, err) assert.False(t, manager.RequestCloseConnections(ctx, nil)) - assert.False(t, manager.RequestCloseConnections(ctx, map[string]string{})) + assert.False(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{})) mockSessionClient.LookupSessionFn = func(_ context.Context, req *pbs.LookupSessionRequest) (*pbs.LookupSessionResponse, error) { return &pbs.LookupSessionResponse{ @@ -98,7 +98,7 @@ func TestManager_RequestCloseConnections(t *testing.T) { mockSessionClient.CloseConnectionFn = func(context.Context, *pbs.CloseConnectionRequest) (*pbs.CloseConnectionResponse, error) { return nil, errors.New("error") } - assert.False(t, manager.RequestCloseConnections(ctx, map[string]string{"connection id": session1.GetId()})) + assert.False(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{"connection id": {SessionId: session1.GetId()}})) mockSessionClient.CloseConnectionFn = func(_ context.Context, req *pbs.CloseConnectionRequest) (*pbs.CloseConnectionResponse, error) { var data []*pbs.CloseConnectionResponseData for _, r := range req.GetCloseRequestData() { @@ -110,8 +110,8 @@ func TestManager_RequestCloseConnections(t *testing.T) { return &pbs.CloseConnectionResponse{CloseResponseData: data}, nil } // There is no connection information yet for this connection. - assert.False(t, manager.RequestCloseConnections(ctx, map[string]string{ - "random_connection_id": session1.GetId(), + assert.False(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{ + "random_connection_id": {SessionId: session1.GetId()}, })) // Load the connection info into the local storage @@ -126,17 +126,17 @@ func TestManager_RequestCloseConnections(t *testing.T) { c1, _, err := session1.RequestAuthorizeConnection(ctx, "worker id", cancelFn) require.NoError(t, err) - assert.True(t, manager.RequestCloseConnections(ctx, map[string]string{ - c1.Id: session1.GetId(), + assert.True(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{ + c1.Id: {SessionId: session1.GetId()}, })) c2, _, err := session2.RequestAuthorizeConnection(ctx, "worker id", cancelFn) require.NoError(t, err) c3, _, err := session3.RequestAuthorizeConnection(ctx, "worker id", cancelFn) require.NoError(t, err) - assert.True(t, manager.RequestCloseConnections(ctx, map[string]string{ - c2.Id: session2.GetId(), - c3.Id: session3.GetId(), + assert.True(t, manager.RequestCloseConnections(ctx, map[string]*ConnectionCloseData{ + c2.Id: {SessionId: session2.GetId()}, + c3.Id: {SessionId: session3.GetId()}, })) } diff --git a/internal/daemon/worker/session/session.go b/internal/daemon/worker/session/session.go index 7562f7b4a..31369037b 100644 --- a/internal/daemon/worker/session/session.go +++ b/internal/daemon/worker/session/session.go @@ -30,11 +30,24 @@ type ConnInfo struct { // closes the proxy connection. connCtxCancelFunc context.CancelFunc + // The number of bytes uploaded from the client. + BytesUp func() uint64 + // The number of bytes downloaded to the client. + BytesDown func() uint64 + // The time the controller has successfully reported that this connection is // closed. CloseTime time.Time } +// ConnectionCloseData encapsulates the data we need to send via CloseConnection +// RPC to the controller. +type ConnectionCloseData struct { + SessionId string + BytesUp uint64 + BytesDown uint64 +} + // Session is the local representation of a session. After initial loading // the only values that will change will be the status (readable from // GetStatus()) and the Connections (GetLocalConnections()). @@ -86,6 +99,12 @@ type Session interface { // call to RequestActivate, subsequent calls will fail. RequestActivate(ctx context.Context, tofu string) error + // ApplyConnectionCounterCallbacks sets a connection's bytes up and bytes + // down callbacks to the provided functions. Both functions must be safe for + // concurrent use. If there is no connection with the provided id, an error + // is returned. + ApplyConnectionCounterCallbacks(connId string, bytesUp func() uint64, bytesDown func() uint64) error + // RequestAuthorizeConnection sends an AuthorizeConnection request to // the controller. // It is called by the worker handler after a connection has been received by @@ -179,6 +198,8 @@ func (s *sess) GetLocalConnections() map[string]ConnInfo { Id: v.Id, Status: v.Status, CloseTime: v.CloseTime, + BytesUp: v.BytesUp, + BytesDown: v.BytesDown, } } return res @@ -283,11 +304,19 @@ func (s *sess) RequestAuthorizeConnection(ctx context.Context, workerId string, case workerId == "": return ConnInfo{}, 0, errors.New("worker id is empty") } + ci, connsLeft, err := authorizeConnection(ctx, s.client, workerId, s.GetId()) if err != nil { return ConnInfo{}, connsLeft, err } ci.connCtxCancelFunc = connCancel + + // Install safe callbacks before connection has been established. These + // should be replaced when `ApplyConnectionCounterCallbacks` gets called on + // the `sess` object. + ci.BytesUp = func() uint64 { return 0 } + ci.BytesDown = func() uint64 { return 0 } + s.lock.Lock() defer s.lock.Unlock() s.connInfoMap[ci.Id] = ci @@ -295,6 +324,8 @@ func (s *sess) RequestAuthorizeConnection(ctx context.Context, workerId string, Id: ci.Id, Status: ci.Status, CloseTime: ci.CloseTime, + BytesUp: ci.BytesUp, + BytesDown: ci.BytesDown, }, connsLeft, err } @@ -355,6 +386,19 @@ func (s *sess) CancelAllLocalConnections() []string { return closedIds } +func (s *sess) ApplyConnectionCounterCallbacks(connId string, bytesUp func() uint64, bytesDown func() uint64) error { + s.lock.Lock() + defer s.lock.Unlock() + + ci, ok := s.connInfoMap[connId] + if !ok { + return fmt.Errorf("failed to find connection info for connection id %q", connId) + } + ci.BytesUp = bytesUp + ci.BytesDown = bytesDown + return nil +} + // activate is a helper worker function that sends session activation request to the // controller. func activate(ctx context.Context, sessClient pbs.SessionServiceClient, sessionId, tofuToken string, version uint32) (pbs.SESSIONSTATUS, error) { @@ -437,8 +481,8 @@ func closeConnection(ctx context.Context, sessClient pbs.SessionServiceClient, r // The boolean indicates whether the function was successful, e.g. had any // errors. Individual events will be sent for the errors if there are any. // -// closeInfo is a map of connections mapped to their individual session. -func closeConnections(ctx context.Context, sessClient pbs.SessionServiceClient, sManager Manager, closeInfo map[string]string) bool { +// closeInfo is a map of connection id to connection metadata. +func closeConnections(ctx context.Context, sessClient pbs.SessionServiceClient, sManager Manager, closeInfo map[string]*ConnectionCloseData) bool { const op = "session.closeConnections" if len(closeInfo) == 0 { // This should not happen, but it's a no-op if it does. Just @@ -496,12 +540,14 @@ func closeConnections(ctx context.Context, sessClient pbs.SessionServiceClient, // sessions IDs that those connections belong to. The values are // ignored; the parameter is expected as such just for convenience of // its caller. -func makeCloseConnectionRequest(closeInfo map[string]string) *pbs.CloseConnectionRequest { +func makeCloseConnectionRequest(closeInfo map[string]*ConnectionCloseData) *pbs.CloseConnectionRequest { closeData := make([]*pbs.CloseConnectionRequestData, 0, len(closeInfo)) - for connId := range closeInfo { + for connId, data := range closeInfo { closeData = append(closeData, &pbs.CloseConnectionRequestData{ ConnectionId: connId, Reason: session.UnknownReason.String(), + BytesUp: data.BytesUp, + BytesDown: data.BytesDown, }) } @@ -516,7 +562,7 @@ func makeCloseConnectionRequest(closeInfo map[string]string) *pbs.CloseConnectio // easily lock on session once for all connections in // setCloseTimeForResponse. func makeSessionCloseInfo( - closeInfo map[string]string, + closeInfo map[string]*ConnectionCloseData, response *pbs.CloseConnectionResponse, ) (map[string][]*pbs.CloseConnectionResponseData, error) { if closeInfo == nil { @@ -525,7 +571,8 @@ func makeSessionCloseInfo( result := make(map[string][]*pbs.CloseConnectionResponseData) for _, v := range response.GetCloseResponseData() { - result[closeInfo[v.GetConnectionId()]] = append(result[closeInfo[v.GetConnectionId()]], v) + sessionId := closeInfo[v.GetConnectionId()].SessionId + result[sessionId] = append(result[sessionId], v) } return result, nil @@ -534,14 +581,15 @@ func makeSessionCloseInfo( // makeFakeSessionCloseInfo makes a "fake" makeFakeSessionCloseInfo, intended // for use when we can't contact the controller. func makeFakeSessionCloseInfo( - closeInfo map[string]string, + closeInfo map[string]*ConnectionCloseData, ) (map[string][]*pbs.CloseConnectionResponseData, error) { if closeInfo == nil { return nil, errMakeSessionCloseInfoNilCloseInfo } result := make(map[string][]*pbs.CloseConnectionResponseData) - for connectionId, sessionId := range closeInfo { + for connectionId, closeData := range closeInfo { + sessionId := closeData.SessionId result[sessionId] = append(result[sessionId], &pbs.CloseConnectionResponseData{ ConnectionId: connectionId, Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED, diff --git a/internal/daemon/worker/session/session_test.go b/internal/daemon/worker/session/session_test.go index 406f342cf..fb61bdf65 100644 --- a/internal/daemon/worker/session/session_test.go +++ b/internal/daemon/worker/session/session_test.go @@ -170,17 +170,26 @@ func TestSession_RequestAuthorizeConnection(t *testing.T) { } conn, left, err := sess.RequestAuthorizeConnection(context.Background(), "workerid", cancel) require.NoError(t, err) - assert.Equal(t, ConnInfo{Id: "conn1", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED}, conn) + + assert.Equal(t, "conn1", conn.Id) + assert.Equal(t, pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED, conn.Status) + assert.NotNil(t, conn.BytesUp) + assert.NotNil(t, conn.BytesDown) + assert.Zero(t, conn.BytesUp()) + assert.Zero(t, conn.BytesDown()) assert.Equal(t, int32(-1), left) } func TestWorkerMakeCloseConnectionRequest(t *testing.T) { require := require.New(t) - in := map[string]string{"foo": "one", "bar": "two"} + in := map[string]*ConnectionCloseData{ + "foo": {SessionId: "one", BytesUp: 1000, BytesDown: 2000}, + "bar": {SessionId: "two", BytesUp: 1000, BytesDown: 2000}, + } expected := &pbs.CloseConnectionRequest{ CloseRequestData: []*pbs.CloseConnectionRequestData{ - {ConnectionId: "foo", Reason: session.UnknownReason.String()}, - {ConnectionId: "bar", Reason: session.UnknownReason.String()}, + {ConnectionId: "foo", Reason: session.UnknownReason.String(), BytesUp: 1000, BytesDown: 2000}, + {ConnectionId: "bar", Reason: session.UnknownReason.String(), BytesUp: 1000, BytesDown: 2000}, }, } actual := makeCloseConnectionRequest(in) @@ -189,7 +198,7 @@ func TestWorkerMakeCloseConnectionRequest(t *testing.T) { func TestMakeSessionCloseInfo(t *testing.T) { require := require.New(t) - closeInfo := map[string]string{"foo": "one", "bar": "two"} + closeInfo := map[string]*ConnectionCloseData{"foo": {SessionId: "one"}, "bar": {SessionId: "two"}} response := &pbs.CloseConnectionResponse{ CloseResponseData: []*pbs.CloseConnectionResponseData{ {ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED}, @@ -218,7 +227,7 @@ func TestMakeSessionCloseInfoErrorIfCloseInfoNil(t *testing.T) { func TestMakeSessionCloseInfoEmpty(t *testing.T) { require := require.New(t) - actual, err := makeSessionCloseInfo(make(map[string]string), nil) + actual, err := makeSessionCloseInfo(make(map[string]*ConnectionCloseData), nil) require.NoError(err) require.Equal( make(map[string][]*pbs.CloseConnectionResponseData), @@ -228,7 +237,7 @@ func TestMakeSessionCloseInfoEmpty(t *testing.T) { func TestMakeFakeSessionCloseInfo(t *testing.T) { require := require.New(t) - closeInfo := map[string]string{"foo": "one", "bar": "two"} + closeInfo := map[string]*ConnectionCloseData{"foo": {SessionId: "one"}, "bar": {SessionId: "two"}} expected := map[string][]*pbs.CloseConnectionResponseData{ "one": { {ConnectionId: "foo", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED}, @@ -251,10 +260,30 @@ func TestMakeFakeSessionCloseInfoErrorIfCloseInfoNil(t *testing.T) { func TestMakeFakeSessionCloseInfoEmpty(t *testing.T) { require := require.New(t) - actual, err := makeFakeSessionCloseInfo(make(map[string]string)) + actual, err := makeFakeSessionCloseInfo(make(map[string]*ConnectionCloseData)) require.NoError(err) require.Equal( make(map[string][]*pbs.CloseConnectionResponseData), actual, ) } + +func TestApplyConnectionCounterCallbacks(t *testing.T) { + s := &sess{connInfoMap: make(map[string]*ConnInfo)} + + connId := "conn1" + bytesUpFn := func() uint64 { return 10 } + bytesDnFn := func() uint64 { return 20 } + err := s.ApplyConnectionCounterCallbacks(connId, bytesUpFn, bytesDnFn) + require.EqualError(t, err, "failed to find connection info for connection id \"conn1\"") + + s.connInfoMap[connId] = &ConnInfo{} + require.NoError(t, s.ApplyConnectionCounterCallbacks("conn1", bytesUpFn, bytesDnFn)) + + ci, ok := s.connInfoMap[connId] + require.True(t, ok) + require.NotNil(t, ci.BytesUp) + require.NotNil(t, ci.BytesDown) + require.EqualValues(t, ci.BytesUp(), 10) + require.EqualValues(t, ci.BytesDown(), 20) +} diff --git a/internal/daemon/worker/status.go b/internal/daemon/worker/status.go index d12d6c5cd..f2c27da37 100644 --- a/internal/daemon/worker/status.go +++ b/internal/daemon/worker/status.go @@ -330,7 +330,7 @@ func (w *Worker) sendWorkerStatus(cancelCtx context.Context, sessionManager sess // connections, regardless of whether or not the session is still active. func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionState bool, sessionManager session.Manager) { const op = "worker.(Worker).cleanupConnections" - closeInfo := make(map[string]string) + closeInfo := make(map[string]*session.ConnectionCloseData) cleanSessionIds := make([]string, 0) sessionManager.ForEachLocalSession(func(s session.Session) bool { switch { @@ -341,8 +341,19 @@ func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionStat // Cancel connections without regard to individual connection // state. closedIds := s.CancelAllLocalConnections() + localConns := s.GetLocalConnections() for _, connId := range closedIds { - closeInfo[connId] = s.GetId() + var bytesUp, bytesDown uint64 + connInfo, ok := localConns[connId] + if ok { + bytesUp = connInfo.BytesUp() + bytesDown = connInfo.BytesDown() + } + closeInfo[connId] = &session.ConnectionCloseData{ + SessionId: s.GetId(), + BytesUp: bytesUp, + BytesDown: bytesDown, + } event.WriteSysEvent(cancelCtx, op, "terminated connection due to cancellation or expiration", "session_id", s.GetId(), "connection_id", connId) } @@ -357,7 +368,7 @@ func (w *Worker) cleanupConnections(cancelCtx context.Context, ignoreSessionStat // state (ie: only ones that the controller has requested be // terminated). for _, connId := range s.CancelOpenLocalConnections() { - closeInfo[connId] = s.GetId() + closeInfo[connId] = &session.ConnectionCloseData{SessionId: s.GetId()} event.WriteSysEvent(cancelCtx, op, "terminated connection due to cancellation or expiration", "session_id", s.GetId(), "connection_id", connId) } }