From b5c5cda2dad814c6ebe90b47e1cdbe42431866f1 Mon Sep 17 00:00:00 2001
From: Andrea <andrea.raspitzu@acinq.fr>
Date: Mon, 22 Jul 2019 14:20:00 +0200
Subject: [PATCH] Store block counter on Channel-DB for fundee entropy

---
 .../scala/fr/acinq/eclair/db/ChannelsDb.scala |  2 ++
 .../eclair/db/sqlite/SqliteChannelsDb.scala   | 34 +++++++++++++++++--
 .../eclair/db/SqliteChannelsDbSpec.scala      | 15 ++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/eclair-core/src/main/scala/fr/acinq/eclair/db/ChannelsDb.scala b/eclair-core/src/main/scala/fr/acinq/eclair/db/ChannelsDb.scala
index dd9fb9e120..ea58fbdabb 100644
--- a/eclair-core/src/main/scala/fr/acinq/eclair/db/ChannelsDb.scala
+++ b/eclair-core/src/main/scala/fr/acinq/eclair/db/ChannelsDb.scala
@@ -21,6 +21,8 @@ import fr.acinq.eclair.channel.HasCommitments
 
 trait ChannelsDb {
 
+  def getCounterFor(blockHeight: Long): Long
+
   def addOrUpdateChannel(state: HasCommitments)
 
   def removeChannel(channelId: ByteVector32)
diff --git a/eclair-core/src/main/scala/fr/acinq/eclair/db/sqlite/SqliteChannelsDb.scala b/eclair-core/src/main/scala/fr/acinq/eclair/db/sqlite/SqliteChannelsDb.scala
index d747a3b9dd..be1d21eb55 100644
--- a/eclair-core/src/main/scala/fr/acinq/eclair/db/sqlite/SqliteChannelsDb.scala
+++ b/eclair-core/src/main/scala/fr/acinq/eclair/db/sqlite/SqliteChannelsDb.scala
@@ -16,7 +16,7 @@
 
 package fr.acinq.eclair.db.sqlite
 
-import java.sql.{Connection, Statement}
+import java.sql.{Connection, SQLException, Statement}
 
 import fr.acinq.bitcoin.ByteVector32
 import fr.acinq.eclair.channel.HasCommitments
@@ -36,6 +36,7 @@ class SqliteChannelsDb(sqlite: Connection) extends ChannelsDb with Logging {
 
   private def migration12(statement: Statement) = {
     statement.executeUpdate("ALTER TABLE local_channels ADD COLUMN is_closed BOOLEAN NOT NULL DEFAULT 0")
+    statement.executeUpdate("CREATE TABLE IF NOT EXISTS block_key_counter (block_height INTEGER NOT NULL PRIMARY KEY, counter INTEGER NOT NULL)")
   }
 
   using(sqlite.createStatement()) { statement =>
@@ -49,11 +50,40 @@ class SqliteChannelsDb(sqlite: Connection) extends ChannelsDb with Logging {
         statement.executeUpdate("CREATE TABLE IF NOT EXISTS local_channels (channel_id BLOB NOT NULL PRIMARY KEY, data BLOB NOT NULL, is_closed BOOLEAN NOT NULL DEFAULT 0)")
         statement.executeUpdate("CREATE TABLE IF NOT EXISTS htlc_infos (channel_id BLOB NOT NULL, commitment_number BLOB NOT NULL, payment_hash BLOB NOT NULL, cltv_expiry INTEGER NOT NULL, FOREIGN KEY(channel_id) REFERENCES local_channels(channel_id))")
         statement.executeUpdate("CREATE INDEX IF NOT EXISTS htlc_infos_idx ON htlc_infos(channel_id, commitment_number)")
-
+        statement.executeUpdate("CREATE TABLE IF NOT EXISTS block_key_counter (block_height INTEGER NOT NULL PRIMARY KEY, counter INTEGER NOT NULL)")
       case unknownVersion => throw new RuntimeException(s"Unknown version of DB $DB_NAME found, version=$unknownVersion")
     }
   }
 
+  override def getCounterFor(blockHeight: Long): Long = synchronized {
+    val counterOld = getCounter(blockHeight)
+    setCounter(blockHeight, counterOld + 1)
+    counterOld
+  }
+
+  private def getCounter(blockHeight: Long): Long = {
+    using(sqlite.prepareStatement("SELECT counter FROM block_key_counter WHERE block_height=?")) { statement =>
+      statement.setLong(1, blockHeight)
+      val rs = statement.executeQuery()
+      if(rs.next()) rs.getLong("counter") else 0
+    }
+  }
+
+  private def setCounter(blockHeight: Long, counter: Long) = {
+    using(sqlite.prepareStatement("UPDATE block_key_counter SET counter=? WHERE block_height=?")) { statement =>
+      statement.setLong(1, counter)
+      statement.setLong(2, blockHeight)
+      if(statement.executeUpdate() != 1){
+        using(sqlite.prepareStatement("INSERT INTO block_key_counter VALUES(?, ?)")) { insert =>
+          insert.setLong(1, blockHeight)
+          insert.setLong(2, counter)
+          insert.executeUpdate()
+        }
+      }
+    }
+  }
+
+
   override def addOrUpdateChannel(state: HasCommitments): Unit = {
     val data = stateDataCodec.encode(state).require.toByteArray
     using (sqlite.prepareStatement("UPDATE local_channels SET data=? WHERE channel_id=?")) { update =>
diff --git a/eclair-core/src/test/scala/fr/acinq/eclair/db/SqliteChannelsDbSpec.scala b/eclair-core/src/test/scala/fr/acinq/eclair/db/SqliteChannelsDbSpec.scala
index 88052c3158..db637f5ddd 100644
--- a/eclair-core/src/test/scala/fr/acinq/eclair/db/SqliteChannelsDbSpec.scala
+++ b/eclair-core/src/test/scala/fr/acinq/eclair/db/SqliteChannelsDbSpec.scala
@@ -93,5 +93,20 @@ class SqliteChannelsDbSpec extends FunSuite {
       assert(getVersion(statement, "channels", 1) == 2) // version changed from 1 -> 2
     }
     assert(db.listLocalChannels() === List(channel))
+    assert(db.getCounterFor(123) === 0)
   }
+
+  test("channel keypath counter should get and increment") {
+
+    val sqlite = TestConstants.sqliteInMemory()
+    val channelDb = new SqliteChannelsDb(sqlite)
+
+    assert(channelDb.getCounterFor(123) == 0)
+    assert(channelDb.getCounterFor(123) == 1)
+    assert(channelDb.getCounterFor(123) == 2)
+    assert(channelDb.getCounterFor(124) == 0)
+    assert(channelDb.getCounterFor(125) == 0)
+    assert(channelDb.getCounterFor(124) == 1)
+  }
+
 }
\ No newline at end of file