cluster/identd/ident/client.go

package ident

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"strconv"

	"github.com/golang/glog"
)

type DialOption func(d *dialOptions)

type dialOptions struct {
	dialer func(context.Context, string, string) (net.Conn, error)
}

// WithDialer configures a Client to use a given dial function instead of the
// default implementation in net.
func WithDialer(dialer func(context.Context, string, string) (net.Conn, error)) DialOption {
	return func(d *dialOptions) {
		d.dialer = dialer
	}
}

// parseTarget interprets a target string (ie. the target address of the Dial
// function) as an ident service address, either a host:port pair, or a host
// (in which case the default ident port, 113, is used).
func parseTarget(s string) (string, uint16, error) {
	host, portStr, err := net.SplitHostPort(s)
	if err == nil {
		port, err := strconv.ParseUint(portStr, 10, 16)
		if err != nil {
			return "", 0, fmt.Errorf("can't parse port %q: %w", portStr, err)
		}
		return host, uint16(port), nil
	}

	// Doesn't look like a host:port pair? Default to port 113.
	return s, 113, nil
}

// Dial sets up an ident protocol Client that will connect to the given target.
// Target can be either a host:port pair, or just a host (in which case the
// default ident port, 113, is used).
// This does not actually connect to identd over TCP - that will be done, as
// necessary, as requests are processed (including reconnections if multiple
// requests are processed on a Client which connects to a server that does not
// support long-standing ident donnections).
func Dial(target string, options ...DialOption) (*Client, error) {
	host, port, err := parseTarget(target)
	if err != nil {
		return nil, fmt.Errorf("invalid target: %v", err)
	}

	dialer := net.Dialer{}
	opts := dialOptions{
		dialer: dialer.DialContext,
	}
	for _, opt := range options {
		opt(&opts)
	}

	return &Client{
		opts:    opts,
		target:  net.JoinHostPort(host, fmt.Sprintf("%d", port)),
		conn:    nil,
		scanner: nil,
	}, nil
}

// Client is an ident protocol client. It maintains a connection to the ident
// server that it's been configured for, reconnecting as necessary. It is not
// safe to be used by multiple goroutines.
type Client struct {
	// opts are the dialOptions with which the client has been constructed.
	opts dialOptions
	// target is the full host:port pair that the client should connect to.
	target string
	// conn is either nil or an active TCP connection to the ident server.
	conn net.Conn
	// scannner is either nil or a line-scanner attached to the receive side of
	// conn.
	scanner *bufio.Scanner
}

func (c *Client) connect(ctx context.Context) error {
	glog.V(1).Infof("Dialing IDENT at %q", c.target)
	conn, err := c.opts.dialer(ctx, "tcp", c.target)
	if err != nil {
		return fmt.Errorf("connecting: %w", err)
	}
	c.conn = conn
	c.scanner = bufio.NewScanner(conn)
	return nil
}

func (c *Client) disconnect() {
	if c.conn == nil {
		return
	}
	c.conn.Close()
	c.conn = nil
}

// Do executes the given Request against the server to which the Client is
// connected.
func (c *Client) Do(ctx context.Context, r *Request) (*Response, error) {
	glog.V(1).Infof("Do(%+v)", r)

	// Connect if needed.
	if c.conn == nil {
		if err := c.connect(ctx); err != nil {
			return nil, err
		}
	}

	// Start a goroutine that will perform the actual request/response
	// processing to the server. A successful response will land in resC, while
	// any protocl-level error will land in errC.
	// We make both channels buffered, because if the context expires without a
	// response, we want the goroutine to be able to write to them even though
	// we're not receiving anymore. The channel will then be garbage collected.
	resC := make(chan *Response, 1)
	errC := make(chan error, 1)
	go func() {
		data := r.encode()
		glog.V(3).Infof(" -> %q", data)
		_, err := c.conn.Write(data)
		if err != nil {
			errC <- fmt.Errorf("Write: %w", err)
			return
		}
		if !c.scanner.Scan() {
			// scanner.Err() returns nil on EOF. We want that EOF, as the ident
			// protocol has special meaning for EOF sent by the server
			// (indicating either a lack of support for multiple requests per
			// connection, or a refusal to serve at an early stage of the
			// connection).
			if err := c.scanner.Err(); err != nil {
				errC <- fmt.Errorf("Read: %w", err)
			} else {
				errC <- fmt.Errorf("Read: %w", io.EOF)
			}
		}
		data = c.scanner.Bytes()
		glog.V(3).Infof(" <- %q", data)
		resp, err := decodeResponse(data)
		if err != nil {
			errC <- err
		} else {
			resC <- resp
		}
	}()

	select {
	case <-ctx.Done():
		// If the context is closed, fail with the context error and kill the
		// connection. The running goroutine will error out on any pending
		// network I/O and fail at some later point.
		// TODO(q3k): make the communication goroutine long-lived and don't
		// kill it here, just let it finish whatever it's doing and ignore the
		// result.
		c.disconnect()
		return nil, ctx.Err()
	case res := <-resC:
		return res, nil
	case err := <-errC:
		// TODO(q3k): interpret EOF, which can mean different things at
		// different times according to the RFC.
		if c.conn != nil {
			c.conn.Close()
			c.conn = nil
		}
		return nil, err
	}
}

// Close closes the Client, closing any underlying TCP connection.
func (c *Client) Close() error {
	if c.conn == nil {
		return nil
	}
	return c.conn.Close()
}

// Query performs a single ident protocol request to a server running at target
// and returns the received ident response. If the ident server cannot be
// queries, or the ident server returns an ident error, an error is returned.
//
// This a convenience wrapper around Dial/Do. The given target must be either a
// host or host:port pair. If not given, the port defaults to 113.
//
// Returned ident server error resposes are *IdentError, and can be tested for
// using errors.Is/errors.As. See the IdentError type documentation for more
// information.
//
// The given context will be used to time out the request, either at the
// connection or request stage. If the context is canceled/times out, the
// context error will be returned and the query aborted.
//
// Since Query opens a connection to the ident server just for a single query,
// it should not be used if a single server is going to be queries about
// multiple addresses, and instead Dial/Do should be used to keep a
// long-standing connection if possible.
func Query(ctx context.Context, target string, client, server uint16) (*IdentResponse, error) {
	cl, err := Dial(target)
	if err != nil {
		return nil, fmt.Errorf("could not dial: %w", err)
	}
	defer cl.Close()
	resp, err := cl.Do(ctx, &Request{
		ClientPort: client,
		ServerPort: server,
	})
	if err != nil {
		return nil, fmt.Errorf("could not query: %w", err)
	}
	if resp.Ident != nil {
		return resp.Ident, nil
	}
	return nil, &IdentError{Inner: resp.Error}
}