<commit_before>#include "WSService.h"

#include <chrono>

#include "Config.h"
#include "HTTPRequest.h"
#include "JSON.h"

namespace rustla2 {

WSService::WSService(std::shared_ptr<DB> db, uWS::Hub* hub)
    : db_(db),
      hub_(hub),
      stream_broadcast_timer_(hub->getLoop()),
      rustler_broadcast_timer_(hub->getLoop()),
      stream_observer_(db_->GetStreams()->CreateObserver()) {
  // Run separate stream/rustler broadcast loops for each hub so we don't
  // need to use thread safe queues.
  stream_broadcast_timer_.setData(this);
  stream_broadcast_timer_.start(
      [](Timer* timer) {
        auto* websocket = static_cast<WSService*>(timer->getData());
        websocket->BroadcastStreams();
      },
      0, Config::Get().GetStreamBroadcastInterval());

  rustler_broadcast_timer_.setData(this);
  rustler_broadcast_timer_.start(
      [](Timer* timer) {
        auto* websocket = static_cast<WSService*>(timer->getData());
        websocket->BroadcastRustlers();
      },
      0, Config::Get().GetRustlerBroadcastInterval());

  hub->onConnection(
      [&](uWS::WebSocket<uWS::SERVER>* ws, uWS::HttpRequest uws_req) {
        HTTPRequest req(uws_req);

        ws->send(last_streams_json_.data(), last_streams_json_.size(),
                 uWS::OpCode::TEXT);

        auto state = new WSState();
        state->id = boost::uuids::random_generator()();
        state->user_id = req.GetSessionID();

        db_->GetViewerStates()->IncrViewerStream(state->user_id, 0);

        LOG(INFO) << "CONN_OPEN ws_id:" << state->id
                  << " user_id:" << state->user_id;

        ws->setUserData(reinterpret_cast<void*>(state));

        if (db_->GetBannedIPs()->Contains(req.GetClientIPHeader())) {
          ws->terminate();
        }
      });

  hub->onMessage([&](uWS::WebSocket<uWS::SERVER>* ws, char* message,
                     size_t length, uWS::OpCode opCode) {
    if (length == 0 || opCode != uWS::OpCode::TEXT) {
      return;
    }

    rapidjson::Document input;
    input.Parse(message, length);
    if (input.HasParseError() || !input.IsArray() || input.Size() == 0) {
      return;
    }

    const auto& command = input.GetArray();
    if (!command[0].IsString()) {
      return;
    }

    const json::StringRef method(command[0]);

    if (method == "setAfk") {
      SetAFK(ws, input);
    } else if (method == "setStream") {
      SetStream(ws, input);
    } else if (method == "getStream") {
      GetStream(ws, input);
    }
  });

  hub->onDisconnection([&](uWS::WebSocket<uWS::SERVER>* ws, int code,
                           char* message, size_t length) {
    UnsetStream(ws);

    auto state = reinterpret_cast<WSState*>(ws->getUserData());
    LOG(INFO) << "CONN_CLOSE ws_id:" << state->id
              << " user_id:" << state->user_id;
    delete state;
  });
}

WSService::~WSService() {
  stream_broadcast_timer_.stop();
  stream_broadcast_timer_.close();
  rustler_broadcast_timer_.stop();
  rustler_broadcast_timer_.close();
}

/**
 * Update rustler afk state.
 */
void WSService::SetAFK(uWS::WebSocket<uWS::SERVER>* ws,
                       const rapidjson::Document& input) {
  buf_.Clear();
  rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);
  writer.StartArray();

  const auto& command = input.GetArray();
  if (command.Size() == 2 && command[1].IsBool()) {
    SetAFK(ws, command[1].GetBool(), &writer);
  } else {
    writer.String("ERR");
    writer.String("Invalid command");
  }

  writer.EndArray();
  ws->send(buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);
}

void WSService::SetAFK(uWS::WebSocket<uWS::SERVER>* ws, const bool afk,
                       rapidjson::Writer<rapidjson::StringBuffer>* writer) {
  auto ws_state = GetWSState(ws);

  if (ws_state->afk == afk) {
    writer->String("ERR");
    writer->String("AFK state unchanged");
    return;
  }

  auto stream = GetWSStream(ws);
  if (stream == nullptr) {
    writer->String("ERR");
    writer->String("Not viewing stream");
    return;
  }

  if (afk) {
    stream->IncrAFKCount();
  } else {
    stream->DecrAFKCount();
  }
  ws_state->afk = afk;

  writer->String("AFK_SET");
  writer->Bool(afk);
}

/**
 * Get a single stream definition for the supplied stream id... this (and all
 * non-broadcast sync operations) should be avoided but it's included for
 * compatability.
 */
void WSService::GetStream(uWS::WebSocket<uWS::SERVER>* ws,
                          const rapidjson::Document& input) {
  buf_.Clear();
  rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);
  writer.StartArray();

  const auto& command = input.GetArray();
  if (command.Size() == 2 && command[1].IsUint64()) {
    GetStreamByID(command[1].GetUint64(), &writer);
  } else {
    writer.String("ERR");
    writer.String("Invalid command");
  }

  writer.EndArray();
  ws->send(buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);
}

/**
 * Handle requests for streams by id
 */
void WSService::GetStreamByID(
    const uint64_t stream_id,
    rapidjson::Writer<rapidjson::StringBuffer>* writer) {
  auto stream = db_->GetStreams()->GetByID(stream_id);
  if (stream == nullptr) {
    writer->String("ERR");
    writer->String("Invalid stream id");
    return;
  }

  writer->String("STREAM_GET");
  stream->WriteJSON(writer);
}

/**
 * Clear the current stream defined for the client. Dispatch the command to
 * the appropriate handler based on the json payload shape.
 *
 * If two strings are given treat them as channel and service names.
 * ex: ["setStream", "angelthump", "dariusirl"]
 *
 * If one string is given treat it as an overrustle user id.
 * ex: ["setStream", "dariusirl"]
 *
 * If a null literal is given ack without setting a stream
 * ex: ["setStream", null]
 */
void WSService::SetStream(uWS::WebSocket<uWS::SERVER>* ws,
                          const rapidjson::Document& input) {
  UnsetStream(ws);

  buf_.Clear();
  rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);
  writer.StartArray();

  const auto& command = input.GetArray();
  uint64_t stream_id = 0;

  if (command.Size() == 3 && command[1].IsString() && command[2].IsString()) {
    // handle ["setStream", "channel", "service"]

    const json::StringRef channel(command[1]);
    const json::StringRef service(command[2]);
    SetStreamToChannel(channel, service, &writer, &stream_id);
  } else if (command.Size() == 2 && command[1].IsString()) {
    // handle ["setStream", "stream_path"]

    const json::StringRef stream_path(command[1]);
    SetStreamToStreamPath(stream_path, &writer, &stream_id);
  } else if (command.Size() == 2 && command[1].IsNull()) {
    // handle ["setStream", null]

    SetStreamToNull(&writer, &stream_id);
  } else {
    writer.String("ERR");
    writer.String("Invalid command");
  }

  writer.EndArray();
  ws->send(buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);

  auto state = GetWSState(ws);
  state->stream_id = stream_id;

  db_->GetViewerStates()->IncrViewerStream(state->user_id, stream_id);

  if (stream_id != 0) {
    const auto stream = db_->GetStreams()->GetByID(stream_id);
    const auto stream_channel = stream->GetChannel();
    LOG(INFO) << "STREAM_OPEN ws_id:" << state->id
              << " user_id:" << state->user_id
              << " stream_service:" << stream_channel->GetService()
              << " stream_channel:" << stream_channel->GetChannel();
  } else {
    LOG(INFO) << "STREAM_CLOSE ws_id:" << state->id
              << " user_id:" << state->user_id;
  }
}

/**
 * Handle request for stream by channel/service
 */
inline void WSService::SetStreamToChannel(
    const std::string& channel, const std::string& service,
    rapidjson::Writer<rapidjson::StringBuffer>* writer, uint64_t* stream_id) {
  Status status;
  auto stream_channel = Channel::Create(channel, service, &status);
  if (!status.Ok()) {
    writer->String("ERR");
    writer->String(status.GetErrorMessage().c_str());
    return;
  }

  SetStreamToChannel(stream_channel, writer, stream_id);
}

/**
 * Handle request for stream by overrustle vanity url
 */
void WSService::SetStreamToStreamPath(
    const std::string& stream_path,
    rapidjson::Writer<rapidjson::StringBuffer>* writer, uint64_t* stream_id) {
  auto user = db_->GetUsers()->GetByStreamPath(stream_path);
  if (user == nullptr) {
    writer->String("ERR");
    writer->String("Invalid OverRustle ID");
    return;
  }

  auto channel = user->GetChannel();
  if (channel->IsEmpty()) {
    writer->String("ERR");
    writer->String("OverRustle user has no channel configured");
    return;
  }

  SetStreamToChannel(*channel, writer, stream_id);
}

/**
 * Handle stream upsert for all stream requests.
 *
 * TODO: this should probably be the model's responsibility...
 */
void WSService::SetStreamToChannel(
    const Channel& channel, rapidjson::Writer<rapidjson::StringBuffer>* writer,
    uint64_t* stream_id) {
  if (db_->GetBannedStreams()->Contains(channel)) {
    writer->String("STREAM_BANNED");
    writer->Null();
    return;
  }

  auto stream = db_->GetStreams()->GetByChannel(channel);
  if (stream == nullptr) {
    stream = db_->GetStreams()->Emplace(channel);
  }

  writer->String("STREAM_SET");
  stream->WriteJSON(writer);
  *stream_id = stream->GetID();
  stream->IncrRustlerCount();
}

/**
 * Handle return to stream index
 */
void WSService::SetStreamToNull(
    rapidjson::Writer<rapidjson::StringBuffer>* writer, uint64_t* stream_id) {
  writer->String("STREAM_SET");
  writer->Null();
}

/**
 * Clear the stream id associated with the supplied client
 */
void WSService::UnsetStream(uWS::WebSocket<uWS::SERVER>* ws) {
  auto stream = GetWSStream(ws);
  auto state = GetWSState(ws);

  if (stream != nullptr) {
    stream->DecrRustlerCount(state->afk);
  }

  db_->GetViewerStates()->DecrViewerStream(state->user_id, state->stream_id);
  state->stream_id = 0;
  state->afk = false;
}

/**
 * Generate STREAMS_SET broadcasts. Syncs updates from upstream services
 * ie. liveness, thumbnail, and viewer count.
 */
void WSService::BroadcastStreams() {
  buf_.Clear();
  rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);

  writer.StartArray();
  writer.String("STREAMS_SET");
  db_->GetStreams()->WriteStreamsJSON(&writer);
  writer.EndArray();

  // if the list hasn't changed don't rebroadcast it
  if (last_streams_json_.compare(buf_.GetString()) != 0) {
    hub_->getDefaultGroup<uWS::SERVER>().broadcast(
        buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);

    last_streams_json_.assign(buf_.GetString(), buf_.GetSize());
  }
}

/**
 * Generate RUSTLERS_SET broadcasts for any streams whose rustler counts
 * changed recently. Handling this in a polling loop reduces the cost of
 * syncing clients by debouncing updates and creates a knob for load shedding.
 */
void WSService::BroadcastRustlers() {
  auto now = std::chrono::steady_clock::now();
  auto refresh_ivl =
      std::chrono::milliseconds(Config::Get().GetRustlerBroadcastInterval());
  auto last_rustler_broadcast_time =
      std::chrono::duration_cast<std::chrono::nanoseconds>(
          (now - refresh_ivl).time_since_epoch())
          .count();

  auto streams = db_->GetStreams();
  uint64_t id;
  while (stream_observer_->Next(&id)) {
    auto stream = streams->GetByID(id);

    buf_.Clear();
    rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);
    writer.StartArray();

    // If the stream was reset since the last RUSTLERS_SET broadcast it's
    // a safe bet clients haven't received it via STREAMS_SET. Rather than
    // broadcasting the id and triggering a flood of `getStream` requests
    // broadcast the change as a STREAM_GET.
    if (stream->GetResetTime() >= last_rustler_broadcast_time) {
      writer.String("STREAM_GET");
      stream->WriteJSON(&writer);
    } else {
      writer.String("RUSTLERS_SET");
      writer.Uint64(stream->GetID());
      writer.Uint64(stream->GetRustlerCount());
      writer.Uint64(stream->GetAFKCount());
    }

    writer.EndArray();
    hub_->getDefaultGroup<uWS::SERVER>().broadcast(
        buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);
  }
}

std::shared_ptr<Stream> WSService::GetWSStream(
    uWS::WebSocket<uWS::SERVER>* ws) {
  auto ws_state = GetWSState(ws);
  if (ws_state->stream_id == 0) {
    return nullptr;
  }
  return db_->GetStreams()->GetByID(ws_state->stream_id);
}

WSState* WSService::GetWSState(uWS::WebSocket<uWS::SERVER>* ws) {
  return reinterpret_cast<WSState*>(ws->getUserData());
}

}  // namespace rustla2
<commit_msg>add ip to ws connect log<commit_after>#include "WSService.h"

#include <chrono>

#include "Config.h"
#include "HTTPRequest.h"
#include "JSON.h"

