Adversarial C2 Protocol Implemented in Zig
Simplify tagged union
+9
-15
+9
-15
src/message.zig
+9
-15
src/message.zig
···
1
-
/// Type tag for Message union.
2
-
/// This is the first value in the actual packet sent over the network.
3
-
pub const PacketType = enum(u16) {
4
-
relay = 0x003C,
5
-
connection = 0x00E9,
6
-
_,
7
-
};
8
-
9
1
pub const MessageTypeError = error{
10
2
NotImplementedSaprusType,
11
3
UnknownSaprusType,
···
16
8
17
9
const message = @This();
18
10
19
-
pub const Message = union(PacketType) {
20
-
relay: Message.Relay,
21
-
connection: Message.Connection,
11
+
pub const Message = union(enum(u16)) {
12
+
relay: Message.Relay = 0x003C,
13
+
connection: Message.Connection = 0x00E9,
14
+
_,
22
15
23
16
pub const Relay = message.Relay;
24
17
pub const Connection = message.Connection;
25
18
26
19
pub fn toBytes(self: message.Message, buf: []u8) []u8 {
27
20
return switch (self) {
28
-
inline else => |m| m.toBytes(buf),
21
+
inline .relay, .connection => |m| m.toBytes(buf),
22
+
else => unreachable,
29
23
};
30
24
}
31
25
···
36
30
37
31
pub fn parse(bytes: []const u8) MessageParseError!Message {
38
32
var in: Reader = .fixed(bytes);
39
-
const @"type" = in.takeEnum(PacketType, .big) catch |err| switch (err) {
33
+
const @"type" = in.takeEnum(std.meta.Tag(Message), .big) catch |err| switch (err) {
40
34
error.InvalidEnumTag => return error.UnknownSaprusType,
41
35
else => return error.InvalidMessage,
42
36
};
···
124
118
/// Asserts that buf is large enough to fit the relay message.
125
119
pub fn toBytes(self: Relay, buf: []u8) []u8 {
126
120
var out: Writer = .fixed(buf);
127
-
out.writeInt(u16, @intFromEnum(PacketType.relay), .big) catch unreachable;
121
+
out.writeInt(u16, @intFromEnum(Message.relay), .big) catch unreachable;
128
122
out.writeInt(u16, @intCast(self.payload.len + 4), .big) catch unreachable; // Length field, but unread. Will switch to checksum
129
123
out.writeAll(&self.dest.bytes) catch unreachable;
130
124
out.writeAll(self.payload) catch unreachable;
···
178
172
/// Asserts that buf is large enough to fit the connection message.
179
173
pub fn toBytes(self: Connection, buf: []u8) []u8 {
180
174
var out: Writer = .fixed(buf);
181
-
out.writeInt(u16, @intFromEnum(PacketType.connection), .big) catch unreachable;
175
+
out.writeInt(u16, @intFromEnum(Message.connection), .big) catch unreachable;
182
176
out.writeInt(u16, @intCast(self.payload.len + 14), .big) catch unreachable; // Saprus length field, unread.
183
177
out.writeInt(u16, self.src, .big) catch unreachable;
184
178
out.writeInt(u16, self.dest, .big) catch unreachable;