7200959802

Committed 13 Dec 2023 08:57PM UTC coverage: 64.045% (+0.02%) from 64.026%

Build # 7200959802

Build Type

Pull #1256

github

Committed by

web-flow

Commit Message

Merge 1b2c60279 into 2e6937052

Pull Request Pull Request #1256: Fix issue with drop count calculations

Run Details

4395 of 8958 branches covered (0.0%)

Branch coverage included in aggregate %.

28 of 28 new or added lines in 2 files covered. (100.0%)

2 existing lines in 1 file now uncovered.

15288 of 21775 relevant lines covered (70.21%)

16028.33 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

73.2

/Source/santad/Logs/EndpointSecurity/Writers/Spool.mm

/// Copyright 2022 Google LLC
///
/// 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
///
///     https://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.

#import "Source/santad/Logs/EndpointSecurity/Writers/Spool.h"

#import "Source/common/SNTLogging.h"
#include "Source/common/santa_proto_include_wrapper.h"
#include "absl/strings/string_view.h"

static const char *kTypeGoogleApisComPrefix = "type.googleapis.com/";

namespace santa::santad::logs::endpoint_security::writers {

std::shared_ptr<Spool> Spool::Create(std::string_view base_dir, size_t max_spool_disk_size,
                                     size_t max_spool_batch_size, uint64_t flush_timeout_ms) {
  dispatch_queue_t q = dispatch_queue_create("com.google.santa.daemon.file_base_q",
                                             DISPATCH_QUEUE_SERIAL_WITH_AUTORELEASE_POOL);
  dispatch_source_t timer_source = dispatch_source_create(DISPATCH_SOURCE_TYPE_TIMER, 0, 0, q);
  dispatch_source_set_timer(timer_source, dispatch_time(DISPATCH_TIME_NOW, 0),
                            NSEC_PER_MSEC * flush_timeout_ms, 0);

  auto spool_writer =
    std::make_shared<Spool>(q, timer_source, base_dir, max_spool_disk_size, max_spool_batch_size);

  spool_writer->BeginFlushTask();

  return spool_writer;
}

// Note: The `log_batch_writer_` has the batch size set to SIZE_T_MAX. This is because
// the decision on whether or not to flush is controlled by the Spool class here based
// on a "size of bytes" threshold, not "count of records" threshold used by the
// FsSpoolLogBatchWriter. As such, calling `FsSpoolLogBatchWriter::WriteMessage`
// should never flush.
Spool::Spool(dispatch_queue_t q, dispatch_source_t timer_source, std::string_view base_dir,
             size_t max_spool_disk_size, size_t max_spool_file_size, void (^write_complete_f)(void),
             void (^flush_task_complete_f)(void))
    : q_(q),
      timer_source_(timer_source),
      spool_writer_(absl::string_view(base_dir.data(), base_dir.length()), max_spool_disk_size),
      log_batch_writer_(&spool_writer_, SIZE_T_MAX),
      spool_file_size_threshold_(max_spool_file_size),
      spool_file_size_threshold_leniency_(spool_file_size_threshold_ *
                                          spool_file_size_threshold_leniency_factor_),
      write_complete_f_(write_complete_f),
      flush_task_complete_f_(flush_task_complete_f) {
  type_url_ = kTypeGoogleApisComPrefix + ::santa::pb::v1::SantaMessage::descriptor()->full_name();
}

Spool::~Spool() {
  // Note: `log_batch_writer_` is automatically flushed when destroyed
  if (!flush_task_started_) {
    // The timer_source_ must be resumed to ensure it has a proper retain count before being
    // destroyed. Additionally, it should first be cancelled to ensure the timer isn't ever fired
    // (see man page for `dispatch_source_cancel(3)`).
    dispatch_source_cancel(timer_source_);
    dispatch_resume(timer_source_);
  }
}

void Spool::BeginFlushTask() {
  if (flush_task_started_) {
    return;
  }

  std::weak_ptr<Spool> weak_writer = weak_from_this();
  dispatch_source_set_event_handler(timer_source_, ^{
    std::shared_ptr<Spool> shared_writer = weak_writer.lock();
    if (!shared_writer) {
      return;
    }

    if (!shared_writer->FlushLocked()) {
      LOGE(@"Spool writer: periodic flush failed.");
    }

    if (shared_writer->flush_task_complete_f_) {
      shared_writer->flush_task_complete_f_();
    }
  });

  dispatch_resume(timer_source_);
  flush_task_started_ = true;
}

void Spool::Flush() {
  dispatch_sync(q_, ^{
    FlushLocked();
  });
}

bool Spool::FlushLocked() {
  if (log_batch_writer_.Flush().ok()) {
    accumulated_bytes_ = 0;
    return true;
  } else {
    return false;
  }
}

void Spool::Write(std::vector<uint8_t> &&bytes) {
  auto shared_this = shared_from_this();

  // Workaround to move `bytes` into the block without a copy
  __block std::vector<uint8_t> temp_bytes = std::move(bytes);

  dispatch_async(q_, ^{
    std::vector<uint8_t> moved_bytes = std::move(temp_bytes);

    // Manually pack an `Any` with a pre-serialized SantaMessage
    google::protobuf::Any any;
#if SANTA_OPEN_SOURCE
    any.set_value(moved_bytes.data(), moved_bytes.size());
#else
    any.set_value(absl::string_view((const char*)moved_bytes.data(), moved_bytes.size()));
#endif
    any.set_type_url(type_url_);

    if (shared_this->accumulated_bytes_ >= shared_this->spool_file_size_threshold_) {
      shared_this->FlushLocked();
    }

    // Only write the new message if we have room left.
    // This will account for Flush failing above.
    // Use the more lenient threshold here in case the Flush failures are transitory.
    if (shared_this->accumulated_bytes_ < shared_this->spool_file_size_threshold_leniency_) {
      auto status = shared_this->log_batch_writer_.WriteMessage(any);
      if (!status.ok()) {
        LOGE(@"ProtoEventLogger::LogProto failed with: %s", status.ToString().c_str());
      }

      shared_this->accumulated_bytes_ += moved_bytes.size();
    }

    if (shared_this->write_complete_f_) {
      shared_this->write_complete_f_();
    }
  });
}

}  // namespace santa::santad::logs::endpoint_security::writers

