Browse Source

write by slice when buffer is too large

Darien Raymond 8 năm trước cách đây
mục cha
commit
31a0951bd2
3 tập tin đã thay đổi với 32 bổ sung5 xóa
  1. 1 1
      app/proxyman/mux/reader.go
  2. 16 4
      app/proxyman/mux/writer.go
  3. 15 0
      common/buf/multi_buffer.go

+ 1 - 1
app/proxyman/mux/reader.go

@@ -28,7 +28,7 @@ func (r *Reader) ReadMetadata() (*FrameMetadata, error) {
 	}
 	metaLen := serial.BytesToUint16(b.Bytes())
 	if metaLen > 512 {
-		return nil, newError("invalid metalen ", metaLen)
+		return nil, newError("invalid metalen ", metaLen).AtWarning()
 	}
 	b.Clear()
 	if err := b.AppendSupplier(buf.ReadFullFrom(r.reader, int(metaLen))); err != nil {

+ 16 - 4
app/proxyman/mux/writer.go

@@ -15,9 +15,10 @@ type Writer struct {
 
 func NewWriter(id uint16, dest net.Destination, writer buf.Writer) *Writer {
 	return &Writer{
-		id:     id,
-		dest:   dest,
-		writer: writer,
+		id:       id,
+		dest:     dest,
+		writer:   writer,
+		followup: false,
 	}
 }
 
@@ -29,7 +30,7 @@ func NewResponseWriter(id uint16, writer buf.Writer) *Writer {
 	}
 }
 
-func (w *Writer) Write(mb buf.MultiBuffer) error {
+func (w *Writer) writeInternal(mb buf.MultiBuffer) error {
 	meta := FrameMetadata{
 		SessionID: w.id,
 		Target:    w.dest,
@@ -58,6 +59,17 @@ func (w *Writer) Write(mb buf.MultiBuffer) error {
 	return w.writer.Write(mb2)
 }
 
+func (w *Writer) Write(mb buf.MultiBuffer) error {
+	const chunkSize = 8 * 1024
+	for !mb.IsEmpty() {
+		slice := mb.SliceBySize(chunkSize)
+		if err := w.writeInternal(slice); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (w *Writer) Close() {
 	meta := FrameMetadata{
 		SessionID:     w.id,

+ 15 - 0
common/buf/multi_buffer.go

@@ -86,3 +86,18 @@ func (mb MultiBuffer) ToNetBuffers() net.Buffers {
 	}
 	return bs
 }
+
+func (mb *MultiBuffer) SliceBySize(size int) MultiBuffer {
+	slice := NewMultiBuffer()
+	sliceSize := 0
+	endIndex := len(*mb)
+	for i, b := range *mb {
+		if b.Len()+sliceSize > size {
+			endIndex = i
+			break
+		}
+		slice.Append(b)
+	}
+	*mb = (*mb)[endIndex:]
+	return slice
+}