impleemnt auth / connection commands

tuhuynh27 · tuhuynh27 · commit 7de2d99791ba · 2021-12-10T07:17:01.000+08:00
diff --git a/docs/src/guide/overview/commands.md b/docs/src/guide/overview/commands.md
@@ -4,12 +4,24 @@ Follows [Redis's commands](https://redis.io/commands).
 
 Implemented commands:
 
+<details>
+    <summary>Server</summary>
+
+- INFO
+- FLUSHDB
+- TIME
+
+</details>
+
 <details>
     <summary>Connection</summary>
 
+- AUTH
 - ECHO
 - PING
 - QUIT
+- CLIENT ID
+- CLIENT INFO
 
 </details>
 
diff --git a/docs/src/guide/overview/install.md b/docs/src/guide/overview/install.md
@@ -30,6 +30,7 @@ Parameters:
     -aof <Boolean>  enable append-only-file (default: false)
     -ai <Integer>   define append-only interval in ms (default: 1000)
     -dir <String>   working directory (default: ./)
+    -pw <String>    authenticate password (default: none - no password)
 
 ## Docker
 
diff --git a/server/src/main/java/dev/keva/server/command/impl/connection/Auth.java b/server/src/main/java/dev/keva/server/command/impl/connection/Auth.java
@@ -0,0 +1,37 @@
+package dev.keva.server.command.impl.connection;
+
+import dev.keva.ioc.annotation.Autowired;
+import dev.keva.ioc.annotation.Component;
+import dev.keva.protocol.resp.reply.ErrorReply;
+import dev.keva.protocol.resp.reply.Reply;
+import dev.keva.protocol.resp.reply.StatusReply;
+import dev.keva.server.command.annotation.CommandImpl;
+import dev.keva.server.command.annotation.Execute;
+import dev.keva.server.command.annotation.ParamLength;
+import dev.keva.server.command.impl.connection.manager.AuthManager;
+import dev.keva.server.config.KevaConfig;
+import io.netty.channel.ChannelHandlerContext;
+
+@Component
+@CommandImpl("auth")
+@ParamLength(1)
+public class Auth {
+    private final KevaConfig kevaConfig;
+    private final AuthManager authManager;
+
+    @Autowired
+    public Auth(KevaConfig kevaConfig, AuthManager authManager) {
+        this.kevaConfig = kevaConfig;
+        this.authManager = authManager;
+    }
+
+    @Execute
+    public Reply<?> execute(byte[] password, ChannelHandlerContext ctx) {
+        String passwordString = new String(password);
+        if (kevaConfig.getPassword().equals(passwordString)) {
+            authManager.authenticate(ctx.channel());
+            return new StatusReply("OK");
+        }
+        return new ErrorReply("ERR WRONGPASS invalid password.");
+    }
+}
diff --git a/server/src/main/java/dev/keva/server/command/impl/connection/Client.java b/server/src/main/java/dev/keva/server/command/impl/connection/Client.java
@@ -0,0 +1,28 @@
+package dev.keva.server.command.impl.connection;
+
+import dev.keva.ioc.annotation.Component;
+import dev.keva.protocol.resp.reply.BulkReply;
+import dev.keva.protocol.resp.reply.ErrorReply;
+import dev.keva.protocol.resp.reply.Reply;
+import dev.keva.server.command.annotation.CommandImpl;
+import dev.keva.server.command.annotation.Execute;
+import dev.keva.server.command.annotation.ParamLength;
+import io.netty.channel.ChannelHandlerContext;
+
+@Component
+@CommandImpl("client")
+@ParamLength(1)
+public class Client {
+    @Execute
+    public Reply<?> execute(byte[] param, ChannelHandlerContext ctx) {
+        String paramStr = new String(param);
+        if (paramStr.equalsIgnoreCase("id")) {
+            return new BulkReply(ctx.channel().id().asShortText());
+        } else if (paramStr.equalsIgnoreCase("info")) {
+            String info = String.format("id=%s, addr=%s\n", ctx.channel().id().asShortText(), ctx.channel().remoteAddress().toString());
+            return new BulkReply(info);
+        } else {
+            return new ErrorReply("ERROR Unsupported query");
+        }
+    }
+}
diff --git a/server/src/main/java/dev/keva/server/command/impl/connection/manager/AuthManager.java b/server/src/main/java/dev/keva/server/command/impl/connection/manager/AuthManager.java
@@ -0,0 +1,24 @@
+package dev.keva.server.command.impl.connection.manager;
+
+import dev.keva.ioc.annotation.Component;
+import io.netty.channel.Channel;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+@Component
+public class AuthManager {
+    private final ConcurrentMap<Channel, Boolean> auths = new ConcurrentHashMap<>();
+
+    public boolean isAuthenticated(Channel channel) {
+        return auths.get(channel) != null && auths.get(channel);
+    }
+
+    public void authenticate(Channel channel) {
+        auths.put(channel, true);
+    }
+
+    public void unAuthenticate(Channel channel) {
+        auths.remove(channel);
+    }
+}
diff --git a/server/src/main/java/dev/keva/server/command/impl/server/FlushDB.java b/server/src/main/java/dev/keva/server/command/impl/server/FlushDB.java
@@ -0,0 +1,27 @@
+package dev.keva.server.command.impl.server;
+
+import dev.keva.ioc.annotation.Autowired;
+import dev.keva.ioc.annotation.Component;
+import dev.keva.protocol.resp.reply.StatusReply;
+import dev.keva.server.command.annotation.CommandImpl;
+import dev.keva.server.command.annotation.Execute;
+import dev.keva.server.command.annotation.ParamLength;
+import dev.keva.store.KevaDatabase;
+
+@Component
+@CommandImpl("flushdb")
+@ParamLength(type = ParamLength.Type.AT_MOST, value = 1)
+public class FlushDB {
+    private final KevaDatabase database;
+
+    @Autowired
+    public FlushDB(KevaDatabase database) {
+        this.database = database;
+    }
+
+    @Execute
+    public StatusReply execute(byte[] ignored) {
+        database.flush();
+        return StatusReply.OK;
+    }
+}
diff --git a/server/src/main/java/dev/keva/server/command/impl/server/Info.java b/server/src/main/java/dev/keva/server/command/impl/server/Info.java
@@ -1,4 +1,4 @@
-package dev.keva.server.command.impl.connection;
+package dev.keva.server.command.impl.server;
 
 import dev.keva.ioc.annotation.Component;
 import dev.keva.protocol.resp.reply.BulkReply;
diff --git a/server/src/main/java/dev/keva/server/command/impl/server/Time.java b/server/src/main/java/dev/keva/server/command/impl/server/Time.java
@@ -0,0 +1,20 @@
+package dev.keva.server.command.impl.server;
+
+import dev.keva.ioc.annotation.Component;
+import dev.keva.protocol.resp.reply.BulkReply;
+import dev.keva.protocol.resp.reply.MultiBulkReply;
+import dev.keva.server.command.annotation.CommandImpl;
+import dev.keva.server.command.annotation.Execute;
+import dev.keva.server.command.annotation.ParamLength;
+
+@Component
+@CommandImpl("time")
+@ParamLength(0)
+public class Time {
+    @Execute
+    public MultiBulkReply execute(byte[] ignored) {
+        BulkReply[] replies = new BulkReply[1];
+        replies[0] = new BulkReply(Long.toString(System.currentTimeMillis() / 1000));
+        return new MultiBulkReply(replies);
+    }
+}
diff --git a/server/src/main/java/dev/keva/server/command/mapping/CommandMapper.java b/server/src/main/java/dev/keva/server/command/mapping/CommandMapper.java
@@ -12,6 +12,7 @@
 import dev.keva.server.command.annotation.Mutate;
 import dev.keva.server.command.annotation.ParamLength;
 import dev.keva.server.aof.AOFContainer;
+import dev.keva.server.command.impl.connection.manager.AuthManager;
 import dev.keva.server.command.impl.transaction.manager.TransactionManager;
 import dev.keva.server.config.KevaConfig;
 import dev.keva.store.KevaDatabase;
@@ -38,6 +39,9 @@ public class CommandMapper {
     @Autowired
     private TransactionManager txManager;
 
+    @Autowired
+    private AuthManager authManager;
+
     @Autowired
     private KevaDatabase database;
 
@@ -60,8 +64,17 @@ public void init() {
                     val paramLengthType = aClass.getAnnotation(ParamLength.class) != null ? aClass.getAnnotation(ParamLength.class).type() : null;
                     val instance = context.getBean(aClass);
                     val isMutate = aClass.getAnnotation(Mutate.class) != null;
+                    val password = kevaConfig.getPassword();
+                    val isAuthEnabled = password != null && password.length() > 0;
 
                     methods.put(new BytesKey(name.getBytes()), (ctx, command) -> {
+                        if (isAuthEnabled && !Arrays.equals(command.getName(), "auth".getBytes())) {
+                            boolean authenticated = authManager.isAuthenticated(ctx.channel());
+                            if (!authenticated) {
+                                return new ErrorReply("ERR NOAUTH Authentication required.");
+                            }
+                        }
+
                         if (ctx != null) {
                             val txContext = txManager.getTransactions().get(ctx.channel());
                             if (txContext != null && txContext.isQueuing()) {
diff --git a/server/src/main/java/dev/keva/server/config/KevaConfig.java b/server/src/main/java/dev/keva/server/config/KevaConfig.java
@@ -40,6 +40,10 @@ public class KevaConfig {
     @CliProp(name = "dir", type = CliPropType.VAL)
     private String workDirectory;
 
+    @ConfigProp(name = "requirepass", defaultVal = "")
+    @CliProp(name = "pw", type = CliPropType.VAL)
+    private String password;
+
     /**
      * @return KevaConfig with sensible defaults
      */
diff --git a/server/src/main/java/dev/keva/server/core/KevaServer.java b/server/src/main/java/dev/keva/server/core/KevaServer.java
@@ -124,6 +124,6 @@ public void run() {
 
     @Override
     public void clear() {
-        database.clear();
+        database.flush();
     }
 }
diff --git a/server/src/test/java/dev/keva/server/core/AbstractServerTest.java b/server/src/test/java/dev/keva/server/core/AbstractServerTest.java
@@ -39,6 +39,20 @@ void info() {
         }
     }
 
+    @Test
+    void flush() {
+        try {
+            val setAbc = jedis.set("abc", "123");
+            assertEquals("OK", setAbc);
+            val flush = jedis.flushDB();
+            assertEquals("OK", flush);
+            val getAbc = jedis.get("abc");
+            assertNull(getAbc);
+        } catch (Exception e) {
+            fail(e);
+        }
+    }
+
     @Test
     void getSetNull() {
         try {
diff --git a/server/src/test/java/dev/keva/server/core/KevaServerTest.java b/server/src/test/java/dev/keva/server/core/KevaServerTest.java
@@ -23,6 +23,7 @@ static void startServer() throws Exception {
                 .aof(false)
                 .hostname(host)
                 .port(port)
+                .password("keva-auth")
                 .build();
         server = KevaServer.of(config);
 
@@ -39,7 +40,9 @@ static void startServer() throws Exception {
         TimeUnit.SECONDS.sleep(2);
 
         jedis = new Jedis(host, port);
+        jedis.auth("keva-auth");
         subscriber = new Jedis(host, port);
+        subscriber.auth("keva-auth");
     }
 
     @AfterAll
diff --git a/store/src/main/java/dev/keva/store/KevaDatabase.java b/store/src/main/java/dev/keva/store/KevaDatabase.java
@@ -5,7 +5,7 @@
 public interface KevaDatabase {
     Lock getLock();
 
-    void clear();
+    void flush();
 
     void put(byte[] key, byte[] val);
 
diff --git a/store/src/main/java/dev/keva/store/impl/OffHeapDatabaseImpl.java b/store/src/main/java/dev/keva/store/impl/OffHeapDatabaseImpl.java
@@ -48,7 +48,7 @@ public OffHeapDatabaseImpl(DatabaseConfig config) {
     }
 
     @Override
-    public void clear() {
+    public void flush() {
         lock.lock();
         try {
             chronicleMap.clear();
diff --git a/store/src/main/java/dev/keva/store/impl/OnHeapDatabaseImpl.java b/store/src/main/java/dev/keva/store/impl/OnHeapDatabaseImpl.java
@@ -19,7 +19,7 @@ public class OnHeapDatabaseImpl implements KevaDatabase {
     private final Map<BytesKey, BytesValue> map = new HashMap<>(100);
 
     @Override
-    public void clear() {
+    public void flush() {
         map.clear();
     }
 

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-package dev.keva.server.command.impl.connection;`
	`1`	`+package dev.keva.server.command.impl.server;`
`2`	`2`
`3`	`3`	`import dev.keva.ioc.annotation.Component;`
`4`	`4`	`import dev.keva.protocol.resp.reply.BulkReply;`