1	/// Copyright 2022 Google LLC
2	///
3	/// Licensed under the Apache License, Version 2.0 (the "License");
4	/// you may not use this file except in compliance with the License.
5	/// You may obtain a copy of the License at
6	///
7	/// https://www.apache.org/licenses/LICENSE-2.0
8	///
9	/// Unless required by applicable law or agreed to in writing, software
10	/// distributed under the License is distributed on an "AS IS" BASIS,
11	/// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	/// See the License for the specific language governing permissions and
13	/// limitations under the License.
14
15	#import "Source/santad/Logs/EndpointSecurity/Writers/Spool.h"
16
17	#import "Source/common/SNTLogging.h"
18	#include "Source/common/santa_proto_include_wrapper.h"
19	#include "absl/strings/string_view.h"
20
21	static const char *kTypeGoogleApisComPrefix = "type.googleapis.com/";
22
23	namespace santa::santad::logs::endpoint_security::writers {
24
25	std::shared_ptr<Spool> Spool::Create(std::string_view base_dir, size_t max_spool_disk_size,
26	size_t max_spool_batch_size, uint64_t flush_timeout_ms) {	1✔
27	dispatch_queue_t q = dispatch_queue_create("com.google.santa.daemon.file_base_q",	1✔
28	DISPATCH_QUEUE_SERIAL_WITH_AUTORELEASE_POOL);	1✔
29	dispatch_source_t timer_source = dispatch_source_create(DISPATCH_SOURCE_TYPE_TIMER, 0, 0, q);	1✔
30	dispatch_source_set_timer(timer_source, dispatch_time(DISPATCH_TIME_NOW, 0),	1✔
31	NSEC_PER_MSEC * flush_timeout_ms, 0);	1✔
32
33	auto spool_writer =	1✔
34	std::make_shared<Spool>(q, timer_source, base_dir, max_spool_disk_size, max_spool_batch_size);	1✔
35
36	spool_writer->BeginFlushTask();	1✔
37
38	return spool_writer;	1✔
39	}	1✔
40
41	// Note: The `log_batch_writer_` has the batch size set to SIZE_T_MAX. This is because
42	// the decision on whether or not to flush is controlled by the Spool class here based
43	// on a "size of bytes" threshold, not "count of records" threshold used by the
44	// FsSpoolLogBatchWriter. As such, calling `FsSpoolLogBatchWriter::WriteMessage`
45	// should never flush.
46	Spool::Spool(dispatch_queue_t q, dispatch_source_t timer_source, std::string_view base_dir,
47	size_t max_spool_disk_size, size_t max_spool_file_size, void (^write_complete_f)(void),
48	void (^flush_task_complete_f)(void))
49	: q_(q),
50	timer_source_(timer_source),
51	spool_writer_(absl::string_view(base_dir.data(), base_dir.length()), max_spool_disk_size),
52	log_batch_writer_(&spool_writer_, SIZE_T_MAX),
53	spool_file_size_threshold_(max_spool_file_size),
54	spool_file_size_threshold_leniency_(spool_file_size_threshold_ *
55	spool_file_size_threshold_leniency_factor_),
56	write_complete_f_(write_complete_f),
57	flush_task_complete_f_(flush_task_complete_f) {	3✔
58	type_url_ = kTypeGoogleApisComPrefix + ::santa::pb::v1::SantaMessage::descriptor()->full_name();	3✔
59	}	3✔
60
61	Spool::~Spool() {	3✔
62	// Note: `log_batch_writer_` is automatically flushed when destroyed
63	if (!flush_task_started_) {	3✔
64	// The timer_source_ must be resumed to ensure it has a proper retain count before being
65	// destroyed. Additionally, it should first be cancelled to ensure the timer isn't ever fired
66	// (see man page for `dispatch_source_cancel(3)`).
67	dispatch_source_cancel(timer_source_);	1✔
68	dispatch_resume(timer_source_);	1✔
69	}	1✔
70	}	3✔
71
72	void Spool::BeginFlushTask() {	2✔
73	if (flush_task_started_) {	2!
74	return;	×
75	}	×
76
77	std::weak_ptr<Spool> weak_writer = weak_from_this();	2✔
78	dispatch_source_set_event_handler(timer_source_, ^{	4✔
79	std::shared_ptr<Spool> shared_writer = weak_writer.lock();	4✔
80	if (!shared_writer) {	4!
UNCOV 81	return;	×
UNCOV 82	}	×
83
84	if (!shared_writer->FlushLocked()) {	4!
85	LOGE(@"Spool writer: periodic flush failed.");	×
86	}	×
87
88	if (shared_writer->flush_task_complete_f_) {	4✔
89	shared_writer->flush_task_complete_f_();	2✔
90	}	2✔
91	});	4✔
92
93	dispatch_resume(timer_source_);	2✔
94	flush_task_started_ = true;	2✔
95	}	2✔
96
97	void Spool::Flush() {	×
98	dispatch_sync(q_, ^{	×
99	FlushLocked();	×
100	});	×
101	}	×
102
103	bool Spool::FlushLocked() {	5✔
104	if (log_batch_writer_.Flush().ok()) {	5!
105	accumulated_bytes_ = 0;	5✔
106	return true;	5✔
107	} else {	×
108	return false;	×
109	}	×
110	}	5✔
111
112	void Spool::Write(std::vector<uint8_t> &&bytes) {	2✔
113	auto shared_this = shared_from_this();	2✔
114
115	// Workaround to move `bytes` into the block without a copy
116	__block std::vector<uint8_t> temp_bytes = std::move(bytes);	2✔
117
118	dispatch_async(q_, ^{	2✔
119	std::vector<uint8_t> moved_bytes = std::move(temp_bytes);	2✔
120
121	// Manually pack an `Any` with a pre-serialized SantaMessage
122	google::protobuf::Any any;	2✔
123	#if SANTA_OPEN_SOURCE	2✔
124	any.set_value(moved_bytes.data(), moved_bytes.size());	2✔
125	#else
126	any.set_value(absl::string_view((const char*)moved_bytes.data(), moved_bytes.size()));
127	#endif
128	any.set_type_url(type_url_);	2✔
129
130	if (shared_this->accumulated_bytes_ >= shared_this->spool_file_size_threshold_) {	2!
131	shared_this->FlushLocked();	×
132	}	×
133
134	// Only write the new message if we have room left.
135	// This will account for Flush failing above.
136	// Use the more lenient threshold here in case the Flush failures are transitory.
137	if (shared_this->accumulated_bytes_ < shared_this->spool_file_size_threshold_leniency_) {	2!
138	auto status = shared_this->log_batch_writer_.WriteMessage(any);	2✔
139	if (!status.ok()) {	2!
140	LOGE(@"ProtoEventLogger::LogProto failed with: %s", status.ToString().c_str());	×
141	}	×
142
143	shared_this->accumulated_bytes_ += moved_bytes.size();	2✔
144	}	2✔
145
146	if (shared_this->write_complete_f_) {	2!
147	shared_this->write_complete_f_();	2✔
148	}	2✔
149	});	2✔
150	}	2✔
151
152	} // namespace santa::santad::logs::endpoint_security::writers

google / santa / 7200959802

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous