mvp with terrible conn handling

src/adapt.go

@@ -18,16 +18,21 @@ func adapt(ctx context.Context, config Config, conn net.Conn) error {
 	for ctx.Err() == nil {
 		if done, err := func() (bool, error) {
 			raw, message, err := readMessage(reader)
-			log.Printf("%q", raw)
 			if err != nil {
 				if err == io.EOF {
 					return true, nil
 				}
 				return true, err
 			}
+			log.Printf("routing: %q (%+v)", raw, message)
 
-			if len(message) > 0 {
-				hashKey := message[max(0, len(message)-1)].(string)
+			if len(message) > 0 && message[0].(string) == "COMMAND" {
+				fmt.Fprintf(conn, "*0\r\n")
+			} else if len(raw) > 0 && len(message) > 0 {
+				hashKey := message[0].(string)
+				if len(message) > 1 {
+					hashKey = message[1].(string)
+				}
 				hash := int(crc32.ChecksumIEEE([]byte(hashKey)))
 				forward := config.forwards[hash%len(config.forwards)]
 				forwardCon := forward.Get()
@@ -35,17 +40,34 @@ func adapt(ctx context.Context, config Config, conn net.Conn) error {
 					return true, io.EOF
 				}
 				forwardConn := forwardCon.(net.Conn)
-				if _, err := forwardConn.Write(raw); err != nil {
-					return true, err
+				log.Printf("forwarding %q", raw)
+				written := 0
+				for written < len(raw) {
+					more, err := forwardConn.Write(raw[written:])
+					if err != nil {
+						forwardConn.Close()
+						return true, err
+					}
+					written += more
 				}
 				replyer := bufio.NewReader(forwardConn)
+				log.Printf("reading reply to %q", raw)
 				raw, _, err := readMessage(replyer)
+				log.Printf("read reply %q", raw)
 				if err != nil {
+					forwardConn.Close()
 					return true, err
 				}
-				log.Printf("%q", raw)
-				if _, err := conn.Write(raw); err != nil {
-					return true, err
+				forward.Put(forwardCon)
+				log.Printf("replying: %v", len(raw))
+				written = 0
+				for written < len(raw) {
+					more, err := conn.Write(raw[written:min(len(raw), written+4096)])
+					if err != nil {
+						return true, err
+					}
+					written += more
+					log.Printf("replied %v of %v...", written, len(raw))
 				}
 			}
 
@@ -61,9 +83,13 @@ func adapt(ctx context.Context, config Config, conn net.Conn) error {
 }
 
 func readMessage(reader *bufio.Reader) ([]byte, []any, error) {
+	b, arr, err := _readMessage(reader)
+	return b, arr, err
+}
+
+func _readMessage(reader *bufio.Reader) ([]byte, []any, error) {
 	w := bytes.NewBuffer(nil)
 	firstLine, _, err := reader.ReadLine()
-	w.Write(firstLine)
 	if err != nil {
 		return w.Bytes(), nil, err
 	}
@@ -71,35 +97,55 @@ func readMessage(reader *bufio.Reader) ([]byte, []any, error) {
 	if len(firstLine) == 0 {
 		return w.Bytes(), nil, nil
 	}
+	fmt.Fprintf(w, "%s\r\n", firstLine)
 
 	switch firstLine[0] {
 	case '+': // simple string, like +OK
 		return w.Bytes(), []any{string(firstLine[1:])}, nil
 	case '-': // simple error, like -message
-		return w.Bytes(), nil, fmt.Errorf("error: %s", firstLine[1:])
+		return w.Bytes(), []any{string(firstLine)}, nil
 	case ':': // number, like /[+-][0-9]+/
 		firstLine = bytes.TrimPrefix(firstLine[1:], []byte("+"))
-		n, err := strconv.Atoi(string(firstLine[1:]))
+		n, err := strconv.Atoi(string(firstLine))
+		if err != nil {
+			return w.Bytes(), nil, fmt.Errorf("num not a num in %q: %q: %w", w.Bytes(), firstLine, err)
+		}
 		return w.Bytes(), []any{n}, err
 	case '$': // long string, like $-1 for nil, like $LEN\r\nSTRING\r\n
 		if firstLine[1] == '-' {
 			return w.Bytes(), []any{nil}, nil
 		}
-		nextLine, _, err := reader.ReadLine()
-		w.Write(nextLine)
-		nextLine = bytes.TrimSuffix(nextLine, []byte("\r\n"))
+		firstLine = bytes.TrimPrefix(firstLine[1:], []byte("+"))
+		n, err := strconv.Atoi(string(firstLine))
+		if err != nil {
+			return w.Bytes(), nil, fmt.Errorf("num not a num in %q: %q: %w", w.Bytes(), firstLine, err)
+		}
+		log.Printf("reading %v+2 bytes for bulk string", n)
+		nextLine := make([]byte, n+2)
+		nAt := 0
+		for nAt < n+2 {
+			nMore, err := reader.Read(nextLine[nAt:])
+			if err != nil {
+				return w.Bytes(), nil, fmt.Errorf("couldnt read %v more/%v bytes for long string: %w", n+2-nAt, n+2, err)
+			}
+			nAt += nMore
+		}
+		fmt.Fprintf(w, "%s", nextLine)
 		return w.Bytes(), []any{string(nextLine)}, err
 	case '*': // array, like *-1 for nil, like *4 for [1,2,3,4]
 		n, err := strconv.Atoi(string(firstLine[1:]))
 		if err != nil {
-			return w.Bytes(), nil, err
+			return w.Bytes(), nil, fmt.Errorf("arr not a num: %q: %w", firstLine, err)
 		} else if n == -1 {
 			return w.Bytes(), nil, nil
 		}
 		var result []any
 		for i := 0; i < n; i++ {
-			moreBytes, more, err := readMessage(reader)
-			w.Write(moreBytes)
+			moreBytes, more, err := _readMessage(reader)
+			moreBytes = bytes.TrimSuffix(moreBytes, []byte("\r\n"))
+			if len(moreBytes) > 0 {
+				fmt.Fprintf(w, "%s\r\n", moreBytes)
+			}
 			if err != nil {
 				return w.Bytes(), nil, err
 			}