cluster/identd/ident/server_test.go

package ident

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"strings"
	"testing"
)

// loopback sets up a net.Listener on any available TCP port and returns it and
// a dialer function that returns open connections to that listener.
func loopback(t *testing.T) (net.Listener, func() net.Conn) {
	t.Helper()

	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("Listen: %v", err)
	}

	return lis, func() net.Conn {
		t.Helper()
		conn, err := net.Dial("tcp", lis.Addr().String())
		if err != nil {
			t.Fatalf("Dial: %v", err)
		}
		return conn
	}
}

// dumbHandler is a handler that returns USERID:UNIX:q3k for every request.
func dumbHandler(ctx context.Context, w ResponseWriter, r *Request) {
	w.SendIdent(&IdentResponse{
		UserID: "q3k",
	})
}

// reqRessps send an ident query to the conn and expects a response with
// USERID:UNIX:q3k on the scanner.
func reqResp(t *testing.T, conn net.Conn, scanner *bufio.Scanner, client, server uint16) {
	t.Helper()
	if _, err := fmt.Fprintf(conn, "%d,%d\r\n", client, server); err != nil {
		t.Fatalf("Write: %v", err)
	}
	if !scanner.Scan() {
		t.Fatalf("Scan: %v", scanner.Err())
	}
	if want, got := fmt.Sprintf("%d,%d:USERID:UNIX:q3k", client, server), scanner.Text(); want != got {
		t.Fatalf("Wanted %q, got %q", want, got)
	}
}

// TestServeSimple exercises the basic Server functionality: responding to
// ident requests.
func TestServeSimple(t *testing.T) {
	lis, dial := loopback(t)
	defer lis.Close()

	isrv := NewServer()
	isrv.HandleFunc(dumbHandler)
	go isrv.Serve(lis)

	conn := dial()
	defer conn.Close()
	scanner := bufio.NewScanner(conn)

	// Send a request, expect response.
	reqResp(t, conn, scanner, 123, 234)
	// Send another request on the same conn, expect response.
	reqResp(t, conn, scanner, 234, 345)

	// Send another request in parallel, expect response.
	conn2 := dial()
	defer conn2.Close()
	scanner2 := bufio.NewScanner(conn2)
	reqResp(t, conn2, scanner2, 345, 456)
}

// TestServeError exercises situations where the server has to deal with
// nasty/broken clients.
func TestServeErrors(t *testing.T) {
	lis, dial := loopback(t)
	defer lis.Close()

	isrv := NewServer()
	isrv.HandleFunc(dumbHandler)
	go isrv.Serve(lis)

	conn := dial()
	defer conn.Close()

	// Send something that's not ident.
	fmt.Fprintf(conn, "GET / HTTP/1.1\r\n\r\n")
	// Expect EOF on read.
	data := make([]byte, 100)
	_, err := conn.Read(data)
	if want, got := io.EOF, err; want != got {
		t.Fatalf("Expected %v, got %v", want, got)
	}

	conn = dial()
	defer conn.Close()

	// Send a very long request line, expect to not be served.
	fmt.Fprintf(conn, "123,%s123\r\n", strings.Repeat(" ", 4096))
	data = make([]byte, 100)
	_, err = conn.Read(data)
	// In a large write, the connection will be closed by the server before
	// we're finished writing. That will cause the connection to be reset, not
	// just EOF'd as above.
	if err == nil {
		t.Fatalf("Read did not fail")
	}
}

// TestServerRestart ensures that the server's serve/stop logic works as expected.
func TestServerRestart(t *testing.T) {
	lis, dial := loopback(t)
	defer lis.Close()

	isrv := NewServer()
	isrv.HandleFunc(dumbHandler)

	// Stop the server before it's even started.
	isrv.Stop()

	// The server should now exit immediately.
	if err := isrv.Serve(lis); err != nil {
		t.Fatalf("Serve: %v", err)
	}

	// On a subsequent run it should, however, start and serve.
	go isrv.Serve(lis)

	conn := dial()
	defer conn.Close()
	scanner := bufio.NewScanner(conn)

	// Send a request, expect response.
	reqResp(t, conn, scanner, 123, 234)

	// Attempting another simultaneous Serve() shoud fail.
	if err := isrv.Serve(lis); err == nil {
		t.Fatal("Serve() returned nil, wanted error")
	}

	// Send a request, expect response.
	reqResp(t, conn, scanner, 234, 345)

	// Stop server, restart server.
	isrv.Stop()
	go isrv.Serve(lis)

	// Send a request, expect response.
	reqResp(t, conn, scanner, 345, 456)
}