Framing in Stream.Write to prevent silent short write

internal/common/copy.go

@@ -77,9 +77,7 @@ func Copy(dst net.Conn, src net.Conn, srcReadTimeout time.Duration) (written int
 		}
 		nr, er := src.Read(buf)
 		if nr > 0 {
-			var offset int
-			for offset < nr {
-				nw, ew := dst.Write(buf[offset:nr])
+			nw, ew := dst.Write(buf[0:nr])
 			if nw > 0 {
 				written += int64(nw)
 			}
@@ -87,7 +85,9 @@ func Copy(dst net.Conn, src net.Conn, srcReadTimeout time.Duration) (written int
 				err = ew
 				break
 			}
-				offset += nw
+			if nr != nw {
+				err = io.ErrShortWrite
+				break
 			}
 		}
 		if er != nil {

internal/multiplex/session.go

@@ -59,6 +59,8 @@ type Session struct {
 	closed uint32
 
 	terminalMsg atomic.Value
+
+	maxStreamUnitWrite int // the max size passed to Write calls before it splits it into multiple frames
 }
 
 func MakeSession(id uint32, config SessionConfig) *Session {
@@ -82,6 +84,7 @@ func MakeSession(id uint32, config SessionConfig) *Session {
 	if config.MaxFrameSize <= 0 {
 		sesh.MaxFrameSize = defaultSendRecvBufSize - 1024
 	}
+	sesh.maxStreamUnitWrite = sesh.MaxFrameSize - HEADER_LEN - sesh.Obfuscator.minOverhead
 
 	sbConfig := switchboardConfig{
 		valve:          sesh.Valve,

internal/multiplex/stream.go

@@ -96,28 +96,29 @@ func (s *Stream) Write(in []byte) (n int, err error) {
 		return 0, ErrBrokenStream
 	}
 
-	var payload []byte
-	maxDataLen := s.session.MaxFrameSize - HEADER_LEN - s.session.minOverhead
-	if len(in) <= maxDataLen {
-		payload = in
+	for n < len(in) {
+		var framePayload []byte
+		if len(in)-n <= s.session.maxStreamUnitWrite {
+			framePayload = in[n:]
 		} else {
-		//TODO: short write isn't the correct behaviour
-		payload = in[:maxDataLen]
+			framePayload = in[n : s.session.maxStreamUnitWrite+n]
 		}
 
 		f := &Frame{
 			StreamID: s.id,
 			Seq:      atomic.AddUint64(&s.nextSendSeq, 1) - 1,
 			Closing:  C_NOOP,
-		Payload:  payload,
+			Payload:  framePayload,
 		}
 
-	i, err := s.session.Obfs(f, s.obfsBuf)
+		var cipherTextLen int
+		cipherTextLen, err = s.session.Obfs(f, s.obfsBuf)
 		if err != nil {
-		return i, err
+			return 0, err
 		}
-	n, err = s.session.sb.send(s.obfsBuf[:i], &s.assignedConnId)
-	log.Tracef("%v sent to remote through stream %v with err %v", len(payload), s.id, err)
+
+		_, err = s.session.sb.send(s.obfsBuf[:cipherTextLen], &s.assignedConnId)
+		log.Tracef("%v sent to remote through stream %v with err %v", len(framePayload), s.id, err)
 		if err != nil {
 			if err == errBrokenSwitchboard {
 				s.session.SetTerminalMsg(err.Error())
@@ -125,8 +126,9 @@ func (s *Stream) Write(in []byte) (n int, err error) {
 			}
 			return
 		}
-	return len(payload), nil
-
+		n += len(framePayload)
+	}
+	return
 }
 
 func (s *Stream) passiveClose() error {