test and fix parsehashkey

src/adapt.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"log"
 	"net"
+	"regexp"
 	"strconv"
 )
 
@@ -15,130 +16,133 @@ func adapt(ctx context.Context, config Config, conn net.Conn) error {
 	reader := bufio.NewReader(conn)
 
 	for ctx.Err() == nil {
-		if done, err := func() (bool, error) {
-			raw, message, err := readMessage(reader)
-			if err != nil {
-				if err == io.EOF {
-					return true, nil
-				}
-				return true, err
-			}
-			log.Printf("routing: %q (%+v)", raw, message)
-
-			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)
-				}
-				if err := config.WithConn(hashKey, func(forwardConn net.Conn) error {
-					//log.Printf("forwarding %q", raw)
-
-					written := 0
-					for written < len(raw) {
-						more, err := forwardConn.Write(raw[written:])
-						if err != nil {
+		if err := adaptOne(ctx, config, conn, reader); err != nil {
 			return err
 		}
-						written += more
-					}
-
-					replyer := bufio.NewReader(forwardConn)
-					//log.Printf("reading reply to %q", raw)
-					_, err = readMessageTo(conn, replyer)
-					//log.Printf("read reply %q", raw)
-					return err
-				}); err != nil {
-					return true, err
-				}
-			}
-
-			return false, nil
-		}(); err != nil {
-			return err
-		} else if done {
-			return nil
-		}
 	}
 
 	return io.EOF
 }
 