namespace rustla2 {

WSService::WSService(std::shared_ptr<DB> db, uWS::Hub* hub)
    : db_(db),
      hub_(hub),
      stream_broadcast_timer_(hub->getLoop()),
      rustler_broadcast_timer_(hub->getLoop()),
      stream_observer_(db_->GetStreams()->CreateObserver()) {
  // Run separate stream/rustler broadcast loops for each hub so we don't
  // need to use thread safe queues.
  stream_broadcast_timer_.setData(this);
  stream_broadcast_timer_.start(
      [](Timer* timer) {
        auto* websocket = static_cast<WSService*>(timer->getData());
        websocket->BroadcastStreams();
      },
      0, Config::Get().GetStreamBroadcastInterval());

  rustler_broadcast_timer_.setData(this);
  rustler_broadcast_timer_.start(
      [](Timer* timer) {
        auto* websocket = static_cast<WSService*>(timer->getData());
        websocket->BroadcastRustlers();
      },
      0, Config::Get().GetRustlerBroadcastInterval());

  hub->onConnection(
      [&](uWS::WebSocket<uWS::SERVER>* ws, uWS::HttpRequest uws_req) {
        HTTPRequest req(uws_req);

        ws->send(last_streams_json_.data(), last_streams_json_.size(),
                 uWS::OpCode::TEXT);

        auto state = new WSState();
        state->id = boost::uuids::random_generator()();
        state->user_id = req.GetSessionID();

        db_->GetViewerStates()->IncrViewerStream(state->user_id, 0);

        LOG(INFO) << "CONN_OPEN ws_id:" << state->id
                  << " user_id:" << state->user_id
                  << " ip:" << req.GetClientIPHeader();

        ws->setUserData(reinterpret_cast<void*>(state));

        if (db_->GetBannedIPs()->Contains(req.GetClientIPHeader())) {
          ws->terminate();
        }
      });

  hub->onMessage([&](uWS::WebSocket<uWS::SERVER>* ws, char* message,
                     size_t length, uWS::OpCode opCode) {
    if (length == 0 || opCode != uWS::OpCode::TEXT) {
      return;
    }

    rapidjson::Document input;
    input.Parse(message, length);
    if (input.HasParseError() || !input.IsArray() || input.Size() == 0) {
      return;
    }

    const auto& command = input.GetArray();
    if (!command[0].IsString()) {
      return;
    }

    const json::StringRef method(command[0]);

    if (method == "setAfk") {
      SetAFK(ws, input);
    } else if (method == "setStream") {
      SetStream(ws, input);
    } else if (method == "getStream") {
      GetStream(ws, input);
    }
  });

  hub->onDisconnection([&](uWS::WebSocket<uWS::SERVER>* ws, int code,
                           char* message, size_t length) {
    UnsetStream(ws);

    auto state = reinterpret_cast<WSState*>(ws->getUserData());
    LOG(INFO) << "CONN_CLOSE ws_id:" << state->id
              << " user_id:" << state->user_id;
    delete state;
  });
}

WSService::~WSService() {
  stream_broadcast_timer_.stop();
  stream_broadcast_timer_.close();
  rustler_broadcast_timer_.stop();
  rustler_broadcast_timer_.close();
}

/**
 * Update rustler afk state.
 */
void WSService::SetAFK(uWS::WebSocket<uWS::SERVER>* ws,
                       const rapidjson::Document& input) {
  buf_.Clear();
  rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);
  writer.StartArray();

  const auto& command = input.GetArray();
  if (command.Size() == 2 && command[1].IsBool()) {
    SetAFK(ws, command[1].GetBool(), &writer);
  } else {
    writer.String("ERR");
    writer.String("Invalid command");
  }

  writer.EndArray();
  ws->send(buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);
}

void WSService::SetAFK(uWS::WebSocket<uWS::SERVER>* ws, const bool afk,
                       rapidjson::Writer<rapidjson::StringBuffer>* writer) {
  auto ws_state = GetWSState(ws);

  if (ws_state->afk == afk) {
    writer->String("ERR");
    writer->String("AFK state unchanged");
    return;
  }

  auto stream = GetWSStream(ws);
  if (stream == nullptr) {
    writer->String("ERR");
    writer->String("Not viewing stream");
    return;
  }

  if (afk) {
    stream->IncrAFKCount();
  } else {
    stream->DecrAFKCount();
  }
  ws_state->afk = afk;

  writer->String("AFK_SET");
  writer->Bool(afk);
}

/**
 * Get a single stream definition for the supplied stream id... this (and all
 * non-broadcast sync operations) should be avoided but it's included for
 * compatability.
 */
void WSService::GetStream(uWS::WebSocket<uWS::SERVER>* ws,
                          const rapidjson::Document& input) {
  buf_.Clear();
  rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);
  writer.StartArray();

  const auto& command = input.GetArray();
  if (command.Size() == 2 && command[1].IsUint64()) {
    GetStreamByID(command[1].GetUint64(), &writer);
  } else {
    writer.String("ERR");
    writer.String("Invalid command");
  }

  writer.EndArray();
  ws->send(buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);
}

/**
 * Handle requests for streams by id
 */
void WSService::GetStreamByID(
    const uint64_t stream_id,
    rapidjson::Writer<rapidjson::StringBuffer>* writer) {
  auto stream = db_->GetStreams()->GetByID(stream_id);
  if (stream == nullptr) {
    writer->String("ERR");
    writer->String("Invalid stream id");
    return;
  }

  writer->String("STREAM_GET");
  stream->WriteJSON(writer);
}

/**
 * Clear the current stream defined for the client. Dispatch the command to
 * the appropriate handler based on the json payload shape.
 *
 * If two strings are given treat them as channel and service names.
 * ex: ["setStream", "angelthump", "dariusirl"]
 *
 * If one string is given treat it as an overrustle user id.
 * ex: ["setStream", "dariusirl"]
 *
 * If a null literal is given ack without setting a stream
 * ex: ["setStream", null]
 */
void WSService::SetStream(uWS::WebSocket<uWS::SERVER>* ws,
                          const rapidjson::Document& input) {
  UnsetStream(ws);

  buf_.Clear();
  rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);
  writer.StartArray();

  const auto& command = input.GetArray();
  uint64_t stream_id = 0;

  if (command.Size() == 3 && command[1].IsString() && command[2].IsString()) {
    // handle ["setStream", "channel", "service"]

    const json::StringRef channel(command[1]);
    const json::StringRef service(command[2]);
    SetStreamToChannel(channel, service, &writer, &stream_id);
  } else if (command.Size() == 2 && command[1].IsString()) {
    // handle ["setStream", "stream_path"]

    const json::StringRef stream_path(command[1]);
    SetStreamToStreamPath(stream_path, &writer, &stream_id);
  } else if (command.Size() == 2 && command[1].IsNull()) {
    // handle ["setStream", null]

    SetStreamToNull(&writer, &stream_id);
  } else {
    writer.String("ERR");
    writer.String("Invalid command");
  }

  writer.EndArray();
  ws->send(buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);

  auto state = GetWSState(ws);
  state->stream_id = stream_id;

  db_->GetViewerStates()->IncrViewerStream(state->user_id, stream_id);

  if (stream_id != 0) {
    const auto stream = db_->GetStreams()->GetByID(stream_id);
    const auto stream_channel = stream->GetChannel();
    LOG(INFO) << "STREAM_OPEN ws_id:" << state->id
              << " user_id:" << state->user_id
              << " stream_service:" << stream_channel->GetService()
              << " stream_channel:" << stream_channel->GetChannel();
  } else {
    LOG(INFO) << "STREAM_CLOSE ws_id:" << state->id
              << " user_id:" << state->user_id;
  }
}

/**
 * Handle request for stream by channel/service
 */
inline void WSService::SetStreamToChannel(
    const std::string& channel, const std::string& service,
    rapidjson::Writer<rapidjson::StringBuffer>* writer, uint64_t* stream_id) {
  Status status;
  auto stream_channel = Channel::Create(channel, service, &status);
  if (!status.Ok()) {
    writer->String("ERR");
    writer->String(status.GetErrorMessage().c_str());
    return;
  }

  SetStreamToChannel(stream_channel, writer, stream_id);
}

/**
 * Handle request for stream by overrustle vanity url
 */
void WSService::SetStreamToStreamPath(
    const std::string& stream_path,
    rapidjson::Writer<rapidjson::StringBuffer>* writer, uint64_t* stream_id) {
  auto user = db_->GetUsers()->GetByStreamPath(stream_path);
  if (user == nullptr) {
    writer->String("ERR");
    writer->String("Invalid OverRustle ID");
    return;
  }

  auto channel = user->GetChannel();
  if (channel->IsEmpty()) {
    writer->String("ERR");
    writer->String("OverRustle user has no channel configured");
    return;
  }

  SetStreamToChannel(*channel, writer, stream_id);
}

/**
 * Handle stream upsert for all stream requests.
 *
 * TODO: this should probably be the model's responsibility...
 */
void WSService::SetStreamToChannel(
    const Channel& channel, rapidjson::Writer<rapidjson::StringBuffer>* writer,
    uint64_t* stream_id) {
  if (db_->GetBannedStreams()->Contains(channel)) {
    writer->String("STREAM_BANNED");
    writer->Null();
    return;
  }

  auto stream = db_->GetStreams()->GetByChannel(channel);
  if (stream == nullptr) {
    stream = db_->GetStreams()->Emplace(channel);
  }

  writer->String("STREAM_SET");
  stream->WriteJSON(writer);
  *stream_id = stream->GetID();
  stream->IncrRustlerCount();
}

/**
 * Handle return to stream index
 */
void WSService::SetStreamToNull(
    rapidjson::Writer<rapidjson::StringBuffer>* writer, uint64_t* stream_id) {
  writer->String("STREAM_SET");
  writer->Null();
}

/**
 * Clear the stream id associated with the supplied client
 */
void WSService::UnsetStream(uWS::WebSocket<uWS::SERVER>* ws) {
  auto stream = GetWSStream(ws);
  auto state = GetWSState(ws);

  if (stream != nullptr) {
    stream->DecrRustlerCount(state->afk);
  }

  db_->GetViewerStates()->DecrViewerStream(state->user_id, state->stream_id);
  state->stream_id = 0;
  state->afk = false;
}

/**
 * Generate STREAMS_SET broadcasts. Syncs updates from upstream services
 * ie. liveness, thumbnail, and viewer count.
 */
void WSService::BroadcastStreams() {
  buf_.Clear();
  rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);

  writer.StartArray();
  writer.String("STREAMS_SET");
  db_->GetStreams()->WriteStreamsJSON(&writer);
  writer.EndArray();

  // if the list hasn't changed don't rebroadcast it
  if (last_streams_json_.compare(buf_.GetString()) != 0) {
    hub_->getDefaultGroup<uWS::SERVER>().broadcast(
        buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);

    last_streams_json_.assign(buf_.GetString(), buf_.GetSize());
  }
}

/**
 * Generate RUSTLERS_SET broadcasts for any streams whose rustler counts
 * changed recently. Handling this in a polling loop reduces the cost of
 * syncing clients by debouncing updates and creates a knob for load shedding.
 */
void WSService::BroadcastRustlers() {
  auto now = std::chrono::steady_clock::now();
  auto refresh_ivl =
      std::chrono::milliseconds(Config::Get().GetRustlerBroadcastInterval());
  auto last_rustler_broadcast_time =
      std::chrono::duration_cast<std::chrono::nanoseconds>(
          (now - refresh_ivl).time_since_epoch())
          .count();

  auto streams = db_->GetStreams();
  uint64_t id;
  while (stream_observer_->Next(&id)) {
    auto stream = streams->GetByID(id);

    buf_.Clear();
    rapidjson::Writer<rapidjson::StringBuffer> writer(buf_);
    writer.StartArray();

    // If the stream was reset since the last RUSTLERS_SET broadcast it's
    // a safe bet clients haven't received it via STREAMS_SET. Rather than
    // broadcasting the id and triggering a flood of `getStream` requests
    // broadcast the change as a STREAM_GET.
    if (stream->GetResetTime() >= last_rustler_broadcast_time) {
      writer.String("STREAM_GET");
      stream->WriteJSON(&writer);
    } else {
      writer.String("RUSTLERS_SET");
      writer.Uint64(stream->GetID());
      writer.Uint64(stream->GetRustlerCount());
      writer.Uint64(stream->GetAFKCount());
    }

    writer.EndArray();
    hub_->getDefaultGroup<uWS::SERVER>().broadcast(
        buf_.GetString(), buf_.GetSize(), uWS::OpCode::TEXT);
  }
}

std::shared_ptr<Stream> WSService::GetWSStream(
    uWS::WebSocket<uWS::SERVER>* ws) {
  auto ws_state = GetWSState(ws);
  if (ws_state->stream_id == 0) {
    return nullptr;
  }
  return db_->GetStreams()->GetByID(ws_state->stream_id);
}

WSState* WSService::GetWSState(uWS::WebSocket<uWS::SERVER>* ws) {
  return reinterpret_cast<WSState*>(ws->getUserData());
}

}  // namespace rustla2
<|endoftext|>


<commit_before>/*=========================================================================
 *
 *  Copyright RTK Consortium
 *
 *  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.txt
 *
 *  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.
 *
 *=========================================================================*/

#ifndef __rtkConjugateGradientConeBeamReconstructionFilter_hxx
#define __rtkConjugateGradientConeBeamReconstructionFilter_hxx

#include "rtkConjugateGradientConeBeamReconstructionFilter.h"

