From 47d561e20da97ca30fb9aa96d0ff6b45279131b3 Mon Sep 17 00:00:00 2001
From: "STeve (Xin) Huang" <huangin52@gmail.com>
Date: Wed, 8 Mar 2023 09:18:33 -0500
Subject: [PATCH] Reapply: Fix an issue status type (+) is not read/write
 properly  (#37)

* Fix an issue status type (+) is not read/write properly (#22)

* add an option for using StatusString vs string
---
 internal/proto/reader.go      | 27 ++++++++++++++++++++++++---
 internal/proto/writer.go      | 14 ++++++++++++++
 internal/proto/writer_test.go | 27 +++++++++++++++++++++++++++
 proto.go                      |  2 +-
 4 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/internal/proto/reader.go b/internal/proto/reader.go
index da3f5af65..d1a833ffa 100644
--- a/internal/proto/reader.go
+++ b/internal/proto/reader.go
@@ -31,6 +31,9 @@ const (
 	RespPush      = '>' // ><len>\r\n... (same as Array)
 )
 
+// StatusString is the read golang type of the RespStatus.
+type StatusString string
+
 // Not used temporarily.
 // Redis has not used these two data types for the time being, and will implement them later.
 // Streamed           = "EOF:"
@@ -52,13 +55,28 @@ func ParseErrorReply(line []byte) error {
 
 //------------------------------------------------------------------------------
 
+type readerOpt struct {
+	// useStatusStringType uses `StatusString` type instead of string for RespStatus.
+	useStatusStringType bool
+}
+
+func ReaderOptUseStatusStringType(o *readerOpt) {
+	o.useStatusStringType = true
+}
+
 type Reader struct {
-	rd *bufio.Reader
+	rd  *bufio.Reader
+	opt readerOpt
 }
 
-func NewReader(rd io.Reader) *Reader {
+func NewReader(rd io.Reader, opts ...func(o *readerOpt)) *Reader {
+	opt := readerOpt{}
+	for _, applyOpt := range opts {
+		applyOpt(&opt)
+	}
 	return &Reader{
-		rd: bufio.NewReader(rd),
+		rd:  bufio.NewReader(rd),
+		opt: opt,
 	}
 }
 
@@ -159,6 +177,9 @@ func (r *Reader) ReadReply() (interface{}, error) {
 
 	switch line[0] {
 	case RespStatus:
+		if r.opt.useStatusStringType {
+			return StatusString(line[1:]), nil
+		}
 		return string(line[1:]), nil
 	case RespInt:
 		return util.ParseInt(line[1:], 10, 64)
diff --git a/internal/proto/writer.go b/internal/proto/writer.go
index fde9bc7d4..82d358732 100644
--- a/internal/proto/writer.go
+++ b/internal/proto/writer.go
@@ -69,6 +69,8 @@ func (w *Writer) WriteArg(v interface{}) error {
 		return w.string("")
 	case string:
 		return w.string(v)
+	case StatusString:
+		return w.status(v)
 	case []byte:
 		return w.bytes(v)
 	case int:
@@ -135,6 +137,18 @@ func (w *Writer) bytes(b []byte) error {
 	return w.crlf()
 }
 
+func (w *Writer) status(s StatusString) error {
+	if err := w.WriteByte(RespStatus); err != nil {
+		return err
+	}
+
+	if _, err := w.Write([]byte(s)); err != nil {
+		return err
+	}
+
+	return w.crlf()
+}
+
 func (w *Writer) string(s string) error {
 	return w.bytes(util.StringToBytes(s))
 }
diff --git a/internal/proto/writer_test.go b/internal/proto/writer_test.go
index c801f87db..865f89eb1 100644
--- a/internal/proto/writer_test.go
+++ b/internal/proto/writer_test.go
@@ -100,3 +100,30 @@ func BenchmarkWriteBuffer_Append(b *testing.B) {
 		}
 	}
 }
+
+func TestWriteStatus(t *testing.T) {
+	inputStatusBytes := []byte("+status\r\n")
+
+	// Read it.
+	reader := proto.NewReader(bytes.NewReader(inputStatusBytes), proto.ReaderOptUseStatusStringType)
+	readStatus, err := reader.ReadReply()
+	if err != nil {
+		t.Errorf("Failed to ReadReply: %v", err)
+	}
+
+	if readStatus != proto.StatusString("status") {
+		t.Errorf("expect read %v but got %v", "status", readStatus)
+	}
+
+	// Write it.
+	outputStatusBytes := new(bytes.Buffer)
+	writer := proto.NewWriter(outputStatusBytes)
+	err = writer.WriteArg(readStatus)
+	if err != nil {
+		t.Errorf("Failed to WriteArg: %v", err)
+	}
+
+	if string(inputStatusBytes) != outputStatusBytes.String() {
+		t.Errorf("expect written %v but got %v", string(inputStatusBytes), outputStatusBytes.String())
+	}
+}
diff --git a/proto.go b/proto.go
index d942290c7..2b0772297 100644
--- a/proto.go
+++ b/proto.go
@@ -15,7 +15,7 @@ const RespArray = proto.RespArray
 const RespInt = proto.RespInt
 
 func NewReader(rd io.Reader) *Reader {
-	return proto.NewReader(rd)
+	return proto.NewReader(rd, proto.ReaderOptUseStatusStringType)
 }
 
 func NewWriter(wr *bytes.Buffer) *Writer {