Skip to content

Commit 0e2b592

Browse files
1358445256713584452567
committed
refactor: Refactor transports, USB reads, and buffer handling
Rework network and USB I/O to improve correctness and reduce allocations. - TcpTransport: reuse 8-byte buffers for message length read/write to avoid per-call allocations. - UdpTransport: major refactor of send/receive logic to handle sequencing, continuation fragments, retries, and errors; split send/receive into smaller helper methods and validate/assemble payloads robustly. - Linux/Windows/macOS/libusbdotnet USB backends: add ReadInto(byte[], offset, length) implementations that return actual byte counts and make Read(int) wrappers return exact-length arrays; add argument validation and chunked transfer handling to support large reads and avoid extra copies. - InternalDownloadDataBytes: use ArrayPool<byte>.Shared.Rent/Return for transfer buffer to reduce allocations and ensure proper return in finally. - InternalHandleResponse: avoid repeated string reallocations by tracking a pendingOffset and compacting the buffer when needed; update parsing to use offset-based substring operations and prevent incorrect prefix handling. Also includes various bounds checks, timeout handling, and clearer error reporting.
1 parent e659516 commit 0e2b592

File tree

8 files changed

+337
-91
lines changed

8 files changed

+337
-91
lines changed

‎FirmwareKit.Comm.Fastboot/Backend/Network/TcpTransport.cs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ public class TcpTransport : IFastbootBufferedTransport
88
{
99
private const int DefaultIoTimeoutMs = 30000;
1010
private readonly TcpClient _client = new();
11+
private readonly byte[] _readLenBuffer = new byte[8];
12+
private readonly byte[] _writeLenBuffer = new byte[8];
1113
private NetworkStream? _stream;
1214
private long _messageBytesLeft = 0;
1315

@@ -96,12 +98,11 @@ public int ReadInto(byte[] buffer, int offset, int length)
9698

9799
if (_messageBytesLeft == 0)
98100
{
99-
byte[] lenBuffer = new byte[8];
100-
if (ReadFully(lenBuffer, 0, 8) != 8)
101+
if (ReadFully(_readLenBuffer, 0, 8) != 8)
101102
{
102103
throw new Exception("Failed to read message length from TCP stream.");
103104
}
104-
_messageBytesLeft = BinaryPrimitives.ReadInt64BigEndian(lenBuffer);
105+
_messageBytesLeft = BinaryPrimitives.ReadInt64BigEndian(_readLenBuffer);
105106
}
106107

107108
int toRead = (int)Math.Min(length, _messageBytesLeft);
@@ -113,10 +114,9 @@ public int ReadInto(byte[] buffer, int offset, int length)
113114
public long Write(byte[] data, int length)
114115
{
115116
if (_stream == null) throw new InvalidOperationException("Stream not initialized");
116-
byte[] header = new byte[8];
117-
BinaryPrimitives.WriteInt64BigEndian(header, length);
117+
BinaryPrimitives.WriteInt64BigEndian(_writeLenBuffer, length);
118118

119-
_stream.Write(header, 0, 8);
119+
_stream.Write(_writeLenBuffer, 0, 8);
120120
_stream.Write(data, 0, length);
121121
_stream.Flush();
122122
return length;

‎FirmwareKit.Comm.Fastboot/Backend/Network/UdpTransport.cs‎

Lines changed: 179 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -209,24 +209,193 @@ public int ReadInto(byte[] buffer, int offset, int length)
209209
throw new ArgumentOutOfRangeException(nameof(length));
210210
}
211211

212-
byte[] response = SendDataInternal(PacketId.Fastboot, [], 0, _maxTransmissionAttempts);
213-
if (response.Length > length)
212+
int read = ReceiveFastbootInto(buffer, offset, length, _maxTransmissionAttempts, out ushort nextSeq);
213+
_sequence = nextSeq;
214+
return read;
215+
}
216+
217+
private int ReceiveFastbootInto(byte[] output, int outputOffset, int outputLength, int attempts, out ushort nextSeq)
218+
{
219+
PacketId currentId = PacketId.Fastboot;
220+
PacketFlag currentFlag = PacketFlag.None;
221+
ushort currentSeq = (ushort)_sequence;
222+
int written = 0;
223+
224+
while (true)
214225
{
215-
throw new Exception("UDP protocol error: receive overflow, target sent too much fastboot data.");
216-
}
226+
byte[] packet = new byte[HeaderSize];
227+
packet[0] = (byte)currentId;
228+
packet[1] = (byte)currentFlag;
229+
BinaryPrimitives.WriteUInt16BigEndian(packet.AsSpan(2, 2), currentSeq);
230+
231+
bool gotValidResponse = false;
232+
for (int i = 0; i < attempts; i++)
233+
{
234+
_client.Send(packet, packet.Length, _endpoint);
235+
236+
try
237+
{
238+
while (true)
239+
{
240+
IPEndPoint from = new(IPAddress.Any, 0);
241+
byte[] rxPacket = _client.Receive(ref from);
242+
if (rxPacket.Length < HeaderSize) continue;
243+
244+
ushort responseSeq = BinaryPrimitives.ReadUInt16BigEndian(rxPacket.AsSpan(2, 2));
245+
byte responseId = rxPacket[0];
246+
if (responseSeq != currentSeq) continue;
247+
if (responseId != (byte)currentId && responseId != (byte)PacketId.Error) continue;
248+
249+
if (responseId == (byte)PacketId.Error)
250+
{
251+
throw new Exception("Target returned error response.");
252+
}
253+
254+
int payloadLen = rxPacket.Length - HeaderSize;
255+
if (payloadLen > 0)
256+
{
257+
if (written + payloadLen > outputLength)
258+
{
259+
throw new Exception("UDP protocol error: receive overflow, target sent too much fastboot data.");
260+
}
261+
262+
Buffer.BlockCopy(rxPacket, HeaderSize, output, outputOffset + written, payloadLen);
263+
written += payloadLen;
264+
}
265+
266+
gotValidResponse = true;
267+
currentSeq = (ushort)((currentSeq + 1) & 0xFFFF);
268+
269+
bool continuation = (rxPacket[1] & (byte)PacketFlag.Continuation) != 0;
270+
if (!continuation)
271+
{
272+
nextSeq = currentSeq;
273+
return written;
274+
}
275+
276+
// Prompt the target for the next continuation fragment.
277+
currentId = (PacketId)responseId;
278+
currentFlag = PacketFlag.None;
279+
break;
280+
}
281+
}
282+
catch (SocketException ex) when (ex.SocketErrorCode == SocketError.TimedOut)
283+
{
284+
continue;
285+
}
286+
287+
if (gotValidResponse)
288+
{
289+
break;
290+
}
291+
}
217292

218-
Buffer.BlockCopy(response, 0, buffer, offset, response.Length);
219-
return response.Length;
293+
if (!gotValidResponse)
294+
{
295+
throw new Exception($"Failed to receive response after {attempts} attempts.");
296+
}
297+
}
220298
}
221299