namespace rtk
{

template< typename TOutputImage>
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>::ConjugateGradientConeBeamReconstructionFilter()
{
  this->SetNumberOfRequiredInputs(3);

  // Set the default values of member parameters
  m_NumberOfIterations=3;
  m_MeasureExecutionTimes=false;
  m_Preconditioned=false;
  m_Gamma = 0;
  m_Regularized = false;
  m_CudaConjugateGradient = true;

  // Create the filters
#ifdef RTK_USE_CUDA
  m_DisplacedDetectorFilter = rtk::CudaDisplacedDetectorImageFilter::New();
  m_ConstantVolumeSource     = rtk::CudaConstantVolumeSource::New();
#else
  m_DisplacedDetectorFilter = DisplacedDetectorFilterType::New();
  m_ConstantVolumeSource     = ConstantImageSourceType::New();
#endif
  m_CGOperator = CGOperatorFilterType::New();

  m_DivideFilter = DivideFilterType::New();
  m_ConstantProjectionsSource = ConstantImageSourceType::New();

  m_MultiplyVolumeFilter = MultiplyFilterType::New();
  m_MultiplyProjectionsFilter = MultiplyFilterType::New();
  m_MultiplyOutputFilter = MultiplyFilterType::New();

  // Set permanent parameters
  m_ConstantVolumeSource->SetConstant(itk::NumericTraits<typename TOutputImage::PixelType>::ZeroValue());
  m_ConstantProjectionsSource->SetConstant(1.0);
  m_DisplacedDetectorFilter->SetPadOnTruncatedSide(false);
}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::SetForwardProjectionFilter (int _arg)
{
  if( _arg != this->GetForwardProjectionFilter() )
    {
    Superclass::SetForwardProjectionFilter( _arg );
    m_ForwardProjectionFilter = this->InstantiateForwardProjectionFilter( _arg );
    m_CGOperator->SetForwardProjectionFilter( m_ForwardProjectionFilter );
    }
}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::SetBackProjectionFilter (int _arg)
{
  if( _arg != this->GetBackProjectionFilter() )
    {
    Superclass::SetBackProjectionFilter( _arg );
    m_BackProjectionFilter = this->InstantiateBackProjectionFilter( _arg );
    m_BackProjectionFilterForB = this->InstantiateBackProjectionFilter( _arg );
    m_BackProjectionFilterForPreconditioning = this->InstantiateBackProjectionFilter( _arg );
    m_BackProjectionFilterForNormalization = this->InstantiateBackProjectionFilter( _arg );
    m_CGOperator->SetBackProjectionFilter( m_BackProjectionFilter);
    }
}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::GenerateInputRequestedRegion()
{
  // Input 0 is the volume we update
  typename Superclass::InputImagePointer inputPtr0 =
          const_cast< TOutputImage * >( this->GetInput(0) );
  if ( !inputPtr0 )
      return;
  inputPtr0->SetRequestedRegion( this->GetOutput()->GetRequestedRegion() );

  // Input 1 is the stack of projections to backproject
  typename Superclass::InputImagePointer  inputPtr1 =
          const_cast< TOutputImage * >( this->GetInput(1) );
  if ( !inputPtr1 )
      return;
  inputPtr1->SetRequestedRegion( inputPtr1->GetLargestPossibleRegion() );

  // Input 2 is the weights map on projections, either user-defined or filled with ones (default)
  typename Superclass::InputImagePointer  inputPtr2 =
          const_cast< TOutputImage * >( this->GetInput(2) );
  if ( !inputPtr2 )
      return;
  inputPtr2->SetRequestedRegion( inputPtr2->GetLargestPossibleRegion() );

}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::GenerateOutputInformation()
{
  // Choose between cuda or non-cuda conjugate gradient filter
  m_ConjugateGradientFilter = ConjugateGradientFilterType::New();
#ifdef RTK_USE_CUDA
  if (m_CudaConjugateGradient)
    m_ConjugateGradientFilter = rtk::CudaConjugateGradientImageFilter_3f::New();
#endif
  m_ConjugateGradientFilter->SetA(m_CGOperator.GetPointer());

  // Set runtime connections
  m_ConstantVolumeSource->SetInformationFromImage(this->GetInput(0));
  m_CGOperator->SetInput(1, this->GetInput(1));
  m_ConjugateGradientFilter->SetX(this->GetInput(0));
  m_DisplacedDetectorFilter->SetInput(this->GetInput(2));

  // Links with the m_BackProjectionFilter should be set here and not
  // in the constructor, as m_BackProjectionFilter is set at runtime
  m_BackProjectionFilterForB->SetInput(0, m_ConstantVolumeSource->GetOutput());
  m_ConjugateGradientFilter->SetB(m_BackProjectionFilterForB->GetOutput());

  // Multiply the projections by the weights map
  m_MultiplyProjectionsFilter->SetInput1(this->GetInput(1));
  m_MultiplyProjectionsFilter->SetInput2(m_DisplacedDetectorFilter->GetOutput());
  m_CGOperator->SetInput(2, m_DisplacedDetectorFilter->GetOutput());
  m_BackProjectionFilterForB->SetInput(1, m_MultiplyProjectionsFilter->GetOutput());

  if (m_Preconditioned)
    {
    // Set the projections source
    m_ConstantProjectionsSource->SetInformationFromImage(this->GetInput(1));

    // Build the part of the pipeline that generates the preconditioning weights
    m_BackProjectionFilterForNormalization->SetInput(0, m_ConstantVolumeSource->GetOutput());
    m_BackProjectionFilterForNormalization->SetInput(1, m_ConstantProjectionsSource->GetOutput());
    m_BackProjectionFilterForPreconditioning->SetInput(0, m_ConstantVolumeSource->GetOutput());
    m_BackProjectionFilterForPreconditioning->SetInput(1, this->GetInput(2));

    m_DivideFilter->SetInput1(m_BackProjectionFilterForNormalization->GetOutput());
    m_DivideFilter->SetInput2(m_BackProjectionFilterForPreconditioning->GetOutput());

    // Multiply the volume by preconditioning weights, and pass them to the conjugate gradient operator
    m_MultiplyVolumeFilter->SetInput1(m_BackProjectionFilterForB->GetOutput());
    m_MultiplyVolumeFilter->SetInput2(m_DivideFilter->GetOutput());
    m_CGOperator->SetInput(3, m_DivideFilter->GetOutput());
    m_ConjugateGradientFilter->SetB(m_MultiplyVolumeFilter->GetOutput());

    // Divide the output by the preconditioning weights
    m_MultiplyOutputFilter->SetInput1(m_ConjugateGradientFilter->GetOutput());
    m_MultiplyOutputFilter->SetInput2(m_DivideFilter->GetOutput());
    }

  // For the same reason, set geometry now
  m_CGOperator->SetGeometry(this->m_Geometry);
  m_BackProjectionFilterForB->SetGeometry(this->m_Geometry.GetPointer());
  m_DisplacedDetectorFilter->SetGeometry(this->m_Geometry);
  m_BackProjectionFilterForNormalization->SetGeometry(this->m_Geometry.GetPointer());
  m_BackProjectionFilterForPreconditioning->SetGeometry(this->m_Geometry.GetPointer());

  // Set runtime parameters
  m_ConjugateGradientFilter->SetNumberOfIterations(this->m_NumberOfIterations);
  m_CGOperator->SetPreconditioned(m_Preconditioned);
  m_CGOperator->SetRegularized(m_Regularized);
  m_CGOperator->SetGamma(m_Gamma);

  // Set memory management parameters
    m_MultiplyProjectionsFilter->ReleaseDataFlagOn();

  if (m_Preconditioned)
    {
    m_ConstantProjectionsSource->ReleaseDataFlagOn();
    m_BackProjectionFilterForPreconditioning->ReleaseDataFlagOn();
    m_BackProjectionFilterForNormalization->ReleaseDataFlagOn();
    m_MultiplyVolumeFilter->ReleaseDataFlagOn();
    m_MultiplyOutputFilter->ReleaseDataFlagOn();
    }
  m_BackProjectionFilterForB->ReleaseDataFlagOn();

  // Have the last filter calculate its output information
  m_ConjugateGradientFilter->UpdateOutputInformation();

  // Copy it as the output information of the composite filter
  this->GetOutput()->CopyInformation( m_ConjugateGradientFilter->GetOutput() );
}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::GenerateData()
{
  itk::TimeProbe ConjugateGradientTimeProbe;

  if(m_MeasureExecutionTimes)
    {
    std::cout << "Starting ConjugateGradient" << std::endl;
    ConjugateGradientTimeProbe.Start();
    }

  if (m_Preconditioned)
    m_DivideFilter->Update();

  m_ConjugateGradientFilter->Update();

  if (m_Preconditioned)
    m_MultiplyOutputFilter->Update();

  if(m_MeasureExecutionTimes)
    {
    ConjugateGradientTimeProbe.Stop();
    std::cout << "ConjugateGradient took " << ConjugateGradientTimeProbe.GetTotal() << ' ' << ConjugateGradientTimeProbe.GetUnit() << std::endl;
    }

  if (m_Preconditioned)
    this->GraftOutput( m_MultiplyOutputFilter->GetOutput() );
  else
    this->GraftOutput( m_ConjugateGradientFilter->GetOutput() );
}

}// end namespace

#endif
<commit_msg>Modified the conjugate gradient preconditioner to take the displaced detector weights into account.<commit_after>/*=========================================================================
 *
 *  Copyright RTK Consortium
 *
 *  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.txt
 *
 *  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.
 *
 *=========================================================================*/

#ifndef __rtkConjugateGradientConeBeamReconstructionFilter_hxx
#define __rtkConjugateGradientConeBeamReconstructionFilter_hxx

#include "rtkConjugateGradientConeBeamReconstructionFilter.h"

