Add support for reading / writing NIO ByteBuffers (#759)

* Add support for reading / writing NIO ByteBuffers Currently one can transfer data using streams or array, but it would be great to have the opportunity to use NIO buffers. This adds two new method to the File class that accept a NIO ByteBuffer. * Implemented ByteBuffer write using ByteChunkProvider Signed-off-by: Jeroen van Erp <[email protected]> --------- Signed-off-by: Jeroen van Erp <[email protected]> Co-authored-by: Christoph Läubrich <[email protected]> Co-authored-by: Jeroen van Erp <[email protected]>
hierynomus · May 8, 2023 · 01d720e · 01d720e
1 parent 350c73e
commit 01d720e
Show file tree

Hide file tree

Showing 5 changed files with 265 additions and 46 deletions.
diff --git a/src/main/java/com/hierynomus/smbj/io/ByteBufferByteChunkProvider.java b/src/main/java/com/hierynomus/smbj/io/ByteBufferByteChunkProvider.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright (C)2016 - SMBJ Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.hierynomus.smbj.io;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+public class ByteBufferByteChunkProvider extends CachingByteChunkProvider {
+    private ByteBuffer buffer;
+
+    public ByteBufferByteChunkProvider(ByteBuffer buffer) {
+        super();
+        this.buffer = buffer;
+    }
+
+    public ByteBufferByteChunkProvider(ByteBuffer buffer, long fileOffset) {
+        super();
+        this.buffer = buffer;
+        this.offset = fileOffset;
+    }
+
+    @Override
+    int prepareChunk(byte[] chunk, int bytesNeeded) throws IOException {
+        int bytesToRead = Math.min(chunk.length, Math.min(bytesNeeded, buffer.remaining()));
+        if (bytesToRead == 0) {
+            return -1;
+        }
+
+        buffer.get(chunk, 0, bytesToRead);
+        return bytesToRead;
+    }
+
+    @Override
+    public boolean isAvailable() {
+        return super.isAvailable() || buffer.hasRemaining();
+    }
+}
diff --git a/src/main/java/com/hierynomus/smbj/io/CachingByteChunkProvider.java b/src/main/java/com/hierynomus/smbj/io/CachingByteChunkProvider.java
@@ -0,0 +1,83 @@
+/*
+ * Copyright (C)2016 - SMBJ Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.hierynomus.smbj.io;
+
+import java.io.IOException;
+
+import com.hierynomus.protocol.commons.buffer.Buffer;
+import com.hierynomus.protocol.commons.buffer.Endian;
+import com.hierynomus.smbj.common.SMBRuntimeException;
+
+abstract class CachingByteChunkProvider extends ByteChunkProvider {
+    private BufferByteChunkProvider cachingProvider;
+    private Buffer<Buffer.PlainBuffer> buffer;
+
+    CachingByteChunkProvider() {
+        this.buffer = new Buffer.PlainBuffer(Endian.BE);
+        this.cachingProvider = new BufferByteChunkProvider(buffer);
+    }
+
+    @Override
+    public void prepareWrite(int maxBytesToPrepare) {
+        if (buffer == null) {
+            return;
+        }
+
+        byte[] chunk = new byte[1024];
+
+        // Before each prepareWrite, compact the buffer to minimize size growth
+        buffer.compact();
+
+        int bytesNeeded = maxBytesToPrepare - buffer.available();
+        int read;
+        try {
+            while (bytesNeeded > 0) {
+                read = prepareChunk(chunk, bytesNeeded);
+                if (read == -1) {
+                    break;
+                }
+
+                // Write the data to the buffer
+                buffer.putRawBytes(chunk, 0, read);
+                bytesNeeded -= read;
+            }
+        } catch (IOException e) {
+            throw new SMBRuntimeException(e);
+        }
+    }
+
+    abstract int prepareChunk(byte[] chunk, int bytesNeeded) throws IOException;
+
+    @Override
+    protected int getChunk(byte[] chunk) throws IOException {
+        return cachingProvider.getChunk(chunk);
+    }
+
+    @Override
+    public int bytesLeft() {
+        return cachingProvider.bytesLeft();
+    }
+
+    @Override
+    public boolean isAvailable() {
+        return cachingProvider.isAvailable();
+    }
+
+    @Override
+    public void close() throws IOException {
+        cachingProvider.close();
+    }
+}
diff --git a/src/main/java/com/hierynomus/smbj/io/InputStreamByteChunkProvider.java b/src/main/java/com/hierynomus/smbj/io/InputStreamByteChunkProvider.java
@@ -15,82 +15,52 @@
  */
 package com.hierynomus.smbj.io;
 
-import com.hierynomus.protocol.commons.buffer.Buffer;
-import com.hierynomus.protocol.commons.buffer.Endian;
-import com.hierynomus.smbj.common.SMBRuntimeException;
-
 import java.io.BufferedInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 
-public class InputStreamByteChunkProvider extends ByteChunkProvider {
+import com.hierynomus.smbj.common.SMBRuntimeException;
+
+public class InputStreamByteChunkProvider extends CachingByteChunkProvider {
 
     private BufferedInputStream is;
-    private BufferByteChunkProvider cachingProvider;
-    private Buffer<Buffer.PlainBuffer> buffer;
+    private boolean close;
 
     public InputStreamByteChunkProvider(InputStream is) {
-        this.buffer = new Buffer.PlainBuffer(Endian.BE);
-        this.cachingProvider = new BufferByteChunkProvider(buffer);
+        super();
         if (is instanceof BufferedInputStream)
             this.is = (BufferedInputStream) is;
-        else
+        else {
             this.is = new BufferedInputStream(is);
-    }
-
-    @Override
-    public void prepareWrite(int maxBytesToPrepare) {
-        if (is == null) {
-            return;
+            this.close = true; // We control the is, so we close it
         }
 
-        byte[] chunk = new byte[1024];
-
-        // Before each prepareWrite, compact the buffer to minimize size growth
-        buffer.compact();
-
-        int bytesNeeded = maxBytesToPrepare - buffer.available();
-        int read;
-        try {
-            while (bytesNeeded > 0) {
-                read = is.read(chunk, 0, chunk.length);
-                if (read == -1) {
-                    break;
-                }
-
-                // Write the data to the buffer
-                buffer.putRawBytes(chunk, 0, read);
-                bytesNeeded -= read;
-            }
-        } catch (IOException e) {
-            throw new SMBRuntimeException(e);
-        }
     }
 
     @Override
-    protected int getChunk(byte[] chunk) throws IOException {
-        return cachingProvider.getChunk(chunk);
-    }
+    int prepareChunk(byte[] chunk, int bytesNeeded) throws IOException {
+        int toRead = Math.min(bytesNeeded, chunk.length);
+        if (toRead == 0) {
+            return -1;
+        }
 
-    @Override
-    public int bytesLeft() {
-        return cachingProvider.bytesLeft();
+        return is.read(chunk, 0, toRead);
     }
 
     @Override
     public boolean isAvailable() {
         try {
-            return cachingProvider.isAvailable() || is.available() > 0;
+            return super.isAvailable() || (is != null && is.available() > 0);
         } catch (IOException e) {
             throw new SMBRuntimeException(e);
         }
     }
 
     @Override
     public void close() throws IOException {
-        cachingProvider.close();
+        super.close();
 
-        if (is != null) {
+        if (is != null && close) {
             try {
                 is.close();
             } finally {

diff --git a/src/main/java/com/hierynomus/smbj/share/File.java b/src/main/java/com/hierynomus/smbj/share/File.java
@@ -18,6 +18,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -39,6 +40,7 @@
 import com.hierynomus.smbj.common.SMBRuntimeException;
 import com.hierynomus.smbj.common.SmbPath;
 import com.hierynomus.smbj.io.ArrayByteChunkProvider;
+import com.hierynomus.smbj.io.ByteBufferByteChunkProvider;
 import com.hierynomus.smbj.io.ByteChunkProvider;
 
 public class File extends DiskEntry {
@@ -186,6 +188,40 @@ public void read(OutputStream destStream, ProgressListener progressListener) thr
         is.close();
     }
 
+    /**
+     * Write the data in a {@link ByteBuffer} to this file at position fileOffset.
+     *
+     * @param buffer     the data to write
+     * @param fileOffset The offset, in bytes, into the file to which the data should be written
+     * @return the actual number of bytes that was written to the file
+     */
+    public long write(ByteBuffer buffer, long fileOffset) {
+        ByteChunkProvider provider = new ByteBufferByteChunkProvider(buffer, fileOffset);
+        return write(provider);
+    }
+
+
+    /**
+     * Read data from this file starting at position fileOffset into the given {@link ByteBuffer}.
+     *
+     * @param buffer     the {@link ByteBuffer} to write into
+     * @param fileOffset The offset, in bytes, into the file from which the data should be read
+     * @return the actual number of bytes that were read; or -1 if the end of the file was reached
+     */
+    public long read(ByteBuffer buffer, long fileOffset) {
+        int remaining = buffer.remaining();
+
+        SMB2ReadResponse response = share.read(fileId, fileOffset, remaining);
+        if (response.getHeader().getStatusCode() == NtStatus.STATUS_END_OF_FILE.getValue()) {
+            return -1;
+        } else {
+            byte[] data = response.getData();
+            int bytesRead = Math.min(remaining, data.length);
+            buffer.put(data, 0, bytesRead);
+            return bytesRead;
+        }
+    }
+
     /**
      * Performs a remote file copy of this file to the given file.
      * <p>

diff --git a/src/test/groovy/com/hierynomus/smbj/io/ByteBufferByteChunkProviderSpec.groovy b/src/test/groovy/com/hierynomus/smbj/io/ByteBufferByteChunkProviderSpec.groovy
@@ -0,0 +1,80 @@
+/*
+ * Copyright (C)2016 - SMBJ Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.hierynomus.smbj.io
+
+import spock.lang.Specification
+import java.nio.ByteBuffer
+
+class ByteBufferByteChunkProviderSpec extends Specification {
+
+  def "should write 1 chunk to outputStream"() {
+    given:
+    def buffer = getBufferWithRandomData(ByteChunkProvider.CHUNK_SIZE)
+    def checkBuff = buffer.duplicate()
+    def provider = new ByteBufferByteChunkProvider(buffer)
+    def baos = new ByteArrayOutputStream()
+
+    when:
+    provider.prepareWrite(ByteChunkProvider.CHUNK_SIZE)
+    provider.writeChunk(baos)
+
+    then:
+    baos.toByteArray() == checkBuff.array()
+    provider.offset == ByteChunkProvider.CHUNK_SIZE
+    !provider.isAvailable()
+  }
+
+  def "should write part of chunk to outputStream"() {
+    given:
+    def buffer = getBufferWithRandomData(1024)
+    def checkBuff = buffer.duplicate()
+    def provider = new ByteBufferByteChunkProvider(buffer)
+    def baos = new ByteArrayOutputStream()
+
+    when:
+    provider.prepareWrite(ByteChunkProvider.CHUNK_SIZE)
+    provider.writeChunk(baos)
+
+    then:
+    baos.toByteArray() == checkBuff.array()
+    provider.offset == 1024
+    !provider.isAvailable()
+
+  }
+
+  def "should have available after writing first chunk"() {
+    given:
+    def buffer = getBufferWithRandomData(ByteChunkProvider.CHUNK_SIZE + 1)
+    def provider = new ByteBufferByteChunkProvider(buffer)
+    def baos = new ByteArrayOutputStream()
+
+    when:
+    provider.prepareWrite(ByteChunkProvider.CHUNK_SIZE)
+    provider.writeChunk(baos)
+
+    then:
+    provider.offset == ByteChunkProvider.CHUNK_SIZE
+    provider.isAvailable()
+
+  }
+
+  private def getBufferWithRandomData(int size) {
+    def bytes = new byte[size]
+    new Random().nextBytes(bytes)
+    def buffer = ByteBuffer.wrap(bytes)
+    return buffer
+  }
+}