222300
public long Write(byte[] data, int length)
223301
{
224-
byte[] response = SendDataInternal(PacketId.Fastboot, data, length, _maxTransmissionAttempts);
225-
if (response.Length > 0)
302+
SendFastbootNoPayload(PacketId.Fastboot, data, length, _maxTransmissionAttempts);
303+
return length;
304+
}
305+
306+
private void SendFastbootNoPayload(PacketId id, byte[] txData, int txLength, int attempts)
307+
{
308+
int offset = 0;
309+
310+
do
226311
{
227-
throw new Exception("UDP protocol error: target sent fastboot data out-of-turn.");
312+
int chunkLen = Math.Min(txLength - offset, _maxDataLength);
313+
PacketFlag flag = (offset + chunkLen < txLength) ? PacketFlag.Continuation : PacketFlag.None;
314+
315+
SendSinglePacketNoPayload(id, (ushort)_sequence, flag, txData, offset, chunkLen, attempts, out ushort nextSeq);
316+
317+
_sequence = nextSeq;
318+
offset += chunkLen;
319+
} while (offset < txLength);
320+
}
321+
322+
private void SendSinglePacketNoPayload(PacketId id, ushort seq, PacketFlag flag, byte[] txData, int txOffset, int txLen, int attempts, out ushort nextSeq)
323+
{
324+
PacketId currentId = id;
325+
PacketFlag currentFlag = flag;
326+
int currentTxOffset = txOffset;
327+
int currentTxLen = txLen;
328+
ushort currentSeq = seq;
329+
330+
while (true)
331+
{
332+
byte[] packet = new byte[HeaderSize + currentTxLen];
333+
packet[0] = (byte)currentId;
334+
packet[1] = (byte)currentFlag;
335+
BinaryPrimitives.WriteUInt16BigEndian(packet.AsSpan(2, 2), currentSeq);
336+
if (currentTxLen > 0) Array.Copy(txData, currentTxOffset, packet, HeaderSize, currentTxLen);
337+
338+
bool gotValidResponse = false;
339+
for (int i = 0; i < attempts; i++)
340+
{
341+
_client.Send(packet, packet.Length, _endpoint);
342+
343+
try
344+
{
345+
while (true)
346+
{
347+
IPEndPoint from = new(IPAddress.Any, 0);
348+
byte[] rxPacket = _client.Receive(ref from);
349+
if (rxPacket.Length < HeaderSize) continue;
350+
351+
ushort responseSeq = BinaryPrimitives.ReadUInt16BigEndian(rxPacket.AsSpan(2, 2));
352+
byte responseId = rxPacket[0];
353+
if (responseSeq != currentSeq) continue;
354+
if (responseId != (byte)currentId && responseId != (byte)PacketId.Error) continue;
355+
356+
if (responseId == (byte)PacketId.Error)
357+
{
358+
throw new Exception("Target returned error response.");
359+
}
360+
361+
if (rxPacket.Length > HeaderSize)
362+
{
363+
throw new Exception("UDP protocol error: target sent fastboot data out-of-turn.");
364+
}
365+
366+
gotValidResponse = true;
367+
currentSeq = (ushort)((currentSeq + 1) & 0xFFFF);
368+
369+
bool continuation = (rxPacket[1] & (byte)PacketFlag.Continuation) != 0;
370+
if (!continuation)
371+
{
372+
nextSeq = currentSeq;
373+
return;
374+
}
375+
376+
currentId = (PacketId)responseId;
377+
currentFlag = PacketFlag.None;
378+
currentTxOffset = 0;
379+
currentTxLen = 0;
380+
break;
381+
}
382+
}
383+
catch (SocketException ex) when (ex.SocketErrorCode == SocketError.TimedOut)
384+
{
385+
continue;
386+
}
387+
388+
if (gotValidResponse)
389+
{
390+
break;
391+
}
392+
}
393+
394+
if (!gotValidResponse)
395+
{
396+
throw new Exception($"Failed to receive response after {attempts} attempts.");
397+
}
228398
}
229-
return length;
230399
}
231400