namespace rtk
{

template< typename TOutputImage>
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>::ConjugateGradientConeBeamReconstructionFilter()
{
  this->SetNumberOfRequiredInputs(3);

  // Set the default values of member parameters
  m_NumberOfIterations=3;
  m_MeasureExecutionTimes=false;
  m_Preconditioned=false;
  m_Gamma = 0;
  m_Regularized = false;
  m_CudaConjugateGradient = true;

  // Create the filters
#ifdef RTK_USE_CUDA
  m_DisplacedDetectorFilter = rtk::CudaDisplacedDetectorImageFilter::New();
  m_ConstantVolumeSource     = rtk::CudaConstantVolumeSource::New();
#else
  m_DisplacedDetectorFilter = DisplacedDetectorFilterType::New();
  m_ConstantVolumeSource     = ConstantImageSourceType::New();
#endif
  m_CGOperator = CGOperatorFilterType::New();

  m_DivideFilter = DivideFilterType::New();
  m_ConstantProjectionsSource = ConstantImageSourceType::New();

  m_MultiplyVolumeFilter = MultiplyFilterType::New();
  m_MultiplyProjectionsFilter = MultiplyFilterType::New();
  m_MultiplyOutputFilter = MultiplyFilterType::New();

  // Set permanent parameters
  m_ConstantVolumeSource->SetConstant(itk::NumericTraits<typename TOutputImage::PixelType>::ZeroValue());
  m_ConstantProjectionsSource->SetConstant(1.0);
  m_DisplacedDetectorFilter->SetPadOnTruncatedSide(false);
}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::SetForwardProjectionFilter (int _arg)
{
  if( _arg != this->GetForwardProjectionFilter() )
    {
    Superclass::SetForwardProjectionFilter( _arg );
    m_ForwardProjectionFilter = this->InstantiateForwardProjectionFilter( _arg );
    m_CGOperator->SetForwardProjectionFilter( m_ForwardProjectionFilter );
    }
}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::SetBackProjectionFilter (int _arg)
{
  if( _arg != this->GetBackProjectionFilter() )
    {
    Superclass::SetBackProjectionFilter( _arg );
    m_BackProjectionFilter = this->InstantiateBackProjectionFilter( _arg );
    m_BackProjectionFilterForB = this->InstantiateBackProjectionFilter( _arg );
    m_BackProjectionFilterForPreconditioning = this->InstantiateBackProjectionFilter( _arg );
    m_BackProjectionFilterForNormalization = this->InstantiateBackProjectionFilter( _arg );
    m_CGOperator->SetBackProjectionFilter( m_BackProjectionFilter);
    }
}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::GenerateInputRequestedRegion()
{
  // Input 0 is the volume we update
  typename Superclass::InputImagePointer inputPtr0 =
          const_cast< TOutputImage * >( this->GetInput(0) );
  if ( !inputPtr0 )
      return;
  inputPtr0->SetRequestedRegion( this->GetOutput()->GetRequestedRegion() );

  // Input 1 is the stack of projections to backproject
  typename Superclass::InputImagePointer  inputPtr1 =
          const_cast< TOutputImage * >( this->GetInput(1) );
  if ( !inputPtr1 )
      return;
  inputPtr1->SetRequestedRegion( inputPtr1->GetLargestPossibleRegion() );

  // Input 2 is the weights map on projections, either user-defined or filled with ones (default)
  typename Superclass::InputImagePointer  inputPtr2 =
          const_cast< TOutputImage * >( this->GetInput(2) );
  if ( !inputPtr2 )
      return;
  inputPtr2->SetRequestedRegion( inputPtr2->GetLargestPossibleRegion() );

}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::GenerateOutputInformation()
{
  // Choose between cuda or non-cuda conjugate gradient filter
  m_ConjugateGradientFilter = ConjugateGradientFilterType::New();
#ifdef RTK_USE_CUDA
  if (m_CudaConjugateGradient)
    m_ConjugateGradientFilter = rtk::CudaConjugateGradientImageFilter_3f::New();
#endif
  m_ConjugateGradientFilter->SetA(m_CGOperator.GetPointer());

  // Set runtime connections
  m_ConstantVolumeSource->SetInformationFromImage(this->GetInput(0));
  m_CGOperator->SetInput(1, this->GetInput(1));
  m_ConjugateGradientFilter->SetX(this->GetInput(0));
  m_DisplacedDetectorFilter->SetInput(this->GetInput(2));

  // Links with the m_BackProjectionFilter should be set here and not
  // in the constructor, as m_BackProjectionFilter is set at runtime
  m_BackProjectionFilterForB->SetInput(0, m_ConstantVolumeSource->GetOutput());
  m_ConjugateGradientFilter->SetB(m_BackProjectionFilterForB->GetOutput());

  // Multiply the projections by the weights map
  m_MultiplyProjectionsFilter->SetInput1(this->GetInput(1));
  m_MultiplyProjectionsFilter->SetInput2(m_DisplacedDetectorFilter->GetOutput());
  m_CGOperator->SetInput(2, m_DisplacedDetectorFilter->GetOutput());
  m_BackProjectionFilterForB->SetInput(1, m_MultiplyProjectionsFilter->GetOutput());

  if (m_Preconditioned)
    {
    // Set the projections source
    m_ConstantProjectionsSource->SetInformationFromImage(this->GetInput(1));

    // Build the part of the pipeline that generates the preconditioning weights
    m_BackProjectionFilterForNormalization->SetInput(0, m_ConstantVolumeSource->GetOutput());
    m_BackProjectionFilterForNormalization->SetInput(1, m_ConstantProjectionsSource->GetOutput());
    m_BackProjectionFilterForPreconditioning->SetInput(0, m_ConstantVolumeSource->GetOutput());
    m_BackProjectionFilterForPreconditioning->SetInput(1, m_DisplacedDetectorFilter->GetOutput());

    m_DivideFilter->SetInput1(m_BackProjectionFilterForNormalization->GetOutput());
    m_DivideFilter->SetInput2(m_BackProjectionFilterForPreconditioning->GetOutput());

    // Multiply the volume by preconditioning weights, and pass them to the conjugate gradient operator
    m_MultiplyVolumeFilter->SetInput1(m_BackProjectionFilterForB->GetOutput());
    m_MultiplyVolumeFilter->SetInput2(m_DivideFilter->GetOutput());
    m_CGOperator->SetInput(3, m_DivideFilter->GetOutput());
    m_ConjugateGradientFilter->SetB(m_MultiplyVolumeFilter->GetOutput());

    // Divide the output by the preconditioning weights
    m_MultiplyOutputFilter->SetInput1(m_ConjugateGradientFilter->GetOutput());
    m_MultiplyOutputFilter->SetInput2(m_DivideFilter->GetOutput());
    }

  // For the same reason, set geometry now
  m_CGOperator->SetGeometry(this->m_Geometry);
  m_BackProjectionFilterForB->SetGeometry(this->m_Geometry.GetPointer());
  m_DisplacedDetectorFilter->SetGeometry(this->m_Geometry);
  m_BackProjectionFilterForNormalization->SetGeometry(this->m_Geometry.GetPointer());
  m_BackProjectionFilterForPreconditioning->SetGeometry(this->m_Geometry.GetPointer());

  // Set runtime parameters
  m_ConjugateGradientFilter->SetNumberOfIterations(this->m_NumberOfIterations);
  m_CGOperator->SetPreconditioned(m_Preconditioned);
  m_CGOperator->SetRegularized(m_Regularized);
  m_CGOperator->SetGamma(m_Gamma);

  // Set memory management parameters
    m_MultiplyProjectionsFilter->ReleaseDataFlagOn();

  if (m_Preconditioned)
    {
    m_ConstantProjectionsSource->ReleaseDataFlagOn();
    m_BackProjectionFilterForPreconditioning->ReleaseDataFlagOn();
    m_BackProjectionFilterForNormalization->ReleaseDataFlagOn();
    m_MultiplyVolumeFilter->ReleaseDataFlagOn();
    m_MultiplyOutputFilter->ReleaseDataFlagOn();
    }
  m_BackProjectionFilterForB->ReleaseDataFlagOn();

  // Have the last filter calculate its output information
  m_ConjugateGradientFilter->UpdateOutputInformation();

  // Copy it as the output information of the composite filter
  this->GetOutput()->CopyInformation( m_ConjugateGradientFilter->GetOutput() );
}

template< typename TOutputImage>
void
ConjugateGradientConeBeamReconstructionFilter<TOutputImage>
::GenerateData()
{
  itk::TimeProbe ConjugateGradientTimeProbe;

  if(m_MeasureExecutionTimes)
    {
    std::cout << "Starting ConjugateGradient" << std::endl;
    ConjugateGradientTimeProbe.Start();
    }

  if (m_Preconditioned)
    m_DivideFilter->Update();

  m_ConjugateGradientFilter->Update();

  if (m_Preconditioned)
    m_MultiplyOutputFilter->Update();

  if(m_MeasureExecutionTimes)
    {
    ConjugateGradientTimeProbe.Stop();
    std::cout << "ConjugateGradient took " << ConjugateGradientTimeProbe.GetTotal() << ' ' << ConjugateGradientTimeProbe.GetUnit() << std::endl;
    }

  if (m_Preconditioned)
    this->GraftOutput( m_MultiplyOutputFilter->GetOutput() );
  else
    this->GraftOutput( m_ConjugateGradientFilter->GetOutput() );
}

}// end namespace

#endif
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2004-2014 ZNC, see the NOTICE file for details.
 *
 * 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.
 */

#include <znc/IRCNetwork.h>
#include <znc/IRCSock.h>

#ifdef HAVE_LIBSSL
#define HAVE_SASL_MECHANISM
#endif

static const struct {
	const char *szName;
	const char *szDescription;
	const bool  bDefault;
} SupportedMechanisms[] = {
	{ "EXTERNAL",           "TLS certificate, for use with the *cert module", false },
#ifdef HAVE_SASL_MECHANISM
	{ "DH-BLOWFISH",        "Negotiation using the DH-BLOWFISH mechanism", false },
	{ "DH-AES",		"Negotiation using the DH-AES mechanism", false },
#endif
	{ "PLAIN",              "Plain text negotiation, this should work always if the network supports SASL", true },
	{ NULL, NULL, false }
};

#define NV_REQUIRE_AUTH     "require_auth"
#define NV_MECHANISMS       "mechanisms"

class Mechanisms : public VCString {
public:
	void SetIndex(unsigned int uiIndex) {
		m_uiIndex = uiIndex;
	}

	unsigned int GetIndex() const {
		return m_uiIndex;
	}

	bool HasNext() const {
		return size() > (m_uiIndex + 1);
	}

	void IncrementIndex() {
		m_uiIndex++;
	}

	CString GetCurrent() const {
		return at(m_uiIndex);
	}

	CString GetNext() const {
		if (HasNext()) {
			return at(m_uiIndex + 1);
		}

		return "";
	}

private:
	unsigned int m_uiIndex;
};

#ifdef HAVE_SASL_MECHANISM
class DHCommon {
public:
	DH *dh;
	unsigned char *secret;
	int key_size;

	DHCommon() {
		dh = DH_new();
		secret = NULL;
		key_size = 0;
	}

	~DHCommon() {
		if (dh)
			DH_free(dh);
		if (secret)
			free(secret);
	}

	bool ParseDH(const CString &sLine) {
		/*
		 * sLine contains the prime, generator and public key of the server.
		 * We first extract this information and then we pass this to OpenSSL.
		 * OpenSSL will generate our own public and private key. Which we then
		 * use to encrypt our password
		 *
		 * sLine will look something like:
		 *
		 *   base64(
		 *     prime length (2 bytes)
		 *     prime
		 *     generator length (2 bytes)
		 *     generator
		 *     servers public key length (2 bytes)
		 *     servers public key
		 *   )
		 */

		/* Decode base64 into (data, length) */
		CString sData = sLine.Base64Decode_n();
		const unsigned char *data = (const unsigned char*)sData.c_str();
		CString::size_type length = sLine.size();

		if (length < 2) {
			DEBUG("sasl: No prime number");
			return false;
		}

		/* Prime number */
		uint16_t size16;
		memcpy(&size16, data, sizeof(size16));
		unsigned int size = ntohs(size16);
		data += 2;
		length -= 2;

		if (size > length) {
			DEBUG("sasl: Extracting prime number. Invalid length");
			return false;
		}

		dh->p = BN_bin2bn(data, size, NULL);
		data += size;

		/* Generator */
		if (length < 2) {
			DEBUG("sasl: No generator");
			return false;
		}

		memcpy(&size16, data, sizeof(size16));
		size = ntohs(size16);
		data += 2;
		length -= 2;

		if (size > length) {
			DEBUG("sasl: Extracting generator. Invalid length");
			return false;
		}

		dh->g = BN_bin2bn(data, size, NULL);
		data += size;

		/* Server public key */
		if (length < 2) {
			DEBUG("sasl: No public key");
			return false;
		}

		memcpy(&size16, data, sizeof(size16));
		size = ntohs(size16);
		data += 2;
		length -= 2;

		if (size > length) {
			DEBUG("sasl: Extracting server public key. Invalid length");
			return false;
		}

		BIGNUM *server_pub_key = BN_bin2bn(data, size, NULL);

		/* Generate our own public/private keys */
		if (!DH_generate_key(dh)) {
			DEBUG("sasl: Failed to generate keys");
			return false;
		}

		/* Compute shared secret */
		secret = (unsigned char*)malloc(DH_size(dh));
		if ((key_size = DH_compute_key(secret, server_pub_key, dh)) == -1) {
			DEBUG("sasl: Failed to compute shared secret");
			return false;
		}

		return true;
	}
};
#endif

class CSASLMod : public CModule {
public:
	MODCONSTRUCTOR(CSASLMod) {
		AddCommand("Help",        static_cast<CModCommand::ModCmdFunc>(&CSASLMod::PrintHelp),
			"search", "Generate this output");
		AddCommand("Set",         static_cast<CModCommand::ModCmdFunc>(&CSASLMod::Set),
			"username password", "Set the password for DH-BLOWFISH/DH-AES/PLAIN");
		AddCommand("Mechanism",   static_cast<CModCommand::ModCmdFunc>(&CSASLMod::SetMechanismCommand),
			"[mechanism[ ...]]", "Set the mechanisms to be attempted (in order)");
		AddCommand("RequireAuth", static_cast<CModCommand::ModCmdFunc>(&CSASLMod::RequireAuthCommand),
			"[yes|no]", "Don't connect if SASL cannot be authenticated");

		m_bAuthenticated = false;
	}

	void PrintHelp(const CString& sLine) {
		HandleHelpCommand(sLine);

		CTable Mechanisms;
		Mechanisms.AddColumn("Mechanism");
		Mechanisms.AddColumn("Description");

		for (size_t i = 0; SupportedMechanisms[i].szName != NULL; i++) {
			Mechanisms.AddRow();
			Mechanisms.SetCell("Mechanism",   SupportedMechanisms[i].szName);
			Mechanisms.SetCell("Description", SupportedMechanisms[i].szDescription);
		}

		PutModule("The following mechanisms are available:");
		PutModule(Mechanisms);
	}

	void Set(const CString& sLine) {
		SetNV("username", sLine.Token(1));
		SetNV("password", sLine.Token(2));

		PutModule("Username has been set to [" + GetNV("username") + "]");
		PutModule("Password has been set to [" + GetNV("password") + "]");
	}

	void SetMechanismCommand(const CString& sLine) {
		CString sMechanisms = sLine.Token(1, true).AsUpper();

		if (!sMechanisms.empty()) {
			VCString vsMechanisms;
			sMechanisms.Split(" ", vsMechanisms);

			for (VCString::const_iterator it = vsMechanisms.begin(); it != vsMechanisms.end(); ++it) {
				if (!SupportsMechanism(*it)) {
					PutModule("Unsupported mechanism: " + *it);
					return;
				}
			}

			SetNV(NV_MECHANISMS, sMechanisms);
		}

		PutModule("Current mechanisms set: " + GetMechanismsString());
	}

	void RequireAuthCommand(const CString& sLine) {
		if (!sLine.Token(1).empty()) {
			SetNV(NV_REQUIRE_AUTH, sLine.Token(1));
		}

		if (GetNV(NV_REQUIRE_AUTH).ToBool()) {
			PutModule("We require SASL negotiation to connect");
		} else {
			PutModule("We will connect even if SASL fails");
		}
	}

	bool SupportsMechanism(const CString& sMechanism) const {
		for (size_t i = 0; SupportedMechanisms[i].szName != NULL; i++) {
			if (sMechanism.Equals(SupportedMechanisms[i].szName)) {
				return true;
			}
		}

		return false;
	}

	CString GetMechanismsString() const {
		if (GetNV(NV_MECHANISMS).empty()) {
			CString sDefaults = "";

			for (size_t i = 0; SupportedMechanisms[i].szName != NULL; i++) {
				if (SupportedMechanisms[i].bDefault) {
					if (!sDefaults.empty()) {
						sDefaults += " ";
					}

					sDefaults += SupportedMechanisms[i].szName;
				}
			}

			return sDefaults;
		}

		return GetNV(NV_MECHANISMS);
	}

