From b66daaaac931711977e8098f30d5590f5fe3cd09 Mon Sep 17 00:00:00 2001
From: Rob Austin <rob.austin@boundedbuffer>
Date: Thu, 23 Oct 2014 11:11:24 +0100
Subject: [PATCH] HCOLL-167 improved the buffer overflow support

---
 .../chronicle/map/StatelessChronicleMap.java  |  7 +-
 .../openhft/chronicle/map/TcpReplicator.java  | 74 ++++++++++++++-----
 2 files changed, 61 insertions(+), 20 deletions(-)

diff --git a/src/main/java/net/openhft/chronicle/map/StatelessChronicleMap.java b/src/main/java/net/openhft/chronicle/map/StatelessChronicleMap.java
index aeb21bf1d..56ecbe3be 100644
--- a/src/main/java/net/openhft/chronicle/map/StatelessChronicleMap.java
+++ b/src/main/java/net/openhft/chronicle/map/StatelessChronicleMap.java
@@ -35,7 +35,7 @@ class StatelessChronicleMap<K, V> implements ChronicleMap<K, V>, Closeable {
 
     private final ByteBuffer outBuffer = allocateDirect(1024);
     private final ByteBufferBytes out = new ByteBufferBytes(outBuffer.slice());
-    private final ByteBufferBytes in = new ByteBufferBytes(allocateDirect(1024));
+    private ByteBufferBytes in = new ByteBufferBytes(allocateDirect(1024));
 
     private final KeyValueSerializer<K, V> keyValueSerializer;
 
@@ -327,9 +327,7 @@ class StatelessChronicleMap<K, V> implements ChronicleMap<K, V>, Closeable {
         writeValue(value);
 
         // get the data back from the server
-
         return readKey(sizeLocation);
-
     }
 
     public synchronized V remove(Object key) {
@@ -527,6 +525,9 @@ class StatelessChronicleMap<K, V> implements ChronicleMap<K, V>, Closeable {
                 // the number of bytes in the response
                 int size = receive(2, timeoutTime).readUnsignedShort();
 
+                if (in.capacity() < size)
+                    in = new ByteBufferBytes(allocateDirect(size));
+
                 receive(size, timeoutTime);
 
                 boolean isException = in.readBoolean();
diff --git a/src/main/java/net/openhft/chronicle/map/TcpReplicator.java b/src/main/java/net/openhft/chronicle/map/TcpReplicator.java
index 09fb0b73f..4de7a43cc 100644
--- a/src/main/java/net/openhft/chronicle/map/TcpReplicator.java
+++ b/src/main/java/net/openhft/chronicle/map/TcpReplicator.java
@@ -1345,8 +1345,6 @@ class StatelessServerConnector<K, V> {
                 str :
                 str.substring(0, (int) (remaining - 4)) + "...";
 
-        out.writeObject(str);
-
         out.writeObject(result);
         writeSizeAndFlags(sizeLocation, false, out);
         return null;
@@ -1467,7 +1465,13 @@ class StatelessServerConnector<K, V> {
     private Work keySet(Bytes in, final Bytes out) {
         final long sizeLocation = reflectTransactionId(in, out);
 
-        final Set<K> ks = map.keySet();
+        final Set<K> ks;
+        try {
+            ks = map.keySet();
+        } catch (RuntimeException e) {
+            return sendException(out, sizeLocation, e);
+        }
+
         out.writeStopBit(ks.size());
 
         final Iterator<K> iterator = ks.iterator();
@@ -1495,39 +1499,75 @@ class StatelessServerConnector<K, V> {
     }
 
     private Work values(Bytes in, Bytes out) {
-        long sizeLocation = reflectTransactionId(in, out);
+        final long sizeLocation = reflectTransactionId(in, out);
         final Collection<V> values;
+
         try {
             values = map.values();
         } catch (RuntimeException e) {
             return sendException(out, sizeLocation, e);
         }
 
-
         out.writeStopBit(values.size());
-        for (final V value : values) {
-            writeValue(value, out);
-        }
-        writeSizeAndFlags(sizeLocation, false, out);
-        return null;
+
+        final Iterator<V> iterator = values.iterator();
+
+        // this allows us to write more data than the buffer will allow
+        return new Work() {
+
+            @Override
+            public boolean doWork(Bytes out) {
+
+                while (iterator.hasNext()) {
+
+                    // we've filled up the buffer, so lets give another channel a chance to send
+                    // some data, we don't know the max key size, we will use the entrySize instead
+                    if (out.remaining() <= maxEntrySizeBytes)
+                        return false;
+
+                    writeValue(iterator.next(), out);
+                }
+
+                writeSizeAndFlags(sizeLocation, false, out);
+                return true;
+            }
+        };
     }
 
     private Work entrySet(Bytes in, Bytes out) {
-        long sizeLocation = reflectTransactionId(in, out);
+        final long sizeLocation = reflectTransactionId(in, out);
 
         final Set<Map.Entry<K, V>> entries;
         try {
-
             entries = map.entrySet();
         } catch (RuntimeException e) {
             return sendException(out, sizeLocation, e);
         }
         out.writeStopBit(entries.size());
-        for (Map.Entry<K, V> e : entries) {
-            writeKey(e.getKey(), out);
-            writeValue(e.getValue(), out);
-        }
-        return null;
+
+        final Iterator<Map.Entry<K, V>> iterator = entries.iterator();
+
+        // this allows us to write more data than the buffer will allow
+        return new Work() {
+
+            @Override
+            public boolean doWork(Bytes out) {
+
+                while (iterator.hasNext()) {
+
+                    // we've filled up the buffer, so lets give another channel a chance to send
+                    // some data, we don't know the max key size, we will use the entrySize instead
+                    if (out.remaining() <= maxEntrySizeBytes)
+                        return false;
+                    Map.Entry<K, V> next = iterator.next();
+                    writeKey(next.getKey(), out);
+                    writeValue(next.getValue(), out);
+                }
+
+                writeSizeAndFlags(sizeLocation, false, out);
+                return true;
+            }
+        };
     }
 
     private Work putIfAbsent(Bytes in, Bytes out) {