232401
public void Dispose()

‎FirmwareKit.Comm.Fastboot/Backend/Usb/Linux/LinuxUsbDevice.cs‎

Lines changed: 30 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -114,24 +114,43 @@ public override int GetSerialNumber()
114114
}
115115

116116
public override byte[] Read(int length)
117+
{
118+
if (length <= 0) return Array.Empty<byte>();
119+
120+
byte[] buffer = new byte[length];
121+
int count = ReadInto(buffer, 0, length);
122+
if (count == length) return buffer;
123+
if (count == 0) return Array.Empty<byte>();
124+
125+
byte[] result = new byte[count];
126+
Buffer.BlockCopy(buffer, 0, result, 0, count);
127+
return result;
128+
}
129+
130+
public override int ReadInto(byte[] buffer, int offset, int length)
117131
{
118132
const uint MAX_USBFS_BULK_SIZE = 16384;
119133
const int MAX_RETRIES = 5;
120-
byte[] buffer = new byte[length];
121-
int count = 0;
122134

123-
while (count < length)
135+
if (length <= 0) return 0;
136+
if (offset < 0 || length < 0 || offset + length > buffer.Length)
124137
{
125-
int xfer = (length - count > (int)MAX_USBFS_BULK_SIZE) ? (int)MAX_USBFS_BULK_SIZE : (length - count);
126-
GCHandle handle = GCHandle.Alloc(buffer, GCHandleType.Pinned);
127-
try
138+
throw new ArgumentOutOfRangeException(nameof(length));
139+
}
140+
141+
int count = 0;
142+
GCHandle handle = GCHandle.Alloc(buffer, GCHandleType.Pinned);
143+
try
144+
{
145+
while (count < length)
128146
{
147+
int xfer = Math.Min(length - count, (int)MAX_USBFS_BULK_SIZE);
129148
usbdevfs_bulktransfer bulk = new usbdevfs_bulktransfer
130149
{
131150
ep = ep_in,
132151
len = (uint)xfer,
133152
timeout = 5000,
134-
data = new IntPtr(handle.AddrOfPinnedObject().ToInt64() + count)
153+
data = new IntPtr(handle.AddrOfPinnedObject().ToInt64() + offset + count)
135154
};
136155

137156
uint bulkCode = (IntPtr.Size == 8) ? USBDEVFS_BULK_X86_64 : USBDEVFS_BULK_X86;
@@ -159,18 +178,13 @@ public override byte[] Read(int length)
159178
count += n;
160179
if (n < xfer) break;
161180
}
162-
finally
163-
{
164-
handle.Free();
165-
}
166181
}
167-
if (count < length)
182+
finally
168183
{
169-
byte[] result = new byte[count];
170-
Array.Copy(buffer, result, count);
171-
return result;
184+
handle.Free();
172185
}
173-
return buffer;
186+
187+
return count;
174188
}
175189

