justme
2e146d4da7
ci/woodpecker/push/main-amd64-build Pipeline was successful
Namespaces, assemblies, solution, binaries, images, systemd units, config paths, environment variables, protocol methods and the Linux user are all renamed. Breaking: existing databases, node identities and installed units must be recreated.
346 lines
11 KiB
C#
346 lines
11 KiB
C#
using System.Buffers.Binary;
|
|
using System.IO.Compression;
|
|
using System.Text.Json;
|
|
using Fleet.Core;
|
|
|
|
namespace MoShiro.Protocol;
|
|
|
|
/// <summary>
|
|
/// RawChunk 的编解码（Probe、Agent、Backend、离线工具共用同一份代码）。
|
|
/// 布局：<c>"MSRC" | u8 container_version | u32 crc32c(压缩数据) | Brotli( u32 header_len | header JSON | 各列数据 )</c>。
|
|
/// 列数据：可空列先写有效位图，再写有效值；plain 为无符号 varint（有符号类型用 zigzag），delta 为与上一个有效值之差的 zigzag varint。
|
|
/// </summary>
|
|
public static class RawChunkCodec
|
|
{
|
|
private const int PrefixBytes = 4 + 1 + 4;
|
|
|
|
private static ReadOnlySpan<byte> Magic => "MSRC"u8;
|
|
|
|
/// <summary>压缩之前的正文：<c>u32 header_len | header JSON | 各列数据</c>。布局变化时黄金字节测试会失败，提醒升级版本。</summary>
|
|
public static byte[] EncodeBody(RawChunk chunk)
|
|
{
|
|
RawChunkHeader header = chunk.Header;
|
|
header.ContainerVersion = ProtocolVersion.Container;
|
|
header.Streams = [];
|
|
foreach (RawStream stream in chunk.Streams)
|
|
{
|
|
var streamHeader = new StreamHeader { Name = stream.Name, Count = stream.Count };
|
|
foreach (RawColumn column in stream.Columns)
|
|
{
|
|
streamHeader.Columns.Add(new ColumnHeader
|
|
{
|
|
Name = column.Name,
|
|
Type = TypeName(column.Type),
|
|
Enc = column.Enc == ColumnEnc.Delta ? "delta" : "plain",
|
|
Nullable = column.Nullable,
|
|
Role = RoleName(column.Role),
|
|
});
|
|
}
|
|
|
|
header.Streams.Add(streamHeader);
|
|
}
|
|
|
|
byte[] headerJson = JsonSerializer.SerializeToUtf8Bytes(header, MoShiroJsonContext.Default.RawChunkHeader);
|
|
|
|
using var body = new MemoryStream();
|
|
Span<byte> length = stackalloc byte[4];
|
|
BinaryPrimitives.WriteUInt32LittleEndian(length, (uint)headerJson.Length);
|
|
body.Write(length);
|
|
body.Write(headerJson);
|
|
foreach (RawStream stream in chunk.Streams)
|
|
{
|
|
foreach (RawColumn column in stream.Columns)
|
|
{
|
|
WriteColumn(body, column);
|
|
}
|
|
}
|
|
|
|
return body.ToArray();
|
|
}
|
|
|
|
public static byte[] Encode(RawChunk chunk)
|
|
{
|
|
byte[] plain = EncodeBody(chunk);
|
|
using var compressed = new MemoryStream();
|
|
using (var brotli = new BrotliStream(compressed, CompressionLevel.Optimal, leaveOpen: true))
|
|
{
|
|
brotli.Write(plain);
|
|
}
|
|
|
|
byte[] payload = compressed.ToArray();
|
|
byte[] result = new byte[PrefixBytes + payload.Length];
|
|
Magic.CopyTo(result);
|
|
result[4] = ProtocolVersion.Container;
|
|
BinaryPrimitives.WriteUInt32LittleEndian(result.AsSpan(5), Crc32c.Compute(payload));
|
|
payload.CopyTo(result, PrefixBytes);
|
|
return result;
|
|
}
|
|
|
|
public static Result<RawChunk> Decode(ReadOnlySpan<byte> data)
|
|
{
|
|
if (data.Length < PrefixBytes || !data[..4].SequenceEqual(Magic))
|
|
{
|
|
return Corrupt("bad magic");
|
|
}
|
|
|
|
if (data[4] != ProtocolVersion.Container)
|
|
{
|
|
return Corrupt($"unsupported container version {data[4]}");
|
|
}
|
|
|
|
ReadOnlySpan<byte> payload = data[PrefixBytes..];
|
|
if (BinaryPrimitives.ReadUInt32LittleEndian(data[5..]) != Crc32c.Compute(payload))
|
|
{
|
|
return Corrupt("crc mismatch");
|
|
}
|
|
|
|
byte[]? plain = Inflate(payload);
|
|
if (plain is null)
|
|
{
|
|
return Corrupt("invalid or oversized compressed data");
|
|
}
|
|
|
|
if (plain.Length < 4)
|
|
{
|
|
return Corrupt("truncated header length");
|
|
}
|
|
|
|
uint headerLength = BinaryPrimitives.ReadUInt32LittleEndian(plain);
|
|
if (headerLength > plain.Length - 4)
|
|
{
|
|
return Corrupt("header length exceeds data");
|
|
}
|
|
|
|
RawChunkHeader? header;
|
|
try
|
|
{
|
|
header = JsonSerializer.Deserialize(plain.AsSpan(4, (int)headerLength), MoShiroJsonContext.Default.RawChunkHeader);
|
|
}
|
|
catch (JsonException)
|
|
{
|
|
return Corrupt("header is not valid json");
|
|
}
|
|
|
|
if (header is null || header.Streams.Count > ProtocolLimits.MaxStreamsPerChunk)
|
|
{
|
|
return Corrupt("header is missing or has too many streams");
|
|
}
|
|
|
|
ReadOnlySpan<byte> columns = plain.AsSpan(4 + (int)headerLength);
|
|
var streams = new List<RawStream>(header.Streams.Count);
|
|
foreach (StreamHeader streamHeader in header.Streams)
|
|
{
|
|
if (streamHeader.Count is < 0 or > ProtocolLimits.MaxRecordsPerChunk || streamHeader.Columns.Count > ProtocolLimits.MaxColumnsPerStream)
|
|
{
|
|
return Corrupt($"stream {streamHeader.Name} has an invalid size");
|
|
}
|
|
|
|
var cols = new List<RawColumn>(streamHeader.Columns.Count);
|
|
foreach (ColumnHeader columnHeader in streamHeader.Columns)
|
|
{
|
|
if (!TryParseType(columnHeader.Type, out ColumnType type) || !TryParseEnc(columnHeader.Enc, out ColumnEnc enc))
|
|
{
|
|
return Corrupt($"column {columnHeader.Name} has an unknown type or encoding");
|
|
}
|
|
|
|
RawColumn? column = ReadColumn(ref columns, columnHeader, type, enc, streamHeader.Count);
|
|
if (column is null)
|
|
{
|
|
return Corrupt($"column {streamHeader.Name}.{columnHeader.Name} is truncated or out of range");
|
|
}
|
|
|
|
cols.Add(column);
|
|
}
|
|
|
|
streams.Add(new RawStream(streamHeader.Name, streamHeader.Count, cols));
|
|
}
|
|
|
|
return columns.Length == 0 ? Result<RawChunk>.Ok(new RawChunk(header, streams)) : Corrupt("trailing bytes");
|
|
}
|
|
|
|
private static Result<RawChunk> Corrupt(string message) =>
|
|
Result<RawChunk>.Fail(MoShiroErrors.RawCorrupt, message);
|
|
|
|
private static byte[]? Inflate(ReadOnlySpan<byte> payload)
|
|
{
|
|
using var input = new MemoryStream(payload.ToArray());
|
|
using var brotli = new BrotliStream(input, CompressionMode.Decompress);
|
|
using var output = new MemoryStream();
|
|
byte[] buffer = new byte[16 * 1024];
|
|
try
|
|
{
|
|
int read;
|
|
while ((read = brotli.Read(buffer, 0, buffer.Length)) > 0)
|
|
{
|
|
output.Write(buffer, 0, read);
|
|
if (output.Length > ProtocolLimits.MaxDecodedChunkBytes)
|
|
{
|
|
return null;
|
|
}
|
|
}
|
|
}
|
|
catch (InvalidDataException)
|
|
{
|
|
return null;
|
|
}
|
|
|
|
return output.ToArray();
|
|
}
|
|
|
|
private static void WriteColumn(MemoryStream body, RawColumn column)
|
|
{
|
|
int count = column.Count;
|
|
if (column.Nullable)
|
|
{
|
|
byte[] bitmap = new byte[(count + 7) / 8];
|
|
for (int i = 0; i < count; i++)
|
|
{
|
|
if (column.IsValid(i))
|
|
{
|
|
bitmap[i >> 3] |= (byte)(1 << (i & 7));
|
|
}
|
|
}
|
|
|
|
body.Write(bitmap);
|
|
}
|
|
|
|
bool signed = column.Type is ColumnType.I32 or ColumnType.I64;
|
|
long previous = 0;
|
|
for (int i = 0; i < count; i++)
|
|
{
|
|
if (!column.IsValid(i))
|
|
{
|
|
continue;
|
|
}
|
|
|
|
long value = column.Values[i];
|
|
if (column.Enc == ColumnEnc.Delta)
|
|
{
|
|
VarInt.Write(body, VarInt.ZigZag(value - previous));
|
|
previous = value;
|
|
}
|
|
else
|
|
{
|
|
VarInt.Write(body, signed ? VarInt.ZigZag(value) : (ulong)value);
|
|
}
|
|
}
|
|
}
|
|
|
|
private static RawColumn? ReadColumn(ref ReadOnlySpan<byte> data, ColumnHeader header, ColumnType type, ColumnEnc enc, int count)
|
|
{
|
|
bool[]? valid = null;
|
|
if (header.Nullable)
|
|
{
|
|
int bitmapBytes = (count + 7) / 8;
|
|
if (data.Length < bitmapBytes)
|
|
{
|
|
return null;
|
|
}
|
|
|
|
valid = new bool[count];
|
|
for (int i = 0; i < count; i++)
|
|
{
|
|
valid[i] = (data[i >> 3] & (1 << (i & 7))) != 0;
|
|
}
|
|
|
|
data = data[bitmapBytes..];
|
|
}
|
|
|
|
bool signed = type is ColumnType.I32 or ColumnType.I64;
|
|
long[] values = new long[count];
|
|
long previous = 0;
|
|
for (int i = 0; i < count; i++)
|
|
{
|
|
if (valid is not null && !valid[i])
|
|
{
|
|
continue;
|
|
}
|
|
|
|
if (!VarInt.TryRead(data, out ulong raw, out int consumed))
|
|
{
|
|
return null;
|
|
}
|
|
|
|
data = data[consumed..];
|
|
long value;
|
|
if (enc == ColumnEnc.Delta)
|
|
{
|
|
value = previous + VarInt.UnZigZag(raw);
|
|
previous = value;
|
|
}
|
|
else
|
|
{
|
|
value = signed ? VarInt.UnZigZag(raw) : unchecked((long)raw);
|
|
}
|
|
|
|
if (!InRange(type, value))
|
|
{
|
|
return null;
|
|
}
|
|
|
|
values[i] = value;
|
|
}
|
|
|
|
return new RawColumn(header.Name, type, enc, header.Nullable, ParseRole(header.Role), values, valid);
|
|
}
|
|
|
|
private static bool InRange(ColumnType type, long value) => type switch
|
|
{
|
|
ColumnType.U8 => value is >= 0 and <= byte.MaxValue,
|
|
ColumnType.U16 => value is >= 0 and <= ushort.MaxValue,
|
|
ColumnType.U32 => value is >= 0 and <= uint.MaxValue,
|
|
ColumnType.I32 => value is >= int.MinValue and <= int.MaxValue,
|
|
_ => true,
|
|
};
|
|
|
|
public static string TypeName(ColumnType type) => type switch
|
|
{
|
|
ColumnType.U8 => "u8",
|
|
ColumnType.U16 => "u16",
|
|
ColumnType.U32 => "u32",
|
|
ColumnType.U64 => "u64",
|
|
ColumnType.I32 => "i32",
|
|
_ => "i64",
|
|
};
|
|
|
|
private static bool TryParseType(string text, out ColumnType type)
|
|
{
|
|
switch (text)
|
|
{
|
|
case "u8": type = ColumnType.U8; return true;
|
|
case "u16": type = ColumnType.U16; return true;
|
|
case "u32": type = ColumnType.U32; return true;
|
|
case "u64": type = ColumnType.U64; return true;
|
|
case "i32": type = ColumnType.I32; return true;
|
|
case "i64": type = ColumnType.I64; return true;
|
|
default: type = default; return false;
|
|
}
|
|
}
|
|
|
|
private static bool TryParseEnc(string text, out ColumnEnc enc)
|
|
{
|
|
switch (text)
|
|
{
|
|
case "plain": enc = ColumnEnc.Plain; return true;
|
|
case "delta": enc = ColumnEnc.Delta; return true;
|
|
default: enc = default; return false;
|
|
}
|
|
}
|
|
|
|
private static string? RoleName(ColumnRole role) => role switch
|
|
{
|
|
ColumnRole.TsProbe => "ts_probe",
|
|
ColumnRole.TsAgent => "ts_agent",
|
|
ColumnRole.TsLocal => "ts_local",
|
|
_ => null,
|
|
};
|
|
|
|
private static ColumnRole ParseRole(string? text) => text switch
|
|
{
|
|
"ts_probe" => ColumnRole.TsProbe,
|
|
"ts_agent" => ColumnRole.TsAgent,
|
|
"ts_local" => ColumnRole.TsLocal,
|
|
_ => ColumnRole.None,
|
|
};
|
|
}
|