feat(protocol): add FTP, FTPS and FTPES downloads (#1569)

internal/protocol/ftp/testserver_test.go

@@ -0,0 +1,587 @@
+package ftp
+
+import (
+	"bufio"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/sha256"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/hex"
+	"errors"
+	"io"
+	"math/big"
+	"net"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	ftpserver "github.com/fclairamb/ftpserverlib"
+	"github.com/spf13/afero"
+
+	pftp "github.com/GopeedLab/gopeed/pkg/protocol/ftp"
+)
+
+// ServerOptions are the knobs of the in-process FTP test server.
+type ServerOptions struct {
+	// TLS is pftp.TLSNone (or empty), pftp.TLSExplicit or pftp.TLSImplicit.
+	// Explicit mode refuses plain logins, so a passing test proves TLS was used.
+	TLS string
+	// MaxLogins caps concurrent control connections. Beyond it the server
+	// greets with "421 Too many connections" and hangs up. Zero means no cap.
+	MaxLogins int
+	// NoRest answers "502" to REST. It filters the plain control stream, so it
+	// works with TLSNone and TLSImplicit only.
+	NoRest bool
+	// NoMLST disables MLST and MLSD, so clients fall back to SIZE and LIST.
+	NoMLST bool
+	// User and Pass are the only accepted credentials. An empty User accepts
+	// any login, anonymous included.
+	User string
+	Pass string
+	// RefuseDelay holds every 421 refusal back this long, like a server
+	// that answers slowly when it is full.
+	RefuseDelay time.Duration
+	// RequireSessionReuse makes every data connection resume the TLS
+	// session of the control connection that opened it, as ProFTPD's
+	// mod_tls does by default. A data connection that does not is refused.
+	// It needs TLSImplicit, where this test's wrapper sees the plain
+	// control stream and so learns which passive port belongs to whom.
+	RequireSessionReuse bool
+}
+
+// FTPServer is an ftpserverlib server over a temp directory.
+type FTPServer struct {
+	Addr string // host:port
+	Root string // directory served as "/"
+	// CA trusts the server's self-signed certificate.
+	CA *x509.CertPool
+
+	opts    ServerOptions
+	tlsConf *tls.Config
+	server  *ftpserver.FtpServer
+
+	// ReadDelay is slept before every read of a transferred file, in
+	// nanoseconds, to slow downloads down. Tests may change it at any time.
+	ReadDelay atomic.Int64
+
+	// DropAfter, when positive, makes the next transfer fail once after it
+	// has sent that many bytes, as if the data connection were killed.
+	DropAfter atomic.Int64
+
+	mu      sync.Mutex
+	active  int
+	retrs   []int64 // offset of every RETR, 0 when no REST came first
+	refused int
+	logins  int
+
+	// Passive port to the control connection that announced it, and the
+	// outcome of every data connection's session check.
+	pasvOwner map[int]string
+	reuseOK   int
+	reuseBad  int
+
+	// Stall state, see StallAfter. stallMu is held across file reads.
+	stallMu     sync.Mutex
+	stallAt     int64
+	sent        int64
+	stallCh     chan struct{}
+	stalled     chan struct{}
+	stalledOnce *sync.Once
+}
+
+// StartFTPServer starts a server and stops it when the test ends.
+func StartFTPServer(t *testing.T, opts ServerOptions) *FTPServer {
+	t.Helper()
+	s := &FTPServer{Root: t.TempDir(), opts: opts}
+	s.tlsConf, s.CA = selfSignedTLS(t)
+
+	base, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	var ln net.Listener = base
+	if opts.RequireSessionReuse {
+		if opts.TLS != pftp.TLSImplicit {
+			t.Fatal("RequireSessionReuse needs TLSImplicit")
+		}
+		s.requireSessionReuse(t, base.Addr().(*net.TCPAddr).Port)
+	}
+	if opts.TLS == pftp.TLSImplicit {
+		ln = tls.NewListener(ln, s.tlsConf)
+	}
+	ln = &limitListener{Listener: ln, srv: s}
+
+	s.Addr = base.Addr().String()
+	s.server = ftpserver.NewFtpServer(&mainDriver{srv: s, listener: ln})
+	if err := s.server.Listen(); err != nil {
+		t.Fatal(err)
+	}
+	go s.server.Serve()
+	t.Cleanup(func() { _ = s.server.Stop() })
+	return s
+}
+
+// WriteFile creates a file under the served root and returns its SHA-256.
+func (s *FTPServer) WriteFile(t *testing.T, rel string, data []byte) string {
+	t.Helper()
+	name := filepath.Join(s.Root, filepath.FromSlash(rel))
+	if err := os.MkdirAll(filepath.Dir(name), 0o755); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(name, data, 0o644); err != nil {
+		t.Fatal(err)
+	}
+	return SHA256Hex(data)
+}
+
+// RetrOffsets returns the start offset of every RETR the server received.
+func (s *FTPServer) RetrOffsets() []int64 {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return append([]int64(nil), s.retrs...)
+}
+
+// RestOffsets returns the non-zero RETR offsets, which only a REST sets.
+func (s *FTPServer) RestOffsets() []int64 {
+	var out []int64
+	for _, off := range s.RetrOffsets() {
+		if off != 0 {
+			out = append(out, off)
+		}
+	}
+	return out
+}
+
+// Refused returns how many connections were turned away with 421.
+func (s *FTPServer) Refused() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.refused
+}
+
+// Logins returns the number of successful logins.
+func (s *FTPServer) Logins() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.logins
+}
+
+// SetMaxLogins changes the login cap while the server runs, as when other
+// clients take slots.
+func (s *FTPServer) SetMaxLogins(n int) {
+	s.mu.Lock()
+	s.opts.MaxLogins = n
+	s.mu.Unlock()
+}
+
+// SessionReuse returns how many data connections resumed their own control
+// connection's TLS session, and how many did not.
+func (s *FTPServer) SessionReuse() (ok, bad int) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.reuseOK, s.reuseBad
+}
+
+// StallAfter makes every transfer stop sending once the server has sent n
+// bytes in total, until Release. Stalled is closed when that happens.
+func (s *FTPServer) StallAfter(n int64) {
+	s.stallMu.Lock()
+	defer s.stallMu.Unlock()
+	s.stallAt, s.sent = n, 0
+	s.stallCh = make(chan struct{})
+	s.stalled = make(chan struct{})
+	s.stalledOnce = &sync.Once{}
+}
+
+// Stalled is closed once a transfer has stopped at the StallAfter limit.
+func (s *FTPServer) Stalled() <-chan struct{} {
+	s.stallMu.Lock()
+	defer s.stallMu.Unlock()
+	return s.stalled
+}
+
+// Sent returns the bytes sent since StallAfter.
+func (s *FTPServer) Sent() int64 {
+	s.stallMu.Lock()
+	defer s.stallMu.Unlock()
+	return s.sent
+}
+
+// Release ends the stall. The stalled transfers fail.
+func (s *FTPServer) Release() {
+	s.stallMu.Lock()
+	defer s.stallMu.Unlock()
+	if s.stallAt > 0 {
+		s.stallAt = 0
+		close(s.stallCh)
+	}
+}
+
+var errStalled = errors.New("test server released a stalled transfer")
+
+// stall limits a read to the StallAfter budget, and blocks once the budget
+// is spent.
+func (s *FTPServer) stall(read func([]byte) (int, error), p []byte) (int, error) {
+	s.stallMu.Lock()
+	if s.stallAt <= 0 {
+		s.stallMu.Unlock()
+		return read(p)
+	}
+	allowed := s.stallAt - s.sent
+	if allowed <= 0 {
+		ch := s.stallCh
+		s.stalledOnce.Do(func() { close(s.stalled) })
+		s.stallMu.Unlock()
+		<-ch
+		return 0, errStalled
+	}
+	if int64(len(p)) > allowed {
+		p = p[:allowed]
+	}
+	n, err := read(p)
+	s.sent += int64(n)
+	s.stallMu.Unlock()
+	return n, err
+}
+
+var pasvReply = regexp.MustCompile(`^(?:229 [^(]*\(\|\|\|(\d+)\|\)|227 [^(]*\(\d+,\d+,\d+,\d+,(\d+),(\d+)\))`)
+
+// notePassive records the passive port a control connection announced.
+func (s *FTPServer) notePassive(owner string, reply []byte) {
+	m := pasvReply.FindSubmatch(reply)
+	if m == nil {
+		return
+	}
+	var port int
+	if len(m[1]) > 0 {
+		port, _ = strconv.Atoi(string(m[1]))
+	} else {
+		hi, _ := strconv.Atoi(string(m[2]))
+		lo, _ := strconv.Atoi(string(m[3]))
+		port = hi<<8 | lo
+	}
+	s.mu.Lock()
+	s.pasvOwner[port] = owner
+	s.mu.Unlock()
+}
+
+// requireSessionReuse stamps each control connection's session tickets with
+// its address, and checks that each data connection resumed a ticket of the
+// control connection that announced its passive port.
+func (s *FTPServer) requireSessionReuse(t *testing.T, controlPort int) {
+	base := s.tlsConf
+	var key [32]byte
+	if _, err := rand.Read(key[:]); err != nil {
+		t.Fatal(err)
+	}
+	base.SetSessionTicketKeys([][32]byte{key})
+	s.pasvOwner = map[int]string{}
+	stamp := func(owner string) func(tls.ConnectionState, *tls.SessionState) ([]byte, error) {
+		return func(cs tls.ConnectionState, ss *tls.SessionState) ([]byte, error) {
+			ss.Extra = [][]byte{[]byte(owner)}
+			return base.EncryptTicket(cs, ss)
+		}
+	}
+	base.GetConfigForClient = func(hello *tls.ClientHelloInfo) (*tls.Config, error) {
+		cfg := base.Clone()
+		cfg.GetConfigForClient = nil
+		local, _ := hello.Conn.LocalAddr().(*net.TCPAddr)
+		if local == nil || local.Port == controlPort {
+			cfg.WrapSession = stamp(hello.Conn.RemoteAddr().String())
+			return cfg, nil
+		}
+		s.mu.Lock()
+		owner := s.pasvOwner[local.Port]
+		s.mu.Unlock()
+		var resumedFrom string
+		cfg.UnwrapSession = func(identity []byte, cs tls.ConnectionState) (*tls.SessionState, error) {
+			ss, err := base.DecryptTicket(identity, cs)
+			if err != nil || ss == nil {
+				return nil, nil
+			}
+			if len(ss.Extra) > 0 {
+				resumedFrom = string(ss.Extra[0])
+			}
+			return ss, nil
+		}
+		cfg.WrapSession = stamp(owner)
+		cfg.VerifyConnection = func(cs tls.ConnectionState) error {
+			ok := owner != "" && cs.DidResume && resumedFrom == owner
+			s.mu.Lock()
+			if ok {
+				s.reuseOK++
+			} else {
+				s.reuseBad++
+			}
+			s.mu.Unlock()
+			if !ok {
+				return errors.New("data connection did not resume its control connection's TLS session")
+			}
+			return nil
+		}
+		return cfg, nil
+	}
+}
+
+func (s *FTPServer) acquire() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.opts.MaxLogins > 0 && s.active >= s.opts.MaxLogins {
+		s.refused++
+		return false
+	}
+	s.active++
+	return true
+}
+
+func (s *FTPServer) release() {
+	s.mu.Lock()
+	s.active--
+	s.mu.Unlock()
+}
+
+// limitListener counts control connections and refuses those beyond
+// MaxLogins with a 421 greeting, like vsftpd's max_clients and max_per_ip.
+type limitListener struct {
+	net.Listener
+	srv *FTPServer
+}
+
+func (l *limitListener) Accept() (net.Conn, error) {
+	for {
+		c, err := l.Listener.Accept()
+		if err != nil {
+			return nil, err
+		}
+		if !l.srv.acquire() {
+			go func() {
+				time.Sleep(l.srv.opts.RefuseDelay)
+				_ = c.SetDeadline(time.Now().Add(5 * time.Second))
+				_, _ = io.WriteString(c, "421 Too many connections\r\n")
+				_ = c.Close()
+			}()
+			continue
+		}
+		cc := &countedConn{Conn: c, release: l.srv.release}
+		if l.srv.opts.RequireSessionReuse {
+			owner := c.RemoteAddr().String()
+			cc.sniff = func(b []byte) { l.srv.notePassive(owner, b) }
+		}
+		var conn net.Conn = cc
+		if l.srv.opts.NoRest {
+			conn = &noRestConn{countedConn: conn.(*countedConn), in: bufio.NewReader(c)}
+		}
+		return conn, nil
+	}
+}
+
+type countedConn struct {
+	net.Conn
+	once    sync.Once
+	release func()
+	sniff   func([]byte) // sees every reply the server writes
+}
+
+func (c *countedConn) Write(b []byte) (int, error) {
+	if c.sniff != nil {
+		c.sniff(b)
+	}
+	return c.Conn.Write(b)
+}
+
+func (c *countedConn) Close() error {
+	c.once.Do(c.release)
+	return c.Conn.Close()
+}
+
+// noRestConn answers REST itself with 502 and hides it from the server.
+type noRestConn struct {
+	*countedConn
+	in      *bufio.Reader
+	pending []byte
+}
+
+func (c *noRestConn) Read(p []byte) (int, error) {
+	for len(c.pending) == 0 {
+		line, err := c.in.ReadBytes('\n')
+		if len(line) > 0 {
+			if strings.HasPrefix(strings.ToUpper(string(line)), "REST ") {
+				if _, werr := io.WriteString(c.Conn, "502 REST not implemented\r\n"); werr != nil {
+					return 0, werr
+				}
+			} else {
+				c.pending = line
+			}
+		}
+		if err != nil {
+			if len(c.pending) > 0 {
+				break
+			}
+			return 0, err
+		}
+	}
+	n := copy(p, c.pending)
+	c.pending = c.pending[n:]
+	return n, nil
+}
+
+type mainDriver struct {
+	srv      *FTPServer
+	listener net.Listener
+}
+
+func (d *mainDriver) GetSettings() (*ftpserver.Settings, error) {
+	settings := &ftpserver.Settings{
+		Listener:    d.listener,
+		DisableMLSD: d.srv.opts.NoMLST,
+		DisableMLST: d.srv.opts.NoMLST,
+	}
+	switch d.srv.opts.TLS {
+	case pftp.TLSExplicit:
+		settings.TLSRequired = ftpserver.MandatoryEncryption
+	case pftp.TLSImplicit:
+		settings.TLSRequired = ftpserver.ImplicitEncryption
+	}
+	return settings, nil
+}
+
+func (d *mainDriver) ClientConnected(ftpserver.ClientContext) (string, error) {
+	return "gopeed test server", nil
+}
+
+func (d *mainDriver) ClientDisconnected(ftpserver.ClientContext) {}
+
+func (d *mainDriver) AuthUser(_ ftpserver.ClientContext, user, pass string) (ftpserver.ClientDriver, error) {
+	if d.srv.opts.User != "" && (user != d.srv.opts.User || pass != d.srv.opts.Pass) {
+		return nil, errors.New("bad credentials")
+	}
+	d.srv.mu.Lock()
+	d.srv.logins++
+	d.srv.mu.Unlock()
+	return &clientDriver{Fs: afero.NewBasePathFs(afero.NewOsFs(), d.srv.Root), srv: d.srv}, nil
+}
+
+func (d *mainDriver) GetTLSConfig() (*tls.Config, error) {
+	return d.srv.tlsConf, nil
+}
+
+type clientDriver struct {
+	afero.Fs
+	srv *FTPServer
+}
+
+// GetHandle sees the REST offset of every transfer. The server seeks the
+// returned file to that offset itself.
+func (d *clientDriver) GetHandle(name string, flags int, offset int64) (ftpserver.FileTransfer, error) {
+	f, err := d.Fs.OpenFile(name, flags, 0o644)
+	if err != nil {
+		return nil, err
+	}
+	if flags == os.O_RDONLY {
+		d.srv.mu.Lock()
+		d.srv.retrs = append(d.srv.retrs, offset)
+		d.srv.mu.Unlock()
+	}
+	return &slowFile{File: f, srv: d.srv, dropAfter: d.srv.DropAfter.Swap(0)}, nil
+}
+
+var errDropped = errors.New("test server dropped the transfer")
+
+type slowFile struct {
+	afero.File
+	srv       *FTPServer
+	sent      int64
+	dropAfter int64
+}
+
+func (f *slowFile) Read(p []byte) (int, error) {
+	if delay := time.Duration(f.srv.ReadDelay.Load()); delay > 0 {
+		time.Sleep(delay)
+		if len(p) > 16<<10 {
+			p = p[:16<<10]
+		}
+	}
+	if f.dropAfter > 0 {
+		if f.sent >= f.dropAfter {
+			return 0, errDropped
+		}
+		if rest := f.dropAfter - f.sent; int64(len(p)) > rest {
+			p = p[:rest]
+		}
+	}
+	n, err := f.srv.stall(f.File.Read, p)
+	f.sent += int64(n)
+	return n, err
+}
+
+// selfSignedTLS returns a server config with a fresh certificate for
+// 127.0.0.1 and localhost, and a pool that trusts it.
+func selfSignedTLS(t *testing.T) (*tls.Config, *x509.CertPool) {
+	t.Helper()
+	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	if err != nil {
+		t.Fatal(err)
+	}
+	tmpl := &x509.Certificate{
+		SerialNumber:          big.NewInt(time.Now().UnixNano()),
+		Subject:               pkix.Name{CommonName: "gopeed ftp test"},
+		NotBefore:             time.Now().Add(-time.Hour),
+		NotAfter:              time.Now().Add(time.Hour),
+		IPAddresses:           []net.IP{net.ParseIP("127.0.0.1")},
+		DNSNames:              []string{"localhost"},
+		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
+		IsCA:                  true,
+		BasicConstraintsValid: true,
+	}
+	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
+	if err != nil {
+		t.Fatal(err)
+	}
+	cert, err := x509.ParseCertificate(der)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pool := x509.NewCertPool()
+	pool.AddCert(cert)
+	return &tls.Config{
+		Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key, Leaf: cert}},
+		MinVersion:   tls.VersionTLS12,
+	}, pool
+}
+
+func SHA256Hex(data []byte) string {
+	sum := sha256.Sum256(data)
+	return hex.EncodeToString(sum[:])
+}
+
+func FileSHA256(t *testing.T, name string) string {
+	t.Helper()
+	data, err := os.ReadFile(name)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return SHA256Hex(data)
+}
+
+// RandomBytes returns n reproducible pseudo-random bytes.
+func RandomBytes(n int, seed byte) []byte {
+	out := make([]byte, n)
+	x := uint32(seed) + 1
+	for i := range out {
+		x ^= x << 13
+		x ^= x >> 17
+		x ^= x << 5
+		out[i] = byte(x)
+	}
+	return out
+}