diff --git a/src/main/java/com/hierynomus/smbj/io/ByteBufferByteChunkProvider.java b/src/main/java/com/hierynomus/smbj/io/ByteBufferByteChunkProvider.java new file mode 100644 index 00000000..4851e230 --- /dev/null +++ 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 new file mode 100644 index 00000000..55e977e6 --- /dev/null +++ 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; + + 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 index a0d79620..23441f40 100644 --- a/src/main/java/com/hierynomus/smbj/io/InputStreamByteChunkProvider.java +++ b/src/main/java/com/hierynomus/smbj/io/InputStreamByteChunkProvider.java @@ -15,72 +15,42 @@ */ 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; + 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); } @@ -88,9 +58,9 @@ public boolean isAvailable() { @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 index ab06cbe8..3580d2d9 100644 --- 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. *

diff --git a/src/test/groovy/com/hierynomus/smbj/io/ByteBufferByteChunkProviderSpec.groovy b/src/test/groovy/com/hierynomus/smbj/io/ByteBufferByteChunkProviderSpec.groovy new file mode 100644 index 00000000..62a2a8fb --- /dev/null +++ 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 + } +}