Lazy spool large messages to disk

Add support for wasm-donut options in windows `execute-shellcode`

server/c2/envelope_reader.go

@@ -0,0 +1,69 @@
+package c2
+
+import (
+	"io"
+	"os"
+)
+
+const (
+	// socketEnvelopeDiskSpoolThreshold is the point where envelope payloads are
+	// streamed to disk first to avoid large attacker-triggered allocations.
+	socketEnvelopeDiskSpoolThreshold = 64 * 1024 * 1024
+
+	// socketEnvelopeReadChunkSize bounds peak in-memory usage while reading large
+	// envelopes from the network.
+	socketEnvelopeReadChunkSize = 1 * 1024 * 1024
+)
+
+func readSocketEnvelopeData(reader io.Reader, dataLength int, inMemoryLimit int) ([]byte, error) {
+	if dataLength <= inMemoryLimit {
+		dataBuf := make([]byte, dataLength)
+		if _, err := io.ReadFull(reader, dataBuf); err != nil {
+			return nil, err
+		}
+		return dataBuf, nil
+	}
+
+	tmpFile, err := os.CreateTemp("", "sliver-envelope-*")
+	if err != nil {
+		return nil, err
+	}
+	tmpPath := tmpFile.Name()
+	defer func() {
+		_ = tmpFile.Close()
+		_ = os.Remove(tmpPath)
+	}()
+
+	chunkSize := socketEnvelopeReadChunkSize
+	if dataLength < chunkSize {
+		chunkSize = dataLength
+	}
+	chunkBuf := make([]byte, chunkSize)
+
+	remaining := dataLength
+	for remaining > 0 {
+		nextRead := chunkSize
+		if remaining < nextRead {
+			nextRead = remaining
+		}
+
+		if _, err := io.ReadFull(reader, chunkBuf[:nextRead]); err != nil {
+			return nil, err
+		}
+		if _, err := tmpFile.Write(chunkBuf[:nextRead]); err != nil {
+			return nil, err
+		}
+
+		remaining -= nextRead
+	}
+
+	if _, err := tmpFile.Seek(0, io.SeekStart); err != nil {
+		return nil, err
+	}
+
+	dataBuf := make([]byte, dataLength)
+	if _, err := io.ReadFull(tmpFile, dataBuf); err != nil {
+		return nil, err
+	}
+	return dataBuf, nil
+}

server/c2/envelope_reader_test.go

