From f7e48ec83633a8d177e36c885d763c44a8c554ce Mon Sep 17 00:00:00 2001 From: Jeff Mitchell Date: Sun, 20 Sep 2020 16:28:00 -0400 Subject: [PATCH] Plumb connection limit to proxy and output it (#384) --- internal/cmd/commands/proxy/funcs.go | 29 ++++++- internal/cmd/commands/proxy/proxy.go | 41 ++++++++-- .../servers/services/session_service.pb.go | 79 +++++++++++-------- .../servers/services/v1/session_service.proto | 1 + .../proto/local/worker/proxy/v1/proxy.proto | 1 + internal/proxy/proxy.pb.go | 23 ++++-- .../handlers/workers/worker_service.go | 16 ++-- internal/servers/worker/handler.go | 6 +- internal/servers/worker/session.go | 12 +-- 9 files changed, 149 insertions(+), 59 deletions(-) diff --git a/internal/cmd/commands/proxy/funcs.go b/internal/cmd/commands/proxy/funcs.go index e447535b9e..59b41ef99c 100644 --- a/internal/cmd/commands/proxy/funcs.go +++ b/internal/cmd/commands/proxy/funcs.go @@ -1,10 +1,12 @@ package proxy import ( + "time" + "github.com/hashicorp/boundary/internal/cmd/base" ) -func generateConnectionInfoTableOutput(in ConnectionInfo) string { +func generateSessionInfoTableOutput(in SessionInfo) string { var ret []string nonAttributeMap := map[string]interface{}{ @@ -29,3 +31,28 @@ func generateConnectionInfoTableOutput(in ConnectionInfo) string { return base.WrapForHelpText(ret) } + +func generateConnectionInfoTableOutput(in ConnectionInfo) string { + var ret []string + + nonAttributeMap := map[string]interface{}{ + "Connections Left": in.ConnectionsLeft, + "Expiration": in.Expiration.Local().Format(time.RFC1123), + } + + maxLength := 0 + for k := range nonAttributeMap { + if len(k) > maxLength { + maxLength = len(k) + } + } + + ret = append(ret, "", "Connection information:") + + ret = append(ret, + // We do +2 because there is another +2 offset for host sets below + base.WrapMap(2, maxLength+2, nonAttributeMap), + ) + + return base.WrapForHelpText(ret) +} diff --git a/internal/cmd/commands/proxy/proxy.go b/internal/cmd/commands/proxy/proxy.go index 7ec7a1504e..86b8807f31 100644 --- a/internal/cmd/commands/proxy/proxy.go +++ b/internal/cmd/commands/proxy/proxy.go @@ -15,6 +15,7 @@ import ( "os" "strings" "sync" + "time" "github.com/btcsuite/btcutil/base58" "github.com/hashicorp/boundary/api/targets" @@ -27,20 +28,29 @@ import ( "github.com/kr/pretty" "github.com/mitchellh/cli" "github.com/posener/complete" + "go.uber.org/atomic" "google.golang.org/protobuf/proto" "nhooyr.io/websocket" "nhooyr.io/websocket/wspb" ) -type ConnectionInfo struct { +type SessionInfo struct { Address string `json:"address"` Port int `json:"port"` Protocol string `json:"protocol"` } +type ConnectionInfo struct { + Expiration time.Time `json:"expiration"` + ConnectionsLeft int32 `json:"connections_left"` +} + var _ cli.Command = (*Command)(nil) var _ cli.CommandAutocomplete = (*Command)(nil) +var connectionsLeft atomic.Int32 +var expiration time.Time + type Command struct { *base.Command @@ -272,6 +282,8 @@ func (c *Command) Run(args []string) (retCode int) { return 1 } + expiration = parsedCert.NotAfter + certPool := x509.NewCertPool() certPool.AddCert(parsedCert) @@ -324,7 +336,7 @@ func (c *Command) Run(args []string) (retCode int) { listenerAddr := listener.Addr().(*net.TCPAddr) - connInfo := ConnectionInfo{ + sessInfo := SessionInfo{ Protocol: "tcp", Address: listenerAddr.IP.String(), Port: listenerAddr.Port, @@ -332,11 +344,11 @@ func (c *Command) Run(args []string) (retCode int) { switch base.Format(c.UI) { case "table": - c.UI.Output(generateConnectionInfoTableOutput(connInfo)) + c.UI.Output(generateSessionInfoTableOutput(sessInfo)) case "json": - out, err := json.Marshal(&connInfo) + out, err := json.Marshal(&sessInfo) if err != nil { - c.UI.Error(fmt.Errorf("error marshaling connection information: %w", err).Error()) + c.UI.Error(fmt.Errorf("error marshaling session information: %w", err).Error()) return 1 } c.UI.Output(string(out)) @@ -363,6 +375,7 @@ AcceptLoop: defer listeningConn.Close() if err := handleConnection( c.Context, + c.UI, connWg, listeningConn, workerAddr, @@ -379,6 +392,7 @@ AcceptLoop: func handleConnection( ctx context.Context, + ui cli.Ui, connWg *sync.WaitGroup, listeningConn *net.TCPConn, workerAddr string, @@ -427,6 +441,23 @@ func handleConnection( if err := wspb.Read(ctx, conn, &handshakeResult); err != nil { return fmt.Errorf("error reading handshake result: %w", err) } + connectionsLeft.Store(handshakeResult.ConnectionsLeft) + + connInfo := ConnectionInfo{ + Expiration: expiration, + ConnectionsLeft: handshakeResult.GetConnectionsLeft(), + } + + switch base.Format(ui) { + case "table": + ui.Output(generateConnectionInfoTableOutput(connInfo)) + case "json": + out, err := json.Marshal(&connInfo) + if err != nil { + ui.Error(fmt.Errorf("error marshaling connection information: %w", err).Error()) + } + ui.Output(string(out)) + } // We don't _rely_ on client-side timeout verification but this prevents us // seeming to be ready for a connection that will immediately fail when we diff --git a/internal/gen/controller/servers/services/session_service.pb.go b/internal/gen/controller/servers/services/session_service.pb.go index 1cb14751bf..c93ee7fd03 100644 --- a/internal/gen/controller/servers/services/session_service.pb.go +++ b/internal/gen/controller/servers/services/session_service.pb.go @@ -355,8 +355,9 @@ type AuthorizeConnectionResponse struct { sizeCache protoimpl.SizeCache unknownFields protoimpl.UnknownFields - ConnectionId string `protobuf:"bytes,10,opt,name=connection_id,json=connectionId,proto3" json:"connection_id,omitempty"` - Status CONNECTIONSTATUS `protobuf:"varint,20,opt,name=status,proto3,enum=controller.servers.services.v1.CONNECTIONSTATUS" json:"status,omitempty"` + ConnectionId string `protobuf:"bytes,10,opt,name=connection_id,json=connectionId,proto3" json:"connection_id,omitempty"` + Status CONNECTIONSTATUS `protobuf:"varint,20,opt,name=status,proto3,enum=controller.servers.services.v1.CONNECTIONSTATUS" json:"status,omitempty"` + ConnectionsLeft int32 `protobuf:"varint,30,opt,name=connections_left,json=connectionsLeft,proto3" json:"connections_left,omitempty"` } func (x *AuthorizeConnectionResponse) Reset() { @@ -405,6 +406,13 @@ func (x *AuthorizeConnectionResponse) GetStatus() CONNECTIONSTATUS { return CONNECTIONSTATUS_CONNECTIONSTATUS_UNSPECIFIED } +func (x *AuthorizeConnectionResponse) GetConnectionsLeft() int32 { + if x != nil { + return x.ConnectionsLeft + } + return 0 +} + var File_controller_servers_services_v1_session_service_proto protoreflect.FileDescriptor var file_controller_servers_services_v1_session_service_proto_rawDesc = []byte{ @@ -475,7 +483,7 @@ var file_controller_servers_services_v1_session_service_proto_rawDesc = []byte{ 0x69, 0x7a, 0x65, 0x43, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x0a, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, - 0x6e, 0x49, 0x64, 0x22, 0x8c, 0x01, 0x0a, 0x1b, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, + 0x6e, 0x49, 0x64, 0x22, 0xb7, 0x01, 0x0a, 0x1b, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x43, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x23, 0x0a, 0x0d, 0x63, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x0a, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, 0x6f, 0x6e, 0x6e, @@ -484,39 +492,42 @@ var file_controller_servers_services_v1_session_service_proto_rawDesc = []byte{ 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x4f, 0x4e, 0x4e, 0x45, 0x43, 0x54, 0x49, 0x4f, 0x4e, 0x53, 0x54, 0x41, 0x54, 0x55, 0x53, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, - 0x75, 0x73, 0x32, 0xaa, 0x03, 0x0a, 0x0e, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x53, 0x65, - 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x7e, 0x0a, 0x0d, 0x4c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x53, - 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x34, 0x2e, 0x63, 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, - 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x2e, 0x73, 0x65, 0x72, 0x76, - 0x69, 0x63, 0x65, 0x73, 0x2e, 0x76, 0x31, 0x2e, 0x4c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x53, 0x65, - 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x35, 0x2e, 0x63, - 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, - 0x73, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x2e, 0x76, 0x31, 0x2e, 0x4c, 0x6f, - 0x6f, 0x6b, 0x75, 0x70, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x84, 0x01, 0x0a, 0x0f, 0x41, 0x63, 0x74, 0x69, 0x76, 0x61, - 0x74, 0x65, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x36, 0x2e, 0x63, 0x6f, 0x6e, 0x74, - 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x2e, 0x73, - 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x2e, 0x76, 0x31, 0x2e, 0x41, 0x63, 0x74, 0x69, 0x76, - 0x61, 0x74, 0x65, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, - 0x74, 0x1a, 0x37, 0x2e, 0x63, 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, + 0x75, 0x73, 0x12, 0x29, 0x0a, 0x10, 0x63, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, + 0x73, 0x5f, 0x6c, 0x65, 0x66, 0x74, 0x18, 0x1e, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0f, 0x63, 0x6f, + 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x4c, 0x65, 0x66, 0x74, 0x32, 0xaa, 0x03, + 0x0a, 0x0e, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, + 0x12, 0x7e, 0x0a, 0x0d, 0x4c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, + 0x6e, 0x12, 0x34, 0x2e, 0x63, 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x2e, - 0x76, 0x31, 0x2e, 0x41, 0x63, 0x74, 0x69, 0x76, 0x61, 0x74, 0x65, 0x53, 0x65, 0x73, 0x73, 0x69, - 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x90, 0x01, 0x0a, - 0x13, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x43, 0x6f, 0x6e, 0x6e, 0x65, 0x63, - 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x3a, 0x2e, 0x63, 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, + 0x76, 0x31, 0x2e, 0x4c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, + 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x35, 0x2e, 0x63, 0x6f, 0x6e, 0x74, 0x72, 0x6f, + 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x2e, 0x73, 0x65, 0x72, + 0x76, 0x69, 0x63, 0x65, 0x73, 0x2e, 0x76, 0x31, 0x2e, 0x4c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x53, + 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, + 0x12, 0x84, 0x01, 0x0a, 0x0f, 0x41, 0x63, 0x74, 0x69, 0x76, 0x61, 0x74, 0x65, 0x53, 0x65, 0x73, + 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x36, 0x2e, 0x63, 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, - 0x65, 0x73, 0x2e, 0x76, 0x31, 0x2e, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x43, - 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, - 0x1a, 0x3b, 0x2e, 0x63, 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, - 0x72, 0x76, 0x65, 0x72, 0x73, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x2e, 0x76, - 0x31, 0x2e, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x43, 0x6f, 0x6e, 0x6e, 0x65, - 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, - 0x51, 0x5a, 0x4f, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x68, 0x61, - 0x73, 0x68, 0x69, 0x63, 0x6f, 0x72, 0x70, 0x2f, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x61, 0x72, 0x79, - 0x2f, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x2f, 0x67, 0x65, 0x6e, 0x2f, 0x63, 0x6f, - 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2f, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, - 0x2f, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x3b, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, - 0x65, 0x73, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x65, 0x73, 0x2e, 0x76, 0x31, 0x2e, 0x41, 0x63, 0x74, 0x69, 0x76, 0x61, 0x74, 0x65, 0x53, 0x65, + 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x37, 0x2e, 0x63, + 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, + 0x73, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x2e, 0x76, 0x31, 0x2e, 0x41, 0x63, + 0x74, 0x69, 0x76, 0x61, 0x74, 0x65, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, + 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x90, 0x01, 0x0a, 0x13, 0x41, 0x75, 0x74, 0x68, + 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x43, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x12, + 0x3a, 0x2e, 0x63, 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, + 0x76, 0x65, 0x72, 0x73, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x2e, 0x76, 0x31, + 0x2e, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x43, 0x6f, 0x6e, 0x6e, 0x65, 0x63, + 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x3b, 0x2e, 0x63, 0x6f, + 0x6e, 0x74, 0x72, 0x6f, 0x6c, 0x6c, 0x65, 0x72, 0x2e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, + 0x2e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x2e, 0x76, 0x31, 0x2e, 0x41, 0x75, 0x74, + 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x43, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, + 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x51, 0x5a, 0x4f, 0x67, 0x69, + 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x68, 0x61, 0x73, 0x68, 0x69, 0x63, 0x6f, + 0x72, 0x70, 0x2f, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x61, 0x72, 0x79, 0x2f, 0x69, 0x6e, 0x74, 0x65, + 0x72, 0x6e, 0x61, 0x6c, 0x2f, 0x67, 0x65, 0x6e, 0x2f, 0x63, 0x6f, 0x6e, 0x74, 0x72, 0x6f, 0x6c, + 0x6c, 0x65, 0x72, 0x2f, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x2f, 0x73, 0x65, 0x72, 0x76, + 0x69, 0x63, 0x65, 0x73, 0x3b, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x73, 0x62, 0x06, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x33, } var ( diff --git a/internal/proto/local/controller/servers/services/v1/session_service.proto b/internal/proto/local/controller/servers/services/v1/session_service.proto index 8cbe1d72cf..d7247a6bc7 100644 --- a/internal/proto/local/controller/servers/services/v1/session_service.proto +++ b/internal/proto/local/controller/servers/services/v1/session_service.proto @@ -57,4 +57,5 @@ message AuthorizeConnectionRequest{ message AuthorizeConnectionResponse{ string connection_id = 10; controller.servers.services.v1.CONNECTIONSTATUS status = 20; + int32 connections_left = 30; } \ No newline at end of file diff --git a/internal/proto/local/worker/proxy/v1/proxy.proto b/internal/proto/local/worker/proxy/v1/proxy.proto index b563be8da6..1015ccf9e2 100644 --- a/internal/proto/local/worker/proxy/v1/proxy.proto +++ b/internal/proto/local/worker/proxy/v1/proxy.proto @@ -11,4 +11,5 @@ message ClientHandshake { message HandshakeResult { google.protobuf.Timestamp expiration = 10; + int32 connections_left = 20; } \ No newline at end of file diff --git a/internal/proxy/proxy.pb.go b/internal/proxy/proxy.pb.go index 7465412b69..cb918cc374 100644 --- a/internal/proxy/proxy.pb.go +++ b/internal/proxy/proxy.pb.go @@ -78,7 +78,8 @@ type HandshakeResult struct { sizeCache protoimpl.SizeCache unknownFields protoimpl.UnknownFields - Expiration *timestamp.Timestamp `protobuf:"bytes,10,opt,name=expiration,proto3" json:"expiration,omitempty"` + Expiration *timestamp.Timestamp `protobuf:"bytes,10,opt,name=expiration,proto3" json:"expiration,omitempty"` + ConnectionsLeft int32 `protobuf:"varint,20,opt,name=connections_left,json=connectionsLeft,proto3" json:"connections_left,omitempty"` } func (x *HandshakeResult) Reset() { @@ -120,6 +121,13 @@ func (x *HandshakeResult) GetExpiration() *timestamp.Timestamp { return nil } +func (x *HandshakeResult) GetConnectionsLeft() int32 { + if x != nil { + return x.ConnectionsLeft + } + return 0 +} + var File_worker_proxy_v1_proxy_proto protoreflect.FileDescriptor var file_worker_proxy_v1_proxy_proto_rawDesc = []byte{ @@ -131,15 +139,18 @@ var file_worker_proxy_v1_proxy_proto_rawDesc = []byte{ 0x30, 0x0a, 0x0f, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x48, 0x61, 0x6e, 0x64, 0x73, 0x68, 0x61, 0x6b, 0x65, 0x12, 0x1d, 0x0a, 0x0a, 0x74, 0x6f, 0x66, 0x75, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x18, 0x0a, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x74, 0x6f, 0x66, 0x75, 0x54, 0x6f, 0x6b, 0x65, - 0x6e, 0x22, 0x4d, 0x0a, 0x0f, 0x48, 0x61, 0x6e, 0x64, 0x73, 0x68, 0x61, 0x6b, 0x65, 0x52, 0x65, + 0x6e, 0x22, 0x78, 0x0a, 0x0f, 0x48, 0x61, 0x6e, 0x64, 0x73, 0x68, 0x61, 0x6b, 0x65, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x12, 0x3a, 0x0a, 0x0a, 0x65, 0x78, 0x70, 0x69, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x0a, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x0a, 0x65, 0x78, 0x70, 0x69, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, - 0x42, 0x34, 0x5a, 0x32, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x68, - 0x61, 0x73, 0x68, 0x69, 0x63, 0x6f, 0x72, 0x70, 0x2f, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x61, 0x72, - 0x79, 0x2f, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x2f, 0x70, 0x72, 0x6f, 0x78, 0x79, - 0x3b, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x12, 0x29, 0x0a, 0x10, 0x63, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x5f, + 0x6c, 0x65, 0x66, 0x74, 0x18, 0x14, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0f, 0x63, 0x6f, 0x6e, 0x6e, + 0x65, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x4c, 0x65, 0x66, 0x74, 0x42, 0x34, 0x5a, 0x32, 0x67, + 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x68, 0x61, 0x73, 0x68, 0x69, 0x63, + 0x6f, 0x72, 0x70, 0x2f, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x61, 0x72, 0x79, 0x2f, 0x69, 0x6e, 0x74, + 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x2f, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x3b, 0x70, 0x72, 0x6f, 0x78, + 0x79, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, } var ( diff --git a/internal/servers/controller/handlers/workers/worker_service.go b/internal/servers/controller/handlers/workers/worker_service.go index a410e8466d..d131623f3f 100644 --- a/internal/servers/controller/handlers/workers/worker_service.go +++ b/internal/servers/controller/handlers/workers/worker_service.go @@ -206,7 +206,7 @@ func (ws *workerServiceServer) AuthorizeConnection(ctx context.Context, req *pbs return nil, status.Errorf(codes.Internal, "error getting session repo: %v", err) } - connectionInfo, connStates, _, err := sessRepo.AuthorizeConnection(ctx, req.GetSessionId()) + connectionInfo, connStates, authzSummary, err := sessRepo.AuthorizeConnection(ctx, req.GetSessionId()) if err != nil { return nil, err } @@ -217,8 +217,14 @@ func (ws *workerServiceServer) AuthorizeConnection(ctx context.Context, req *pbs return nil, status.Error(codes.Internal, "Invalid connection state in authorize response.") } - return &pbs.AuthorizeConnectionResponse{ - ConnectionId: connectionInfo.GetPublicId(), - Status: connStates[0].Status.ProtoVal(), - }, nil + ret := &pbs.AuthorizeConnectionResponse{ + ConnectionId: connectionInfo.GetPublicId(), + Status: connStates[0].Status.ProtoVal(), + ConnectionsLeft: authzSummary.ConnectionLimit, + } + if ret.ConnectionsLeft != -1 { + ret.ConnectionsLeft -= int32(authzSummary.CurrentConnectionCount) + } + + return ret, nil } diff --git a/internal/servers/worker/handler.go b/internal/servers/worker/handler.go index 143fcfc8c0..bd1775652c 100644 --- a/internal/servers/worker/handler.go +++ b/internal/servers/worker/handler.go @@ -110,7 +110,8 @@ func (w *Worker) handleProxy() http.HandlerFunc { } var ci *connInfo - ci, err = w.authorizeConnection(r.Context(), sessionId) + var connsLeft int32 + ci, connsLeft, err = w.authorizeConnection(r.Context(), sessionId) if err != nil { w.logger.Error("unable to authorize conneciton", "error", err) conn.Close(websocket.StatusInternalError, "unable to authorize connection") @@ -127,7 +128,8 @@ func (w *Worker) handleProxy() http.HandlerFunc { w.logger.Trace("authorized connection", "connection_id", ci.id) handshakeResult := &proxy.HandshakeResult{ - Expiration: expiration, + Expiration: expiration, + ConnectionsLeft: connsLeft, } if err := wspb.Write(connCtx, conn, handshakeResult); err != nil { w.logger.Error("error sending handshake result to client", "error", err) diff --git a/internal/servers/worker/session.go b/internal/servers/worker/session.go index 96a6886d83..c1f6836809 100644 --- a/internal/servers/worker/session.go +++ b/internal/servers/worker/session.go @@ -148,28 +148,28 @@ func (w *Worker) activateSession(ctx context.Context, sessionId, tofuToken strin return resp.GetStatus(), nil } -func (w *Worker) authorizeConnection(ctx context.Context, sessionId string) (*connInfo, error) { +func (w *Worker) authorizeConnection(ctx context.Context, sessionId string) (*connInfo, int32, error) { rawConn := w.controllerSessionConn.Load() if rawConn == nil { - return nil, errors.New("could not get a controller client") + return nil, 0, errors.New("could not get a controller client") } conn, ok := rawConn.(pbs.SessionServiceClient) if !ok { - return nil, errors.New("could not cast atomic controller client to the real thing") + return nil, 0, errors.New("could not cast atomic controller client to the real thing") } if conn == nil { - return nil, errors.New("controller client is nil") + return nil, 0, errors.New("controller client is nil") } resp, err := conn.AuthorizeConnection(ctx, &pbs.AuthorizeConnectionRequest{ SessionId: sessionId, }) if err != nil { - return nil, fmt.Errorf("error authorizing connection: %w", err) + return nil, 0, fmt.Errorf("error authorizing connection: %w", err) } return &connInfo{ id: resp.ConnectionId, status: resp.GetStatus(), - }, nil + }, resp.GetConnectionsLeft(), nil }