Blockchain implementation in Java 90% complete

Author: rampatra

src/main/java/com/rampatra/blockchain/Block.java

@@ -1,10 +1,12 @@
 package com.rampatra.blockchain;
 
+import java.io.Serializable;
+
 /**
  * @author rampatra
  * @since 2019-03-05
  */
-public class Block {
+public class Block implements Serializable {
 
     private int index;
     private String previousHash;

src/main/java/com/rampatra/blockchain/Blockchain.java

@@ -3,43 +3,70 @@
 import java.nio.charset.StandardCharsets;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
-import java.util.ArrayList;
 import java.util.Date;
 import java.util.List;
+import java.util.ListIterator;
 
 /**
  * @author rampatra
  * @since 2019-03-05
  */
-public class Blockchain { // TODO: P2P
+public class Blockchain {
 
-    private List<Block> blockchain;
+    private List<Block> blocks;
     private int difficulty;
 
-    public Blockchain(List<Block> blockchain, int difficulty) {
-        this.blockchain = blockchain;
+    public Blockchain(List<Block> blocks, int difficulty) {
+        this.blocks = blocks;
         this.difficulty = difficulty;
-        this.blockchain.add(getGenesisBlock("Blockchain in Java"));
+        this.blocks.add(getGenesisBlock());
     }
 
-    public List<Block> getBlockchain() {
-        return blockchain;
+    public List<Block> getBlocks() {
+        return blocks;
     }
 
-    public void mine(String data) {
-        Block previousBlock = blockchain.get(blockchain.size() - 1);
+    public int getSize() {
+        return blocks.size();
+    }
+
+    public Block getLatestBlock() {
+        if (blocks.isEmpty()) return null;
+
+        return blocks.get(blocks.size() - 1);
+    }
+
+    public void addBlock(Block block) {
+        blocks.add(block);
+    }
+
+    public Block mine(String data) {
+        Block previousBlock = getLatestBlock();
         Block nextBlock = getNextBlock(previousBlock, data);
-        
+
         if (isValidNextBlock(previousBlock, nextBlock)) {
-            blockchain.add(nextBlock);
+            blocks.add(nextBlock);
+            return nextBlock;
         } else {
             throw new RuntimeException("Invalid block");
         }
     }
 
-    private Block getGenesisBlock(String data) {
+    public boolean isValidChain() {
+        ListIterator<Block> listIterator = blocks.listIterator();
+        listIterator.next();
+        while (listIterator.hasPrevious() && listIterator.hasNext()) {
+            if (!isValidNextBlock(listIterator.previous(), listIterator.next())) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    private Block getGenesisBlock() {
         final long timestamp = new Date().getTime();
         int nonce = 0;
+        String data = "Blockchain in Java";
         String hash;
         while (!isValidHashDifficulty(hash = calculateHashForBlock(0, "0", timestamp, data, nonce))) {
             nonce++;
@@ -108,11 +135,4 @@ private static String bytesToHex(byte[] hash) {
         }
         return hexString.toString();
     }
-
-    public static void main(String[] args) {
-        Blockchain blockchain = new Blockchain(new ArrayList<>(), 3);
-        blockchain.mine("12");
-        blockchain.mine("26");
-        System.out.println("Blockchain: " + blockchain.getBlockchain());
-    }
 }

src/main/java/com/rampatra/blockchain/Message.java

@@ -0,0 +1,70 @@
+package com.rampatra.blockchain;
+
+import java.io.Serializable;
+
+/**
+ * @author rampatra
+ * @since 2019-03-10
+ */
+public class Message implements Serializable {
+
+    enum MessageType {
+        REQUEST_LATEST_BLOCK,
+        RECEIVE_LATEST_BLOCK,
+        REQUEST_BLOCKCHAIN,
+        RECEIVE_BLOCKCHAIN
+    }
+
+    private MessageType messageType;
+    private Block latestBlock;
+    private Blockchain blockchain;
+
+
+    public MessageType getMessageType() {
+        return messageType;
+    }
+
+    public Block getLatestBlock() {
+        return latestBlock;
+    }
+
+    public Blockchain getBlockchain() {
+        return blockchain;
+    }
+
+    public static final class MessageBuilder {
+        private MessageType messageType;
+        private Block latestBlock;
+        private Blockchain blockchain;
+
+        private MessageBuilder() {
+        }
+
+        public static MessageBuilder aMessage() {
+            return new MessageBuilder();
+        }
+
+        public MessageBuilder withMessageType(MessageType messageType) {
+            this.messageType = messageType;
+            return this;
+        }
+
+        public MessageBuilder withLatestBlock(Block latestBlock) {
+            this.latestBlock = latestBlock;
+            return this;
+        }
+
+        public MessageBuilder withBlockchain(Blockchain blockchain) {
+            this.blockchain = blockchain;
+            return this;
+        }
+
+        public Message build() {
+            Message message = new Message();
+            message.latestBlock = this.latestBlock;
+            message.blockchain = this.blockchain;
+            message.messageType = this.messageType;
+            return message;
+        }
+    }
+}

src/main/java/com/rampatra/blockchain/MessageHandler.java

@@ -0,0 +1,96 @@
+package com.rampatra.blockchain;
+
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.net.Socket;
+
+/**
+ * @author rampatra
+ * @since 2019-03-10
+ */
+public class MessageHandler implements Runnable {
+
+    ObjectInputStream in;
+    ObjectOutputStream out;
+    Peer peer;
+
+    MessageHandler(Socket client, Peer peer) {
+        try {
+            in = new ObjectInputStream(client.getInputStream());
+            out = new ObjectOutputStream(client.getOutputStream());
+            this.peer = peer;
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void run() {
+        Message message;
+        try {
+            while ((message = (Message) in.readObject()) != null) {
+                handleMessage(message);
+            }
+            in.close();
+            out.close();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private void handleMessage(Message message) {
+        Message.MessageType type = message.getMessageType();
+        switch (type) {
+            case REQUEST_LATEST_BLOCK:
+                sendMessage(Message.MessageBuilder
+                        .aMessage()
+                        .withMessageType(Message.MessageType.RECEIVE_LATEST_BLOCK)
+                        .withLatestBlock(peer.getBlockchain().getLatestBlock())
+                        .build());
+                break;
+            case RECEIVE_LATEST_BLOCK:
+                handleLatestBlock(message.getLatestBlock());
+                break;
+            case REQUEST_BLOCKCHAIN:
+                sendMessage(Message.MessageBuilder
+                        .aMessage()
+                        .withMessageType(Message.MessageType.RECEIVE_BLOCKCHAIN)
+                        .withBlockchain(peer.getBlockchain())
+                        .build());
+                break;
+            case RECEIVE_BLOCKCHAIN:
+                handleBlockchain(message.getBlockchain());
+                break;
+            default:
+                throw new RuntimeException("Unknown message type");
+        }
+    }
+
+    private void sendMessage(Message message) {
+        try {
+            out.writeObject(message);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private void handleLatestBlock(Block latestBlock) {
+        Block latestBlockWithCurrentPeer = peer.getBlockchain().getLatestBlock();
+
+        if (latestBlock.getPreviousHash().equals(latestBlockWithCurrentPeer.getHash())) {
+            peer.getBlockchain().addBlock(latestBlock);
+        } else if (latestBlock.getIndex() > latestBlockWithCurrentPeer.getIndex()) {
+            sendMessage(Message.MessageBuilder
+                    .aMessage()
+                    .withMessageType(Message.MessageType.REQUEST_BLOCKCHAIN)
+                    .build());
+        }
+    }
+
+    private void handleBlockchain(Blockchain blockchain) {
+        if (blockchain.isValidChain() && blockchain.getSize() > peer.getBlockchain().getSize()) {
+            peer.setBlockchain(blockchain);
+        }
+    }
+}

src/main/java/com/rampatra/blockchain/P2P.java

@@ -0,0 +1,36 @@
+package com.rampatra.blockchain;
+
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+/**
+ * @author rampatra
+ * @since 2019-03-09
+ */
+public class P2P {
+
+    private static List<Peer> peers = new ArrayList<>();
+    private static int lastUsedPort = 4000;
+
+    public static Peer addPeer(Blockchain blockchain) {
+        Peer peer = new Peer(blockchain, lastUsedPort++);
+        peers.add(peer);
+        return peer;
+    }
+
+    public static void removePeer(Peer peer) {
+        Iterator<Peer> iterator = peers.iterator();
+        while (iterator.hasNext()) {
+            if (iterator.next().equals(peer)) {
+                peer.stopServer();
+                iterator.remove();
+                break;
+            }
+        }
+    }
+
+    public static List<Peer> getPeers() {
+        return peers;
+    }
+}