2222import java .io .*;
2323import java .nio .file .Files ;
2424import java .security .NoSuchAlgorithmException ;
25- import java .time .Duration ;
2625import java .util .ArrayList ;
2726import java .util .Arrays ;
2827import java .util .List ;
@@ -218,6 +217,7 @@ public ObjectInfo get(String objectName, OutputStream out) throws IOException, J
218217 Digester digester = new Digester ();
219218 long totalBytes = 0 ;
220219 long totalChunks = 0 ;
220+ long expectedChunks = oi .getChunks ();
221221
222222 // if there is one chunk, just go get the message directly and we're done.
223223 if (oi .getChunks () == 1 ) {
@@ -236,30 +236,42 @@ public ObjectInfo get(String objectName, OutputStream out) throws IOException, J
236236 JetStreamSubscription sub = js .subscribe (rawChunkSubject (oi .getNuid ()),
237237 PushSubscribeOptions .builder ().stream (streamName ).ordered (true ).build ());
238238
239- Message m = sub .nextMessage (Duration . ofSeconds ( 1 ));
239+ Message m = sub .nextMessage (jsm . getTimeout ( ));
240240 while (m != null ) {
241+ // track the byte count and chunks
242+ long pending = m .metaData ().pendingCount ();
243+ if (expectedChunks != pending + (++totalChunks )) {
244+ throw OsGetChunksMismatch .instance (); // short circuit, we already know there are not enough chunks.
245+ }
246+
241247 byte [] data = m .getData ();
248+ totalBytes += data .length ;
242249
243- // track the byte count and chunks
244250 // update the digest
245- // write the bytes to the output file
246- totalBytes += data .length ;
247- totalChunks ++;
248251 digester .update (data );
252+
253+ // write the bytes to the output file
249254 out .write (data );
250255
251256 // read until the subject is complete
252- m = sub .nextMessage (Duration .ofSeconds (1 ));
257+ if (pending == 0 ) {
258+ break ;
259+ }
260+ m = sub .nextMessage (jsm .getTimeout ());
253261 }
254262
255- sub .unsubscribe ();
263+ try {
264+ sub .unsubscribe ();
265+ }
266+ catch (RuntimeException ignore ) {}
256267 }
257- out .flush ();
258268
259- if (totalBytes != oi .getSize ()) { throw OsGetSizeMismatch .instance (); }
260269 if (totalChunks != oi .getChunks ()) { throw OsGetChunksMismatch .instance (); }
270+ if (totalBytes != oi .getSize ()) { throw OsGetSizeMismatch .instance (); }
261271 if (!digester .matches (oi .getDigest ())) { throw OsGetDigestMismatch .instance (); }
262272
273+ out .flush (); // moved after validation, no need if invalid
274+
263275 return oi ;
264276 }
265277
0 commit comments