2828
2929public class TDF {
3030
31+ private final long maximumSize ;
32+
33+ public TDF () {
34+ this (MAX_TDF_INPUT_SIZE );
35+ }
36+
37+ // constructor for tests so that we can set a maximum size that's tractable for tests
38+ TDF (long maximumInputSize ) {
39+ this .maximumSize = maximumInputSize ;
40+ }
41+
3142 public static Logger logger = LoggerFactory .getLogger (TDF .class );
3243
3344 private static final long MAX_TDF_INPUT_SIZE = 68719476736L ;
@@ -66,12 +77,6 @@ public KasPublicKeyMissing(String errorMessage) {
6677 }
6778 }
6879
69- public static class InputStreamReadFailed extends Exception {
70- public InputStreamReadFailed (String errorMessage ) {
71- super (errorMessage );
72- }
73- }
74-
7580 public static class FailedToCreateGMAC extends Exception {
7681 public FailedToCreateGMAC (String errorMessage ) {
7782 super (errorMessage );
@@ -300,16 +305,10 @@ private static String calculateSignature(byte[] data, byte[] secret, Config.Inte
300305 return Hex .encodeHexString (gmacPayload );
301306 }
302307
303- public TDFObject createTDF (InputStream inputStream ,
304- long inputSize ,
308+ public TDFObject createTDF (InputStream payload ,
305309 OutputStream outputStream ,
306310 Config .TDFConfig tdfConfig , SDK .KAS kas ) throws Exception {
307- if (inputSize > MAX_TDF_INPUT_SIZE ) {
308- throw new DataSizeNotSupported ("can't create tdf larger than 64gb" );
309- }
310-
311311 if (tdfConfig .kasInfoList .isEmpty ()) {
312-
313312 throw new KasInfoMissing ("kas information is missing" );
314313 }
315314
@@ -318,51 +317,41 @@ public TDFObject createTDF(InputStream inputStream,
318317 TDFObject tdfObject = new TDFObject ();
319318 tdfObject .prepareManifest (tdfConfig );
320319
321- int segmentSize = tdfConfig .defaultSegmentSize ;
322- long totalSegments = inputSize / segmentSize ;
323- if (inputSize % segmentSize != 0 ) {
324- totalSegments += 1 ;
325- }
326-
327- // Empty payload we still want to create a payload
328- if (totalSegments == 0 ) {
329- totalSegments = 1 ;
330- }
331-
332- long encryptedSegmentSize = segmentSize + kGcmIvSize + kAesBlockSize ;
320+ long encryptedSegmentSize = tdfConfig .defaultSegmentSize + kGcmIvSize + kAesBlockSize ;
333321 TDFWriter tdfWriter = new TDFWriter (outputStream );
334322
335- long readPos = 0 ;
336323 StringBuilder aggregateHash = new StringBuilder ();
337324 byte [] readBuf = new byte [tdfConfig .defaultSegmentSize ];
338325
339326 tdfObject .manifest .encryptionInformation .integrityInformation .segments = new ArrayList <>();
340- while (totalSegments != 0 ) {
341- long readSize = segmentSize ;
342- if ((inputSize - readPos ) < segmentSize ) {
343- readSize = inputSize - readPos ;
344- }
345-
346- long n = inputStream .read (readBuf , 0 , (int ) readSize );
347- if (n != readSize ) {
348- throw new InputStreamReadFailed ("Input stream read miss match" );
349- }
350-
351- byte [] cipherData = tdfObject .aesGcm .encrypt (readBuf , 0 , (int ) readSize );
352- tdfWriter .appendPayload (cipherData );
327+ long totalSize = 0 ;
328+ boolean finished ;
329+ try (var payloadOutput = tdfWriter .payload ()) {
330+ do {
331+ int nRead = 0 ;
332+ int readThisLoop = 0 ;
333+ while (readThisLoop < readBuf .length && (nRead = payload .read (readBuf , readThisLoop , readBuf .length - readThisLoop )) > 0 ) {
334+ readThisLoop += nRead ;
335+ }
336+ finished = nRead < 0 ;
337+ totalSize += readThisLoop ;
353338
354- String segmentSig = calculateSignature (cipherData , tdfObject .payloadKey , tdfConfig .segmentIntegrityAlgorithm );
339+ if (totalSize > maximumSize ) {
340+ throw new DataSizeNotSupported ("can't create tdf larger than 64gb" );
341+ }
342+ byte [] cipherData = tdfObject .aesGcm .encrypt (readBuf , 0 , readThisLoop );
343+ payloadOutput .write (cipherData );
355344
356- aggregateHash .append (segmentSig );
357- Manifest .Segment segmentInfo = new Manifest .Segment ();
358- segmentInfo .hash = Base64 .getEncoder ().encodeToString (segmentSig .getBytes (StandardCharsets .UTF_8 ));
359- segmentInfo .segmentSize = readSize ;
360- segmentInfo .encryptedSegmentSize = cipherData .length ;
345+ String segmentSig = calculateSignature (cipherData , tdfObject .payloadKey , tdfConfig .segmentIntegrityAlgorithm );
361346
362- tdfObject .manifest .encryptionInformation .integrityInformation .segments .add (segmentInfo );
347+ aggregateHash .append (segmentSig );
348+ Manifest .Segment segmentInfo = new Manifest .Segment ();
349+ segmentInfo .hash = Base64 .getEncoder ().encodeToString (segmentSig .getBytes (StandardCharsets .UTF_8 ));
350+ segmentInfo .segmentSize = readThisLoop ;
351+ segmentInfo .encryptedSegmentSize = cipherData .length ;
363352
364- totalSegments -= 1 ;
365- readPos += readSize ;
353+ tdfObject . manifest . encryptionInformation . integrityInformation . segments . add ( segmentInfo ) ;
354+ } while (! finished ) ;
366355 }
367356
368357 Manifest .RootSignature rootSignature = new Manifest .RootSignature ();
@@ -377,7 +366,7 @@ public TDFObject createTDF(InputStream inputStream,
377366 rootSignature .algorithm = alg ;
378367 tdfObject .manifest .encryptionInformation .integrityInformation .rootSignature = rootSignature ;
379368
380- tdfObject .manifest .encryptionInformation .integrityInformation .segmentSizeDefault = segmentSize ;
369+ tdfObject .manifest .encryptionInformation .integrityInformation .segmentSizeDefault = tdfConfig . defaultSegmentSize ;
381370 tdfObject .manifest .encryptionInformation .integrityInformation .encryptedSegmentSizeDefault = (int )encryptedSegmentSize ;
382371
383372 tdfObject .manifest .encryptionInformation .integrityInformation .segmentHashAlg = kGmacIntegrityAlgorithm ;
0 commit comments