bucketeer/internal/probe/protocols_test.go

package probe

import (
	"context"
	"encoding/binary"
	"io"
	"net"
	"testing"
	"time"

	"tea.karlcom.de/karl/bucketeer/internal/model"
)

func listenFake(t *testing.T, handle func(net.Conn)) (string, int) {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })
	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		handle(conn)
		conn.Close()
	}()
	host, port := splitAddr(t, ln.Addr().String())
	return host, port
}

func splitAddr(t *testing.T, addr string) (string, int) {
	t.Helper()
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		t.Fatal(err)
	}
	return host, mustAtoi(portStr)
}

func TestProbeRedisReadable(t *testing.T) {
	host, port := listenFake(t, func(c net.Conn) {
		buf := make([]byte, 6)
		_, _ = io.ReadFull(c, buf)
		_, _ = c.Write([]byte("+PONG\r\n"))
	})
	p := New(Config{Workers: 1, Timeout: time.Second})
	state, detail, ok := p.probeRedis(context.Background(), host, port)
	if !ok || state != model.ExposurePublicRead {
		t.Fatalf("got %s (%s) ok=%v, want public_read", state, detail, ok)
	}
}

func TestProbeRedisAuthRequired(t *testing.T) {
	host, port := listenFake(t, func(c net.Conn) {
		buf := make([]byte, 6)
		_, _ = io.ReadFull(c, buf)
		_, _ = c.Write([]byte("-NOAUTH Authentication required.\r\n"))
	})
	p := New(Config{Workers: 1, Timeout: time.Second})
	state, _, _ := p.probeRedis(context.Background(), host, port)
	if state != model.ExposureAuthRequired {
		t.Fatalf("got %s, want auth_required", state)
	}
}

func TestProbePostgresTrust(t *testing.T) {
	host, port := listenFake(t, func(c net.Conn) {
		ssl := make([]byte, 8)
		if _, err := io.ReadFull(c, ssl); err != nil {
			return
		}
		_, _ = c.Write([]byte{'N'})
		var l [4]byte
		if _, err := io.ReadFull(c, l[:]); err != nil {
			return
		}
		n := int(binary.BigEndian.Uint32(l[:])) - 4
		_, _ = io.ReadFull(c, make([]byte, n))
		// AuthenticationOk: 'R', length 8, auth type 0
		_, _ = c.Write([]byte{'R', 0, 0, 0, 8, 0, 0, 0, 0})
	})
	p := New(Config{Workers: 1, Timeout: time.Second})
	state, detail, _ := p.probePostgres(context.Background(), host, port)
	if state != model.ExposurePublicRead {
		t.Fatalf("got %s (%s), want public_read", state, detail)
	}
}

func TestProbePostgresPassword(t *testing.T) {
	host, port := listenFake(t, func(c net.Conn) {
		ssl := make([]byte, 8)
		if _, err := io.ReadFull(c, ssl); err != nil {
			return
		}
		_, _ = c.Write([]byte{'N'})
		var l [4]byte
		if _, err := io.ReadFull(c, l[:]); err != nil {
			return
		}
		n := int(binary.BigEndian.Uint32(l[:])) - 4
		_, _ = io.ReadFull(c, make([]byte, n))
		// AuthenticationMD5Password: 'R', length 8, auth type 5
		_, _ = c.Write([]byte{'R', 0, 0, 0, 8, 0, 0, 0, 5})
	})
	p := New(Config{Workers: 1, Timeout: time.Second})
	state, _, _ := p.probePostgres(context.Background(), host, port)
	if state != model.ExposureAuthRequired {
		t.Fatalf("got %s, want auth_required", state)
	}
}

func TestParseBSONFields(t *testing.T) {
	var doc []byte
	addDouble := func(name string, v uint64) {
		doc = append(doc, 0x01)
		doc = append(doc, appendCString(nil, name)...)
		var tmp [8]byte
		binary.LittleEndian.PutUint64(tmp[:], v)
		doc = append(doc, tmp[:]...)
	}
	addString := func(name, val string) {
		doc = append(doc, 0x02)
		doc = append(doc, appendCString(nil, name)...)
		var tmp [4]byte
		binary.LittleEndian.PutUint32(tmp[:], uint32(len(val)+1))
		doc = append(doc, tmp[:]...)
		doc = append(doc, appendCString(nil, val)...)
	}
	addDouble("ok", 0x3FF0000000000000) // 1.0
	addString("errmsg", "requires authentication")
	body := make([]byte, 0, 5+len(doc)+5)
	body = append(body, 0, 0, 0, 0, 0) // flags + section kind 0
	full := make([]byte, 4, 4+len(doc)+1)
	binary.LittleEndian.PutUint32(full, uint32(4+len(doc)+1))
	full = append(full, doc...)
	full = append(full, 0)
	body = append(body, full...)

	ok, errmsg := parseBSONFields(body)
	if ok != 1 || errmsg != "requires authentication" {
		t.Fatalf("got ok=%v errmsg=%q", ok, errmsg)
	}
}