V13

core/internal/wire/protocol.go

@@ -0,0 +1,177 @@
+package wire
+
+import (
+	"crypto/sha256"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+)
+
+const (
+	RequestHeaderSize  = 29
+	ResponseHeaderSize = 5
+	MaxPayload         = 2 * 1024 * 1024
+
+	ModeProbe    byte = 0
+	ModeOpen     byte = 1
+	ModeUpload   byte = 2
+	ModeDownload byte = 3
+	ModeClose    byte = 4
+
+	StatusOK    byte = 0
+	StatusError byte = 1
+	StatusData  byte = 2
+	StatusWait  byte = 3
+	StatusEOF   byte = 4
+
+	ProbeUpload    byte = 1
+	ProbeDownload  byte = 2
+	ProbeKeepalive byte = 3
+	ProbeBatch     byte = 4
+)
+
+var ProbeMagic = [4]byte{'D', 'T', 'P', '2'}
+
+type SessionID [16]byte
+
+type Request struct {
+	Mode    byte
+	Session SessionID
+	Seq     uint64
+	Payload []byte
+}
+
+func MaskInPlace(data []byte, sid SessionID, mode byte, seq uint64, response bool) {
+	if len(data) == 0 {
+		return
+	}
+
+	var seed [30]byte
+	copy(seed[:16], sid[:])
+	seed[16] = mode
+	binary.BigEndian.PutUint64(seed[17:25], seq)
+	if response {
+		seed[25] = 1
+	}
+
+	var counter uint32
+	for off := 0; off < len(data); {
+		binary.BigEndian.PutUint32(seed[26:30], counter)
+		block := sha256.Sum256(seed[:])
+		n := len(data) - off
+		if n > len(block) {
+			n = len(block)
+		}
+		for i := 0; i < n; i++ {
+			data[off+i] ^= block[i]
+		}
+		off += n
+		counter++
+	}
+}
+
+func WriteRequest(w io.Writer, mode byte, sid SessionID, seq uint64, plaintext []byte) error {
+	if len(plaintext) > MaxPayload {
+		return fmt.Errorf("request payload too large: %d", len(plaintext))
+	}
+
+	packet := make([]byte, RequestHeaderSize+len(plaintext))
+	packet[0] = mode
+	copy(packet[1:17], sid[:])
+	binary.BigEndian.PutUint64(packet[17:25], seq)
+	binary.BigEndian.PutUint32(packet[25:29], uint32(len(plaintext)))
+	copy(packet[29:], plaintext)
+	MaskInPlace(packet[29:], sid, mode, seq, false)
+	return writeAll(w, packet)
+}
+
+func ReadRequest(r io.Reader) (Request, error) {
+	var req Request
+	var header [RequestHeaderSize]byte
+	if _, err := io.ReadFull(r, header[:]); err != nil {
+		return req, err
+	}
+
+	req.Mode = header[0]
+	copy(req.Session[:], header[1:17])
+	req.Seq = binary.BigEndian.Uint64(header[17:25])
+	n := binary.BigEndian.Uint32(header[25:29])
+	if n > MaxPayload {
+		return req, errors.New("request payload too large")
+	}
+
+	if n > 0 {
+		req.Payload = make([]byte, int(n))
+		if _, err := io.ReadFull(r, req.Payload); err != nil {
+			return req, err
+		}
+		MaskInPlace(req.Payload, req.Session, req.Mode, req.Seq, false)
+	}
+	return req, nil
+}
+
+func WriteResponse(w io.Writer, status byte, body []byte) error {
+	if len(body) > MaxPayload {
+		return fmt.Errorf("response body too large: %d", len(body))
+	}
+	packet := make([]byte, ResponseHeaderSize+len(body))
+	packet[0] = status
+	binary.BigEndian.PutUint32(packet[1:5], uint32(len(body)))
+	copy(packet[5:], body)
+	return writeAll(w, packet)
+}
+
+func WriteMaskedResponse(w io.Writer, status byte, body []byte, sid SessionID, mode byte, seq uint64) error {
+	if len(body) > MaxPayload {
+		return fmt.Errorf("response body too large: %d", len(body))
+	}
+	packet := make([]byte, ResponseHeaderSize+len(body))
+	packet[0] = status
+	binary.BigEndian.PutUint32(packet[1:5], uint32(len(body)))
+	copy(packet[5:], body)
+	MaskInPlace(packet[5:], sid, mode, seq, true)
+	return writeAll(w, packet)
+}
+
+func ReadResponse(r io.Reader) (byte, []byte, error) {
+	var header [ResponseHeaderSize]byte
+	if _, err := io.ReadFull(r, header[:]); err != nil {
+		return 0, nil, err
+	}
+	n := binary.BigEndian.Uint32(header[1:5])
+	if n > MaxPayload {
+		return 0, nil, errors.New("response body too large")
+	}
+	var body []byte
+	if n > 0 {
+		body = make([]byte, int(n))
+		if _, err := io.ReadFull(r, body); err != nil {
+			return 0, nil, err
+		}
+	}
+	return header[0], body, nil
+}
+
+func DecodeMaskedResponse(status byte, body []byte, sid SessionID, mode byte, seq uint64) []byte {
+	if len(body) == 0 || status == StatusError {
+		return body
+	}
+	out := append([]byte(nil), body...)
+	MaskInPlace(out, sid, mode, seq, true)
+	return out
+}
+
+func writeAll(w io.Writer, b []byte) error {
+	for len(b) > 0 {
+		n, err := w.Write(b)
+		if err != nil {
+			return err
+		}
+		if n <= 0 {
+			return io.ErrShortWrite
+		}
+		b = b[n:]
+	}
+	return nil
+}