66
77import java .lang .foreign .Arena ;
88import java .lang .foreign .MemorySegment ;
9+ import java .util .Arrays ;
910import java .util .Random ;
1011import org .assertj .core .api .ThrowableAssert .ThrowingCallable ;
1112import org .junit .jupiter .api .Test ;
@@ -76,7 +77,7 @@ void prefixIsAppliedAndRequiredToDecode() {
7677 try (ZstdDecompressCtx dctx = new ZstdDecompressCtx ()) {
7778 MemorySegment out = arena .allocate (random .length );
7879 long m = dctx .decompress (out , copy (arena , frame ));
79- reproduced = java . util . Arrays .equals (bytes (out , (int ) m ), random );
80+ reproduced = Arrays .equals (bytes (out , (int ) m ), random );
8081 } catch (ZstdException e ) {
8182 reproduced = false ;
8283 }
@@ -111,7 +112,39 @@ void clearingPrefixWithNullCompressesPlainly() {
111112 }
112113
113114 @ Test
114- void refPrefixRejectsAHeapSegment () {
115+ void prefixIsSingleUseAndDoesNotStickAcrossFrames () {
116+ // Given — random data identical to the prefix (see prefixIsAppliedAndRequiredToDecode),
117+ // and one context kept open across two compressions with the prefix set once.
118+ byte [] random = randomBytes (0xCAFE , 16384 );
119+ try (Arena arena = Arena .ofConfined ();
120+ ZstdCompressCtx cctx = new ZstdCompressCtx ().level (19 )) {
121+ MemorySegment prefix = copy (arena , random );
122+ MemorySegment src = copy (arena , random );
123+ MemorySegment first = arena .allocate (Zstd .compressBound (random .length ));
124+ MemorySegment second = arena .allocate (Zstd .compressBound (random .length ));
125+
126+ // When — first frame consumes the prefix; second is compressed with no re-set
127+ cctx .refPrefix (prefix );
128+ long n1 = cctx .compress (first , src );
129+ long n2 = cctx .compress (second , src );
130+
131+ // Then — the prefix shrank only the first frame; the second got no prefix
132+ // (incompressible random with no prefix stays ~full size)
133+ assertThat (n1 ).isLessThan (n2 / 10 );
134+
135+ // And — the second frame carries no prefix, so a plain decoder decodes it
136+ byte [] restored ;
137+ try (ZstdDecompressCtx dctx = new ZstdDecompressCtx ()) {
138+ MemorySegment out = arena .allocate (random .length );
139+ long m = dctx .decompress (out , second .asSlice (0 , n2 ));
140+ restored = bytes (out , (int ) m );
141+ }
142+ assertThat (restored ).isEqualTo (random );
143+ }
144+ }
145+
146+ @ Test
147+ void compressRefPrefixRejectsAHeapSegment () {
115148 // Given
116149 MemorySegment heap = MemorySegment .ofArray ("prefix" .getBytes ());
117150
@@ -128,6 +161,24 @@ void refPrefixRejectsAHeapSegment() {
128161 .hasMessageContaining ("prefix" );
129162 }
130163
164+ @ Test
165+ void decompressRefPrefixRejectsAHeapSegment () {
166+ // Given
167+ MemorySegment heap = MemorySegment .ofArray ("prefix" .getBytes ());
168+
169+ // When
170+ ThrowingCallable result = () -> {
171+ try (ZstdDecompressCtx dctx = new ZstdDecompressCtx ()) {
172+ dctx .refPrefix (heap );
173+ }
174+ };
175+
176+ // Then
177+ assertThatThrownBy (result )
178+ .isInstanceOf (IllegalArgumentException .class )
179+ .hasMessageContaining ("prefix" );
180+ }
181+
131182 private static MemorySegment copy (Arena arena , byte [] src ) {
132183 MemorySegment seg = arena .allocate (src .length );
133184 MemorySegment .copy (src , 0 , seg , JAVA_BYTE , 0 , src .length );
0 commit comments