	bool CheckRequireAuth() {
		if (!m_bAuthenticated && GetNV(NV_REQUIRE_AUTH).ToBool()) {
			GetNetwork()->SetIRCConnectEnabled(false);
			PutModule("Disabling network, we require authentication.");
			PutModule("Use 'RequireAuth no' to disable.");
			return true;
		}

		return false;
	}

#ifdef HAVE_SASL_MECHANISM
	bool AuthenticateAES(const CString& sLine) {
		CString::size_type length;

		DHCommon dh;
		if (!dh.ParseDH(sLine))
			return false;

		const int len = GetNV("username").size() + GetNV("password").size() + 2;
		const int padlen = 16 - (len % 16);
		CString::size_type userpass_length = len + padlen;
		unsigned char *encrypted_userpass = (unsigned char *)malloc(userpass_length);
		unsigned char *plaintext_userpass = (unsigned char *)malloc(userpass_length);

		memset(encrypted_userpass, 0, userpass_length);

		/* Create plaintext message */
		unsigned char *ptr = plaintext_userpass;
		memcpy(ptr, GetNV("username").c_str(), GetNV("username").size() + 1);
		ptr += GetNV("username").size() + 1;
		memcpy(ptr, GetNV("password").c_str(), GetNV("password").size() + 1);
		ptr += GetNV("password").size() + 1;
		if (padlen)
		{
			/* Padding */
			unsigned char randbytes[16];
			if (!RAND_bytes(randbytes, padlen)) {
				DEBUG("sasl: DH-AES: Unable to pad");
				free(encrypted_userpass);
				free(plaintext_userpass);
				return false;
			}
			memcpy(ptr, randbytes, padlen);
		}

		/* Create the IV
		 * It is changed during encryption for some reason - so we need to keep a copy.
		 */
		unsigned char iv[16], iv_copy[16];
		if (!RAND_bytes(iv, sizeof (iv))) {
			DEBUG("sasl: DH-AES: Unable to create IV");
			free(encrypted_userpass);
			free(plaintext_userpass);
			return false;
		}
		memcpy(iv_copy, iv, sizeof(iv));

		/* Encrypt */
		AES_KEY key;
		AES_set_encrypt_key(dh.secret, dh.key_size * 8, &key);
		AES_cbc_encrypt(plaintext_userpass, encrypted_userpass, userpass_length,
				&key, iv_copy, AES_ENCRYPT);

		free(plaintext_userpass);

		/* Build our response */
		length = 2 + dh.key_size + sizeof(iv) + userpass_length;
		char *response = (char *)malloc(length);
		char *out_ptr = response;

		/* Size of the key + key */
		uint16_t size16 = htons((uint16_t)dh.key_size);
		memcpy(out_ptr, &size16, sizeof(size16));
		out_ptr += 2;
		BN_bn2bin(dh.dh->pub_key, (unsigned char *)out_ptr);
		out_ptr += dh.key_size;

		/* Add the IV */
		memcpy(out_ptr, iv, sizeof(iv));
		out_ptr += sizeof(iv);

		/* Add encrypted userpass to the response */
		memcpy(out_ptr, encrypted_userpass, userpass_length);
		free(encrypted_userpass);

		PutIRC("AUTHENTICATE " + CString((const char *)response, length).Base64Encode_n());

		DEBUG(CString((const char *)response, length).Base64Encode_n());

		free(response);
		return true;
	}

	bool AuthenticateBlowfish(const CString& sLine) {
		/* Encrypt our sasl password with blowfish
		 *
		 * Our response should look something like:
		 *
		 *   base64(
		 *     our public key length (2 bytes)
		 *     our public key
		 *     sasl username + \0
		 *     blowfish(
		 *       sasl password
		 *     )
		 *   )
		 */
		CString::size_type length;

		/* Our DH params */
		DHCommon dh;
		if (!dh.ParseDH(sLine))
			return false;

		// TODO for passwords with length 8, 16, 24, 32, etc. this will have 8 additional zero bytes at the end...
		// But it works when treated as null-terminated string anyway, and if it works I don't want to touch it right now.
		CString::size_type password_length = GetNV("password").size() + (8 - (GetNV("password").size() % 8));
		unsigned char *encrypted_password = (unsigned char *)malloc(password_length);
		char *plaintext_password = (char *)malloc(password_length);

		memset(encrypted_password, 0, password_length);
		memset(plaintext_password, 0, password_length);
		memcpy(plaintext_password, GetNV("password").c_str(), GetNV("password").size());

		BF_KEY key;
		BF_set_key(&key, dh.key_size, dh.secret);

		char *out_ptr = (char *)encrypted_password;
		char *in_ptr = (char *)plaintext_password;
		for (length = password_length; length; length -= 8, in_ptr += 8, out_ptr += 8) {
			BF_ecb_encrypt((unsigned char *)in_ptr, (unsigned char *)out_ptr, &key, BF_ENCRYPT);
		}

		free(plaintext_password);

		/* Build our response */
		length = 2 + BN_num_bytes(dh.dh->pub_key) + password_length + GetNV("username").size() + 1;
		char *response = (char *)malloc(length);
		out_ptr = response;

		/* Add our key to the response */
		uint16_t size16 = htons((uint16_t)BN_num_bytes(dh.dh->pub_key));
		memcpy(out_ptr, &size16, sizeof(size16));
		out_ptr += 2;
		BN_bn2bin(dh.dh->pub_key, (unsigned char *)out_ptr);
		out_ptr += BN_num_bytes(dh.dh->pub_key);

		/* Add sasl username to response */
		memcpy(out_ptr, GetNV("username").c_str(), GetNV("username").length() + 1); // +1 for zero byte in the end
		out_ptr += GetNV("username").length() + 1;

		/* Finally add the encrypted password to the response */
		memcpy(out_ptr, encrypted_password, password_length);
		free(encrypted_password);

		/* Base 64 encode and send! */
		PutIRC("AUTHENTICATE " + CString((const char *)response, length).Base64Encode_n());

		free(response);
		return true;
	}
#endif

	void Authenticate(const CString& sLine) {
		if (m_Mechanisms.GetCurrent().Equals("PLAIN") && sLine.Equals("+")) {
			CString sAuthLine = GetNV("username") + '\0' + GetNV("username")  + '\0' + GetNV("password");
			sAuthLine.Base64Encode();
			PutIRC("AUTHENTICATE " + sAuthLine);
#ifdef HAVE_SASL_MECHANISM
		} else if (m_Mechanisms.GetCurrent().Equals("DH-BLOWFISH")) {
			AuthenticateBlowfish(sLine);
		} else if (m_Mechanisms.GetCurrent().Equals("DH-AES")) {
			AuthenticateAES(sLine);
#endif
		} else {
			/* Send blank authenticate for other mechanisms (like EXTERNAL). */
			PutIRC("AUTHENTICATE +");
		}
	}

	virtual bool OnServerCapAvailable(const CString& sCap) override {
		return sCap.Equals("sasl");
	}

	virtual void OnServerCapResult(const CString& sCap, bool bSuccess) override {
		if (sCap.Equals("sasl")) {
			if (bSuccess) {
				GetMechanismsString().Split(" ", m_Mechanisms);

				if (m_Mechanisms.empty()) {
					CheckRequireAuth();
					return;
				}

				GetNetwork()->GetIRCSock()->PauseCap();

				m_Mechanisms.SetIndex(0);
				PutIRC("AUTHENTICATE " + m_Mechanisms.GetCurrent());
			} else {
				CheckRequireAuth();
			}
		}
	}

	virtual EModRet OnRaw(CString &sLine) override {
		if (sLine.Token(0).Equals("AUTHENTICATE")) {
			Authenticate(sLine.Token(1, true));
		} else if (sLine.Token(1).Equals("903")) {
			/* SASL success! */
			GetNetwork()->GetIRCSock()->ResumeCap();
			m_bAuthenticated = true;
			DEBUG("sasl: Authenticated with mechanism [" << m_Mechanisms.GetCurrent() << "]");
		} else if (sLine.Token(1).Equals("904") || sLine.Token(1).Equals("905")) {
			DEBUG("sasl: Mechanism [" << m_Mechanisms.GetCurrent() << "] failed.");
			PutModule(m_Mechanisms.GetCurrent() + " mechanism failed.");

			if (m_Mechanisms.HasNext()) {
				m_Mechanisms.IncrementIndex();
				PutIRC("AUTHENTICATE " + m_Mechanisms.GetCurrent());
			} else {
				CheckRequireAuth();
				GetNetwork()->GetIRCSock()->ResumeCap();
			}
		} else if (sLine.Token(1).Equals("906")) {
			/* CAP wasn't paused? */
			DEBUG("sasl: Reached 906.");
			CheckRequireAuth();
		} else if (sLine.Token(1).Equals("907")) {
			m_bAuthenticated = true;
			GetNetwork()->GetIRCSock()->ResumeCap();
			DEBUG("sasl: Received 907 -- We are already registered");
		} else {
			return CONTINUE;
		}

		return HALT;
	}

	virtual void OnIRCConnected() override {
		/* Just incase something slipped through, perhaps the server doesn't
		 * respond to our CAP negotiation. */

		CheckRequireAuth();
	}

	virtual void OnIRCDisconnected() override {
		m_bAuthenticated = false;
	}
private:
	Mechanisms m_Mechanisms;
	bool m_bAuthenticated;
};

template<> void TModInfo<CSASLMod>(CModInfo& Info) {
	Info.SetWikiPage("sasl");
}

NETWORKMODULEDEFS(CSASLMod, "Adds support for sasl authentication capability to authenticate to an IRC server")
<commit_msg>Remove DH-BLOWFISH and DH-AES from the sasl module<commit_after>/*
 * Copyright (C) 2004-2014 ZNC, see the NOTICE file for details.
 *
 * 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.
 */

#include <znc/IRCNetwork.h>
#include <znc/IRCSock.h>

static const struct {
	const char *szName;
	const char *szDescription;
	const bool  bDefault;
} SupportedMechanisms[] = {
	{ "EXTERNAL",           "TLS certificate, for use with the *cert module", false },
	{ "PLAIN",              "Plain text negotiation, this should work always if the network supports SASL", true },
	{ NULL, NULL, false }
};

#define NV_REQUIRE_AUTH     "require_auth"
#define NV_MECHANISMS       "mechanisms"

class Mechanisms : public VCString {
public:
	void SetIndex(unsigned int uiIndex) {
		m_uiIndex = uiIndex;
	}

	unsigned int GetIndex() const {
		return m_uiIndex;
	}

	bool HasNext() const {
		return size() > (m_uiIndex + 1);
	}

	void IncrementIndex() {
		m_uiIndex++;
	}

	CString GetCurrent() const {
		return at(m_uiIndex);
	}

	CString GetNext() const {
		if (HasNext()) {
			return at(m_uiIndex + 1);
		}

		return "";
	}

private:
	unsigned int m_uiIndex;
};

class CSASLMod : public CModule {
public:
	MODCONSTRUCTOR(CSASLMod) {
		AddCommand("Help",        static_cast<CModCommand::ModCmdFunc>(&CSASLMod::PrintHelp),
			"search", "Generate this output");
		AddCommand("Set",         static_cast<CModCommand::ModCmdFunc>(&CSASLMod::Set),
			"username password", "Set the password for PLAIN");
		AddCommand("Mechanism",   static_cast<CModCommand::ModCmdFunc>(&CSASLMod::SetMechanismCommand),
			"[mechanism[ ...]]", "Set the mechanisms to be attempted (in order)");
		AddCommand("RequireAuth", static_cast<CModCommand::ModCmdFunc>(&CSASLMod::RequireAuthCommand),
			"[yes|no]", "Don't connect if SASL cannot be authenticated");

		m_bAuthenticated = false;
	}

	void PrintHelp(const CString& sLine) {
		HandleHelpCommand(sLine);

		CTable Mechanisms;
		Mechanisms.AddColumn("Mechanism");
		Mechanisms.AddColumn("Description");

		for (size_t i = 0; SupportedMechanisms[i].szName != NULL; i++) {
			Mechanisms.AddRow();
			Mechanisms.SetCell("Mechanism",   SupportedMechanisms[i].szName);
			Mechanisms.SetCell("Description", SupportedMechanisms[i].szDescription);
		}

		PutModule("The following mechanisms are available:");
		PutModule(Mechanisms);
	}

	void Set(const CString& sLine) {
		SetNV("username", sLine.Token(1));
		SetNV("password", sLine.Token(2));

		PutModule("Username has been set to [" + GetNV("username") + "]");
		PutModule("Password has been set to [" + GetNV("password") + "]");
	}

	void SetMechanismCommand(const CString& sLine) {
		CString sMechanisms = sLine.Token(1, true).AsUpper();

		if (!sMechanisms.empty()) {
			VCString vsMechanisms;
			sMechanisms.Split(" ", vsMechanisms);

			for (VCString::const_iterator it = vsMechanisms.begin(); it != vsMechanisms.end(); ++it) {
				if (!SupportsMechanism(*it)) {
					PutModule("Unsupported mechanism: " + *it);
					return;
				}
			}

			SetNV(NV_MECHANISMS, sMechanisms);
		}

		PutModule("Current mechanisms set: " + GetMechanismsString());
	}

	void RequireAuthCommand(const CString& sLine) {
		if (!sLine.Token(1).empty()) {
			SetNV(NV_REQUIRE_AUTH, sLine.Token(1));
		}

		if (GetNV(NV_REQUIRE_AUTH).ToBool()) {
			PutModule("We require SASL negotiation to connect");
		} else {
			PutModule("We will connect even if SASL fails");
		}
	}

	bool SupportsMechanism(const CString& sMechanism) const {
		for (size_t i = 0; SupportedMechanisms[i].szName != NULL; i++) {
			if (sMechanism.Equals(SupportedMechanisms[i].szName)) {
				return true;
			}
		}

		return false;
	}

	CString GetMechanismsString() const {
		if (GetNV(NV_MECHANISMS).empty()) {
			CString sDefaults = "";

			for (size_t i = 0; SupportedMechanisms[i].szName != NULL; i++) {
				if (SupportedMechanisms[i].bDefault) {
					if (!sDefaults.empty()) {
						sDefaults += " ";
					}

					sDefaults += SupportedMechanisms[i].szName;
				}
			}

			return sDefaults;
		}

		return GetNV(NV_MECHANISMS);
	}

	bool CheckRequireAuth() {
		if (!m_bAuthenticated && GetNV(NV_REQUIRE_AUTH).ToBool()) {
			GetNetwork()->SetIRCConnectEnabled(false);
			PutModule("Disabling network, we require authentication.");
			PutModule("Use 'RequireAuth no' to disable.");
			return true;
		}

		return false;
	}

	void Authenticate(const CString& sLine) {
		if (m_Mechanisms.GetCurrent().Equals("PLAIN") && sLine.Equals("+")) {
			CString sAuthLine = GetNV("username") + '\0' + GetNV("username")  + '\0' + GetNV("password");
			sAuthLine.Base64Encode();
			PutIRC("AUTHENTICATE " + sAuthLine);
		} else {
			/* Send blank authenticate for other mechanisms (like EXTERNAL). */
			PutIRC("AUTHENTICATE +");
		}
	}

	virtual bool OnServerCapAvailable(const CString& sCap) override {
		return sCap.Equals("sasl");
	}

	virtual void OnServerCapResult(const CString& sCap, bool bSuccess) override {
		if (sCap.Equals("sasl")) {
			if (bSuccess) {
				GetMechanismsString().Split(" ", m_Mechanisms);

				if (m_Mechanisms.empty()) {
					CheckRequireAuth();
					return;
				}

				GetNetwork()->GetIRCSock()->PauseCap();

				m_Mechanisms.SetIndex(0);
				PutIRC("AUTHENTICATE " + m_Mechanisms.GetCurrent());
			} else {
				CheckRequireAuth();
			}
		}
	}

	virtual EModRet OnRaw(CString &sLine) override {
		if (sLine.Token(0).Equals("AUTHENTICATE")) {
			Authenticate(sLine.Token(1, true));
		} else if (sLine.Token(1).Equals("903")) {
			/* SASL success! */
			GetNetwork()->GetIRCSock()->ResumeCap();
			m_bAuthenticated = true;
			DEBUG("sasl: Authenticated with mechanism [" << m_Mechanisms.GetCurrent() << "]");
		} else if (sLine.Token(1).Equals("904") || sLine.Token(1).Equals("905")) {
			DEBUG("sasl: Mechanism [" << m_Mechanisms.GetCurrent() << "] failed.");
			PutModule(m_Mechanisms.GetCurrent() + " mechanism failed.");

			if (m_Mechanisms.HasNext()) {
				m_Mechanisms.IncrementIndex();
				PutIRC("AUTHENTICATE " + m_Mechanisms.GetCurrent());
			} else {
				CheckRequireAuth();
				GetNetwork()->GetIRCSock()->ResumeCap();
			}
		} else if (sLine.Token(1).Equals("906")) {
			/* CAP wasn't paused? */
			DEBUG("sasl: Reached 906.");
			CheckRequireAuth();
		} else if (sLine.Token(1).Equals("907")) {
			m_bAuthenticated = true;
			GetNetwork()->GetIRCSock()->ResumeCap();
			DEBUG("sasl: Received 907 -- We are already registered");
		} else {
			return CONTINUE;
		}

		return HALT;
	}

	virtual void OnIRCConnected() override {
		/* Just incase something slipped through, perhaps the server doesn't
		 * respond to our CAP negotiation. */

		CheckRequireAuth();
	}

	virtual void OnIRCDisconnected() override {
		m_bAuthenticated = false;
	}
private:
	Mechanisms m_Mechanisms;
	bool m_bAuthenticated;
};

template<> void TModInfo<CSASLMod>(CModInfo& Info) {
	Info.SetWikiPage("sasl");
}

NETWORKMODULEDEFS(CSASLMod, "Adds support for sasl authentication capability to authenticate to an IRC server")
<|endoftext|>


<commit_before>/*
  This file is part of KDE Kontact.

  Copyright (C) 2003 Cornelius Schumacher <schumacher@kde.org>
  Copyright (C) 2008 Rafael Fernández López <ereslibre@kde.org>

  This program is free software; you can redistribute it and/or
  modify it under the terms of the GNU General Public
  License as published by the Free Software Foundation; either
  version 2 of the License, or (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; see the file COPYING.  If not, write to
  the Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
  Boston, MA 02110-1301, USA.
*/

#include "iconsidepane.h"
#include "mainwindow.h"
#include "plugin.h"
#include "prefs.h"

#include <QtGui/QStringListModel>
#include <QtGui/QSortFilterProxyModel>
#include <QtGui/QDragEnterEvent>
#include <QtGui/QDragMoveEvent>
#include <QtGui/QStyledItemDelegate>
#include <QtGui/QScrollBar>

#include <KLocalizedString>
#include <KDialog>
#include <KIcon>

namespace Kontact
{

class Navigator;

class Model : public QStringListModel
{
  public:
    enum SpecialRoles {
      PluginName = Qt::UserRole
    };

    Model( Navigator *parentNavigator = 0 )
      : QStringListModel( parentNavigator ), mNavigator(parentNavigator)
    {
    }

    void setPluginList( const QList<Kontact::Plugin*> &list ) {
      pluginList = list;
    }

    virtual Qt::ItemFlags flags( const QModelIndex &index ) const
    {
      Qt::ItemFlags flags = QStringListModel::flags( index );

      if ( index.isValid() ) {
        if ( static_cast<Kontact::Plugin*>( index.internalPointer() )->disabled() ) {
          flags &= ~Qt::ItemIsEnabled;
          flags &= ~Qt::ItemIsSelectable;
          flags &= ~Qt::ItemIsDropEnabled;
        } else {
          flags |= Qt::ItemIsDropEnabled;
        }
      } else {
        flags &= ~Qt::ItemIsDropEnabled;
      }

      return flags;
    }

    virtual QModelIndex index( int row, int column,
                               const QModelIndex &parent = QModelIndex() ) const
    {
      Q_UNUSED( parent );
      if ( row < 0 || row >= pluginList.count() ) {
        return QModelIndex();
      }
      return createIndex( row, column, pluginList[row] );
    }

    virtual QVariant data( const QModelIndex &index, int role = Qt::DisplayRole ) const
    {
      if ( !index.isValid() || !index.internalPointer() ) {
        return QVariant();
      }

      if ( role == Qt::DisplayRole ) {
        if ( !mNavigator->showText() ) {
          return QVariant();
        }
        return static_cast<Kontact::Plugin*>( index.internalPointer() )->title();
      } else if ( role == Qt::DecorationRole ) {
        if ( !mNavigator->showIcons() ) {
          return QVariant();
        }
        return KIcon( static_cast<Kontact::Plugin*>( index.internalPointer() )->icon() );
      } else if ( role == Qt::TextAlignmentRole ) {
        return Qt::AlignCenter;
      } else if ( role == Qt::ToolTipRole ) {
        if ( !mNavigator->showText() ) {
          return static_cast<Kontact::Plugin*>( index.internalPointer() )->title();
        }
        return QVariant();
      } else if ( role == PluginName ) {
        return static_cast<Kontact::Plugin*>( index.internalPointer() )->identifier();
      }
      return QStringListModel::data( index, role );
    }

  private:
    QList<Kontact::Plugin*> pluginList;
    Navigator *mNavigator;
};

class SortFilterProxyModel
  : public QSortFilterProxyModel
{
  public:
    SortFilterProxyModel( QObject *parent = 0 ): QSortFilterProxyModel( parent )
    {
      setDynamicSortFilter( true );
      sort ( 0 );
    }
  protected:
    bool lessThan( const QModelIndex &left, const QModelIndex &right ) const
    {
      Kontact::Plugin *leftPlugin = static_cast<Kontact::Plugin*>( left.internalPointer() );
      Kontact::Plugin *rightPlugin = static_cast<Kontact::Plugin*>( right.internalPointer() );
      return leftPlugin->weight() < rightPlugin->weight();
    }
};

class Delegate : public QStyledItemDelegate
{
  public:
    Delegate( Navigator *parentNavigator = 0 )
      : QStyledItemDelegate( parentNavigator ), mNavigator( parentNavigator )
    {
    }

    void paint( QPainter *painter, const QStyleOptionViewItem &option,
                const QModelIndex &index ) const
    {
      if ( !index.isValid() || !index.internalPointer() ) {
        return;
      }

      QStyleOptionViewItemV4 optionCopy( *static_cast<const QStyleOptionViewItemV4*>( &option ) );
      optionCopy.decorationPosition = QStyleOptionViewItem::Top;
      optionCopy.decorationSize = QSize( mNavigator->iconSize(), mNavigator->iconSize() );
      optionCopy.textElideMode = Qt::ElideNone;
      QStyledItemDelegate::paint( painter, optionCopy, index );
    }

    QSize sizeHint( const QStyleOptionViewItem &option, const QModelIndex &index ) const
    {
      if ( !index.isValid() || !index.internalPointer() ) {
        return QSize();
      }

      QStyleOptionViewItemV4 optionCopy( *static_cast<const QStyleOptionViewItemV4*>( &option ) );
      optionCopy.decorationPosition = QStyleOptionViewItem::Top;
      optionCopy.decorationSize =
        mNavigator->showIcons() ? QSize( mNavigator->iconSize(), mNavigator->iconSize() ) : QSize();
      optionCopy.textElideMode = Qt::ElideNone;
      return QStyledItemDelegate::sizeHint( optionCopy, index );
    }

  private:
    Navigator *mNavigator;
};

Navigator::Navigator( SidePaneBase *parent )
  : QListView( parent ), mSidePane( parent )
{
  setVerticalScrollMode( ScrollPerPixel );
  setHorizontalScrollMode( ScrollPerPixel );

  mIconSize = Prefs::self()->sidePaneIconSize();
  mShowIcons = Prefs::self()->sidePaneShowIcons();
  mShowText = Prefs::self()->sidePaneShowText();

  QActionGroup *viewMode = new QActionGroup( this );

  mShowIconsAction = new KAction( i18n( "Show Icons Only" ), this );
  mShowIconsAction->setCheckable( true );
  mShowIconsAction->setActionGroup( viewMode );
  mShowIconsAction->setChecked( !mShowText && mShowIcons );
  connect( mShowIconsAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  mShowTextAction = new KAction( i18n( "Show Text Only" ), this );
  mShowTextAction->setCheckable( true );
  mShowTextAction->setActionGroup( viewMode );
  mShowTextAction->setChecked( mShowText && !mShowIcons );
  connect( mShowTextAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  mShowBothAction = new KAction( i18n( "Show Icons and Text" ), this );
  mShowBothAction->setCheckable( true );
  mShowBothAction->setActionGroup( viewMode );
  mShowBothAction->setChecked( mShowText && mShowIcons );
  connect( mShowBothAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  KAction *sep = new KAction( this );
  sep->setSeparator( true );

  QActionGroup *iconSize = new QActionGroup( this );

  mBigIconsAction = new KAction( i18n( "Big Icons" ), this );
  mBigIconsAction->setCheckable( iconSize );
  mBigIconsAction->setActionGroup( iconSize );
  mBigIconsAction->setChecked( mIconSize == KIconLoader::SizeLarge );
  connect( mBigIconsAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  mNormalIconsAction = new KAction( i18n( "Normal Icons" ), this );
  mNormalIconsAction->setCheckable( true );
  mNormalIconsAction->setActionGroup( iconSize );
  mNormalIconsAction->setChecked( mIconSize == KIconLoader::SizeMedium );
  connect( mNormalIconsAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  mSmallIconsAction = new KAction( i18n( "Small Icons" ), this );
  mSmallIconsAction->setCheckable( true );
  mSmallIconsAction->setActionGroup( iconSize );
  mSmallIconsAction->setChecked( mIconSize == KIconLoader::SizeSmall );
  connect( mSmallIconsAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  QList<QAction*> actionList;
  actionList << mShowIconsAction << mShowTextAction << mShowBothAction << sep
             << mBigIconsAction << mNormalIconsAction << mSmallIconsAction;

  insertActions( 0, actionList );

  setContextMenuPolicy( Qt::ActionsContextMenu );
  setViewMode( ListMode );
  setItemDelegate( new Delegate( this ) );
  mModel = new Model( this );
  SortFilterProxyModel *sortFilterProxyModel = new SortFilterProxyModel;
  sortFilterProxyModel->setSourceModel( mModel );
  setModel( sortFilterProxyModel );

  setDragDropMode( DropOnly );
  viewport()->setAcceptDrops( true );
  setDropIndicatorShown( true );

  connect( selectionModel(), SIGNAL(currentChanged(QModelIndex,QModelIndex)),
           this, SLOT(slotCurrentChanged(QModelIndex)) );
}

void Navigator::updatePlugins( QList<Kontact::Plugin*> plugins_ )
{
  QList<Kontact::Plugin*> pluginsToShow;
  foreach ( Kontact::Plugin *plugin, plugins_ ) {
    if ( plugin->showInSideBar() ) {
      pluginsToShow << plugin;
    }
  }

  mModel->setPluginList( pluginsToShow );

  mModel->removeRows( 0, mModel->rowCount() );
  mModel->insertRows( 0, pluginsToShow.count() );
}

void Navigator::setCurrentPlugin( const QString &plugin )
{
  for ( int i = 0; i < model()->rowCount(); i++ ) {
    const QModelIndex index = model()->index( i, 0 );
    const QString pluginName = model()->data( index, Model::PluginName ).toString();

    if ( plugin == pluginName ) {
      selectionModel()->setCurrentIndex( index, QItemSelectionModel::SelectCurrent );
      break;
    }
  }
}

QSize Navigator::sizeHint() const
{
  //### TODO: We can cache this value, so this reply is faster. Since here we won't
  //          have too many elements, it is not that important. When caching this value
  //          make sure it is updated correctly when new rows have been added or
  //          removed. (ereslibre)

  int maxWidth = 0;
  for ( int i = 0; i < model()->rowCount(); i++ ) {
    const QModelIndex index = model()->index( i, 0 );
    maxWidth = qMax( maxWidth, sizeHintForIndex( index ).width() );
  }

  if ( !verticalScrollBar()->isVisible() ) {
    maxWidth += style()->pixelMetric( QStyle::PM_ScrollBarExtent ) * 2;
  }

  int viewHeight = QListView::sizeHint().height();

  return QSize( maxWidth + rect().width() - contentsRect().width(), viewHeight );
}

void Navigator::dragEnterEvent( QDragEnterEvent *event )
{
  if ( event->proposedAction() == Qt::IgnoreAction ) {
    return;
  }
  event->acceptProposedAction();
}

void Navigator::dragMoveEvent( QDragMoveEvent *event )
{
  if ( event->proposedAction() == Qt::IgnoreAction ) {
    return;
  }

  const QModelIndex dropIndex = indexAt( event->pos() );

  if ( !dropIndex.isValid() ||
       !( dropIndex.model()->flags( dropIndex ) & Qt::ItemIsEnabled ) ) {
    event->setAccepted( false );
    return;
  } else {
    const QModelIndex sourceIndex =
      static_cast<const QSortFilterProxyModel*>( model() )->mapToSource( dropIndex );
    Kontact::Plugin *plugin =
      static_cast<Kontact::Plugin*>( sourceIndex.internalPointer() );
    if ( !plugin->canDecodeMimeData( event->mimeData() ) ) {
      event->setAccepted( false );
      return;
    }
  }

  event->acceptProposedAction();
}

void Navigator::dropEvent( QDropEvent *event )
{
  if ( event->proposedAction() == Qt::IgnoreAction ) {
    return;
  }

  const QModelIndex dropIndex = indexAt( event->pos() );

  if ( !dropIndex.isValid() ) {
    return;
  } else {
    const QModelIndex sourceIndex =
      static_cast<const QSortFilterProxyModel*>( model() )->mapToSource( dropIndex );
    Kontact::Plugin *plugin =
      static_cast<Kontact::Plugin*>( sourceIndex.internalPointer() );
    plugin->processDropEvent( event );
  }
}

void Navigator::slotCurrentChanged( const QModelIndex &current )
{
  if ( !current.isValid() || !current.internalPointer() ||
       !( current.model()->flags( current ) & Qt::ItemIsEnabled ) ) {
    return;
  }

  QModelIndex source =
    static_cast<const QSortFilterProxyModel*>( current.model() )->mapToSource( current );

  emit pluginActivated( static_cast<Kontact::Plugin*>( source.internalPointer() ) );
}

void Navigator::slotActionTriggered( bool checked )
{
  QObject *object = sender();

  if ( object == mShowIconsAction ) {
    mShowIcons = checked;
    mShowText = !checked;
  } else if ( object == mShowTextAction ) {
    mShowIcons = !checked;
    mShowText = checked;
  } else if ( object == mShowBothAction ) {
    mShowIcons = checked;
    mShowText = checked;
  } else if ( object == mBigIconsAction ) {
    mIconSize = KIconLoader::SizeLarge;
  } else if ( object == mNormalIconsAction ) {
    mIconSize = KIconLoader::SizeMedium;
  } else if ( object == mSmallIconsAction ) {
    mIconSize = KIconLoader::SizeSmallMedium;
  }

  Prefs::self()->setSidePaneIconSize( mIconSize );
  Prefs::self()->setSidePaneShowIcons( mShowIcons );
  Prefs::self()->setSidePaneShowText( mShowText );

  scheduleDelayedItemsLayout();

  parentWidget()->setMaximumWidth( sizeHint().width() );
  parentWidget()->setMinimumWidth( sizeHint().width() );
}

IconSidePane::IconSidePane( Core *core, QWidget *parent )
  : SidePaneBase( core, parent )
{
  mNavigator = new Navigator( this );
  connect( mNavigator, SIGNAL(pluginActivated(Kontact::Plugin*)),
           SIGNAL(pluginSelected(Kontact::Plugin*)) );
}

IconSidePane::~IconSidePane()
{
}

void IconSidePane::setCurrentPlugin( const QString &plugin )
{
  mNavigator->setCurrentPlugin( plugin );
}

void IconSidePane::updatePlugins()
{
  mNavigator->updatePlugins( core()->pluginList() );
}

}

#include "iconsidepane.moc"

// vim: sw=2 sts=2 et tw=80
<commit_msg>This sets the background color as the old sidebar (the base one). Please revert this patch if you want the other solution.<commit_after>/*
  This file is part of KDE Kontact.

  Copyright (C) 2003 Cornelius Schumacher <schumacher@kde.org>
  Copyright (C) 2008 Rafael Fernández López <ereslibre@kde.org>

  This program is free software; you can redistribute it and/or
  modify it under the terms of the GNU General Public
  License as published by the Free Software Foundation; either
  version 2 of the License, or (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; see the file COPYING.  If not, write to
  the Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
  Boston, MA 02110-1301, USA.
*/

#include "iconsidepane.h"
#include "mainwindow.h"
#include "plugin.h"
#include "prefs.h"

#include <QtGui/QStringListModel>
#include <QtGui/QSortFilterProxyModel>
#include <QtGui/QDragEnterEvent>
#include <QtGui/QDragMoveEvent>
#include <QtGui/QStyledItemDelegate>
#include <QtGui/QScrollBar>

#include <KLocalizedString>
#include <KDialog>
#include <KIcon>

namespace Kontact
{

class Navigator;

class Model : public QStringListModel
{
  public:
    enum SpecialRoles {
      PluginName = Qt::UserRole
    };

    Model( Navigator *parentNavigator = 0 )
      : QStringListModel( parentNavigator ), mNavigator(parentNavigator)
    {
    }

    void setPluginList( const QList<Kontact::Plugin*> &list ) {
      pluginList = list;
    }

    virtual Qt::ItemFlags flags( const QModelIndex &index ) const
    {
      Qt::ItemFlags flags = QStringListModel::flags( index );

      if ( index.isValid() ) {
        if ( static_cast<Kontact::Plugin*>( index.internalPointer() )->disabled() ) {
          flags &= ~Qt::ItemIsEnabled;
          flags &= ~Qt::ItemIsSelectable;
          flags &= ~Qt::ItemIsDropEnabled;
        } else {
          flags |= Qt::ItemIsDropEnabled;
        }
      } else {
        flags &= ~Qt::ItemIsDropEnabled;
      }

      return flags;
    }

    virtual QModelIndex index( int row, int column,
                               const QModelIndex &parent = QModelIndex() ) const
    {
      Q_UNUSED( parent );
      if ( row < 0 || row >= pluginList.count() ) {
        return QModelIndex();
      }
      return createIndex( row, column, pluginList[row] );
    }

    virtual QVariant data( const QModelIndex &index, int role = Qt::DisplayRole ) const
    {
      if ( !index.isValid() || !index.internalPointer() ) {
        return QVariant();
      }

      if ( role == Qt::DisplayRole ) {
        if ( !mNavigator->showText() ) {
          return QVariant();
        }
        return static_cast<Kontact::Plugin*>( index.internalPointer() )->title();
      } else if ( role == Qt::DecorationRole ) {
        if ( !mNavigator->showIcons() ) {
          return QVariant();
        }
        return KIcon( static_cast<Kontact::Plugin*>( index.internalPointer() )->icon() );
      } else if ( role == Qt::TextAlignmentRole ) {
        return Qt::AlignCenter;
      } else if ( role == Qt::ToolTipRole ) {
        if ( !mNavigator->showText() ) {
          return static_cast<Kontact::Plugin*>( index.internalPointer() )->title();
        }
        return QVariant();
      } else if ( role == PluginName ) {
        return static_cast<Kontact::Plugin*>( index.internalPointer() )->identifier();
      }
      return QStringListModel::data( index, role );
    }

  private:
    QList<Kontact::Plugin*> pluginList;
    Navigator *mNavigator;
};

class SortFilterProxyModel
  : public QSortFilterProxyModel
{
  public:
    SortFilterProxyModel( QObject *parent = 0 ): QSortFilterProxyModel( parent )
    {
      setDynamicSortFilter( true );
      sort ( 0 );
    }
  protected:
    bool lessThan( const QModelIndex &left, const QModelIndex &right ) const
    {
      Kontact::Plugin *leftPlugin = static_cast<Kontact::Plugin*>( left.internalPointer() );
      Kontact::Plugin *rightPlugin = static_cast<Kontact::Plugin*>( right.internalPointer() );
      return leftPlugin->weight() < rightPlugin->weight();
    }
};

class Delegate : public QStyledItemDelegate
{
  public:
    Delegate( Navigator *parentNavigator = 0 )
      : QStyledItemDelegate( parentNavigator ), mNavigator( parentNavigator )
    {
    }

    void paint( QPainter *painter, const QStyleOptionViewItem &option,
                const QModelIndex &index ) const
    {
      if ( !index.isValid() || !index.internalPointer() ) {
        return;
      }

      QStyleOptionViewItemV4 optionCopy( *static_cast<const QStyleOptionViewItemV4*>( &option ) );
      optionCopy.decorationPosition = QStyleOptionViewItem::Top;
      optionCopy.decorationSize = QSize( mNavigator->iconSize(), mNavigator->iconSize() );
      optionCopy.textElideMode = Qt::ElideNone;
      QStyledItemDelegate::paint( painter, optionCopy, index );
    }

    QSize sizeHint( const QStyleOptionViewItem &option, const QModelIndex &index ) const
    {
      if ( !index.isValid() || !index.internalPointer() ) {
        return QSize();
      }

      QStyleOptionViewItemV4 optionCopy( *static_cast<const QStyleOptionViewItemV4*>( &option ) );
      optionCopy.decorationPosition = QStyleOptionViewItem::Top;
      optionCopy.decorationSize =
        mNavigator->showIcons() ? QSize( mNavigator->iconSize(), mNavigator->iconSize() ) : QSize();
      optionCopy.textElideMode = Qt::ElideNone;
      return QStyledItemDelegate::sizeHint( optionCopy, index );
    }

  private:
    Navigator *mNavigator;
};

Navigator::Navigator( SidePaneBase *parent )
  : QListView( parent ), mSidePane( parent )
{
  setViewport( new QWidget( this ) );

  setVerticalScrollMode( ScrollPerPixel );
  setHorizontalScrollMode( ScrollPerPixel );

  mIconSize = Prefs::self()->sidePaneIconSize();
  mShowIcons = Prefs::self()->sidePaneShowIcons();
  mShowText = Prefs::self()->sidePaneShowText();

  QActionGroup *viewMode = new QActionGroup( this );

  mShowIconsAction = new KAction( i18n( "Show Icons Only" ), this );
  mShowIconsAction->setCheckable( true );
  mShowIconsAction->setActionGroup( viewMode );
  mShowIconsAction->setChecked( !mShowText && mShowIcons );
  connect( mShowIconsAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  mShowTextAction = new KAction( i18n( "Show Text Only" ), this );
  mShowTextAction->setCheckable( true );
  mShowTextAction->setActionGroup( viewMode );
  mShowTextAction->setChecked( mShowText && !mShowIcons );
  connect( mShowTextAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  mShowBothAction = new KAction( i18n( "Show Icons and Text" ), this );
  mShowBothAction->setCheckable( true );
  mShowBothAction->setActionGroup( viewMode );
  mShowBothAction->setChecked( mShowText && mShowIcons );
  connect( mShowBothAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  KAction *sep = new KAction( this );
  sep->setSeparator( true );

  QActionGroup *iconSize = new QActionGroup( this );

  mBigIconsAction = new KAction( i18n( "Big Icons" ), this );
  mBigIconsAction->setCheckable( iconSize );
  mBigIconsAction->setActionGroup( iconSize );
  mBigIconsAction->setChecked( mIconSize == KIconLoader::SizeLarge );
  connect( mBigIconsAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  mNormalIconsAction = new KAction( i18n( "Normal Icons" ), this );
  mNormalIconsAction->setCheckable( true );
  mNormalIconsAction->setActionGroup( iconSize );
  mNormalIconsAction->setChecked( mIconSize == KIconLoader::SizeMedium );
  connect( mNormalIconsAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  mSmallIconsAction = new KAction( i18n( "Small Icons" ), this );
  mSmallIconsAction->setCheckable( true );
  mSmallIconsAction->setActionGroup( iconSize );
  mSmallIconsAction->setChecked( mIconSize == KIconLoader::SizeSmall );
  connect( mSmallIconsAction, SIGNAL(triggered(bool)), this, SLOT(slotActionTriggered(bool)) );

  QList<QAction*> actionList;
  actionList << mShowIconsAction << mShowTextAction << mShowBothAction << sep
             << mBigIconsAction << mNormalIconsAction << mSmallIconsAction;

  insertActions( 0, actionList );

  setContextMenuPolicy( Qt::ActionsContextMenu );
  setViewMode( ListMode );
  setItemDelegate( new Delegate( this ) );
  mModel = new Model( this );
  SortFilterProxyModel *sortFilterProxyModel = new SortFilterProxyModel;
  sortFilterProxyModel->setSourceModel( mModel );
  setModel( sortFilterProxyModel );

  setDragDropMode( DropOnly );
  viewport()->setAcceptDrops( true );
  setDropIndicatorShown( true );

  connect( selectionModel(), SIGNAL(currentChanged(QModelIndex,QModelIndex)),
           this, SLOT(slotCurrentChanged(QModelIndex)) );
}

void Navigator::updatePlugins( QList<Kontact::Plugin*> plugins_ )
{
  QList<Kontact::Plugin*> pluginsToShow;
  foreach ( Kontact::Plugin *plugin, plugins_ ) {
    if ( plugin->showInSideBar() ) {
      pluginsToShow << plugin;
    }
  }

  mModel->setPluginList( pluginsToShow );

  mModel->removeRows( 0, mModel->rowCount() );
  mModel->insertRows( 0, pluginsToShow.count() );
}

void Navigator::setCurrentPlugin( const QString &plugin )
{
  for ( int i = 0; i < model()->rowCount(); i++ ) {
    const QModelIndex index = model()->index( i, 0 );
    const QString pluginName = model()->data( index, Model::PluginName ).toString();

    if ( plugin == pluginName ) {
      selectionModel()->setCurrentIndex( index, QItemSelectionModel::SelectCurrent );
      break;
    }
  }
}

QSize Navigator::sizeHint() const
{
  //### TODO: We can cache this value, so this reply is faster. Since here we won't
  //          have too many elements, it is not that important. When caching this value
  //          make sure it is updated correctly when new rows have been added or
  //          removed. (ereslibre)

  int maxWidth = 0;
  for ( int i = 0; i < model()->rowCount(); i++ ) {
    const QModelIndex index = model()->index( i, 0 );
    maxWidth = qMax( maxWidth, sizeHintForIndex( index ).width() );
  }

  if ( !verticalScrollBar()->isVisible() ) {
    maxWidth += style()->pixelMetric( QStyle::PM_ScrollBarExtent ) * 2;
  }

  int viewHeight = QListView::sizeHint().height();

  return QSize( maxWidth + rect().width() - contentsRect().width(), viewHeight );
}

void Navigator::dragEnterEvent( QDragEnterEvent *event )
{
  if ( event->proposedAction() == Qt::IgnoreAction ) {
    return;
  }
  event->acceptProposedAction();
}

void Navigator::dragMoveEvent( QDragMoveEvent *event )
{
  if ( event->proposedAction() == Qt::IgnoreAction ) {
    return;
  }

  const QModelIndex dropIndex = indexAt( event->pos() );

  if ( !dropIndex.isValid() ||
       !( dropIndex.model()->flags( dropIndex ) & Qt::ItemIsEnabled ) ) {
    event->setAccepted( false );
    return;
  } else {
    const QModelIndex sourceIndex =
      static_cast<const QSortFilterProxyModel*>( model() )->mapToSource( dropIndex );
    Kontact::Plugin *plugin =
      static_cast<Kontact::Plugin*>( sourceIndex.internalPointer() );
    if ( !plugin->canDecodeMimeData( event->mimeData() ) ) {
      event->setAccepted( false );
      return;
    }
  }

  event->acceptProposedAction();
}

void Navigator::dropEvent( QDropEvent *event )
{
  if ( event->proposedAction() == Qt::IgnoreAction ) {
    return;
  }

  const QModelIndex dropIndex = indexAt( event->pos() );

  if ( !dropIndex.isValid() ) {
    return;
  } else {
    const QModelIndex sourceIndex =
      static_cast<const QSortFilterProxyModel*>( model() )->mapToSource( dropIndex );
    Kontact::Plugin *plugin =
      static_cast<Kontact::Plugin*>( sourceIndex.internalPointer() );
    plugin->processDropEvent( event );
  }
}

void Navigator::slotCurrentChanged( const QModelIndex &current )
{
  if ( !current.isValid() || !current.internalPointer() ||
       !( current.model()->flags( current ) & Qt::ItemIsEnabled ) ) {
    return;
  }

  QModelIndex source =
    static_cast<const QSortFilterProxyModel*>( current.model() )->mapToSource( current );

  emit pluginActivated( static_cast<Kontact::Plugin*>( source.internalPointer() ) );
}

void Navigator::slotActionTriggered( bool checked )
{
  QObject *object = sender();

  if ( object == mShowIconsAction ) {
    mShowIcons = checked;
    mShowText = !checked;
  } else if ( object == mShowTextAction ) {
    mShowIcons = !checked;
    mShowText = checked;
  } else if ( object == mShowBothAction ) {
    mShowIcons = checked;
    mShowText = checked;
  } else if ( object == mBigIconsAction ) {
    mIconSize = KIconLoader::SizeLarge;
  } else if ( object == mNormalIconsAction ) {
    mIconSize = KIconLoader::SizeMedium;
  } else if ( object == mSmallIconsAction ) {
    mIconSize = KIconLoader::SizeSmallMedium;
  }

  Prefs::self()->setSidePaneIconSize( mIconSize );
  Prefs::self()->setSidePaneShowIcons( mShowIcons );
  Prefs::self()->setSidePaneShowText( mShowText );

  scheduleDelayedItemsLayout();

  parentWidget()->setMaximumWidth( sizeHint().width() );
  parentWidget()->setMinimumWidth( sizeHint().width() );
}

IconSidePane::IconSidePane( Core *core, QWidget *parent )
  : SidePaneBase( core, parent )
{
  mNavigator = new Navigator( this );
  connect( mNavigator, SIGNAL(pluginActivated(Kontact::Plugin*)),
           SIGNAL(pluginSelected(Kontact::Plugin*)) );
}

IconSidePane::~IconSidePane()
{
}

void IconSidePane::setCurrentPlugin( const QString &plugin )
{
  mNavigator->setCurrentPlugin( plugin );
}

void IconSidePane::updatePlugins()
{
  mNavigator->updatePlugins( core()->pluginList() );
}

}

#include "iconsidepane.moc"

// vim: sw=2 sts=2 et tw=80
<|endoftext|>


<commit_before>
#include "NNApplication.h"
#include "GameScene.h"
#include "StoryScene.h"
//#include "MainMenuScene.h"
#include "LoadScene.h"

#ifdef _DEBUG
#pragma comment(lib, "NNGameFramework_DEBUG.lib")
#else
#pragma comment(lib, "NNGameFramework.lib")
#endif
#define _CRT_SECURE_NO_WARNINGS
int WINAPI WinMain( HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpszCmdParam, int nShowCmd )
{
#ifdef _DEBUG
	/*_CrtSetDbgFlag(_CRTDBG_ALLOC_MEM_DF|_CRTDBG_LEAK_CHECK_DF);
	_CrtSetBreakAlloc( 69253 );

	AllocConsole();
	FILE* console;
	freopen_s( &console, "CONOUT$", "wt", stdout );

	printf_s("Console Open \n");*/
#endif

	char* param = lpszCmdParam;
	char* param1 = nullptr;
	int roomNum, playerID;
	roomNum = playerID = 0;
	int i=0;

	while(param[i++] != ' ');
	param1 = param + i;
	param[i-1] = NULL;
	roomNum = atoi(param);
	playerID = atoi(param1);
	printf("\n%d\n%d\n",roomNum, playerID);

	NNApplication* Application = NNApplication::GetInstance();
	Application->Init( L"CSM Project", 800, 600, D2D );

	//NNSceneDirector::GetInstance()->ChangeScene( MainMenuScene::Create() );
	NNSceneDirector::GetInstance()->ChangeScene( LoadScene::Create(L"resource/map/44.csm",roomNum,playerID));
	//NNSceneDirector::GetInstance()->ChangeScene(CGameScene::Create(L"resource/map/44.csm"));

	Application->Run();
	Application->Release();

#ifdef _DEBUG
	//FreeConsole();
#endif

	return 0;
}<commit_msg># Ban incorrect access<commit_after>
#include "NNApplication.h"
#include "GameScene.h"
#include "StoryScene.h"
//#include "MainMenuScene.h"
#include "LoadScene.h"

#ifdef _DEBUG
#pragma comment(lib, "NNGameFramework_DEBUG.lib")
#else
#pragma comment(lib, "NNGameFramework.lib")
#endif
#define _CRT_SECURE_NO_WARNINGS
int WINAPI WinMain( HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpszCmdParam, int nShowCmd )
{
#ifdef _DEBUG
	_CrtSetDbgFlag(_CRTDBG_ALLOC_MEM_DF|_CRTDBG_LEAK_CHECK_DF);
	_CrtSetBreakAlloc( 69253 );

	AllocConsole();
	FILE* console;
	freopen_s( &console, "CONOUT$", "wt", stdout );

	printf_s("Console Open \n");
#endif

	char* param = lpszCmdParam;
	char* param1 = nullptr;
	int roomNum, playerID;
	roomNum = playerID = 0;
	int i=0;

	while(param[i++] != ' ');
	param1 = param + i;
	param[i-1] = NULL;
	roomNum = atoi(param);
	playerID = atoi(param1);
	printf("\n%d\n%d\n",roomNum, playerID);

	NNApplication* Application = NNApplication::GetInstance();
	Application->Init( L"CSM Project", 800, 600, D2D );

	if (roomNum == 0 || playerID == 0)
	{
		MessageBox(NNApplication::GetInstance()->GetHWND(), L"߸ Դϴ!", L"ERROR", 0);
		exit(0);
	}

	//NNSceneDirector::GetInstance()->ChangeScene( MainMenuScene::Create() );
	NNSceneDirector::GetInstance()->ChangeScene( LoadScene::Create(L"resource/map/44.csm",roomNum,playerID));
	//NNSceneDirector::GetInstance()->ChangeScene(CGameScene::Create(L"resource/map/44.csm"));

	Application->Run();
	Application->Release();

#ifdef _DEBUG
	//FreeConsole();
#endif

	return 0;
}<|endoftext|>


<commit_before>// Copyright (c) 2009 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/extensions/extension_apitest.h"
#include "chrome/browser/geolocation/location_arbitrator.h"
#include "chrome/browser/geolocation/mock_location_provider.h"
#include "chrome/common/chrome_switches.h"

class GeolocationApiTest : public ExtensionApiTest {
public:
  // InProcessBrowserTest
  void SetUpInProcessBrowserTestFixture() {
    ExtensionApiTest::SetUpInProcessBrowserTestFixture();
    GeolocationArbitrator::SetProviderFactoryForTest(
        &NewAutoSuccessMockLocationProvider);
  }
};

IN_PROC_BROWSER_TEST_F(GeolocationApiTest, ExtensionGeolocationAccessFail) {
  // Test that geolocation cannot be accessed from extension without permission.
  ASSERT_TRUE(RunExtensionTest("geolocation/no_permission")) << message_;
}

IN_PROC_BROWSER_TEST_F(GeolocationApiTest, ExtensionGeolocationAccessPass) {
  // Test that geolocation can be accessed from extension with permission.
  ASSERT_TRUE(RunExtensionTest("geolocation/has_permission")) << message_;
}
<commit_msg>GeolocationApiTest.ExtensionGeolocationAccessFail is flaky on Mac<commit_after>// Copyright (c) 2009 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/extensions/extension_apitest.h"
#include "chrome/browser/geolocation/location_arbitrator.h"
#include "chrome/browser/geolocation/mock_location_provider.h"
#include "chrome/common/chrome_switches.h"

class GeolocationApiTest : public ExtensionApiTest {
public:
  // InProcessBrowserTest
  void SetUpInProcessBrowserTestFixture() {
    ExtensionApiTest::SetUpInProcessBrowserTestFixture();
    GeolocationArbitrator::SetProviderFactoryForTest(
        &NewAutoSuccessMockLocationProvider);
  }
};

IN_PROC_BROWSER_TEST_F(GeolocationApiTest, FLAKY_ExtensionGeolocationAccessFail) {
  // Test that geolocation cannot be accessed from extension without permission.
  ASSERT_TRUE(RunExtensionTest("geolocation/no_permission")) << message_;
}

IN_PROC_BROWSER_TEST_F(GeolocationApiTest, ExtensionGeolocationAccessPass) {
  // Test that geolocation can be accessed from extension with permission.
  ASSERT_TRUE(RunExtensionTest("geolocation/has_permission")) << message_;
}
<|endoftext|>