@@ -0,0 +1,135 @@
+package c2
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"io"
+	"math"
+	"net"
+	"runtime"
+	"testing"
+
+	"github.com/bishopfox/sliver/protobuf/sliverpb"
+	"github.com/bishopfox/sliver/util/minisign"
+)
+
+type readSizeRecorder struct {
+	reader  io.Reader
+	maxRead int
+}
+
+func (r *readSizeRecorder) Read(p []byte) (int, error) {
+	if len(p) > r.maxRead {
+		r.maxRead = len(p)
+	}
+	return r.reader.Read(p)
+}
+
+func envelopeFrameWithDeclaredLength(dataLength uint32, payload []byte) []byte {
+	frame := make([]byte, minisign.RawSigSize+4+len(payload))
+	binary.LittleEndian.PutUint16(frame[:2], minisign.EdDSA)
+	binary.LittleEndian.PutUint32(frame[minisign.RawSigSize:minisign.RawSigSize+4], dataLength)
+	copy(frame[minisign.RawSigSize+4:], payload)
+	return frame
+}
+
+func measureEnvelopeReadAlloc(t *testing.T, readFunc func(net.Conn) (*sliverpb.Envelope, error), frame []byte) (uint64, error) {
+	t.Helper()
+
+	srv, cli := net.Pipe()
+	defer srv.Close()
+	defer cli.Close()
+
+	errCh := make(chan error, 1)
+	go func() {
+		_, err := readFunc(srv)
+		errCh <- err
+	}()
+
+	runtime.GC()
+	var before runtime.MemStats
+	runtime.ReadMemStats(&before)
+
+	go func() {
+		_, _ = cli.Write(frame)
+		_ = cli.Close()
+	}()
+
+	err := <-errCh
+
+	runtime.GC()
+	var after runtime.MemStats
+	runtime.ReadMemStats(&after)
+
+	if after.TotalAlloc < before.TotalAlloc {
+		return 0, err
+	}
+	return after.TotalAlloc - before.TotalAlloc, err
+}
+
+func assertLargeTruncatedFrameHasBoundedAlloc(t *testing.T, readFunc func(net.Conn) (*sliverpb.Envelope, error)) {
+	t.Helper()
+
+	declaredLength := uint32(socketEnvelopeDiskSpoolThreshold + (8 * 1024 * 1024))
+	frame := envelopeFrameWithDeclaredLength(declaredLength, []byte{0x01})
+
+	minAllocDelta := uint64(math.MaxUint64)
+	for i := 0; i < 3; i++ {
+		allocDelta, err := measureEnvelopeReadAlloc(t, readFunc, frame)
+		if err == nil {
+			t.Fatalf("expected read failure for truncated frame")
+		}
+		if !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) {
+			t.Fatalf("expected EOF-like error for truncated frame, got %v", err)
+		}
+		if allocDelta < minAllocDelta {
+			minAllocDelta = allocDelta
+		}
+	}
+
+	const maxAllowedAlloc = 24 * 1024 * 1024
+	if minAllocDelta > maxAllowedAlloc {
+		t.Fatalf("expected bounded allocation <= %d bytes, got %d", maxAllowedAlloc, minAllocDelta)
+	}
+}
+
+func TestReadSocketEnvelopeDataKeepsSmallPayloadsInMemory(t *testing.T) {
+	payload := bytes.Repeat([]byte{0x7a}, 4096)
+	recorder := &readSizeRecorder{reader: bytes.NewReader(payload)}
+
+	got, err := readSocketEnvelopeData(recorder, len(payload), 8192)
+	if err != nil {
+		t.Fatalf("readSocketEnvelopeData returned err: %v", err)
+	}
+	if !bytes.Equal(got, payload) {
+		t.Fatalf("unexpected payload contents")
+	}
+	if recorder.maxRead != len(payload) {
+		t.Fatalf("expected in-memory read size %d, got %d", len(payload), recorder.maxRead)
+	}
+}
+
+func TestReadSocketEnvelopeDataSpoolsLargePayloadsInChunks(t *testing.T) {
+	payload := bytes.Repeat([]byte{0x42}, (2*socketEnvelopeReadChunkSize)+123)
+	recorder := &readSizeRecorder{reader: bytes.NewReader(payload)}
+
+	got, err := readSocketEnvelopeData(recorder, len(payload), socketEnvelopeReadChunkSize/2)
+	if err != nil {
+		t.Fatalf("readSocketEnvelopeData returned err: %v", err)
+	}
+	if !bytes.Equal(got, payload) {
+		t.Fatalf("unexpected payload contents")
+	}
+	if recorder.maxRead > socketEnvelopeReadChunkSize {
+		t.Fatalf("expected max read chunk <= %d, got %d", socketEnvelopeReadChunkSize, recorder.maxRead)
+	}
+}
+
+func TestMTLSSocketReadEnvelopeLargeTruncatedMessageBoundsAllocation(t *testing.T) {
+	assertLargeTruncatedFrameHasBoundedAlloc(t, socketReadEnvelope)
+}
+
+func TestWGSocketReadEnvelopeLargeTruncatedMessageBoundsAllocation(t *testing.T) {
+	assertLargeTruncatedFrameHasBoundedAlloc(t, socketWGReadEnvelope)
+}

server/c2/mtls.go

@@ -358,10 +358,8 @@ func socketReadEnvelope(connection net.Conn) (*sliverpb.Envelope, error) {
 		return nil, errors.New("[pivot] invalid data length")
 	}
 
-	dataBuf := make([]byte, dataLength)
-
-	n, err = io.ReadFull(connection, dataBuf)
-	if err != nil || n != dataLength {
+	dataBuf, err := readSocketEnvelopeData(connection, dataLength, socketEnvelopeDiskSpoolThreshold)
+	if err != nil {
 		mtlsLog.Errorf("Socket error (read data): %v", err)
 		return nil, err
 	}

server/c2/wireguard.go

@@ -452,11 +452,8 @@ func socketWGReadEnvelope(connection net.Conn) (*sliverpb.Envelope, error) {
 		return nil, errors.New("[wireguard] zero data length")
 	}
 
-	dataBuf := make([]byte, dataLength)
-
-	n, err = io.ReadFull(connection, dataBuf)
-
-	if err != nil || n != dataLength {
+	dataBuf, err := readSocketEnvelopeData(connection, dataLength, socketEnvelopeDiskSpoolThreshold)
+	if err != nil {
 		wgLog.Errorf("Socket error (read data): %v", err)
 		return nil, err
 	}