websocket_handler.h

Go to the documentation of this file.

/******************************************************************************
 * Copyright 2017 The Apollo Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/

#pragma once

#include <memory>
#include <mutex>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

#include "CivetServer.h"
#include "nlohmann/json.hpp"

namespace apollo {
namespace dreamview {

class WebSocketHandler : public CivetWebSocketHandler {
  // In case of receiving fragmented message,
  // websocket opcode and accumulated data are stored.
  thread_local static unsigned char current_opcode_;
  thread_local static std::stringstream data_;

 public:
  using Json = nlohmann::json;
  using Connection = struct mg_connection;
  using MessageHandler = std::function<void(const Json &, Connection *)>;
  using ConnectionReadyHandler = std::function<void(Connection *)>;

  explicit WebSocketHandler(const std::string &name) : name_(name) {}

  bool handleConnection(CivetServer *server, const Connection *conn) override {
    return true;
  }

  void handleReadyState(CivetServer *server, Connection *conn) override;

  bool handleData(CivetServer *server, Connection *conn, int bits, char *data,
                  size_t data_len) override;

  bool handleJsonData(Connection *conn, const std::string &data);
  bool handleBinaryData(Connection *conn, const std::string &data);

  void handleClose(CivetServer *server, const Connection *conn) override;

  bool BroadcastData(const std::string &data, bool skippable = false);

  bool SendData(Connection *conn, const std::string &data,
                bool skippable = false, int op_code = MG_WEBSOCKET_OPCODE_TEXT);

  bool SendBinaryData(Connection *conn, const std::string &data,
                      bool skippable = false);

  void RegisterMessageHandler(std::string type, MessageHandler handler) {
    message_handlers_[type] = handler;
  }

  void RegisterConnectionReadyHandler(ConnectionReadyHandler handler) {
    connection_ready_handlers_.emplace_back(handler);
  }

 private:
  const std::string name_;

  // Message handlers keyed by message type.
  std::unordered_map<std::string, MessageHandler> message_handlers_;
  // New connection ready handlers.
  std::vector<ConnectionReadyHandler> connection_ready_handlers_;

  // The mutex guarding the connection set. We are not using read
  // write lock, as the server is not expected to get many clients
  // (connections).
  // CAVEAT: Execution section while holding this global lock should be as
  // brief as possible.
  mutable std::mutex mutex_;

  // The pool of all maintained connections. Each connection has a lock to
  // guard against simultaneous write.
  std::unordered_map<Connection *, std::shared_ptr<std::mutex>> connections_;
};

}  // namespace dreamview
}  // namespace apollo