chore: use range over func in go1.23 for iteration cipherstream.Frame

nange · Nov 15, 2024 · 9ce91cd · 9ce91cd
1 parent 85abe77
commit 9ce91cd
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/cipherstream/cipherstream.go b/cipherstream/cipherstream.go
@@ -183,8 +183,7 @@ func (cs *CipherStream) Read(b []byte) (int, error) {
 	}
 
 	var frame *Frame
-	for {
-		frame = cs.frameIter.Next()
+	for frame = range cs.frameIter.Iter() {
 		if cs.frameIter.Error() != nil {
 			if timeout(cs.frameIter.Error()) {
 				return 0, errors.Join(cs.frameIter.Error(), ErrTimeout)
@@ -208,6 +207,9 @@ func (cs *CipherStream) Read(b []byte) (int, error) {
 		break
 	}
 
+	if frame == nil {
+		panic("frame must not be nil")
+	}
 	rawData := frame.RawDataPayload()
 	cn := copy(b, rawData)
 	if cn < len(rawData) {
@@ -218,8 +220,10 @@ func (cs *CipherStream) Read(b []byte) (int, error) {
 }
 
 func (cs *CipherStream) ReadFrame() (*Frame, error) {
-	frame := cs.frameIter.Next()
-	return frame, cs.frameIter.Error()
+	for frame := range cs.frameIter.Iter() {
+		return frame, cs.frameIter.Error()
+	}
+	return nil, cs.frameIter.Error()
 }
 
 func (cs *CipherStream) Release() {

diff --git a/cipherstream/frame.go b/cipherstream/frame.go
@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
+	"iter"
 	mr "math/rand"
 
 	"github.com/nange/easyss/v2/util"
@@ -303,6 +304,16 @@ func NewFrameIter(r io.Reader, cipher AEADCipher) *FrameIter {
 	}
 }
 
+func (fi *FrameIter) Iter() iter.Seq[*Frame] {
+	return func(yield func(*Frame) bool) {
+		for {
+			if !yield(fi.Next()) {
+				return
+			}
+		}
+	}
+}
+
 func (fi *FrameIter) Next() *Frame {
 	hBuf := fi.buf[:Http2HeaderLen+fi.cipher.NonceSize()+fi.cipher.Overhead()]
 	_, err := io.ReadFull(fi.r, hBuf)