fakeserver: CompressShape shapes the COMPRESS DEFLATE output stream

fakeserver/compress.go

@@ -2,13 +2,13 @@ package fakeserver
 
 import (
 	"bufio"
-	"bytes"
-	"compress/flate"
 	"errors"
 	"io"
 	"net"
 	"regexp"
+	"slices"
 	"strings"
+	"time"
 
 	"github.com/go-while/nntp-testkit/internal/inflate"
 )
@@ -27,15 +27,20 @@ import (
 // handlers return and the connection is closed, as RFC 8054 2.2.2 requires. Raw returns the
 // compressing conn afterwards; keep Raw() from before the call to write bytes past the
 // compressor (corrupt output for a client test). It fails when compression is already on.
-func (c *Conn) StartCompress() error { return c.startCompress(false) }
+func (c *Conn) StartCompress() error { return c.startCompress(false, CompressShape{}) }
+
+// StartCompressShaped is StartCompress with the output stream shaped (CompressShape): flush
+// kind, splits, a final block. EmptyFlushWith206 and With206 have no effect here, since the 206
+// is already written; the handlers of the package write them with the 206.
+func (c *Conn) StartCompressShaped(shape CompressShape) error { return c.startCompress(false, shape) }
 
 // StartCompressOutputOnly is a broken COMPRESS for client tests: after the 206 only what the
 // server writes is deflated, while the client's bytes are read as they arrive, without
 // inflating them, as a server that compresses its responses but never decompresses its
 // input. A client that compresses its commands, as RFC 8054 requires, then gets no answer
 // (its deflated bytes rarely hold a line end) or "500" answers to binary lines, and Lines
 // records those raw bytes.
-func (c *Conn) StartCompressOutputOnly() error { return c.startCompress(true) }
+func (c *Conn) StartCompressOutputOnly() error { return c.startCompress(true, CompressShape{}) }
 
 // Compressed reports whether StartCompress or StartCompressOutputOnly switched the connection
 // to compression.
@@ -44,16 +49,19 @@ func (c *Conn) Compressed() bool {
 	return ok
 }
 
-func (c *Conn) startCompress(outputOnly bool) error {
+func (c *Conn) startCompress(outputOnly bool, shape CompressShape) error {
 	if c.Compressed() {
 		return errors.New("fakeserver: compression is already active")
 	}
-	d := &deflateConn{Conn: c.conn, r: c.r}
-	zw, err := flate.NewWriter(&d.out, flate.DefaultCompression)
-	if err != nil {
-		return err
+	enc := c.enc // the one that wrote the pipelined bytes of the 206, when it did
+	c.enc = nil
+	if enc == nil {
+		var err error
+		if enc, err = newEncoder(shape); err != nil {
+			return err
+		}
 	}
-	d.zw = zw
+	d := &deflateConn{Conn: c.conn, r: c.r, enc: enc, sleep: c.sleep, deadline: c.s.deadline}
 	if !outputOnly {
 		pr, pw := io.Pipe()
 		d.r, d.pr = pr, pr
@@ -74,11 +82,25 @@ func (c *Conn) startCompress(outputOnly bool) error {
 // addresses and Close are those of the conn underneath.
 type deflateConn struct {
 	net.Conn
-	r    io.Reader
-	pr   *io.PipeReader // the inflater's output; nil for StartCompressOutputOnly
-	zw   *flate.Writer  // writes into out
-	out  bytes.Buffer   // the deflated bytes of one Write, sent in one write
-	werr error          // the first failed write: the stream is broken from there on
+	r        io.Reader
+	pr       *io.PipeReader // the inflater's output; nil for StartCompressOutputOnly
+	enc      *encoder
+	sleep    func(time.Duration) bool // Conn.sleep: bounded by the Deadline, false when the server closed
+	deadline time.Duration
+	werr     error     // the first failed write: the stream is broken from there on
+	wdl      time.Time // the write deadline last set through this conn (zero: none)
+}
+
+// SetWriteDeadline and SetDeadline remember the write deadline, so that send can put back the
+// one writeKind set (Drip, LineDelay, Split, Endless) instead of clearing it.
+func (d *deflateConn) SetWriteDeadline(t time.Time) error {
+	d.wdl = t
+	return d.Conn.SetWriteDeadline(t)
+}
+
+func (d *deflateConn) SetDeadline(t time.Time) error {
+	d.wdl = t
+	return d.Conn.SetDeadline(t)
 }
 
 func (d *deflateConn) Read(p []byte) (int, error) { return d.r.Read(p) }
@@ -91,29 +113,56 @@ func (d *deflateConn) Close() error {
 	return d.Conn.Close()
 }
 
-// Write deflates p, flushes the compressor (sync flush) and writes the result to the conn
-// underneath in one write, as a server that flushes its output buffer, so the client can
-// inflate all of p at once (compress/flate's Writer on its own writes a short input in five
-// pieces). After a failed write the deflate stream is broken, so every later Write returns
-// that error.
+// Write deflates p, flushes the compressor as the shape says (sync flush by default) and
+// writes the result to the conn underneath in one write (or the pieces of the shape's Split),
+// as a server that flushes its output buffer, so the client can inflate all of p at once
+// (compress/flate's Writer on its own writes a short input in five pieces). After a failed
+// write the deflate stream is broken, so every later Write returns that error.
 func (d *deflateConn) Write(p []byte) (int, error) {
 	if d.werr != nil {
 		return 0, d.werr
 	}
-	d.out.Reset()
-	if _, err := d.zw.Write(p); err != nil { // into d.out: cannot fail
-		return 0, err
-	}
-	if err := d.zw.Flush(); err != nil {
+	b, err := d.enc.write(p)
+	if err != nil {
 		return 0, err
 	}
-	if _, err := d.Conn.Write(d.out.Bytes()); err != nil {
+	if err := d.send(b); err != nil {
 		d.werr = err
 		return 0, err
 	}
 	return len(p), nil
 }
 
+// send writes b to the conn underneath, cut by the shape's Split. The pauses between the
+// pieces end at the server's Deadline from the call, with a write deadline set meanwhile.
+func (d *deflateConn) send(b []byte) error {
+	if len(b) == 0 {
+		return nil
+	}
+	sp := d.enc.shape.Split
+	cuts := sp.cuts(len(b))
+	if len(cuts) == 0 {
+		_, err := d.Conn.Write(b)
+		return err
+	}
+	end := time.Now().Add(d.deadline)
+	if sp.Gap > 0 && (d.wdl.IsZero() || end.Before(d.wdl)) {
+		_ = d.Conn.SetWriteDeadline(end)
+		defer d.Conn.SetWriteDeadline(d.wdl) // what the caller had set, not "none"
+	}
+	prev := 0
+	for _, k := range append(cuts, len(b)) {
+		if _, err := d.Conn.Write(b[prev:k]); err != nil {
+			return err
+		}
+		prev = k
+		if k < len(b) && sp.Gap > 0 && (time.Now().After(end) || !d.sleep(sp.Gap)) {
+			return errors.New("fakeserver: compressed split ran past the Deadline or the server closed")
+		}
+	}
+	return nil
+}
+
 // algorithm is the syntax of an RFC 8054 algorithm name (5.3): 1 to 20 upper-case letters,
 // digits, "-" or "_"; case-sensitive.
 var algorithm = regexp.MustCompile(`^[A-Z0-9_-]{1,20}$`)
@@ -139,5 +188,25 @@ func (c *Conn) compressAfter(cmd, status string, outputOnly bool) bool {
 	if cmd != "COMPRESS" || (status != "206" && !strings.HasPrefix(status, "206 ")) || c.Compressed() {
 		return true
 	}
-	return c.startCompress(outputOnly) == nil
+	return c.startCompress(outputOnly, c.shape) == nil
+}
+
+// pipeline206 returns the compressed bytes the shape writes in the same socket write as the
+// 206 reply to a command (nil for none), and keeps the encoder for startCompress. writeKind
+// calls it with the reply's status line; armed is set for the command that starts compression
+// (COMPRESS, or any command of a Script step with Compress).
+func (c *Conn) pipeline206(armed bool, status string) []byte {
+	if !armed || !c.shape.pipelined() || c.Compressed() || (status != "206" && !strings.HasPrefix(status, "206 ")) {
+		return nil
+	}
+	enc, err := newEncoder(c.shape)
+	if err != nil {
+		return nil
+	}
+	b, err := enc.flushData(c.shape.With206)
+	if err != nil {
+		return nil
+	}
+	c.enc = enc
+	return slices.Clone(b)
 }

fakeserver/compressshape.go

@@ -0,0 +1,203 @@
+package fakeserver
+
+import (
+	"bytes"
+	"compress/flate"
+	"math/bits"
+	"slices"
+	"sort"
+)
+
+// CompressFlush is how the compressed output stream is flushed at the end of every write.
+type CompressFlush int
+
+const (
+	// FlushSync ends every write with a sync flush (zlib's Z_SYNC_FLUSH): an empty stored
+	// block, the marker 00 00 FF FF on a byte boundary. The default.
+	FlushSync CompressFlush = iota
+	// FlushPartial ends every write with zlib's Z_PARTIAL_FLUSH, as INN's nnrpd does: an empty
+	// fixed-Huffman block (10 bits), of which only the whole bytes are sent. The last 0 to 7
+	// bits stay in the compressor until the next write; the data before them is complete.
+	// The data before them is complete: the held bits are only the tail of the empty block's
+	// end-of-block code, which is zeros. There is no 00 00 FF FF marker in the stream. The data are literals in fixed-Huffman
+	// blocks, one per write (no matches), so the output is larger than FlushSync's. A client
+	// that waits for a stored block, or for the end of the block it is in, hangs on this
+	// stream, as INN's clients must not.
+	FlushPartial
+	// FlushFinish ends the deflate stream (a final block) with the first write after the 206:
+	// that write is deflated and the final empty block follows in the same socket write. Later
+	// writes are plain bytes on the socket, after the end of the stream (the kit still inflates
+	// what the client sends). CompressShape.AfterEnd adds bytes behind the final block.
+	FlushFinish
+)
+
+// CompressShape shapes what the fake server sends on a COMPRESS DEFLATE connection (RFC 8054),
+// for client tests: StoreOptions.CompressShape, PeerOptions.CompressShape, Step.CompressShape,
+// Conn.StartCompressShaped. The zero value is the default: every write sync-flushed and sent
+// in one socket write.
+type CompressShape struct {
+	// Flush is the flush at the end of each write, see CompressFlush.
+	Flush CompressFlush
+
+	// Split cuts the compressed bytes of each write (reply, Faults piece) into separate socket
+	// writes: At holds offsets from the start of the write, a negative offset counts from its
+	// end (with FlushSync -1 .. -3 cut inside the sync marker 00 00 FF FF, -4 before it, and
+	// with FlushPartial -1 cuts the empty block's last byte), Size cuts a piece every Size
+	// bytes when At is empty, Gap pauses after each piece but the last, bounded by the
+	// server's Deadline like Faults.Split. A client that reads between the pieces sees
+	// compressed bytes that end mid-block or mid-marker.
+	Split Split
+
+	// EmptyFlushWith206 writes an empty flush block (the marker of Flush) right after the CRLF
+	// of the "206" in the same socket write as the 206 line, as a server that flushes its
+	// compressor when it switches it on. A client must keep those bytes for its inflater. The
+	// 206 is written by the kit's handlers (Store, Peer with Compress or CompressOutputOnly,
+	// Script step with Compress); a Handler that writes its own 206 cannot use this. Faults on
+	// the COMPRESS reply (DropAfterBytes, Split offsets) count these bytes too.
+	EmptyFlushWith206 bool
+	// With206 is deflated and flushed like a reply and written in the same socket write as the
+	// 206 (it ends with its own flush, so EmptyFlushWith206 adds nothing then): compressed
+	// payload pipelined behind the 206, for a client that reads past the 206 line.
+	With206 []byte
+
+	// AfterEnd is written in the same socket write as the final block, with FlushFinish:
+	// trailing bytes after the end of the stream (junk, or plain text).
+	AfterEnd []byte
+}
+
+func (s CompressShape) pipelined() bool { return s.EmptyFlushWith206 || len(s.With206) > 0 }
+
+// cuts returns the sorted offsets in (0, n) where a write of n bytes is cut.
+func (s Split) cuts(n int) []int {
+	var out []int
+	for _, a := range s.At {
+		if a < 0 {
+			a += n
+		}
+		if a > 0 && a < n {
+			out = append(out, a)
+		}
+	}
+	sort.Ints(out)
+	out = slices.Compact(out)
+	if len(s.At) == 0 && s.Size > 0 {
+		for k := s.Size; k < n; k += s.Size {
+			out = append(out, k)
+		}
+	}
+	return out
+}
+
+// encoder deflates the output of a compressed connection as its shape says. Not safe for
+// concurrent use (one connection goroutine writes).
+type encoder struct {
+	shape CompressShape
+	zw    *flate.Writer // FlushSync and FlushFinish; writes into out
+	out   bytes.Buffer
+	bw    bitWriter // FlushPartial
+	ended bool      // FlushFinish: the stream is over
+}
+
+func newEncoder(s CompressShape) (*encoder, error) {
+	e := &encoder{shape: s}
+	if s.Flush != FlushPartial {
+		zw, err := flate.NewWriter(&e.out, flate.DefaultCompression)
+		if err != nil {
+			return nil, err
+		}
+		e.zw = zw
+	}
+	return e, nil
+}
+
+// flushData deflates p and flushes (sync, or partial for FlushPartial; FlushFinish counts as
+// sync here). An empty p is an empty flush block. The result is valid until the next call.
+func (e *encoder) flushData(p []byte) ([]byte, error) {
+	if e.shape.Flush == FlushPartial {
+		return e.partial(p), nil
+	}
+	e.out.Reset()
+	if _, err := e.zw.Write(p); err != nil { // into e.out: cannot fail
+		return nil, err
+	}
+	if err := e.zw.Flush(); err != nil {
+		return nil, err
+	}
+	return e.out.Bytes(), nil
+}
+
+// write returns the bytes to send for a write of p: deflated and flushed, or with FlushFinish
+// the end of the stream (once; then p itself).
+func (e *encoder) write(p []byte) ([]byte, error) {
+	if e.ended {
+		return p, nil
+	}
+	if e.shape.Flush != FlushFinish {
+		return e.flushData(p)
+	}
+	e.out.Reset()
+	if _, err := e.zw.Write(p); err != nil {
+		return nil, err
+	}
+	if err := e.zw.Close(); err != nil {
+		return nil, err
+	}
+	e.ended = true
+	return append(e.out.Bytes(), e.shape.AfterEnd...), nil
+}
+
+// partial is zlib's Z_PARTIAL_FLUSH (deflate.c: FLUSH_BLOCK, then _tr_align): the data as one
+// fixed-Huffman block of literals when there are any, then an empty fixed-Huffman block (3
+// header bits, the 7-bit end-of-block code), and only the whole bytes are sent. Since zlib
+// 1.2.5.2 that is all: older versions added a second empty block when fewer than 9 bits of
+// the last block's end code had been sent.
+func (e *encoder) partial(p []byte) []byte {
+	w := &e.bw
+	if len(p) > 0 {
+		w.block(p)
+	}
+	w.block(nil)
+	return w.take()
+}
+
+// bitWriter writes DEFLATE bits, the first one in the lowest bit of a byte (RFC 1951 3.1.1).
+type bitWriter struct {
+	out []byte
+	acc uint32
+	n   uint // bits in acc, fewer than 8 after put
+}
+
+func (w *bitWriter) put(v uint32, n uint) {
+	w.acc |= v << w.n
+	w.n += n
+	for w.n >= 8 {
+		w.out = append(w.out, byte(w.acc))
+		w.acc >>= 8
+		w.n -= 8
+	}
+}
+
+// code writes an n-bit Huffman code, which goes most significant bit first.
+func (w *bitWriter) code(v uint32, n uint) { w.put(bits.Reverse32(v)>>(32-n), n) }
+
+// block writes a non-final fixed-Huffman block (BFINAL 0, BTYPE 01) of the literals p and the
+// end-of-block code (RFC 1951 3.2.6: literals 0-143 are 8 bits from 00110000, 144-255 are 9
+// bits from 110010000, the end of block is 7 zero bits).
+func (w *bitWriter) block(p []byte) {
+	w.put(2, 3) // BFINAL 0 in the lowest bit, BTYPE 01 above it
+	for _, b := range p {
+		if b < 144 {
+			w.code(0x30+uint32(b), 8)
+		} else {
+			w.code(0x190+uint32(b)-144, 9)
+		}
+	}
+	w.code(0, 7)
+}
+
+// take returns the whole bytes written so far; the last bits stay.
+func (w *bitWriter) take() []byte {
+	b := w.out
+	w.out = nil
+	return b
+}