Some changes needed for POW server and some general improvements

domain/src/main/java/ch/dissem/bitmessage/BitmessageContext.java

@@ -66,6 +66,8 @@ public class BitmessageContext {
     private final Listener listener;
     private final NetworkHandler.MessageListener networkListener;
 
+    private final boolean sendPubkeyOnIdentityCreation;
+
     private BitmessageContext(Builder builder) {
         ctx = new InternalContext(builder);
         listener = builder.listener;
@@ -75,6 +77,8 @@ public class BitmessageContext {
         // one should be executed at any time.
         pool = Executors.newFixedThreadPool(1);
 
+        sendPubkeyOnIdentityCreation = builder.sendPubkeyOnIdentityCreation;
+
         new Timer().schedule(new TimerTask() {
             @Override
             public void run() {
@@ -100,12 +104,14 @@ public class BitmessageContext {
                 features
         ));
         ctx.getAddressRepository().save(identity);
+        if (sendPubkeyOnIdentityCreation) {
             pool.submit(new Runnable() {
                 @Override
                 public void run() {
                     ctx.sendPubkey(identity, identity.getStream());
                 }
             });
+        }
         return identity;
     }
 
@@ -325,6 +331,8 @@ public class BitmessageContext {
         Listener listener;
         int connectionLimit = 150;
         long connectionTTL = 12 * HOUR;
+        boolean sendPubkeyOnIdentityCreation = true;
+        long pubkeyTTL = 28;
 
         public Builder() {
         }
@@ -399,6 +407,30 @@ public class BitmessageContext {
             return this;
         }
 
+        /**
+         * By default a client will send the public key when an identity is being created. On weaker devices
+         * this behaviour might not be desirable.
+         */
+        public Builder doNotSendPubkeyOnIdentityCreation() {
+            this.sendPubkeyOnIdentityCreation = false;
+            return this;
+        }
+
+        /**
+         * Time to live in seconds for public keys the client sends. Defaults to the maximum of 28 days,
+         * but on weak devices smaller values might be desirable.
+         * <p>
+         *     Please be aware that this might cause some problems where you can't receive a message (the
+         *     sender can't receive your public key) in some special situations. Also note that it's probably
+         *     not a good idea to set it too low.
+         * </p>
+         */
+        public Builder pubkeyTTL(long days) {
+            if (days < 0 || days > 28 * DAY) throw new IllegalArgumentException("TTL must be between 1 and 28 days");
+            this.pubkeyTTL = days;
+            return this;
+        }
+
         public BitmessageContext build() {
             nonNull("inventory", inventory);
             nonNull("nodeRegistry", nodeRegistry);

domain/src/main/java/ch/dissem/bitmessage/InternalContext.java

@@ -16,7 +16,9 @@
 
 package ch.dissem.bitmessage;
 
-import ch.dissem.bitmessage.entity.*;
+import ch.dissem.bitmessage.entity.BitmessageAddress;
+import ch.dissem.bitmessage.entity.Encrypted;
+import ch.dissem.bitmessage.entity.ObjectMessage;
 import ch.dissem.bitmessage.entity.payload.Broadcast;
 import ch.dissem.bitmessage.entity.payload.GetPubkey;
 import ch.dissem.bitmessage.entity.payload.ObjectPayload;
@@ -29,8 +31,6 @@ import org.slf4j.LoggerFactory;
 import java.io.IOException;
 import java.util.TreeSet;
 
-import static ch.dissem.bitmessage.utils.UnixTime.DAY;
-
 /**
  * The internal context should normally only be used for port implementations. If you need it in your client
  * implementation, you're either doing something wrong, something very weird, or the BitmessageContext should
@@ -59,6 +59,7 @@ public class InternalContext {
     private final long clientNonce;
     private final long networkNonceTrialsPerByte = 1000;
     private final long networkExtraBytes = 1000;
+    private final long pubkeyTTL;
     private long connectionTTL;
     private int connectionLimit;
 
@@ -78,6 +79,7 @@ public class InternalContext {
         this.port = builder.port;
         this.connectionLimit = builder.connectionLimit;
         this.connectionTTL = builder.connectionTTL;
+        this.pubkeyTTL = builder.pubkeyTTL;
 
         Singleton.initialize(security);
 
@@ -193,7 +195,7 @@ public class InternalContext {
 
     public void sendPubkey(final BitmessageAddress identity, final long targetStream) {
         try {
-            long expires = UnixTime.now(+28 * DAY);
+            long expires = UnixTime.now(pubkeyTTL);
             LOG.info("Expires at " + expires);
             final ObjectMessage response = new ObjectMessage.Builder()
                     .stream(targetStream)
@@ -211,7 +213,7 @@ public class InternalContext {
     }
 
     public void requestPubkey(final BitmessageAddress contact) {
-        long expires = UnixTime.now(+2 * DAY);
+        long expires = UnixTime.now(+pubkeyTTL);
         LOG.info("Expires at " + expires);
         final ObjectMessage response = new ObjectMessage.Builder()
                 .stream(contact.getStream())

domain/src/main/java/ch/dissem/bitmessage/ProofOfWorkService.java

@@ -8,6 +8,10 @@ import ch.dissem.bitmessage.ports.MessageRepository;
 import ch.dissem.bitmessage.ports.ProofOfWorkEngine;
 import ch.dissem.bitmessage.ports.ProofOfWorkRepository;
 import ch.dissem.bitmessage.ports.Security;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.List;
 
 import static ch.dissem.bitmessage.utils.Singleton.security;
 
@@ -15,13 +19,19 @@ import static ch.dissem.bitmessage.utils.Singleton.security;
  * @author Christian Basler
  */
 public class ProofOfWorkService implements ProofOfWorkEngine.Callback, InternalContext.ContextHolder {
+    private final static Logger LOG = LoggerFactory.getLogger(ProofOfWorkService.class);
+
     private Security security;
     private InternalContext ctx;
     private ProofOfWorkRepository powRepo;
     private MessageRepository messageRepo;
 
     public void doMissingProofOfWork() {
-        for (byte[] initialHash : powRepo.getItems()) {
+        List<byte[]> items = powRepo.getItems();
+        if (items.isEmpty()) return;
+
+        LOG.info("Doing POW for " + items.size() + " tasks.");
+        for (byte[] initialHash : items) {
             ProofOfWorkRepository.Item item = powRepo.getItem(initialHash);
             security.doProofOfWork(item.object, item.nonceTrialsPerByte, item.extraBytes, this);
         }
@@ -32,8 +42,10 @@ public class ProofOfWorkService implements ProofOfWorkEngine.Callback, InternalC
     }
 
     public void doProofOfWork(BitmessageAddress recipient, ObjectMessage object) {
-        long nonceTrialsPerByte = recipient == null ? 0 : recipient.getPubkey().getNonceTrialsPerByte();
+        long nonceTrialsPerByte = recipient == null ?
-        long extraBytes = recipient == null ? 0 : recipient.getPubkey().getExtraBytes();
+                ctx.getNetworkNonceTrialsPerByte() : recipient.getPubkey().getNonceTrialsPerByte();
+        long extraBytes = recipient == null ?
+                ctx.getNetworkExtraBytes() : recipient.getPubkey().getExtraBytes();
 
         powRepo.putObject(object, nonceTrialsPerByte, extraBytes);
         if (object.getPayload() instanceof PlaintextHolder) {

domain/src/main/java/ch/dissem/bitmessage/entity/CustomMessage.java

@@ -43,7 +43,7 @@ public class CustomMessage implements MessagePayload {
         this.data = data;
     }
 
-    public static MessagePayload read(InputStream in, int length) throws IOException {
+    public static CustomMessage read(InputStream in, int length) throws IOException {
         AccessCounter counter = new AccessCounter();
         return new CustomMessage(varString(in, counter), bytes(in, length - counter.length()));
     }

domain/src/main/java/ch/dissem/bitmessage/ports/AbstractSecurity.java

@@ -43,6 +43,8 @@ public abstract class AbstractSecurity implements Security, InternalContext.Cont
     public static final Logger LOG = LoggerFactory.getLogger(Security.class);
     private static final SecureRandom RANDOM = new SecureRandom();
     private static final BigInteger TWO = BigInteger.valueOf(2);
+    private static final BigInteger TWO_POW_64 = TWO.pow(64);
+    private static final BigInteger TWO_POW_16 = TWO.pow(16);
 
     private final String provider;
     private InternalContext context;
@@ -96,7 +98,6 @@ public abstract class AbstractSecurity implements Security, InternalContext.Cont
 
     public void doProofOfWork(ObjectMessage object, long nonceTrialsPerByte,
                               long extraBytes, ProofOfWorkEngine.Callback callback) {
-        try {
         nonceTrialsPerByte = max(nonceTrialsPerByte, context.getNetworkNonceTrialsPerByte());
         extraBytes = max(extraBytes, context.getNetworkExtraBytes());
 
@@ -105,9 +106,6 @@ public abstract class AbstractSecurity implements Security, InternalContext.Cont
         byte[] target = getProofOfWorkTarget(object, nonceTrialsPerByte, extraBytes);
 
         context.getProofOfWorkEngine().calculateNonce(initialHash, target, callback);
-        } catch (IOException e) {
-            throw new RuntimeException(e);
-        }
     }
 
     public void checkProofOfWork(ObjectMessage object, long nonceTrialsPerByte, long extraBytes)
@@ -124,11 +122,20 @@ public abstract class AbstractSecurity implements Security, InternalContext.Cont
         return sha512(object.getPayloadBytesWithoutNonce());
     }
 
-    private byte[] getProofOfWorkTarget(ObjectMessage object, long nonceTrialsPerByte, long extraBytes) throws IOException {
+    @Override
+    public byte[] getProofOfWorkTarget(ObjectMessage object, long nonceTrialsPerByte, long extraBytes) {
+        if (nonceTrialsPerByte == 0) nonceTrialsPerByte = context.getNetworkNonceTrialsPerByte();
+        if (extraBytes == 0) extraBytes = context.getNetworkExtraBytes();
+
         BigInteger TTL = BigInteger.valueOf(object.getExpiresTime() - UnixTime.now());
-        BigInteger numerator = TWO.pow(64);
+        BigInteger numerator = TWO_POW_64;
         BigInteger powLength = BigInteger.valueOf(object.getPayloadBytesWithoutNonce().length + extraBytes);
-        BigInteger denominator = BigInteger.valueOf(nonceTrialsPerByte).multiply(powLength.add(powLength.multiply(TTL).divide(BigInteger.valueOf(2).pow(16))));
+        BigInteger denominator = BigInteger.valueOf(nonceTrialsPerByte)
+                .multiply(
+                        powLength.add(
+                                powLength.multiply(TTL).divide(TWO_POW_16)
+                        )
+                );
         return Bytes.expand(numerator.divide(denominator).toByteArray(), 8);
     }
 

domain/src/main/java/ch/dissem/bitmessage/ports/Security.java

@@ -136,6 +136,8 @@ public interface Security {
 
     byte[] getInitialHash(ObjectMessage object);
 
+    byte[] getProofOfWorkTarget(ObjectMessage object, long nonceTrialsPerByte, long extraBytes);
+
     /**
      * Calculates the MAC for a message (data)
      *

extensions/src/main/java/ch/dissem/bitmessage/extensions/CryptoCustomMessage.java

@@ -54,8 +54,8 @@ public class CryptoCustomMessage<T extends Streamable> extends CustomMessage {
         this.dataReader = dataReader;
     }
 
-    public static <T extends Streamable> CryptoCustomMessage<T> read(byte[] data, Reader<T> dataReader) throws IOException {
+    public static <T extends Streamable> CryptoCustomMessage<T> read(CustomMessage data, Reader<T> dataReader) throws IOException {
-        CryptoBox cryptoBox = CryptoBox.read(new ByteArrayInputStream(data), data.length);
+        CryptoBox cryptoBox = CryptoBox.read(new ByteArrayInputStream(data.getData()), data.getData().length);
         return new CryptoCustomMessage<>(cryptoBox, dataReader);
     }
 
@@ -111,6 +111,7 @@ public class CryptoCustomMessage<T extends Streamable> extends CustomMessage {
 
     @Override
     public void write(OutputStream out) throws IOException {
+        Encode.varString(COMMAND, out);
         container.write(out);
     }
 

extensions/src/main/java/ch/dissem/bitmessage/extensions/pow/ProofOfWorkRequest.java

@@ -24,6 +24,7 @@ import ch.dissem.bitmessage.utils.Encode;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.util.Arrays;
 
 import static ch.dissem.bitmessage.utils.Decode.*;
 
@@ -80,6 +81,28 @@ public class ProofOfWorkRequest implements Streamable {
         Encode.varBytes(data, out);
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        ProofOfWorkRequest other = (ProofOfWorkRequest) o;
+
+        if (!sender.equals(other.sender)) return false;
+        if (!Arrays.equals(initialHash, other.initialHash)) return false;
+        if (request != other.request) return false;
+        return Arrays.equals(data, other.data);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = sender.hashCode();
+        result = 31 * result + Arrays.hashCode(initialHash);
+        result = 31 * result + request.hashCode();
+        result = 31 * result + Arrays.hashCode(data);
+        return result;
+    }
+
     public static class Reader implements CryptoCustomMessage.Reader<ProofOfWorkRequest> {
         private final BitmessageAddress identity;
 
@@ -93,7 +116,6 @@ public class ProofOfWorkRequest implements Streamable {
         }
     }
 
-
     public enum Request {
         CALCULATE,
         CALCULATING,

extensions/src/test/java/ch/dissem/bitmessage/extensions/CryptoCustomMessageTest.java

@@ -17,8 +17,10 @@
 package ch.dissem.bitmessage.extensions;
 
 import ch.dissem.bitmessage.entity.BitmessageAddress;
+import ch.dissem.bitmessage.entity.CustomMessage;
 import ch.dissem.bitmessage.entity.payload.GenericPayload;
 import ch.dissem.bitmessage.entity.valueobject.PrivateKey;
+import ch.dissem.bitmessage.extensions.pow.ProofOfWorkRequest;
 import ch.dissem.bitmessage.utils.TestBase;
 import ch.dissem.bitmessage.utils.TestUtils;
 import org.junit.Test;
@@ -33,7 +35,7 @@ import static org.junit.Assert.assertEquals;
 
 public class CryptoCustomMessageTest extends TestBase {
     @Test
-    public void testEncryptThenDecrypt() throws Exception {
+    public void ensureEncryptThenDecryptYieldsSameObject() throws Exception {
         PrivateKey privateKey = PrivateKey.read(TestUtils.getResource("BM-2cSqjfJ8xK6UUn5Rw3RpdGQ9RsDkBhWnS8.privkey"));
         BitmessageAddress sendingIdentity = new BitmessageAddress(privateKey);
 
@@ -45,7 +47,9 @@ public class CryptoCustomMessageTest extends TestBase {
         messageBefore.write(out);
         ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
 
-        CryptoCustomMessage<GenericPayload> messageAfter = CryptoCustomMessage.read(out.toByteArray(), new CryptoCustomMessage.Reader<GenericPayload>() {
+        CustomMessage customMessage = CustomMessage.read(in, out.size());
+        CryptoCustomMessage<GenericPayload> messageAfter = CryptoCustomMessage.read(customMessage,
+                new CryptoCustomMessage.Reader<GenericPayload>() {
                     @Override
                     public GenericPayload read(BitmessageAddress ignore, InputStream in) throws IOException {
                         return GenericPayload.read(0, in, 1, 100);
@@ -55,4 +59,28 @@ public class CryptoCustomMessageTest extends TestBase {
 
         assertEquals(payloadBefore, payloadAfter);
     }
+
+    @Test
+    public void testWithActualRequest() throws Exception {
+        PrivateKey privateKey = PrivateKey.read(TestUtils.getResource("BM-2cSqjfJ8xK6UUn5Rw3RpdGQ9RsDkBhWnS8.privkey"));
+        final BitmessageAddress sendingIdentity = new BitmessageAddress(privateKey);
+
+        ProofOfWorkRequest requestBefore = new ProofOfWorkRequest(sendingIdentity, security().randomBytes(64),
+                ProofOfWorkRequest.Request.CALCULATE);
+
+        CryptoCustomMessage<ProofOfWorkRequest> messageBefore = new CryptoCustomMessage<>(requestBefore);
+        messageBefore.signAndEncrypt(sendingIdentity, security().createPublicKey(sendingIdentity.getPublicDecryptionKey()));
+
+
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        messageBefore.write(out);
+        ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
+
+        CustomMessage customMessage = CustomMessage.read(in, out.size());
+        CryptoCustomMessage<ProofOfWorkRequest> messageAfter = CryptoCustomMessage.read(customMessage,
+                new ProofOfWorkRequest.Reader(sendingIdentity));
+        ProofOfWorkRequest requestAfter = messageAfter.decrypt(sendingIdentity.getPublicDecryptionKey());
+
+        assertEquals(requestBefore, requestAfter);
+    }
 }