waitfordb/internal/driver/valkey_test.go

package driver

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"git.unkin.net/unkin/waitfordb/internal/config"
	"git.unkin.net/unkin/waitfordb/internal/wait"
)

// fakeValkey is a minimal RESP server. With user/password set it rejects
// PING until a matching AUTH arrives, like a real ACL-enabled server.
type fakeValkey struct {
	ln       net.Listener
	user     string
	password string
	// refuseFirst counts connections to close immediately before serving.
	refuseFirst atomic.Int32
}

func newFakeValkey(t *testing.T, user, password string) *fakeValkey {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	f := &fakeValkey{ln: ln, user: user, password: password}
	t.Cleanup(func() { ln.Close() })
	go f.serve()
	return f
}

func (f *fakeValkey) addr() string { return f.ln.Addr().String() }

func (f *fakeValkey) serve() {
	for {
		conn, err := f.ln.Accept()
		if err != nil {
			return
		}
		if f.refuseFirst.Add(-1) >= 0 {
			conn.Close()
			continue
		}
		go f.handle(conn)
	}
}

func (f *fakeValkey) handle(conn net.Conn) {
	defer conn.Close()
	r := bufio.NewReader(conn)
	authed := f.password == ""
	for {
		cmd, err := readCommand(r)
		if err != nil {
			return
		}
		switch strings.ToUpper(cmd[0]) {
		case "AUTH":
			user, pass := "default", cmd[len(cmd)-1]
			if len(cmd) == 3 {
				user = cmd[1]
			}
			wantUser := f.user
			if wantUser == "" {
				wantUser = "default"
			}
			if user == wantUser && pass == f.password {
				authed = true
				fmt.Fprint(conn, "+OK\r\n")
			} else {
				fmt.Fprint(conn, "-WRONGPASS invalid username-password pair or user is disabled.\r\n")
			}
		case "PING":
			if !authed {
				fmt.Fprint(conn, "-NOAUTH Authentication required.\r\n")
			} else {
				fmt.Fprint(conn, "+PONG\r\n")
			}
		default:
			fmt.Fprint(conn, "-ERR unknown command\r\n")
		}
	}
}

func readCommand(r *bufio.Reader) ([]string, error) {
	head, err := r.ReadString('\n')
	if err != nil {
		return nil, err
	}
	head = strings.TrimRight(head, "\r\n")
	if !strings.HasPrefix(head, "*") {
		return nil, fmt.Errorf("bad command header %q", head)
	}
	n, err := strconv.Atoi(head[1:])
	if err != nil || n < 1 {
		return nil, fmt.Errorf("bad command header %q", head)
	}
	out := make([]string, 0, n)
	for range n {
		size, err := r.ReadString('\n')
		if err != nil {
			return nil, err
		}
		size = strings.TrimRight(size, "\r\n")
		length, err := strconv.Atoi(strings.TrimPrefix(size, "$"))
		if err != nil {
			return nil, fmt.Errorf("bad bulk header %q", size)
		}
		buf := make([]byte, length+2)
		if _, err := io.ReadFull(r, buf); err != nil {
			return nil, err
		}
		out = append(out, string(buf[:length]))
	}
	return out, nil
}

func openValkey(t *testing.T, cfg config.Config) Pinger {
	t.Helper()
	drv, err := Get("valkey")
	if err != nil {
		t.Fatalf("Get(valkey): %v", err)
	}
	p, err := drv.Open(cfg)
	if err != nil {
		t.Fatalf("Open: %v", err)
	}
	t.Cleanup(func() { p.Close() })
	return p
}

func splitAddr(t *testing.T, addr string) (host, port string) {
	t.Helper()
	host, port, err := net.SplitHostPort(addr)
	if err != nil {
		t.Fatalf("split %q: %v", addr, err)
	}
	return host, port
}

func pingCtx(t *testing.T) context.Context {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	t.Cleanup(cancel)
	return ctx
}

func TestValkeyPingNoAuth(t *testing.T) {
	f := newFakeValkey(t, "", "")
	host, port := splitAddr(t, f.addr())
	p := openValkey(t, config.Config{Host: host, Port: port})
	if err := p.Ping(pingCtx(t)); err != nil {
		t.Fatalf("ping: %v", err)
	}
}

func TestValkeyPingDefaultUserAuth(t *testing.T) {
	f := newFakeValkey(t, "", "sekrit")
	host, port := splitAddr(t, f.addr())

	p := openValkey(t, config.Config{Host: host, Port: port, Password: "sekrit"})
	if err := p.Ping(pingCtx(t)); err != nil {
		t.Fatalf("ping with correct password: %v", err)
	}

	unauth := openValkey(t, config.Config{Host: host, Port: port})
	if err := unauth.Ping(pingCtx(t)); err == nil || !strings.Contains(err.Error(), "NOAUTH") {
		t.Fatalf("ping without password: want NOAUTH error, got %v", err)
	}

	wrong := openValkey(t, config.Config{Host: host, Port: port, Password: "nope"})
	if err := wrong.Ping(pingCtx(t)); err == nil || !strings.Contains(err.Error(), "WRONGPASS") {
		t.Fatalf("ping with wrong password: want WRONGPASS error, got %v", err)
	}
}

func TestValkeyPingACLUserAuth(t *testing.T) {
	f := newFakeValkey(t, "_operator", "op-pass")
	host, port := splitAddr(t, f.addr())

	p := openValkey(t, config.Config{Host: host, Port: port, User: "_operator", Password: "op-pass"})
	if err := p.Ping(pingCtx(t)); err != nil {
		t.Fatalf("ping as _operator: %v", err)
	}

	wrong := openValkey(t, config.Config{Host: host, Port: port, User: "other", Password: "op-pass"})
	if err := wrong.Ping(pingCtx(t)); err == nil {
		t.Fatal("ping as wrong user should fail")
	}
}

func TestValkeyDSN(t *testing.T) {
	f := newFakeValkey(t, "app", "pw")
	p := openValkey(t, config.Config{DSN: "redis://app:pw@" + f.addr()})
	if err := p.Ping(pingCtx(t)); err != nil {
		t.Fatalf("ping via DSN: %v", err)
	}
}

func TestValkeyDSNBadScheme(t *testing.T) {
	drv, _ := Get("valkey")
	if _, err := drv.Open(config.Config{DSN: "postgres://h:5432/d"}); err == nil {
		t.Fatal("expected error for non-redis DSN scheme")
	}
}

func TestValkeyWaitRetriesUntilUp(t *testing.T) {
	f := newFakeValkey(t, "", "pw")
	f.refuseFirst.Store(2)
	host, port := splitAddr(t, f.addr())
	p := openValkey(t, config.Config{Host: host, Port: port, Password: "pw"})

	res := wait.Run(context.Background(),
		wait.Params{Timeout: 10 * time.Second, Interval: 10 * time.Millisecond, ConnectTimeout: time.Second},
		p.Ping,
		func(int, error, time.Duration, time.Duration, time.Duration) {},
		wait.RealClock{})

	if !res.OK {
		t.Fatalf("expected readiness, got %+v", res)
	}
	if res.Attempts < 2 {
		t.Errorf("expected retries before success, got %d attempts", res.Attempts)
	}
}