Add argument error handling to server commands

Signed-off-by: Devansh Singh <[email protected]>

server/commands.go

@@ -1,6 +1,11 @@
 package server
 
-import "log"
+import (
+	"errors"
+	"log"
+)
+
+var ErrNotEnoughArgs = errors.New("not enough arguments for command")
 
 func (c *Command) write(msg string) {
 	_, err := c.Conn.Write([]byte(msg + "\n"))

server/generic.go

@@ -7,6 +7,10 @@ import (
 )
 
 func (s *Server) get(cmd Command) {
+	if len(cmd.Args) < 1 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	result, err := s.DB.Get(cmd.Args[0])
 	if err != nil {
 		cmd.error(err)
@@ -16,13 +20,21 @@ func (s *Server) get(cmd Command) {
 }
 
 func (s *Server) set(cmd Command) {
+	if len(cmd.Args) < 2 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	err := s.DB.Set(cmd.Args[0], []byte(cmd.Args[1]))
 	if err != nil {
 		cmd.error(err)
 	}
 }
 
 func (s *Server) setEx(cmd Command) {
+	if len(cmd.Args) < 3 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	ttl, err := strconv.Atoi(cmd.Args[2])
 	if err != nil {
 		cmd.error(err)
@@ -34,6 +46,10 @@ func (s *Server) setEx(cmd Command) {
 }
 
 func (s *Server) del(cmd Command) {
+	if len(cmd.Args) < 1 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	err := s.DB.Del(cmd.Args[0])
 	if err != nil {
 		cmd.error(err)
@@ -48,6 +64,10 @@ func (s *Server) mGet(cmd Command) {
 }
 
 func (s *Server) expire(cmd Command) {
+	if len(cmd.Args) < 2 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	expiration, err := strconv.Atoi(cmd.Args[1])
 	if err != nil {
 		cmd.error(err)
@@ -66,6 +86,10 @@ func (s *Server) keys(cmd Command) {
 }
 
 func (s *Server) exists(cmd Command) {
+	if len(cmd.Args) < 1 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	ok := s.DB.Exists(cmd.Args[0])
 	if !ok {
 		cmd.write("FALSE")
@@ -75,13 +99,21 @@ func (s *Server) exists(cmd Command) {
 }
 
 func (s *Server) persist(cmd Command) {
+	if len(cmd.Args) < 1 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	err := s.DB.Persist(cmd.Args[0])
 	if err != nil {
 		cmd.error(err)
 	}
 }
 
 func (s *Server) expireTime(cmd Command) {
+	if len(cmd.Args) < 1 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	exp, err := s.DB.ExpireTime(cmd.Args[0])
 	if err != nil {
 		cmd.error(err)
@@ -91,6 +123,10 @@ func (s *Server) expireTime(cmd Command) {
 }
 
 func (s *Server) ttl(cmd Command) {
+	if len(cmd.Args) < 1 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	ttl, err := s.DB.TTL(cmd.Args[0])
 	if err != nil {
 		cmd.error(err)

server/set.go

@@ -6,10 +6,18 @@ import (
 )
 
 func (s *Server) sAdd(cmd Command) {
+	if len(cmd.Args) < 2 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	s.DB.SAdd(cmd.Args[0], cmd.Args[1])
 }
 
 func (s *Server) sMembers(cmd Command) {
+	if len(cmd.Args) < 1 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	elements, err := s.DB.SMembers(cmd.Args[0])
 	if err != nil {
 		cmd.error(err)
@@ -21,6 +29,10 @@ func (s *Server) sMembers(cmd Command) {
 }
 
 func (s *Server) sCard(cmd Command) {
+	if len(cmd.Args) < 1 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	size, err := s.DB.SCard(cmd.Args[0])
 	if err != nil {
 		cmd.error(err)
@@ -30,6 +42,10 @@ func (s *Server) sCard(cmd Command) {
 }
 
 func (s *Server) sIsMember(cmd Command) {
+	if len(cmd.Args) < 2 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	ok, err := s.DB.SIsMember(cmd.Args[0], cmd.Args[1])
 	if err != nil {
 		cmd.error(err)
@@ -43,6 +59,10 @@ func (s *Server) sIsMember(cmd Command) {
 }
 
 func (s *Server) sDiff(cmd Command) {
+	if len(cmd.Args) < 2 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	elements, err := s.DB.SDiff(cmd.Args[0], cmd.Args[1])
 	if err != nil {
 		cmd.error(err)
@@ -54,6 +74,10 @@ func (s *Server) sDiff(cmd Command) {
 }
 
 func (s *Server) sDiffStore(cmd Command) {
+	if len(cmd.Args) < 3 {
+		cmd.error(ErrNotEnoughArgs)
+		return
+	}
 	err := s.DB.SDiffStore(cmd.Args[0], cmd.Args[1], cmd.Args[2])
 	if err != nil {
 		cmd.error(err)