176190
public override long Write(byte[] data, int length)

‎FirmwareKit.Comm.Fastboot/Backend/Usb/Windows/WinUSBDevice.cs‎

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -137,21 +137,43 @@ public override int GetSerialNumber()
137137

138138
public override byte[] Read(int length)
139139
{
140-
if (WinUSBHandle == IntPtr.Zero) throw new Exception("Device handle is closed.");
140+
if (length <= 0) return Array.Empty<byte>();
141141

142142
byte[] data = new byte[length];
143-
uint totalBytesRead = 0;
143+
int totalBytesRead = ReadInto(data, 0, length);
144+
if (totalBytesRead == length) return data;
145+
if (totalBytesRead == 0) return Array.Empty<byte>();
146+
147+
byte[] realData = new byte[totalBytesRead];
148+
Buffer.BlockCopy(data, 0, realData, 0, totalBytesRead);
149+
return realData;
150+
}
151+
152+
public override int ReadInto(byte[] buffer, int offset, int length)
153+
{
154+
if (WinUSBHandle == IntPtr.Zero) throw new Exception("Device handle is closed.");
155+
if (length <= 0) return 0;
156+
if (offset < 0 || length < 0 || offset + length > buffer.Length)
157+
{
158+
throw new ArgumentOutOfRangeException(nameof(length));
159+
}
160+
161+
const int MaxChunkSize = 1024 * 1024;
162+
int totalBytesRead = 0;
163+
byte[] chunkBuffer = new byte[Math.Min(length, MaxChunkSize)];
144164

145165
// AOSP style: Read in a loop until requested length is met or a short packet is received.
146166
while (totalBytesRead < length)
147167
{
148-
uint toRead = (uint)Math.Min(length - totalBytesRead, 1024 * 1024);
168+
uint toRead = (uint)Math.Min(length - totalBytesRead, chunkBuffer.Length);
149169
uint bytesRead;
150170

151-
if (WinUsb_ReadPipe(WinUSBHandle, ReadBulkID, data, toRead, out bytesRead, IntPtr.Zero))
171+
if (WinUsb_ReadPipe(WinUSBHandle, ReadBulkID, chunkBuffer, toRead, out bytesRead, IntPtr.Zero))
152172
{
153173
if (bytesRead == 0) break;
154-
totalBytesRead += bytesRead;
174+
175+
Buffer.BlockCopy(chunkBuffer, 0, buffer, offset + totalBytesRead, (int)bytesRead);
176+
totalBytesRead += (int)bytesRead;
155177

156178
// If we got a short packet (less than requested), it signifies end of transfer.
157179
if (bytesRead < toRead) break;
@@ -164,13 +186,7 @@ public override byte[] Read(int length)
164186
}
165187
}
166188

167-
if (totalBytesRead == length) return data;
168-
169-
if (totalBytesRead == 0 && length > 0) return Array.Empty<byte>();
170-
171-
byte[] realData = new byte[totalBytesRead];
172-
Array.Copy(data, realData, (int)totalBytesRead);
173-
return realData;
189+
return totalBytesRead;
174190
}
175191

176192
public override long Write(byte[] data, int length)

0 commit comments

Comments
 (0)