(ctx context.Context, logger slog.Logger, conn net.Conn)
| 143 | } |
| 144 | |
| 145 | func (s *Server) handleConn(ctx context.Context, logger slog.Logger, conn net.Conn) (retErr error) { |
| 146 | defer conn.Close() |
| 147 | s.connectionsTotal.Add(1) |
| 148 | s.connCount.Add(1) |
| 149 | defer s.connCount.Add(-1) |
| 150 | |
| 151 | // This cannot use a JSON decoder, since that can |
| 152 | // buffer additional data that is required for the PTY. |
| 153 | rawLen := make([]byte, 2) |
| 154 | _, err := conn.Read(rawLen) |
| 155 | if err != nil { |
| 156 | // logging at info since a single incident isn't too worrying (the client could just have |
| 157 | // hung up), but if we get a lot of these we'd want to investigate. |
| 158 | logger.Info(ctx, "failed to read AgentReconnectingPTYInit length", slog.Error(err)) |
| 159 | return nil |
| 160 | } |
| 161 | length := binary.LittleEndian.Uint16(rawLen) |
| 162 | data := make([]byte, length) |
| 163 | _, err = conn.Read(data) |
| 164 | if err != nil { |
| 165 | // logging at info since a single incident isn't too worrying (the client could just have |
| 166 | // hung up), but if we get a lot of these we'd want to investigate. |
| 167 | logger.Info(ctx, "failed to read AgentReconnectingPTYInit", slog.Error(err)) |
| 168 | return nil |
| 169 | } |
| 170 | var msg workspacesdk.AgentReconnectingPTYInit |
| 171 | err = json.Unmarshal(data, &msg) |
| 172 | if err != nil { |
| 173 | logger.Warn(ctx, "failed to unmarshal init", slog.F("raw", data)) |
| 174 | return nil |
| 175 | } |
| 176 | |
| 177 | connectionID := uuid.NewString() |
| 178 | connLogger := logger.With(slog.F("message_id", msg.ID), slog.F("connection_id", connectionID), slog.F("container", msg.Container), slog.F("container_user", msg.ContainerUser)) |
| 179 | connLogger.Debug(ctx, "starting handler") |
| 180 | |
| 181 | defer func() { |
| 182 | if err := retErr; err != nil { |
| 183 | // If the context is done, we don't want to log this as an error since it's expected. |
| 184 | if ctx.Err() != nil { |
| 185 | connLogger.Info(ctx, "reconnecting pty failed with attach error (agent closed)", slog.Error(err)) |
| 186 | } else { |
| 187 | connLogger.Error(ctx, "reconnecting pty failed with attach error", slog.Error(err)) |
| 188 | } |
| 189 | } |
| 190 | connLogger.Info(ctx, "reconnecting pty connection closed") |
| 191 | }() |
| 192 | |
| 193 | var rpty ReconnectingPTY |
| 194 | sendConnected := make(chan ReconnectingPTY, 1) |
| 195 | // On store, reserve this ID to prevent multiple concurrent new connections. |
| 196 | waitReady, ok := s.reconnectingPTYs.LoadOrStore(msg.ID, sendConnected) |
| 197 | if ok { |
| 198 | close(sendConnected) // Unused. |
| 199 | connLogger.Debug(ctx, "connecting to existing reconnecting pty") |
| 200 | c, ok := waitReady.(chan ReconnectingPTY) |
| 201 | if !ok { |
| 202 | return xerrors.Errorf("found invalid type in reconnecting pty map: %T", waitReady) |
no test coverage detected