| Alex Gough | 5033b85d | 2024-03-27 21:29:22 | [diff] [blame] | 1 | // Copyright 2024 The Chromium Authors |
| 2 | // Use of this source code is governed by a BSD-style license that can be |
| 3 | // found in the LICENSE file. |
| 4 | |
| 5 | #ifndef MOJO_PUBLIC_CPP_BASE_PROTO_WRAPPER_H_ |
| 6 | #define MOJO_PUBLIC_CPP_BASE_PROTO_WRAPPER_H_ |
| 7 | |
| 8 | #include <optional> |
| 9 | #include <string> |
| 10 | |
| 11 | #include "base/component_export.h" |
| 12 | #include "base/containers/span.h" |
| 13 | #include "base/gtest_prod_util.h" |
| 14 | #include "mojo/public/cpp/base/big_buffer.h" |
| 15 | #include "mojo/public/cpp/base/proto_wrapper_passkeys.h" |
| 16 | #include "mojo/public/cpp/bindings/default_construct_tag.h" |
| 17 | |
| 18 | namespace google::protobuf { |
| 19 | class MessageLite; |
| 20 | } // namespace google::protobuf |
| 21 | |
| 22 | namespace mojo_base { |
| 23 | |
| 24 | namespace mojom { |
| 25 | class ProtoWrapperDataView; |
| 26 | } // namespace mojom |
| 27 | |
| 28 | template <typename T> |
| 29 | concept IsProtoMessage = |
| 30 | !std::is_abstract<T>() && |
| 31 | std::is_base_of<google::protobuf::MessageLite, T>::value; |
| 32 | |
| 33 | class COMPONENT_EXPORT(MOJO_BASE_PROTOBUF_SUPPORT) ProtoWrapper { |
| 34 | public: |
| 35 | ProtoWrapper(ProtoWrapper&& other); |
| 36 | ProtoWrapper& operator=(ProtoWrapper&& other); |
| 37 | ProtoWrapper(const ProtoWrapper&) = delete; |
| 38 | ProtoWrapper& operator=(const ProtoWrapper&) = delete; |
| 39 | ~ProtoWrapper(); |
| 40 | |
| 41 | // Exposed so mojo can create this class. |
| 42 | explicit ProtoWrapper(mojo::DefaultConstruct::Tag passkey); |
| 43 | |
| 44 | // Construct from a protobuf message. May CHECK if the message fails |
| 45 | // serialization. Once constructed the message can be discarded. The typename |
| 46 | // of the message is stored along with the bytes that serialize the message, |
| 47 | // and must match the typename of the message this wrapper is deserialized to. |
| 48 | explicit ProtoWrapper(const google::protobuf::MessageLite& message); |
| 49 | |
| 50 | // Construct from an already serialized proto stream - only use this if the |
| 51 | // stream has come from an external source (e.g. the network, an OS service) |
| 52 | // and you need to get the message into the mojo IPC system with a typename. |
| 53 | // This constructor does not validate that the stream of bytes can populate |
| 54 | // its wrapped protobuf Message until an unwrapping is attempted. |
| 55 | // Makes a copy of the data in the provided span. |
| 56 | explicit ProtoWrapper(base::span<const uint8_t> data, |
| 57 | std::string type_name, |
| 58 | base::PassKey<ProtoWrapperBytes> passkey); |
| 59 | |
| 60 | template <IsProtoMessage ProtoMessage> |
| Joe Mason | a759465 | 2024-09-18 19:12:17 | [diff] [blame] | 61 | std::optional<ProtoMessage> As() const { |
| Alex Gough | 5033b85d | 2024-03-27 21:29:22 | [diff] [blame] | 62 | ProtoMessage message; |
| 63 | if (DeserializeToMessage(message)) { |
| 64 | return message; |
| 65 | } |
| 66 | return std::nullopt; |
| 67 | } |
| 68 | |
| 69 | // Access this to store the bytes somewhere else, or pass to another IPC |
| 70 | // system. The bytes may be mapped from a hostile process so while the size |
| 71 | // cannot change, the contents might. If you want to unpack the contained |
| 72 | // protobuf Message, use As<T>(); |
| 73 | std::optional<base::span<const uint8_t>> byte_span( |
| Joe Mason | a759465 | 2024-09-18 19:12:17 | [diff] [blame] | 74 | base::PassKey<ProtoWrapperBytes> passkey) const { |
| Alex Gough | 5033b85d | 2024-03-27 21:29:22 | [diff] [blame] | 75 | if (!is_valid()) { |
| 76 | return std::nullopt; |
| 77 | } |
| David Benjamin | 37b313c | 2024-04-05 05:00:03 | [diff] [blame] | 78 | return base::span(*bytes_); |
| Alex Gough | 5033b85d | 2024-03-27 21:29:22 | [diff] [blame] | 79 | } |
| 80 | |
| 81 | private: |
| 82 | friend struct mojo::StructTraits<mojo_base::mojom::ProtoWrapperDataView, |
| 83 | mojo_base::ProtoWrapper>; |
| 84 | FRIEND_TEST_ALL_PREFIXES(ProtoWrapperTest, TraitsOk); |
| 85 | FRIEND_TEST_ALL_PREFIXES(ProtoWrapperTest, LargeMessage); |
| 86 | FRIEND_TEST_ALL_PREFIXES(ProtoWrapperTest, TraitsEquivilentMessages); |
| 87 | FRIEND_TEST_ALL_PREFIXES(ProtoWrapperTest, TraitsDistinctMessages); |
| 88 | FRIEND_TEST_ALL_PREFIXES(ProtoWrapperTest, ToFromBytes); |
| 89 | |
| 90 | // Prevent creation of invalid wrappers. |
| 91 | ProtoWrapper(); |
| 92 | |
| 93 | // is_valid() implies the wrapper wraps some data - it does not mean that the |
| 94 | // bytes will deserialize to a valid message. |
| 95 | bool is_valid() const { return bytes_.has_value(); } |
| 96 | |
| Joe Mason | a759465 | 2024-09-18 19:12:17 | [diff] [blame] | 97 | bool DeserializeToMessage(google::protobuf::MessageLite& message) const; |
| Alex Gough | 5033b85d | 2024-03-27 21:29:22 | [diff] [blame] | 98 | |
| 99 | std::string proto_name_; |
| 100 | std::optional<BigBuffer> bytes_; |
| 101 | }; |
| 102 | |
| 103 | } // namespace mojo_base |
| 104 | |
| 105 | #endif // MOJO_PUBLIC_CPP_BASE_PROTO_WRAPPER_H_ |