Avoid unnecessary reallocation of the IV for every packet

This commit is contained in:
Cameron Gutman 2024-01-24 19:03:57 -06:00
parent 3f31400e7f
commit 65851407b1

View File

@ -246,16 +246,13 @@ namespace stream {
// return bytes written on success
// return -1 on error
static inline int
encode_audio(bool encrypted, const audio::buffer_t &plaintext, audio_packet_t &destination, std::uint32_t avRiKeyIv, crypto::cipher::cbc_t &cbc) {
encode_audio(bool encrypted, const audio::buffer_t &plaintext, audio_packet_t &destination, crypto::aes_t &iv, crypto::cipher::cbc_t &cbc) {
// If encryption isn't enabled
if (!encrypted) {
std::copy(std::begin(plaintext), std::end(plaintext), destination->payload());
return plaintext.size();
}
crypto::aes_t iv(16);
*(std::uint32_t *) iv.data() = util::endian::big<std::uint32_t>(avRiKeyIv + destination->rtp.sequenceNumber);
return cbc.encrypt(std::string_view { (char *) std::begin(plaintext), plaintext.size() }, destination->payload(), &iv);
}
@ -401,6 +398,8 @@ namespace stream {
struct {
crypto::cipher::gcm_t cipher;
crypto::aes_t legacy_input_enc_iv; // Only used when the client doesn't support full control stream encryption
crypto::aes_t incoming_iv;
crypto::aes_t outgoing_iv;
std::uint32_t connect_data; // Used for new clients with ML_FF_SESSION_ID_V1
std::string expected_peer_address; // Only used for legacy clients without ML_FF_SESSION_ID_V1
@ -437,7 +436,7 @@ namespace stream {
auto seq = session->control.seq++;
crypto::aes_t iv;
auto &iv = session->control.outgoing_iv;
if (session->config.encryptionFlagsEnabled & SS_ENC_CONTROL_V2) {
// We use the deterministic IV construction algorithm specified in NIST SP 800-38D
// Section 8.2.1. The sequence number is our "invocation" field and the 'CH' in the
@ -977,7 +976,7 @@ namespace stream {
std::string_view tagged_cipher { (char *) header->payload(), (size_t) tagged_cipher_length };
auto &cipher = session->control.cipher;
crypto::aes_t iv;
auto &iv = session->control.incoming_iv;
if (session->config.encryptionFlagsEnabled & SS_ENC_CONTROL_V2) {
// We use the deterministic IV construction algorithm specified in NIST SP 800-38D
// Section 8.2.1. The sequence number is our "invocation" field and the 'CC' in the
@ -1241,6 +1240,7 @@ namespace stream {
platf::adjust_thread_priority(platf::thread_priority_e::high);
stat_trackers::min_max_avg_tracker<uint16_t> frame_processing_latency_tracker;
crypto::aes_t iv(12);
while (auto packet = packets->pop()) {
if (shutdown_event->peek()) {
@ -1413,7 +1413,6 @@ namespace stream {
//
// The IV counter is 64 bits long which allows for 2^64 encrypted video packets
// to be sent to each client before the IV repeats.
crypto::aes_t iv(12);
std::copy_n((uint8_t *) &session->video.gcm_iv_counter, sizeof(session->video.gcm_iv_counter), std::begin(iv));
iv[11] = 'V'; // Video stream
session->video.gcm_iv_counter++;
@ -1486,6 +1485,7 @@ namespace stream {
audio_packet_t audio_packet { (audio_packet_raw_t *) malloc(sizeof(audio_packet_raw_t) + max_block_size) };
fec::rs_t rs { reed_solomon_new(RTPA_DATA_SHARDS, RTPA_FEC_SHARDS) };
crypto::aes_t iv(16);
// For unknown reasons, the RS parity matrix computed by our RS implementation
// doesn't match the one Nvidia uses for audio data. I'm not exactly sure why,
@ -1513,11 +1513,9 @@ namespace stream {
auto sequenceNumber = session->audio.sequenceNumber;
auto timestamp = session->audio.timestamp;
// This will be mapped to big-endianness later
// For now, encode_audio needs it to be the proper sequenceNumber
audio_packet->rtp.sequenceNumber = sequenceNumber;
*(std::uint32_t *) iv.data() = util::endian::big<std::uint32_t>(session->audio.avRiKeyId + sequenceNumber);
auto bytes = encode_audio(session->config.encryptionFlagsEnabled & SS_ENC_AUDIO, packet_data, audio_packet, session->audio.avRiKeyId, session->audio.cipher);
auto bytes = encode_audio(session->config.encryptionFlagsEnabled & SS_ENC_AUDIO, packet_data, audio_packet, iv, session->audio.cipher);
if (bytes < 0) {
BOOST_LOG(error) << "Couldn't encode audio packet"sv;
break;