Fix blocking mux reads and RAW download parsing

internal/server/tcp/downloader.go

@@ -1,6 +1,8 @@
 package tcp
 
 import (
+	"errors"
+	"fmt"
 	"io"
 	"log"
 	"net"
@@ -12,30 +14,23 @@ import (
 	"github.com/NHAS/reverse_ssh/pkg/logger"
 )
 
+const (
+	rawDownloadPrefix        = "RAW"
+	rawDownloadMaxNameLength = 64
+	rawDownloadReadTimeout   = 3 * time.Second
+)
+
 func handleBashConn(conn net.Conn) {
 	defer conn.Close()
 
 	downloadLog := logger.NewLog(conn.RemoteAddr().String())
 
-	conn.SetDeadline(time.Now().Add(3 * time.Second))
-	// RAW header prefix + 64 bytes for file ID
-	fileID := make([]byte, 67)
-
-	n, err := conn.Read(fileID)
+	filename, err := readRawDownloadName(conn)
 	if err != nil {
 		downloadLog.Warning("failed to download file using raw tcp: %s", err)
 		return
 	}
 
-	conn.SetDeadline(time.Time{})
-
-	if n == 0 || n < 3 {
-		downloadLog.Warning("recieved malformed raw download request")
-		return
-	}
-
-	filename := strings.TrimSpace(string(fileID[3:n]))
-
 	f, err := data.GetDownload(filename)
 	if err != nil {
 		downloadLog.Warning("failed to get file %q: err %s", filename, err)
@@ -54,6 +49,51 @@ func handleBashConn(conn net.Conn) {
 	io.Copy(conn, file)
 }
 
+func readRawDownloadName(conn net.Conn) (string, error) {
+	_ = conn.SetReadDeadline(time.Now().Add(rawDownloadReadTimeout))
+	defer conn.SetReadDeadline(time.Time{})
+
+	limit := len(rawDownloadPrefix) + rawDownloadMaxNameLength + 1
+	request := make([]byte, 0, limit)
+	var b [1]byte
+
+	for len(request) < limit {
+		n, err := conn.Read(b[:])
+		if n > 0 {
+			if b[0] == '\n' {
+				break
+			}
+			request = append(request, b[0])
+		}
+
+		if err != nil {
+			if errors.Is(err, io.EOF) && len(request) > 0 {
+				break
+			}
+			return "", err
+		}
+	}
+
+	if len(request) >= limit {
+		return "", fmt.Errorf("raw download request exceeds %d bytes", limit)
+	}
+
+	requestString := strings.TrimSpace(string(request))
+	if !strings.HasPrefix(requestString, rawDownloadPrefix) {
+		return "", fmt.Errorf("malformed raw download request")
+	}
+
+	filename := strings.TrimSpace(strings.TrimPrefix(requestString, rawDownloadPrefix))
+	if filename == "" {
+		return "", fmt.Errorf("empty raw download filename")
+	}
+	if len(filename) > rawDownloadMaxNameLength {
+		return "", fmt.Errorf("raw download filename exceeds %d bytes", rawDownloadMaxNameLength)
+	}
+
+	return filename, nil
+}
+
 func Start(listener net.Listener) {
 
 	log.Println("Started Raw Download Server")

internal/server/tcp/downloader_test.go

@@ -0,0 +1,138 @@
+package tcp
+
+import (
+	"io"
+	"net"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestReadRawDownloadName(t *testing.T) {
+	tests := []struct {
+		name    string
+		chunks  []string
+		want    string
+		wantErr bool
+	}{
+		{
+			name:   "shell script",
+			chunks: []string{"RAWclient.sh\n"},
+			want:   "client.sh",
+		},
+		{
+			name:   "powershell script",
+			chunks: []string{"RAWclient.ps1\r\n"},
+			want:   "client.ps1",
+		},
+		{
+			name:   "raw binary without newline",
+			chunks: []string{"RAWpayload.bin"},
+			want:   "payload.bin",
+		},
+		{
+			name:   "max length filename",
+			chunks: []string{"RAW" + strings.Repeat("a", rawDownloadMaxNameLength) + "\n"},
+			want:   strings.Repeat("a", rawDownloadMaxNameLength),
+		},
+		{
+			name:   "split raw request",
+			chunks: []string{"R", "AW", "payload.bin", "\n"},
+			want:   "payload.bin",
+		},
+		{
+			name:    "missing raw prefix",
+			chunks:  []string{"GET /payload.bin HTTP/1.1\r\n"},
+			wantErr: true,
+		},
+		{
+			name:    "empty filename",
+			chunks:  []string{"RAW\n"},
+			wantErr: true,
+		},
+		{
+			name:    "filename too long",
+			chunks:  []string{"RAW" + strings.Repeat("a", rawDownloadMaxNameLength+1) + "\n"},
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			conn := &chunkedConn{chunks: byteChunks(tt.chunks)}
+
+			got, err := readRawDownloadName(conn)
+			if !conn.readDeadline.IsZero() {
+				t.Fatal("read deadline was not reset")
+			}
+			if tt.wantErr {
+				if err == nil {
+					t.Fatal("expected error")
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("readRawDownloadName failed: %v", err)
+			}
+			if got != tt.want {
+				t.Fatalf("filename = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}
+
+func byteChunks(chunks []string) [][]byte {
+	result := make([][]byte, 0, len(chunks))
+	for _, chunk := range chunks {
+		result = append(result, []byte(chunk))
+	}
+	return result
+}
+
+type chunkedConn struct {
+	chunks       [][]byte
+	readDeadline time.Time
+}
+
+func (c *chunkedConn) Read(b []byte) (int, error) {
+	if len(c.chunks) == 0 {
+		return 0, io.EOF
+	}
+
+	n := copy(b, c.chunks[0])
+	if n == len(c.chunks[0]) {
+		c.chunks = c.chunks[1:]
+	} else {
+		c.chunks[0] = c.chunks[0][n:]
+	}
+	return n, nil
+}
+
+func (c *chunkedConn) Write([]byte) (int, error) {
+	return 0, nil
+}
+
+func (c *chunkedConn) Close() error {
+	return nil
+}
+
+func (c *chunkedConn) LocalAddr() net.Addr {
+	return nil
+}
+
+func (c *chunkedConn) RemoteAddr() net.Addr {
+	return nil
+}
+
+func (c *chunkedConn) SetDeadline(time.Time) error {
+	return nil
+}
+
+func (c *chunkedConn) SetReadDeadline(t time.Time) error {
+	c.readDeadline = t
+	return nil
+}
+
+func (c *chunkedConn) SetWriteDeadline(time.Time) error {
+	return nil
+}

pkg/mux/bufferedConn.go

@@ -13,18 +13,8 @@ type bufferedConn struct {
 func (bc *bufferedConn) Read(b []byte) (n int, err error) {
 	if len(bc.prefix) > 0 {
 		n = copy(b, bc.prefix)
-
 		bc.prefix = bc.prefix[n:]
-
-		var err error
-		if len(b)-n > 0 {
-			// If we havent exhausted the size of b, read some more
-			var actualRead int
-			actualRead, err = bc.conn.Read(b[n:])
-			n += actualRead
-		}
-
-		return n, err
+		return n, nil
 	}
 
 	return bc.conn.Read(b)

pkg/mux/multiplexer_test.go

@@ -0,0 +1,56 @@
+package mux
+
+import (
+	"net"
+	"testing"
+	"time"
+)
+
+func TestBufferedConnReadUsesPrefixWithoutBlocking(t *testing.T) {
+	request := []byte("GET /main.sh HTTP/1.0\r\nHost: example\r\n\r\n")
+	blocking := &blockingConn{}
+	bc := &bufferedConn{prefix: append([]byte(nil), request...), conn: blocking}
+
+	buf := make([]byte, 4096)
+	type readResult struct {
+		n   int
+		err error
+	}
+	done := make(chan readResult, 1)
+	go func() {
+		n, err := bc.Read(buf)
+		done <- readResult{n: n, err: err}
+	}()
+
+	select {
+	case result := <-done:
+		if result.err != nil {
+			t.Fatalf("read failed: %v", result.err)
+		}
+		if result.n != len(request) {
+			t.Fatalf("read %d bytes, want %d", result.n, len(request))
+		}
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("bufferedConn.Read blocked waiting for underlying conn")
+	}
+	if blocking.readCalled {
+		t.Fatal("bufferedConn should not read from underlying conn while prefix has data")
+	}
+}
+
+type blockingConn struct {
+	readCalled bool
+}
+
+func (c *blockingConn) Read([]byte) (int, error) {
+	c.readCalled = true
+	select {}
+}
+
+func (c *blockingConn) Write([]byte) (int, error)        { return 0, nil }
+func (c *blockingConn) Close() error                     { return nil }
+func (c *blockingConn) LocalAddr() net.Addr              { return nil }
+func (c *blockingConn) RemoteAddr() net.Addr             { return nil }
+func (c *blockingConn) SetDeadline(time.Time) error      { return nil }
+func (c *blockingConn) SetReadDeadline(time.Time) error  { return nil }
+func (c *blockingConn) SetWriteDeadline(time.Time) error { return nil }