Kaynağa Gözat

check msg size in writepacket

Darien Raymond 7 yıl önce
ebeveyn
işleme
b4f3106ad6
1 değiştirilmiş dosya ile 12 ekleme ve 6 silme
  1. 12 6
      common/crypto/auth.go

+ 12 - 6
common/crypto/auth.go

@@ -245,6 +245,11 @@ func (w *AuthenticationWriter) seal(b *buf.Buffer) (*buf.Buffer, error) {
 		paddingSize = int(w.padding.NextPaddingLen())
 	}
 
+	totalSize := encryptedSize + paddingSize
+	if totalSize > buf.Size {
+		return nil, newError("size too large: ", totalSize)
+	}
+
 	eb := buf.New()
 	common.Must(eb.Reset(func(bb []byte) (int, error) {
 		w.sizeParser.Encode(uint16(encryptedSize+paddingSize), bb[:0])
@@ -301,17 +306,17 @@ func (w *AuthenticationWriter) writePacket(mb buf.MultiBuffer) error {
 
 	mb2Write := buf.NewMultiBufferCap(int32(len(mb)) + 1)
 
-	for !mb.IsEmpty() {
-		b := mb.SplitFirst()
+	for _, b := range mb {
 		if b.IsEmpty() {
 			continue
 		}
+
 		eb, err := w.seal(b)
-		b.Release()
 		if err != nil {
-			mb2Write.Release()
-			return err
+			newError("failed to encrypt packet").Base(err).AtWarning().WriteToLog()
+			continue
 		}
+
 		mb2Write.Append(eb)
 	}
 
@@ -324,7 +329,8 @@ func (w *AuthenticationWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
 		b := buf.New()
 		defer b.Release()
 
-		eb, _ := w.seal(b)
+		eb, err := w.seal(b)
+		common.Must(err)
 		return w.writer.WriteMultiBuffer(buf.NewMultiBufferValue(eb))
 	}