-func readMessage(reader *bufio.Reader) ([]byte, []any, error) {
-	w := bytes.NewBuffer(nil)
-	arr, err := readMessageTo(w, reader)
-	return w.Bytes(), arr, err
+func adaptOne(ctx context.Context, config Config, conn io.Writer, reader *bufio.Reader) error {
+	raw, err := readMessage(reader)
+	if err != nil {
+		return err
+	} else if len(raw) == 0 {
+		return nil
+	}
+	log.Printf("routing: %q", raw)
+
+	return config.WithConn(parseHashKey(raw), func(forwardConn net.Conn) error {
+		if _, err := io.Copy(forwardConn, bytes.NewReader(raw)); err != nil {
+			return err
+		}
+
+		return readMessageTo(conn, bufio.NewReader(forwardConn))
+	})
 }
 
-func readMessageTo(w io.Writer, reader *bufio.Reader) ([]any, error) {
+func readMessage(reader *bufio.Reader) ([]byte, error) {
+	w := bytes.NewBuffer(nil)
+	err := readMessageTo(w, reader)
+	return w.Bytes(), err
+}
+
+func readMessageTo(w io.Writer, reader *bufio.Reader) error {
+	w2 := bufio.NewWriter(w)
+	defer w2.Flush()
+	err := _readMessageTo(w2, reader)
+	if err != nil {
+		return err
+	}
+	return w2.Flush()
+}
+
+func _readMessageTo(w io.Writer, reader *bufio.Reader) error {
 	firstLine, _, err := reader.ReadLine()
 	if err != nil {
-		return nil, err
+		return err
 	}
 	firstLine = bytes.TrimSuffix(firstLine, []byte("\r\n"))
 	if len(firstLine) == 0 {
-		return nil, nil
+		return nil
 	}
 	fmt.Fprintf(w, "%s\r\n", firstLine)
 
 	switch firstLine[0] {
 	case '+': // simple string, like +OK
-		return []any{string(firstLine[1:])}, nil
+		return nil
 	case '-': // simple error, like -message
-		return []any{string(firstLine)}, nil
+		return nil
 	case ':': // number, like /[+-][0-9]+/
-		firstLine = bytes.TrimPrefix(firstLine[1:], []byte("+"))
-		n, err := strconv.Atoi(string(firstLine))
-		if err != nil {
-			return nil, fmt.Errorf("num not a num: %q: %w", firstLine, err)
-		}
-		return []any{n}, err
+		return err
 	case '$': // long string, like $-1 for nil, like $LEN\r\nSTRING\r\n
 		if firstLine[1] == '-' {
-			return []any{nil}, nil
+			return nil
 		}
 		firstLine = bytes.TrimPrefix(firstLine[1:], []byte("+"))
 		n, err := strconv.Atoi(string(firstLine))
 		if err != nil {
-			return nil, fmt.Errorf("num not a num: %q: %w", firstLine, err)
+			return fmt.Errorf("num not a num: %q: %w", 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 nil, fmt.Errorf("couldnt read %v more/%v bytes for long string: %w", n+2-nAt, n+2, err)
+				return 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 []any{string(nextLine)}, err
-	case '*': // array, like *-1 for nil, like *4 for [1,2,3,4]
+		return err
+	case '*', '~', '%': // *=array ~=set %=map, like *-1 for nil, like *4 for [1,2,3,4]
 		n, err := strconv.Atoi(string(firstLine[1:]))
 		if err != nil {
-			return nil, fmt.Errorf("arr not a num: %q: %w", firstLine, err)
+			return fmt.Errorf("arr not a num: %q: %w", firstLine, err)
 		} else if n == -1 {
-			return nil, nil
+			return nil
+		}
+		if isMap := firstLine[0] == '%'; isMap {
+			n *= 2
 		}
-		var result []any
 		for i := 0; i < n; i++ {
-			more, err := readMessageTo(w, reader)
+			err := _readMessageTo(w, reader)
 			if err != nil {
-				return nil, err
+				return err
 			}
-			result = append(result, more...)
 		}
-		return result, nil
+		return nil
 	case '_': // nil
-		return nil, nil
+		return nil
 	case '#': // boolean, like #t or #f
-		return []any{firstLine[1] == 't'}, nil
+		return nil
 	case ',': // double
-		log.Fatal("not impl")
+		return nil
 	}
 
 	log.Fatalf("not impl: %q", firstLine)
-	return nil, nil
+	return nil
+}
+
+var alpha = regexp.MustCompile(`[a-zA-Z]`)
+var parts = regexp.MustCompile(`(.*?)\r\n*`)
+
+func parseHashKey(raw []byte) string {
+	matches := parts.FindAllSubmatch(raw, 5)
+	if len(matches) == 0 {
+		return ""
+	}
+	for i := range matches {
+		if isTheCommand := alpha.Match(matches[i][0]); isTheCommand {
+			for j := i + 1; j < len(matches); j++ {
+				if matches[j][0][0] == byte('$') {
+					continue
+				}
+				return string(matches[j][1])
+			}
+			return string(matches[i][1])
+		}
+	}
+	return ""
 }

src/adapt_test.go

@@ -0,0 +1,64 @@
+package src
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"strings"
+	"testing"
+)
+
+func TestParseHashKey(t *testing.T) {
+	cases := map[string]string{
+		"+OK\r\n":                        "+OK",
+		"+string\r\n":                    "+string",
+		"*1\r\n$4\r\nping\r\n":           "ping",
+		"*2\r\n$3\r\nget\r\n$1\r\nx\r\n": "x",
+		"*2\r\n$3\r\nset\r\n$1\r\nx\r\n$1\r\ny\r\n": "x",
+	}
+
+	for given, expect := range cases {
+		given, expect := given, expect
+		t.Run(fmt.Sprintf("%q", given), func(t *testing.T) {
+			got := parseHashKey([]byte(given))
+			if string(got) != expect {
+				t.Errorf("expected %q=>%q but got =>%q", given, expect, got)
+			}
+		})
+	}
+}
+
+func TestReadMessage(t *testing.T) {
+	cases := map[string]string{
+		"ok":              "+OK\r\n",
+		"simple string":   "+my string\r\n",
+		"error":           "-my err\r\n",
+		"number":          ":1\r\n",
+		"negative number": ":-1\r\n",
+		"positive number": ":+1\r\n",
+		"nil string":      "$-1\r\n",
+		"long string":     "$100\r\n" + strings.Repeat("a", 100) + "\r\n",
+		"longer string":   "$10000\r\n" + strings.Repeat("a", 10000) + "\r\n",
+		"zero array":      "*0\r\n",
+		"array":           "*3\r\n+hello\r\n:world\r\n$1\r\n!\r\n",
+		"set":             "~3\r\n+hello\r\n:world\r\n$1\r\n!\r\n",
+		"map":             "%1\r\n+foo\r\n+bar\r\n",
+		"nil":             "_\r\n",
+		"boolean":         "#t\r\n",
+		"double":          ",1.5\r\n",
+		// "empty":           "", (eof)
+	}
+
+	for name, payload := range cases {
+		payload := payload
+		t.Run(name, func(t *testing.T) {
+			got, err := readMessage(bufio.NewReader(strings.NewReader(payload)))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !bytes.Equal(got, []byte(payload)) {
+				t.Errorf("expected %q but got %q", payload, got)
+			}
+		})
+	}
+}

src/config.go

@@ -2,6 +2,7 @@ package src
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -17,6 +18,7 @@ import (
 )
 
 type Config struct {
+	ctx      context.Context
 	Listen   string `json:"LISTEN"`
 	Forwards string `json:"FORWARDS"`
 	Hello    string `json:"HELLO_B64"`
@@ -25,14 +27,28 @@ type Config struct {
 
 func NewConfig(ctx context.Context) (Config, error) {
 	config := Config{
+		ctx:      ctx,
 		Listen:   ":10000",
 		Forwards: "",
 		Hello:    base64.StdEncoding.EncodeToString([]byte("*1\r\n$4\r\nping\r\n")),
 	}
-	b, _ := json.Marshal(config)
+
+	if err := config.loadEnv(); err != nil {
+		return config, err
+	}
+
+	if err := config.setForwards(); err != nil {
+		return config, err
+	}
+
+	return config, nil
+}
+
+func (c *Config) loadEnv() error {
+	b, _ := json.Marshal(*c)
 	var m map[string]any
 	if err := json.Unmarshal(b, &m); err != nil {
-		return config, err
+		return err
 	}
 	for k := range m {
 		if v := os.Getenv(k); v != "" {
@@ -40,65 +56,76 @@ func NewConfig(ctx context.Context) (Config, error) {
 		}
 	}
 	b2, _ := json.Marshal(m)
-	if err := json.Unmarshal(b2, &config); err != nil {
-		return config, err
-	}
+	return json.Unmarshal(b2, c)
+}
 
-	hello, err := base64.StdEncoding.DecodeString(config.Hello)
+func (c *Config) setForwards() error {
+	hello, err := base64.StdEncoding.DecodeString(c.Hello)
 	if err != nil {
-		return config, err
+		return err
 	}
 
-	forwards := strings.Split(config.Forwards, ",")
-	forwards = slices.DeleteFunc(forwards, func(s string) bool { return s == "" })
-	if len(forwards) == 0 {
-		return config, fmt.Errorf("at least one $FORWARD required")
+	addresses := strings.Split(c.Forwards, ",")
+	addresses = slices.DeleteFunc(addresses, func(s string) bool { return s == "" })
+	if len(addresses) == 0 {
+		return fmt.Errorf("at least one $FORWARD required")
 	}
-	config.forwards = make([]*sync.Pool, len(forwards))
-	for i := range forwards {
-		forward := forwards[i]
-		config.forwards[i] = &sync.Pool{
+
+	c.forwards = make([]*sync.Pool, len(addresses))
+	for i := range addresses {
+		address := addresses[i]
+		c.forwards[i] = &sync.Pool{
 			New: func() any {
-				if ctx.Err() != nil {
-					return nil
-				}
-				log.Printf("dialing %q with %q", forward, hello)
-				v, err := (&net.Dialer{}).DialContext(ctx, "tcp", forward)
+				conn, err := c.dial(address, hello)
 				if err != nil {
-					log.Printf("! failed dial %q: %v", forward, err)
-					return nil
+					log.Printf("! failed to dial %q: %v", address, err)
 				}
-				if _, err := v.Write([]byte(hello)); err != nil {
-					log.Printf("! failed write hello %q: %v", forward, err)
-					v.Close()
-					return nil
-				}
-				if raw, _, err := readMessage(bufio.NewReader(v)); err != nil {
-					log.Printf("! failed read hello %q: %v", forward, err)
-					v.Close()
-					return nil
-				} else {
-					log.Printf("dial reply: %q", raw)
-				}
-				return v
+				return conn
 			},
 		}
 	}
+	return nil
+}
 
-	return config, nil
+func (c *Config) dial(address string, hello []byte) (net.Conn, error) {
+	if err := c.ctx.Err(); err != nil {
+		return nil, err
+	}
+
+	log.Printf("dialing %q with %q", address, hello)
+	v, err := (&net.Dialer{}).DialContext(c.ctx, "tcp", address)
+	if err != nil {
+		return nil, err
+	}
+
+	if _, err := v.Write([]byte(hello)); err != nil {
+		v.Close()
+		return nil, err
+	}
+
+	if raw, err := readMessage(bufio.NewReader(v)); err != nil {
+		v.Close()
+		return nil, err
+	} else if bytes.HasPrefix(raw, []byte("-")) {
+		v.Close()
+		return nil, fmt.Errorf("failed hello: %s", raw)
+	}
+
+	return v, nil
 }
 
 func (c Config) Close() {
 	for i := range c.forwards {
 		if c.forwards[i] != nil {
 			i := i
-			go func() {
+			func() {
 				c.forwards[i].New = nil
 				for {
 					got := c.forwards[i].Get()
-					if got != nil {
-						go got.(net.Conn).Close()
+					if got == nil {
+						break
 					}
+					got.(net.Conn).Close()
 				}
 			}()
 		}

src/config_test.go

@@ -0,0 +1,72 @@
+package src
+
+import (
+	"fmt"
+	"io"
+	"log"
+	"net"
+	"os"
+	"strconv"
+	"sync"
+	"testing"
+)
+
+func TestConfigLoadEnv(t *testing.T) {
+	c := Config{
+		Listen:   "default",
+		Forwards: "default",
+		Hello:    "default",
+	}
+
+	os.Setenv("LISTEN", "$LISTEN")
+	os.Setenv("FORWARDS", "")
+	os.Setenv("HELLO_B64", "$HELLO_B64")
+
+	if err := c.loadEnv(); err != nil {
+		t.Fatal(err)
+	}
+
+	if c.Listen != "$LISTEN" {
+		t.Error("listen", c.Listen)
+	}
+	if c.Forwards != "default" {
+		t.Error("forwards", c.Forwards)
+	}
+	if c.Hello != "$HELLO_B64" {
+		t.Error("hello", c.Hello)
+	}
+}
+
+func TestConfigWithConn(t *testing.T) {
+	log.SetOutput(io.Discard)
+
+	c := Config{}
+
+	seen0, seen1 := false, false
+	conn0, conn1 := &net.IPConn{}, &net.IPConn{}
+	pool0, pool1 := &sync.Pool{}, &sync.Pool{}
+	pool0.Put(conn0)
+	pool1.Put(conn1)
+	c.forwards = []*sync.Pool{pool0, pool1}
+
+	for i := 0; i < 10; i++ {
+		i := strconv.Itoa(i)
+		t.Run(i, func(t *testing.T) {
+			if err := c.WithConn(i, func(c net.Conn) error {
+				c, _ = c.(*net.IPConn)
+				seen0 = seen0 || c == conn0
+				seen1 = seen1 || c == conn1
+				if c != conn0 && c != conn1 {
+					return fmt.Errorf("got unexepcted conn %v", c)
+				}
+				return nil
+			}); err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+
+	if !seen0 || !seen1 {
+		t.Fatalf("didnt see both conns")
+	}
+}

src/listen.go

@@ -31,7 +31,6 @@ func listen(ctx context.Context, config Config) error {
 		} else {
 			wg.Add(1)
 			go func() {
-				defer log.Println("/handle()")
 				defer wg.Done()
 				handle(ctx, config, conn)
 			}()
@@ -50,5 +49,12 @@ func handle(ctx context.Context, config Config, conn net.Conn) {
 func _handle(ctx context.Context, config Config, conn net.Conn) error {
 	defer conn.Close()
 
+	ctx, can := context.WithCancel(ctx)
+	defer can()
+	go func() {
+		<-ctx.Done()
+		conn.Close()
+	}()
+
 	return adapt(ctx, config, conn)
 }

src/testdata/redises.sh

@@ -1,9 +1,14 @@
 #! /usr/bin/bash
 
+bin=valkey-server
+if ! which valkey-server; then
+   bin=redis-server
+fi
+
 d=$(mktemp -d)
 mkdir -p $d/1
 mkdir -p $d/2
-valkey-server --dir $d/1 --port 60113 --logfile $d/1.log &
+$bin --dir $d/1 --port 60113 --logfile $d/1.log &
 bg=${!}
 cleanup() {
    kill $bg
@@ -13,5 +18,5 @@ trap cleanup EXIT
    sleep 5
    echo FORWARDS=127.0.0.1:60113,127.0.0.1:60114
 ) &
-valkey-server --dir $d/2 --port 60114 --logfile $d/2.log 
+$bin --dir $d/2 --port 60114 --logfile $d/2.log