Normalize packet handling

protocol/handler.go

@@ -0,0 +1,73 @@
+package protocol
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"github.com/bolkedebruin/rdpgw/transport"
+	"io"
+)
+
+type Handler struct {
+	Transport transport.Transport
+}
+
+func NewHandler(t transport.Transport) *Handler {
+	h := &Handler{
+		Transport: t,
+	}
+	return h
+}
+
+func (p *Handler) ReadMessage() (pt int, n int, msg []byte, err error) {
+	fragment := false
+	index := 0
+	buf := make([]byte, 4096)
+
+	for {
+		size, pkt, err := p.Transport.ReadPacket()
+		if err != nil {
+			return 0, 0, []byte{0,0}, err
+		}
+
+		// check for fragments
+		var pt uint16
+		var sz uint32
+		var msg []byte
+
+		if !fragment {
+			pt, sz, msg, err = readHeader(pkt[:size])
+			if err != nil {
+				fragment = true
+				index = copy(buf, pkt[:size])
+				continue
+			}
+			index = 0
+		} else {
+			fragment = false
+			pt, sz, msg, err = readHeader(append(buf[:index], pkt[:size]...))
+			// header is corrupted even after defragmenting
+			if err != nil {
+				return 0, 0, []byte{0,0}, err
+			}
+		}
+		if !fragment {
+			return int(pt), int(sz), msg, nil
+		}
+	}
+}
+
+func readHeader(data []byte) (packetType uint16, size uint32, packet []byte, err error) {
+	// header needs to be 8 min
+	if len(data) < 8 {
+		return 0, 0, nil, errors.New("header too short, fragment likely")
+	}
+	r := bytes.NewReader(data)
+	binary.Read(r, binary.LittleEndian, &packetType)
+	r.Seek(4, io.SeekStart)
+	binary.Read(r, binary.LittleEndian, &size)
+	if len(data) < int(size) {
+		return packetType, size, data[8:], errors.New("data incomplete, fragment received")
+	}
+	return packetType, size, data[8:], nil
+}