Sending correct packet size for multicredit packets (fixes #29) (#31)

* Correctly set the payload size of the write request in correspondence with credits assigned (Fixes #29) * Fixed branch build
hierynomus · Dec 30, 2016 · e6e4667 · e6e4667
1 parent 0d7a6ae
commit e6e4667
Show file tree

Hide file tree

Showing 8 changed files with 58 additions and 40 deletions.
diff --git a/src/main/java/com/hierynomus/mssmb2/SMB2Header.java b/src/main/java/com/hierynomus/mssmb2/SMB2Header.java
@@ -229,4 +229,8 @@ public String toString() {
             nextCommandOffset);
 
     }
+
+    public int getCreditCharge() {
+        return creditCharge;
+    }
 }
diff --git a/src/main/java/com/hierynomus/mssmb2/SMB2MultiCreditPacket.java b/src/main/java/com/hierynomus/mssmb2/SMB2MultiCreditPacket.java
@@ -17,20 +17,25 @@
 
 public class SMB2MultiCreditPacket extends SMB2Packet {
 
-    protected int creditsAssigned = 1;
-    private int payloadSize;
+    private int maxPayloadSize;
 
-    public SMB2MultiCreditPacket(int structureSize, SMB2Dialect dialect, SMB2MessageCommandCode messageType, long sessionId, long treeId, int payloadSize) {
+    public SMB2MultiCreditPacket(int structureSize, SMB2Dialect dialect, SMB2MessageCommandCode messageType, long sessionId, long treeId, int maxPayloadSize) {
         super(structureSize, dialect, messageType, sessionId, treeId);
-        this.payloadSize = payloadSize;
+        this.maxPayloadSize = maxPayloadSize;
+    }
+
+    public int getMaxPayloadSize() {
+        return this.maxPayloadSize;
     }
 
     public int getPayloadSize() {
-        return this.payloadSize;
+        return Math.min(maxPayloadSize, SINGLE_CREDIT_PAYLOAD_SIZE * getCreditsAssigned());
     }
 
+    public int getCreditsAssigned() {
+        return getHeader().getCreditCharge();
+    }
     public void setCreditsAssigned(int creditsAssigned) {
-        this.creditsAssigned = creditsAssigned;
         getHeader().setCreditCharge(creditsAssigned);
     }
 }
diff --git a/src/main/java/com/hierynomus/mssmb2/SMB2Packet.java b/src/main/java/com/hierynomus/mssmb2/SMB2Packet.java
@@ -20,6 +20,7 @@
 import com.hierynomus.smbj.common.SMBBuffer;
 
 public class SMB2Packet implements Packet<SMB2Packet, SMBBuffer> {
+    public static int SINGLE_CREDIT_PAYLOAD_SIZE = 64 * 1024;
     protected final SMB2Header header = new SMB2Header();
     protected int structureSize;
     SMBBuffer buffer;

diff --git a/src/main/java/com/hierynomus/mssmb2/messages/SMB2ReadRequest.java b/src/main/java/com/hierynomus/mssmb2/messages/SMB2ReadRequest.java
@@ -15,11 +15,11 @@
  */
 package com.hierynomus.mssmb2.messages;
 
+import com.hierynomus.mssmb2.SMB2Dialect;
 import com.hierynomus.mssmb2.SMB2FileId;
 import com.hierynomus.mssmb2.SMB2MessageCommandCode;
 import com.hierynomus.mssmb2.SMB2MultiCreditPacket;
 import com.hierynomus.smbj.common.SMBBuffer;
-import com.hierynomus.smbj.connection.NegotiatedProtocol;
 
 /**
  * [MS-SMB2].pdf 2.2.19 SMB2 READ Request
@@ -30,9 +30,9 @@ public class SMB2ReadRequest extends SMB2MultiCreditPacket {
     private final SMB2FileId fileId;
 
     public SMB2ReadRequest(
-        NegotiatedProtocol negotiatedProtocol, SMB2FileId fileId,
-        long sessionId, long treeId, long offset) {
-        super(49, negotiatedProtocol.getDialect(), SMB2MessageCommandCode.SMB2_READ, sessionId, treeId, negotiatedProtocol.getMaxReadSize());
+        SMB2Dialect dialect, SMB2FileId fileId,
+        long sessionId, long treeId, long offset, int maxPayloadSize) {
+        super(49, dialect, SMB2MessageCommandCode.SMB2_READ, sessionId, treeId, maxPayloadSize);
         this.fileId = fileId;
         this.offset = offset;
     }
@@ -42,7 +42,7 @@ protected void writeTo(SMBBuffer buffer) {
         buffer.putUInt16(structureSize); // StructureSize (2 bytes)
         buffer.putByte((byte) 0); // Padding (1 byte)
         buffer.putByte((byte) 0); // Flags (1 byte)
-        buffer.putUInt32(NegotiatedProtocol.SINGLE_CREDIT_PAYLOAD_SIZE * creditsAssigned); // Length (4 bytes)
+        buffer.putUInt32(SINGLE_CREDIT_PAYLOAD_SIZE * getCreditsAssigned()); // Length (4 bytes)
         buffer.putUInt64(offset); // Offset (8 bytes)
         fileId.write(buffer);  // FileId (16 bytes)
         buffer.putUInt32(1); // MinimumCount (4 bytes)

diff --git a/src/main/java/com/hierynomus/mssmb2/messages/SMB2WriteRequest.java b/src/main/java/com/hierynomus/mssmb2/messages/SMB2WriteRequest.java
@@ -48,6 +48,6 @@ protected void writeTo(SMBBuffer buffer) {
         buffer.putUInt16(0); // WriteChannelInfoOffset (2 bytes)
         buffer.putUInt16(0); // WriteChannelInfoLength (2 bytes)
         buffer.putUInt32(0); // Flags (4 bytes)
-        byteProvider.writeChunks(buffer, creditsAssigned);
+        byteProvider.writeChunks(buffer, getCreditsAssigned());
     }
 }
diff --git a/src/main/java/com/hierynomus/smbj/connection/Connection.java b/src/main/java/com/hierynomus/smbj/connection/Connection.java
@@ -59,9 +59,9 @@
 
 import net.engio.mbassy.listener.Handler;
 
+import static com.hierynomus.mssmb2.SMB2Packet.SINGLE_CREDIT_PAYLOAD_SIZE;
 import static com.hierynomus.mssmb2.messages.SMB2SessionSetup.SMB2SecurityMode.SMB2_NEGOTIATE_SIGNING_ENABLED;
 import static com.hierynomus.protocol.commons.EnumWithValue.EnumUtils.isSet;
-import static com.hierynomus.smbj.connection.NegotiatedProtocol.SINGLE_CREDIT_PAYLOAD_SIZE;
 import static java.lang.String.format;
 
 /**
@@ -192,38 +192,23 @@ private <T extends SMB2Packet> T sendAndReceive(SMB2Packet packet) throws Transp
     /**
      * send a packet, potentially signed
      *
-     * @param packet     SMBPacket to send
-     * @param signingKeySpec if null, do not sign the packet.  Otherwise, the signingKey will be used to sign the packet.
+     * @param packet         SMBPacket to send
+     * @param signingKeySpec if null, do not sign the packet. Otherwise, the signingKeySpec will be used to sign the packet.
      * @return a Future to be used to retrieve the response packet
      * @throws TransportException
      */
     public <T extends SMB2Packet> Future<T> send(SMB2Packet packet, SecretKeySpec signingKeySpec) throws TransportException {
         lock.lock();
         try {
             int availableCredits = connectionInfo.getSequenceWindow().available();
-            int grantCredits;
-            if (packet instanceof SMB2MultiCreditPacket) {
-                int payloadSize = ((SMB2MultiCreditPacket) packet).getPayloadSize();
-                int creditsNeeded = creditsNeeded(payloadSize);
-                // Scale the credits granted to the message dynamically.
-                if (availableCredits == 0) {
-                    throw new NoSuchElementException("TODO ([MS-SMB2].pdf 3.2.5.1.4 Granting Message Credits)! No credits left.");
-                } else if (creditsNeeded < availableCredits) {
-                    grantCredits = creditsNeeded;
-                } else if (creditsNeeded > 1 && availableCredits > 1) { // creditsNeeded >= availableCredits
-                    grantCredits = availableCredits - 1; // Keep 1 credit left for a simple request
-                } else {
-                    grantCredits = 1;
-                }
-                long[] messageIds = connectionInfo.getSequenceWindow().get(grantCredits);
-                ((SMB2MultiCreditPacket) packet).setCreditsAssigned(grantCredits);
-                packet.getHeader().setMessageId(messageIds[0]);
-                logger.debug("Granted {} credits to {} with message id << {} >>", grantCredits, packet.getHeader().getMessage(), packet.getHeader().getMessageId());
-            } else {
-                grantCredits = 1;
-                long messageId = connectionInfo.getSequenceWindow().get();
-                packet.getHeader().setMessageId(messageId);
+            if (availableCredits == 0) {
+                throw new NoSuchElementException("TODO ([MS-SMB2].pdf 3.2.5.1.4 Granting Message Credits)! No credits left.");
             }
+
+            int grantCredits = calculateGrantedCredits(packet, availableCredits);
+            long[] messageIds = connectionInfo.getSequenceWindow().get(grantCredits);
+            packet.getHeader().setMessageId(messageIds[0]);
+            logger.debug("Granted {} credits to {} with message id << {} >>", grantCredits, packet.getHeader().getMessage(), packet.getHeader().getMessageId());
             packet.getHeader().setCreditRequest(Math.max(SequenceWindow.PREFERRED_MINIMUM_CREDITS - availableCredits - grantCredits, grantCredits));
 
             Request request = new Request(packet.getHeader().getMessageId(), UUID.randomUUID(), packet);
@@ -239,6 +224,26 @@ public <T extends SMB2Packet> Future<T> send(SMB2Packet packet, SecretKeySpec si
         }
     }
 
+    private int calculateGrantedCredits(final SMB2Packet packet, final int availableCredits) {
+        final int grantCredits;
+        if (packet instanceof SMB2MultiCreditPacket) {
+            int maxPayloadSize = ((SMB2MultiCreditPacket) packet).getMaxPayloadSize();
+            int creditsNeeded = creditsNeeded(maxPayloadSize);
+            // Scale the credits granted to the message dynamically.
+            if (creditsNeeded < availableCredits) {
+                grantCredits = creditsNeeded;
+            } else if (creditsNeeded > 1 && availableCredits > 1) { // creditsNeeded >= availableCredits
+                grantCredits = availableCredits - 1; // Keep 1 credit left for a simple request
+            } else {
+                grantCredits = 1;
+            }
+            ((SMB2MultiCreditPacket) packet).setCreditsAssigned(grantCredits);
+        } else {
+            grantCredits = 1;
+        }
+        return grantCredits;
+    }
+
     private void negotiateDialect() throws TransportException {
         logger.info("Negotiating dialects {} with server {}", config.getSupportedDialects(), getRemoteHostname());
         SMB2Packet negotiatePacket = new SMB2NegotiateRequest(config.getSupportedDialects(), connectionInfo.getClientGuid());

diff --git a/src/main/java/com/hierynomus/smbj/connection/NegotiatedProtocol.java b/src/main/java/com/hierynomus/smbj/connection/NegotiatedProtocol.java
@@ -17,11 +17,12 @@
 
 import com.hierynomus.mssmb2.SMB2Dialect;
 
+import static com.hierynomus.mssmb2.SMB2Packet.SINGLE_CREDIT_PAYLOAD_SIZE;
+
 /**
  * Encapsulates the details of the Protocol Negotiation
  */
 public class NegotiatedProtocol {
-    public static int SINGLE_CREDIT_PAYLOAD_SIZE = 65536;
     private SMB2Dialect dialect;
     private int maxTransactSize;
     private int maxReadSize;

diff --git a/src/main/java/com/hierynomus/smbj/share/FileInputStream.java b/src/main/java/com/hierynomus/smbj/share/FileInputStream.java
@@ -28,6 +28,7 @@
 import com.hierynomus.smbj.ProgressListener;
 import com.hierynomus.smbj.common.SMBApiException;
 import com.hierynomus.smbj.connection.Connection;
+import com.hierynomus.smbj.connection.NegotiatedProtocol;
 import com.hierynomus.smbj.session.Session;
 import com.hierynomus.smbj.transport.TransportException;
 
@@ -118,8 +119,9 @@ private void loadBuffer() throws IOException {
     }
 
     private Future<SMB2ReadResponse> sendRequest() throws IOException {
-        SMB2ReadRequest rreq = new SMB2ReadRequest(connection.getNegotiatedProtocol(), fileId,
-            session.getSessionId(), treeConnect.getTreeId(), offset);
+        NegotiatedProtocol negotiatedProtocol = connection.getNegotiatedProtocol();
+        SMB2ReadRequest rreq = new SMB2ReadRequest(negotiatedProtocol.getDialect(), fileId,
+            session.getSessionId(), treeConnect.getTreeId(), offset, negotiatedProtocol.getMaxReadSize());
         return session.send(rreq);
     }
 }