From a36ea6fb1941de1e927f7733e462bde7b5238645 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maciej=20=C5=BBenczykowski?= <maze@google.com>
Date: Thu, 1 Apr 2021 20:34:57 -0700
Subject: [PATCH] implement insertOrReplace()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Test: atest com.android.networkstack.tethering.BpfMapTest
Signed-off-by: Maciej Żenczykowski <maze@google.com>
Change-Id: I4a40898f03293d6d79b57c35f743271c669a8ea7
---
 .../networkstack/tethering/BpfMap.java        | 30 +++++++++++++++++++
 .../networkstack/tethering/BpfMapTest.java    | 25 +++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/Tethering/src/com/android/networkstack/tethering/BpfMap.java b/Tethering/src/com/android/networkstack/tethering/BpfMap.java
index e9b4ccf2f4..1363dc5150 100644
--- a/Tethering/src/com/android/networkstack/tethering/BpfMap.java
+++ b/Tethering/src/com/android/networkstack/tethering/BpfMap.java
@@ -98,6 +98,7 @@ public class BpfMap<K extends Struct, V extends Struct> implements AutoCloseable
 
     /**
      * Update an existing or create a new key -> value entry in an eBbpf map.
+     * (use insertOrReplaceEntry() if you need to know whether insert or replace happened)
      */
     public void updateEntry(K key, V value) throws ErrnoException {
         writeToMapEntry(mMapFd, key.writeToBytes(), value.writeToBytes(), BPF_ANY);
@@ -133,6 +134,35 @@ public class BpfMap<K extends Struct, V extends Struct> implements AutoCloseable
         }
     }
 
+    /**
+     * Update an existing or create a new key -> value entry in an eBbpf map.
+     * Returns true if inserted, false if replaced.
+     * (use updateEntry() if you don't care whether insert or replace happened)
+     * Note: see inline comment below if running concurrently with delete operations.
+     */
+    public boolean insertOrReplaceEntry(K key, V value)
+            throws ErrnoException {
+        try {
+            writeToMapEntry(mMapFd, key.writeToBytes(), value.writeToBytes(), BPF_NOEXIST);
+            return true;   /* insert succeeded */
+        } catch (ErrnoException e) {
+            if (e.errno != EEXIST) throw e;
+        }
+        try {
+            writeToMapEntry(mMapFd, key.writeToBytes(), value.writeToBytes(), BPF_EXIST);
+            return false;   /* replace succeeded */
+        } catch (ErrnoException e) {
+            if (e.errno != ENOENT) throw e;
+        }
+        /* If we reach here somebody deleted after our insert attempt and before our replace:
+         * this implies a race happened.  The kernel bpf delete interface only takes a key,
+         * and not the value, so we can safely pretend the replace actually succeeded and
+         * was immediately followed by the other thread's delete, since the delete cannot
+         * observe the potential change to the value.
+         */
+        return false;   /* pretend replace succeeded */
+    }
+
     /** Remove existing key from eBpf map. Return false if map was not modified. */
     public boolean deleteEntry(K key) throws ErrnoException {
         return deleteMapEntry(mMapFd, key.writeToBytes());
diff --git a/Tethering/tests/privileged/src/com/android/networkstack/tethering/BpfMapTest.java b/Tethering/tests/privileged/src/com/android/networkstack/tethering/BpfMapTest.java
index 62302c37c8..d24c35a44e 100644
--- a/Tethering/tests/privileged/src/com/android/networkstack/tethering/BpfMapTest.java
+++ b/Tethering/tests/privileged/src/com/android/networkstack/tethering/BpfMapTest.java
@@ -209,7 +209,7 @@ public final class BpfMapTest {
     }
 
     @Test
-    public void testUpdateBpfMap() throws Exception {
+    public void testUpdateEntry() throws Exception {
         final TetherDownstream6Key key = mTestData.keyAt(0);
         final Tether6Value value = mTestData.valueAt(0);
         final Tether6Value value2 = mTestData.valueAt(1);
@@ -231,6 +231,29 @@ public final class BpfMapTest {
         assertFalse(mTestMap.containsKey(key));
     }
 
+    @Test
+    public void testInsertOrReplaceEntry() throws Exception {
+        final TetherDownstream6Key key = mTestData.keyAt(0);
+        final Tether6Value value = mTestData.valueAt(0);
+        final Tether6Value value2 = mTestData.valueAt(1);
+        assertFalse(mTestMap.deleteEntry(key));
+
+        // insertOrReplaceEntry will create an entry if it does not exist already.
+        assertTrue(mTestMap.insertOrReplaceEntry(key, value));
+        assertTrue(mTestMap.containsKey(key));
+        final Tether6Value result = mTestMap.getValue(key);
+        assertEquals(value, result);
+
+        // updateEntry will update an entry that already exists.
+        assertFalse(mTestMap.insertOrReplaceEntry(key, value2));
+        assertTrue(mTestMap.containsKey(key));
+        final Tether6Value result2 = mTestMap.getValue(key);
+        assertEquals(value2, result2);
+
+        assertTrue(mTestMap.deleteEntry(key));
+        assertFalse(mTestMap.containsKey(key));
+    }
+
     @Test
     public void testInsertReplaceEntry() throws Exception {
         final TetherDownstream6Key key = mTestData.keyAt(0);